diff --git a/zipline/monitor.py b/zipline/monitor.py index 450fc444..d5fea061 100644 --- a/zipline/monitor.py +++ b/zipline/monitor.py @@ -8,17 +8,16 @@ class Controller(object): polling = False debug = False - def __init__(self, pull_socket, pub_socket, context=None, logging = None): + def __init__(self, pull_socket, pub_socket, logging = None): + + self._ctx = None self.associated = [] - if not context: - self._ctx = zmq.Context() - else: - self._ctx = context - self.pull_socket = pull_socket + self.push_socket = pull_socket # same port self.pub_socket = pub_socket + self.sub_socket = pub_socket # same port if logging: self.logging = logging @@ -30,18 +29,26 @@ class Controller(object): self.success = 0 self.failed = 0 - def run(self, debug=False): + def run(self, debug=False, context=None): self.polling = True - #if debug: - return self._poll() - #else: - #return self._poll_fast() + if not context: + self._ctx = zmq.Context() + else: + self._ctx = context + + if debug: + return self._poll_fast() # the c loop + else: + return self._poll() # use a python loop def _poll_fast(self): """ C version of the polling forwarder. """ + self.pull = self._ctx.socket(zmq.PULL) + self.pub = self._ctx.socket(zmq.PUB) + zmq.device(zmq.FORWARDER, self.pull, self.pub) def _poll(self): @@ -60,7 +67,9 @@ class Controller(object): while self.polling: try: - self.pub.send(self.pull.recv()) + msg = self.pull.recv() + print msg + self.pub.send(msg) except KeyboardInterrupt: self.polling = False break @@ -75,24 +84,36 @@ class Controller(object): self.failed += 1 continue - def message_sender(self): + # ------------------- + # Hooks for Endpoints + # ------------------- + + def message_sender(self, context=None): """ Spin off a socket used for sending messages to this controller. """ - s = self._ctx.socket(zmq.PUSH) - s.connect(self.pull_socket) + + if not context: + context = zmq.Context() + + s = context.socket(zmq.PUSH) + s.connect(self.push_socket) self.associated.append(s) return s - def message_listener(self): + def message_listener(self, context = None, filters=None): """ Spin off a socket used for receiving messages from this controller. """ - s = self._ctx.socket(zmq.SUB) - s.connect(self.pub_socket) - s.setsockopt(zmq.SUBSCRIBE, '') + + if not context: + context = zmq.Context() + + s = context.socket(zmq.SUB) + s.connect(self.sub_socket) + s.setsockopt(zmq.SUBSCRIBE, filters or '') self.associated.append(s) return s @@ -116,3 +137,13 @@ class Controller(object): return return float(self.success) / (self.success + self.failed) +if __name__ == '__main__': + print 'Running on ',\ + 'tcp://127.0.0.1:5000', \ + 'tcp://127.0.0.1:5001', + + controller = Controller( + 'tcp://127.0.0.1:5000', + 'tcp://127.0.0.1:5001', + ) + controller.run() diff --git a/zipline/test/test_monitor.py b/zipline/test/test_monitor.py new file mode 100644 index 00000000..5835b779 --- /dev/null +++ b/zipline/test/test_monitor.py @@ -0,0 +1,77 @@ +""" +Dummy simulator backported from Qexec for development on Zipline. +""" + +import logging +from multiprocessing import Process +from unittest2 import TestCase + +from zipline.monitor import Controller + +import gevent +from gevent_zeromq import zmq + +ctx = zmq.Context() + +class TestControlProtocol(TestCase): + + def setUpController(self): + self.controller.run() + + def setUp(self): + self.controller = Controller( + 'tcp://127.0.0.1:5000', + 'tcp://127.0.0.1:5001', + ) + self.control_proc = Process(target=self.setUpController) + self.control_proc.start() + + def tearDown(self): + self.control_proc.terminate() + ctx.destroy() + + def asyncMessage(self, socket): + return socket.recv() + + def send_and_receive(self, push, sub, message='\x01'): + msg = gevent.spawn(sub.recv) + push.send(message) + gevent.sleep(0) # explicit gevent yield + msg.join() + self.assertEqual(msg.value, message) + + def test_control_message(self): + + sub = self.controller.message_listener(context=ctx) + message = gevent.spawn(self.asyncMessage, sub) + + push = self.controller.message_sender(context=ctx) + + # Don't like introducing time constants but because of + # the way gevent scheduler works zmq will often send all + # the messages off before the other thread even gets to + # listening. + self.send_and_receive(push, sub) + sub.close() + push.close() + + def test_control_delivery(self): + # Assert that the number of messages sent on the wire is + # the number of messages received, ie we don't drop any. + # This is of course depenendent on the topology of the + # listeners being fixed. Which normally it isn't. + + sub = self.controller.message_listener(context=ctx) + message = gevent.spawn(self.asyncMessage, sub) + + push = self.controller.message_sender(context=ctx) + + # Don't like introducing time constants but because of + # the way gevent scheduler works zmq will often send all + # the messages off before the other thread even gets to + # listening. + for i in xrange(25): + self.send_and_receive(push, sub) + + sub.close() + push.close() diff --git a/zipline/topos.py b/zipline/topos.py new file mode 100644 index 00000000..23e37ff7 --- /dev/null +++ b/zipline/topos.py @@ -0,0 +1,189 @@ +import uuid +import copy +import atexit +import pickle + +from datetime import datetime +from collections import defaultdict + +from UserDict import DictMixin + +class Snapshot(object, DictMixin): + """ + A snapshot in time of a history container. + """ + + def __init__(self, state, version, ts): + self.version = version + self.timestamp = ts + self._state = state + + def keys(self): + return self._state.keys() + + def values(self): + return self._state.values() + + def items(self): + return self._state.items() + + def __getitem__(self, key): + return self._state.__getitem__(key) + + def has_key(self, key): + return self._state.has_key(key) + + def copy(self): + return copy.copy(self._state) + +class History(object, DictMixin): + """ + A duck-typed dictionary that tracks its time evolution. + + Worth noting this not a particuarly high-performance + data structure due to the copious amount of copying going on. + """ + + def __init__(self, default=None): + if default: + initial = defaultdict(default) + else: + initial = {} + + self.version = 0 + self.changeset = [('CREATE', None)] + self.current = Snapshot(initial, version=self.version, ts=datetime.now()) + self._history = [self.current] + + def items(self, version=-1): + return self._history[version].items() + + def keys(self, version=-1): + return self._history[version].keys() + + def rollback(self, version): + pass + + def event(self, tup): + self.changeset.append(tup) + + def __getitem__(self, key, version=-1): + return self._history[version].__getitem__(key) + + def __setitem__(self, key, val): + if self.current.has_key(key): + self.changeset.append(('CHANGE', key)) + else: + self.changeset.append(('ADD', key)) + + state = self.current.copy() + state[key] = val + + self.version += 1 + self.current = Snapshot(state, self.version, datetime.now()) + self._history.append(self.current) + + def __delitem__(self, key): + self.changeset.append(('REMOVE', key)) + + state = self.current.copy() + del state[key] + + self.version += 1 + self.current = Snapshot(state, self.version, datetime.now()) + self._history.append(self.current) + + def history(self): + for change in self.changeset: + print change + + def __repr__(self): + return ':'.join(['historical', self.current._state.__repr__()]) + +SocketHistory = History() +ContextHistory = History() + +def patch_zmq(_zmq=None): + """ + Monkey patch zeromq to allow for socket tracking. + """ + if _zmq: + zmq = _zmq + else: + import zmq + + _Context = zmq.Context + _Socket = zmq.Socket + + class TrackedSocket(zmq.Socket): + + def __init__(self, context, socket_type): + self.context = context + self.uuid = str(uuid.uuid4()) + SocketHistory[self.uuid] = self + _Socket.__init__(self, context, socket_type) + + def connect(self, address): + SocketHistory.event(('CONNECT', self.uuid, address)) + _Socket.connect(self, address) + + def bind(self, address): + SocketHistory.event(('BIND', self.uuid, address)) + _Socket.bind(self, address) + + def close(self, *args, **kwargs): + del SocketHistory[self.uuid] + _Socket.close(self, *args, **kwargs) + + def setsockopt(self, option, optval): + if option == zmq.IDENTITY: + old = SocketHistory[self.uuid] + SocketHistory[optval] = old + del SocketHistory[self.uuid] + self.uuid = optval + + _Socket.setsockopt(self, option, optval) + + class TrackedContext(zmq.Context): + + def __init__(self, *args, **kwargs): + self.sockets = {} + _Context.__init__(self, *args, **kwargs) + self.uuid = str(uuid.uuid4()) + ContextHistory[self.uuid] = self + + def socket(self, socket_type): + sock = TrackedSocket(self, socket_type) + ContextHistory.event(('EMBED', self.uuid, sock.uuid)) + self.sockets[sock.uuid] = sock + return sock + + def name(self, name): + """ + Name the context. Is a superset of the vanilla pyzmq + API. + """ + old = ContextHistory[self.context.uuid] + ContextHistory[name] = old + del ContextHistory[self.context.uuid] + self.uuid = name + + def term(self, *args, **kwargs): + for uid, sock in self.sockets.iteritems(): + if not sock.closed: + del SocketHistory[sock.uuid] + del ContextHistory[self.uuid] + _Context.term(self, *args, **kwargs) + + def destroy(self, *args, **kwargs): + ContextHistory.event(('DESTROY', self.uuid)) + _Context.destroy(self, *args, **kwargs) + + zmq.Context = TrackedContext + zmq.Socket = TrackedSocket + return TrackedContext, TrackedSocket + +def track_to_file(f): + def write_track(): + pickle.dump(SocketHistory.changeset, file(f, 'wb+')) + atexit.register(write_track)