From c2bd2872c1e91ed4598b6a93f02b328ec3b61b9a Mon Sep 17 00:00:00 2001 From: Stephen Diehl Date: Tue, 3 Apr 2012 18:12:16 -0400 Subject: [PATCH] Concurrency agnostic controller. --- zipline/messaging.py | 1 + zipline/monitor.py | 185 ++++++++++++++++++++++--------------------- zipline/protocol.py | 24 +++--- zipline/simulator.py | 5 +- 4 files changed, 108 insertions(+), 107 deletions(-) diff --git a/zipline/messaging.py b/zipline/messaging.py index 552c1115..1bb80dcf 100644 --- a/zipline/messaging.py +++ b/zipline/messaging.py @@ -56,6 +56,7 @@ class ComponentHost(Component): raise Exception("There can be only one!") self.controller = controller + self.controller.zmq_flavor = self.zmq_flavor # Propogate the controller to all the subcomponents for component in self.components.itervalues(): diff --git a/zipline/monitor.py b/zipline/monitor.py index 2e0f2066..378e09e4 100644 --- a/zipline/monitor.py +++ b/zipline/monitor.py @@ -1,11 +1,19 @@ import time import gevent - -from gevent_zeromq import zmq +import itertools import util as qutil +# pyzmq +import zmq +# gevent_zeromq +import gevent_zeromq +# zmq_ctypes +#import zmq_ctypes + from protocol import CONTROL_PROTOCOL, CONTROL_FRAME, \ - CONTROL_UNFRAME, CONTROL_STATES + CONTROL_UNFRAME, CONTROL_STATES, INVALID_CONTROL_FRAME + +from gpoll import _Poller as GeventPoller # TODO: print -> qutil.LOGGER.info @@ -147,7 +155,10 @@ class Controller(object): def __init__(self, pub_socket, route_socket, logging = None): - self._ctx = None + self.context = None + self.zmq = None + self.zmq_poller = None + polling = False self.polling = polling @@ -171,6 +182,31 @@ class Controller(object): self.logging = False self.dologging = False + def init_zmq(self, flavor): + + assert self.zmq_flavor in ['thread', 'mp', 'green'] + + if flavor == 'mp': + self.zmq = zmq + self.context = self.zmq.Context() + self.zmq_poller = self.zmq.Poller + return + if flavor == 'thread': + self.zmq = zmq + self.context = self.zmq.Context.instance() + self.zmq_poller = self.zmq.Poller + return + if flavor == 'green': + self.zmq = gevent_zeromq.zmq + self.context = self.zmq.Context.instance() + self.zmq_poller = GeventPoller + return + if flavor == 'pypy': + self.zmq = zmq + self.context = self.zmq.Context.instance() + self.zmq_poller = self.zmq.Poller + return + def manage(self, topology, states=None, context=None): """ Give the controller a set set of components to manage and @@ -198,65 +234,75 @@ class Controller(object): # Start off in RUNNING, state self.state = self.states[0] - if not context: - self._ctx = zmq.Context.instance() - else: - self._ctx = context + def run(self): + self.init_zmq(self.zmq_flavor) - def run(self, flavor): - """ - Abstracts - """ - assert flavor in ['thread', 'mp', 'green'] - return self._poll() # use a python loop + try: + return self._poll() # use a python loop + except KeyboardInterrupt: + print 'Shutdown event loop' def log_status(self): print "[Controller] Tracking : %s" % ([c for c in self.tracked],) - gevent.spawn_later(5, self.log_status) + + # ----------- + # Event Loops + # ----------- def _poll(self): - self.pub = self._ctx.socket(zmq.PUB) + self.pub = self.context.socket(self.zmq.PUB) self.pub.bind(self.pub_socket) - self.router = self._ctx.socket(zmq.ROUTER) + self.router = self.context.socket(self.zmq.ROUTER) self.router.bind(self.route_socket) self.associated.extend([self.pub, self.router]) - # Spin the coroutines - gevent.spawn(self.log_status) - gevent.spawn_raw(self.recv_hearts_router) + poller = self.zmq.Poller() + poller.register(self.router, self.zmq.POLLIN) + + buffer = [] + + for i in itertools.count(0): + self.log_status() + self.responses = set() + + self.ctime = time.time() + self.pub.send(str(self.ctime)) + + while self.polling: + # Reset the responses for this cycle + + socks = dict(poller.poll(5)) + tic = time.time() + + if tic - self.ctime > 1: + break + + if self.router in socks and socks[self.router] == self.zmq.POLLIN: + rawmessage = self.router.recv() + + if rawmessage: + buffer.append(rawmessage) + + try: + if not self.router.getsockopt(self.zmq.RCVMORE): + self.handle_recv(buffer[:]) + buffer = [] + except INVALID_CONTROL_FRAME: + print 'Invalid frame', rawmessage + pass - while self.polling: self.beat() - gevent.sleep(self.period) - try: - msg = self.pull.recv() - self.pub.send(msg) - except KeyboardInterrupt: - self.polling = False - break - except zmq.ZMQError: - self.polling = False - break - except Exception as e: - # Its common to wrap these in wildcard exceptions so - # that we don't loose messages, ever - if self.logging: - self.logging.error(str(e)) - continue + if self.zmq_flavor == 'green': + gevent.sleep(0) # After loop exits self.terminated = True def beat(self): - toc = time.time() - - self.ctime += toc - self.tic - self.tic = toc - self.message = str(self.ctime) # These the set overloaded operations # A & B ~ set.intersection @@ -279,11 +325,6 @@ class Controller(object): for component in bad: self.fail(component) - # Reset the responses for this cycle - self.responses = set() - - gevent.spawn_raw(self.pub.send, str(self.ctime)) - # ------------------ # Component Handlers # ------------------ @@ -291,7 +332,7 @@ class Controller(object): # The various "states of being that a component can inform us # of def new(self, component): - print '[Controller] Tracking "%s" ' % component + print '[Controller] New Tracked "%s" ' % component if component in self.topology or self.freeform: self.tracked.add(component) @@ -313,36 +354,6 @@ class Controller(object): def exception(self, component, failure): pass - # -------- - # IO Loops - # -------- - - def recv_hearts(self): - buffer = [] - - while True: - message = self.router.recv(zmq.NOBLOCK, copy=False) - - if message: - buffer.append(message) - if not self.router.rcvmore(): - break - - gevent.spawn_later(0.001, self.handle_pong, buffer) - - def recv_hearts_router(self): - buffer = [] - - while True: - message = self.router.recv() - - if message: - buffer.append(message) - - if not self.router.getsockopt(zmq.RCVMORE): - gevent.spawn_raw(self.handle_recv, buffer[:]) - buffer = [] - # ----------------- # Protocol Handling # ----------------- @@ -354,7 +365,6 @@ class Controller(object): be coming over the wire. Which shouldn't happen ... right? """ identity = msg[0] - id, status = CONTROL_UNFRAME(msg[1]) # A component is telling us its alive: @@ -365,7 +375,7 @@ class Controller(object): else: # Otherwise its something weird and we don't know # what to do so just say so - print "Weird stuff happened: %s" % msg[1] + print "Weird stuff happened: %s" % status, self.ctime # A component is telling us it failed, and how if id is CONTROL_PROTOCOL.EXCEPTION: @@ -392,7 +402,7 @@ class Controller(object): """ if not context: - context = zmq.Context.instance() + context = self.zmq.Context.instance() s = context.socket(zmq.DEALER) s.connect(self.route_socket) @@ -408,7 +418,7 @@ class Controller(object): """ if not context: - context = zmq.Context.instance() + context = self.zmq.Context.instance() s = context.socket(zmq.SUB) s.connect(self.pub_socket) @@ -423,7 +433,7 @@ class Controller(object): assert hard or soft, """ Must specify kill hard or soft """ if not context: - context = zmq.Context.instance() + context = self.zmq.Context.instance() if hard: self.state = CONTROL_STATES.SHUTDOWN @@ -441,15 +451,6 @@ class Controller(object): #for asoc in self.associated: #asoc.close() - def destroy(self): - """ - Manual cleanup. - """ - self.shutdown() - - #def __del__(self): - #self.shutdown() - if __name__ == '__main__': print 'Running on ',\ diff --git a/zipline/protocol.py b/zipline/protocol.py index a1341179..4730e8ff 100644 --- a/zipline/protocol.py +++ b/zipline/protocol.py @@ -142,17 +142,17 @@ CONTROL_STATES = Enum( CONTROL_PROTOCOL = Enum( 'HEARTBEAT' , # 0 - req - 'SHUTDOWN' , # 3 - req - 'KILL' , # 4 - req + 'SHUTDOWN' , # 1 - req + 'KILL' , # 2 - req - 'OK' , # 5 - rep - 'DONE' , # 6 - rep - 'EXCEPTION' , # 7 - rep + 'OK' , # 3 - rep + 'DONE' , # 4 - rep + 'EXCEPTION' , # 5 - rep ) -def CONTROL_FRAME(id, status): - assert isinstance(id, basestring,) - assert isinstance(status, int) +def CONTROL_FRAME(event, payload): + assert isinstance(event, int,) + assert isinstance(payload, basestring) return msgpack.dumps(tuple([id, status])) @@ -163,11 +163,11 @@ def CONTROL_UNFRAME(msg): assert isinstance(msg, basestring) try: - id, status = msgpack.loads(msg) - assert isinstance(id, basestring) - assert isinstance(status, int) + event, payload = msgpack.loads(msg) + assert isinstance(event, int) + assert isinstance(payload, basestring) - return id, status + return event, payload except TypeError: raise INVALID_CONTROL_FRAME(msg) except ValueError: diff --git a/zipline/simulator.py b/zipline/simulator.py index 8a3abb27..c43f21fc 100644 --- a/zipline/simulator.py +++ b/zipline/simulator.py @@ -43,7 +43,6 @@ class Simulator(ComponentHost): def launch_controller(self): thread = threading.Thread( target=self.controller.run, - args=('thread',) ) thread.start() @@ -71,8 +70,8 @@ class Simulator(ComponentHost): if not self.running: return - #if self.controller: - #self.controller.shutdown() + if self.controller: + self.controller.shutdown() for component in self.components.itervalues(): component.shutdown()