class WebSocketZMQBridgeHandler(websocket.WebSocketHandler): def open(self, *args, **kwargs): self.currentMessage = [] self.__endpoint = 'tcp://localhost:224' socket = zmq.Context().socket(zmq.DEALER) self.__stream = ZMQStream(socket, IOLoop.current()) self.__stream.on_recv(self.__onReceive) self.__stream.socket.setsockopt(zmq.LINGER, 0) self.__stream.connect(self.__endpoint) def on_close(self, *args, **kwargs): self.__stream.close() def on_message(self, message): hasMore = message[0] self.currentMessage.append(message[1:]) if not hasMore: sendingMessage = self.currentMessage self.currentMessage = [] self.__stream.send_multipart(sendingMessage) def __onReceive(self, msg): for frame in msg[:-1]: self.write_message(b'\x01' + frame, binary=True) self.write_message(b'\x00' + msg[-1], binary=True)
class Subscriber(object): def __init__(self, context, sub_address, sub_topics): self.context = context self.subscriber_address = sub_address self.subscriber_topics = sub_topics socket = self.context.socket(zmq.SUB) ioloop = IOLoop.instance() self.subscriber = ZMQStream(socket, ioloop) self.subscriber.setsockopt(zmq.LINGER, 0) self.subscriber.on_recv(callback=self.subscriber_recv) self.subscriber.setsockopt(zmq.SUBSCRIBE, "") self.subscriber.connect(self.subscriber_address) return def shutdown(self): self.subscriber.on_recv(callback=None) self.subscriber.socket.disconnect(self.subscriber_address) self.subscriber = None return def subscriber_recv(self, msg): return
def launch_scheduler(in_addr, out_addr, mon_addr, not_addr, config=None, logname='ZMQ', log_addr=None, loglevel=logging.DEBUG, scheme='lru', identity=b'task'): from zmq.eventloop import ioloop from zmq.eventloop.zmqstream import ZMQStream if config: # unwrap dict back into Config config = Config(config) ctx = zmq.Context() loop = ioloop.IOLoop() ins = ZMQStream(ctx.socket(zmq.XREP), loop) ins.setsockopt(zmq.IDENTITY, identity) ins.bind(in_addr) outs = ZMQStream(ctx.socket(zmq.XREP), loop) outs.setsockopt(zmq.IDENTITY, identity) outs.bind(out_addr) mons = ZMQStream(ctx.socket(zmq.PUB), loop) mons.connect(mon_addr) nots = ZMQStream(ctx.socket(zmq.SUB), loop) nots.setsockopt(zmq.SUBSCRIBE, '') nots.connect(not_addr) scheme = globals().get(scheme, None) # setup logging if log_addr: connect_logger(logname, ctx, log_addr, root="scheduler", loglevel=loglevel) else: local_logger(logname, loglevel) scheduler = TaskScheduler(client_stream=ins, engine_stream=outs, mon_stream=mons, notifier_stream=nots, scheme=scheme, loop=loop, logname=logname, config=config) scheduler.start() try: loop.start() except KeyboardInterrupt: print("interrupted, exiting...", file=sys.__stderr__)
async def handle_stream(self, stream, address): def onReceive(msg): msg[-2] = b'JSON' doc = msgpack.unpackb(msg[-1], encoding='UTF-8') if doc.__contains__('ResponseID'): doc['ResponseID'] = int(doc['ResponseID']) msg[-1] = json.dumps(doc).encode('UTF-8') # print("ready to send to Arduino. ", msg) stream.write(b'\xAE\xAE\xAE\xAE') for frame in msg[:-1]: stream.write(bytes([len(frame) + 1]) + b'\x01' + frame) stream.write(bytes([len(msg[-1]) + 1]) + b'\x00' + msg[-1]) stream.write(b'\xAF\xAF\xAF\xAF') endpoint = 'tcp://localhost:224' socket = zmq.Context().socket(zmq.DEALER) zmqStream = ZMQStream(socket, IOLoop.current()) zmqStream.on_recv(onReceive) zmqStream.socket.setsockopt(zmq.LINGER, 0) zmqStream.connect(endpoint) def on_message(message): frames = [] position = 0 while True: if len(message) <= position: break size = message[position] if len(message) < position + size + 1: raise RuntimeError("Invalid Frame.") frames.append(message[position + 2:position + size + 1]) position += size + 1 if len(frames) != 7: raise RuntimeError("Bad Frame.") if len(frames[0]) != 0: raise RuntimeError("Bad Frame 0.") if frames[1] != b'IF1': raise RuntimeError("Bad Protocol.") if frames[3] != b'Broker' and frames[3] != b'Service' and frames[ 3] != b'Direct': raise RuntimeError("Bad Distributing Mode.") frames[-2] = b'Msgpack' doc = json.loads(frames[-1]) frames[-1] = msgpack.packb(doc) # print('send a message a ZMQ server: ', frames) zmqStream.send_multipart(frames) errorCount = 0 while True: try: await stream.read_until(b'\xAE\xAE\xAE\xAE') data = await stream.read_until(b'\xAF\xAF\xAF\xAF') data = data[:-4] on_message(data) except StreamClosedError: print('closed') break except RuntimeError as e: print(e)
class Worker(object): def __init__(endpoints): context = zmq.Context() socket = context.socket(zmq.REQ) ioloop = IOLoop.instance() self.endpoints = endpoints self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsocketopt(zmq.LINGER, 0) for endpoint in endpoints: self.stream.connect(endpoint) self._send_ready()
def launch_scheduler(in_addr, out_addr, mon_addr, not_addr, config=None,logname='ZMQ', log_addr=None, loglevel=logging.DEBUG, scheme='lru', identity=b'task'): from zmq.eventloop import ioloop from zmq.eventloop.zmqstream import ZMQStream if config: # unwrap dict back into Config config = Config(config) ctx = zmq.Context() loop = ioloop.IOLoop() ins = ZMQStream(ctx.socket(zmq.XREP),loop) ins.setsockopt(zmq.IDENTITY, identity) ins.bind(in_addr) outs = ZMQStream(ctx.socket(zmq.XREP),loop) outs.setsockopt(zmq.IDENTITY, identity) outs.bind(out_addr) mons = ZMQStream(ctx.socket(zmq.PUB),loop) mons.connect(mon_addr) nots = ZMQStream(ctx.socket(zmq.SUB),loop) nots.setsockopt(zmq.SUBSCRIBE, '') nots.connect(not_addr) scheme = globals().get(scheme, None) # setup logging if log_addr: connect_logger(logname, ctx, log_addr, root="scheduler", loglevel=loglevel) else: local_logger(logname, loglevel) scheduler = TaskScheduler(client_stream=ins, engine_stream=outs, mon_stream=mons, notifier_stream=nots, scheme=scheme, loop=loop, logname=logname, config=config) scheduler.start() try: loop.start() except KeyboardInterrupt: print ("interrupted, exiting...", file=sys.__stderr__)
class LocalRequestProxy: ''' This class is responsible for routing client requests coming from a particular server to the RouterPubSubProxy, which will route them to the workers. ''' def __init__(self, front_end_name, back_end_name, loop): ''' Initializes an instance of LocalRequestProxy @param front_end_name - name of the front end socket. It will be initialized with the Router socket. @param back_end_name - name of the back end socket. It will be initialized with the Dealer socket. @param loop - zmq IOLoop ''' self._loop = loop ctx = zmq.Context.instance() # Create the front end stream front_address = ZmqAddress(chan_name=front_end_name, transport=INPROC) self._front_end = ZMQStream(ctx.socket(zmq.ROUTER), io_loop=loop) self._front_end.setsockopt(zmq.ROUTER_MANDATORY, 1) self._front_end.bind(front_address.zmq_url()) # Create the back end stream back_address = ZmqAddress(chan_name=back_end_name) self._back_end = ZMQStream(ctx.socket(zmq.DEALER), io_loop=loop) self._back_end.connect(back_address.zmq_url()) def callback(from_name, to_name, zmq_stream, msgs): log.debug("Routing from {0} to {1} messages {2}" .format(from_name, to_name, msgs)) zmq_stream.send_multipart(msgs) zmq_stream.flush() self._front_end.on_recv(lambda msgs: callback(front_end_name, back_end_name, self._back_end, msgs)) self._back_end.on_recv(lambda msgs: callback(back_end_name, front_end_name, self._front_end, msgs))
def stream(self, sock_type, sock_addr, sock_bind, callback=None, subscribe=''): assert self.ctx is not None sock_addr = sock_addr % { 'port': random.randint(1024,65535), } s = ZMQStream( self.ctx.socket(sock_type)) if sock_type == zmq.SUB: s.setsockopt(zmq.SUBSCRIBE, subscribe) if sock_bind: s.bind(sock_addr) else: s.connect(sock_addr) if callback: s.on_recv(callback) return (s, sock_addr)
class MNClient(MN_object): """Class for the MN client side. Thin asynchronous encapsulation of a zmq.REQ socket. Provides a :func:`request` method with optional timeout. :param context: the ZeroMQ context to create the socket in. :type context: zmq.Context :param endpoint: the endpoint to connect to. :type endpoint: str :param service: the service the client should use :type service: str """ _proto_version = CLIENT_PROTO def __init__(self, context, endpoint, service): """Initialize the MNClient. """ self.context = context self.service = service self.endpoint = endpoint self.can_send = True self._proto_prefix = ['', CLIENT_PROTO, service] self._tmo = None self.timed_out = False self._create_stream() return def _create_stream(self): """Helper for opening a stream. """ socket = self.context.socket(zmq.DEALER) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) def shutdown(self): """Method to deactivate the client connection completely. Will delete the stream and the underlying socket. .. warning:: The instance MUST not be used after :func:`shutdown` has been called. :rtype: None """ if not self.stream: return self.stream.socket.close() self.stream.close() self.stream = None return def request(self, msg, timeout=None): """Send the given message. :param msg: message parts to send. :type msg: list of str :param timeout: time to wait in milliseconds. :type timeout: int :rtype None: """ if not self.can_send: raise InvalidStateError() if isinstance(msg, bytes): msg = [msg] # prepare full message # to_send = self._proto_prefix[:] # to_send.extend(msg) to_send = msg if self.stream.closed(): self._create_stream() # TODO check this self.stream.send_multipart(to_send) self.can_send = False if timeout: self._start_timeout(timeout) return def _on_timeout(self): """Helper called after timeout. """ self.timed_out = True self._tmo = None self.on_timeout() return def _start_timeout(self, timeout): """Helper for starting the timeout. :param timeout: the time to wait in milliseconds. :type timeout: int """ self._tmo = DelayedCallback(self._on_timeout, timeout) self._tmo.start() return def _on_message(self, msg): """Helper method called on message receive. :param msg: list of message parts. :type msg: list of str """ if self._tmo: # disable timout self._tmo.stop() self._tmo = None # setting state before invoking on_message, so we can request from there self.can_send = True self.on_message(msg) return def on_message(self, msg): """Public method called when a message arrived. .. note:: Does nothing. Should be overloaded! """ pass def on_timeout(self): """Public method called when a timeout occurred. .. note:: Does nothing. Should be overloaded! """ pass
class MNWorker(MN_object): """Class for the MN worker side. Thin encapsulation of a zmq.DEALER socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. :param context: the context to use for socket creation. :type context: zmq.Context :param endpoint: endpoint to connect to. :type endpoint: str :param service: the name of the service we support. :type service: byte-string """ _proto_version = b'MNPW01' # worker protocol version def __init__(self, context, endpoint, service, worker_type, address, protocols): """Initialize the MNWorker. """ self.context = context self.endpoint = endpoint self.service = service self.type = worker_type self.address = address self.protocols = protocols self.envelope = None self.HB_RETRIES = HB_RETRIES self.HB_INTERVAL = HB_INTERVAL self._data = {} self.stream = None self._tmo = None self.timed_out = False self.need_handshake = True self.connected = False self.ticker = None self._delayed_cb = None self._create_stream() _LOG.info("Worker initialized and can be found at '%s'" % endpoint) return def _create_stream(self): """Helper to create the socket and the stream. """ socket = self.context.socket(zmq.DEALER) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() self.ticker.start() return def _send_ready(self): """Helper method to prepare and send the workers READY message. """ _LOG.debug("Informing broker I am ready") ready_msg = [ b'', WORKER_PROTO, MSG_READY, self.service, self.type, self.address, self.protocols ] if self.stream.closed(): self.shutdown() self.stream.send_multipart(ready_msg) self.curr_retries = self.HB_RETRIES return def _tick(self): """Method called every HB_INTERVAL milliseconds. """ self.curr_retries -= 1 self.send_hb() if self.curr_retries >= 0: return # connection seems to be dead self.shutdown() # try to recreate it # self._delayed_cb = IOLoop.call_later(self._create_stream, 5000) # self._delayed_cb = IOLoop.add_timeout(self._create_stream, 5000) self._delayed_cb = DelayedCallback(self._create_stream, self.HB_INTERVAL) self._delayed_cb.start() return def send_hb(self): """Construct and send HB message to broker. """ _LOG.debug("Sending heartbeat") msg = [b'', WORKER_PROTO, MSG_HEARTBEAT] if self.stream.closed(): self.shutdown() self.stream.send_multipart(msg) return def shutdown(self): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. """ if self.ticker: self.ticker.stop() self.ticker = None if not self.stream: return self.stream.socket.close() self.stream.close() self.stream = None self.timed_out = False self.need_handshake = True self.connected = False return def reply(self, msg): """Send the given message. :param msg: full message to send. :type msg: can either be a byte-string or a list of byte-strings """ if self.need_handshake: raise ConnectionNotReadyError() to_send = self.envelope self.envelope = None if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) if self.stream.closed(): self.shutdown() self.stream.send_multipart(to_send) return def _on_message(self, msg): """Helper method called on message receive. :param msg: a list w/ the message parts :type msg: a list of byte-strings """ _LOG.debug("Received: %s." % msg) # 1st part is empty msg.pop(0) # 2nd part is protocol version proto = msg.pop(0) if proto != WORKER_PROTO: # ignore message from not supported protocol pass # 3rd part is message type msg_type = msg.pop(0) # XXX: hardcoded message types! # any message resets the retries counter self.need_handshake = False self.curr_retries = self.HB_RETRIES if msg_type == MSG_DISCONNECT: # disconnect _LOG.info("Broker wants us to disconnect.") self.curr_retries = 0 # reconnect will be triggered by hb timer elif msg_type == MSG_QUERY: # request # remaining parts are the user message _LOG.debug("Received new request: %s." % msg) envelope, msg = split_address(msg) envelope.append(b'') envelope = [b'', WORKER_PROTO, MSG_REPLY] + envelope # reply self.envelope = envelope self.on_request(msg) else: # invalid message # ignored _LOG.debug('ignoring message with invalid id') pass return def on_request(self, msg): """Public method called when a request arrived. :param msg: a list w/ the message parts :type msg: a list of byte-strings Must be overloaded to provide support for various services! """ pass
class MDPWorker(object): """Class for the MDP worker side. Thin encapsulation of a zmq.DEALER socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. """ _proto_version = b'MDPW01' # TODO: integrate that into API HB_INTERVAL = 1000 # in milliseconds HB_LIVENESS = 3 # HBs to miss before connection counts as dead def __init__(self, context, endpoint, service): """Initialize the MDPWorker. context is the zmq context to create the socket from. service is a byte-string with the service name. """ self.context = context self.endpoint = endpoint self.service = service self.stream = None self._tmo = None self.need_handshake = True self.ticker = None self._delayed_cb = None self._create_stream() return def _create_stream(self): """Helper to create the socket and the stream. """ socket = self.context.socket(zmq.DEALER) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() self.ticker.start() return def _send_ready(self): """Helper method to prepare and send the workers READY message. """ ready_msg = [b'', self._proto_version, b'\x01', self.service] self.stream.send_multipart(ready_msg) self.curr_liveness = self.HB_LIVENESS return def _tick(self): """Method called every HB_INTERVAL milliseconds. """ self.curr_liveness -= 1 ## print '%.3f tick - %d' % (time.time(), self.curr_liveness) self.send_hb() if self.curr_liveness >= 0: return ## print '%.3f lost connection' % time.time() # ouch, connection seems to be dead self.shutdown() # try to recreate it self._delayed_cb = DelayedCallback(self._create_stream, 5000) self._delayed_cb.start() return def send_hb(self): """Construct and send HB message to broker. """ msg = [b'', self._proto_version, b'\x04'] self.stream.send_multipart(msg) return def shutdown(self): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. """ if self.ticker: self.ticker.stop() self.ticker = None if not self.stream: return self.stream.socket.close() self.stream.close() self.stream = None self.timed_out = False self.need_handshake = True self.connected = False return def reply(self, msg): """Send the given message. msg can either be a byte-string or a list of byte-strings. """ ## if self.need_handshake: ## raise ConnectionNotReadyError() # prepare full message to_send = self.envelope self.envelope = None if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self.stream.send_multipart(to_send) return def _on_message(self, msg): """Helper method called on message receive. msg is a list w/ the message parts """ # 1st part is empty msg.pop(0) # 2nd part is protocol version # TODO: version check proto = msg.pop(0) # 3rd part is message type msg_type = msg.pop(0) # XXX: hardcoded message types! # any message resets the liveness counter self.need_handshake = False self.curr_liveness = self.HB_LIVENESS if msg_type == b'\x05': # disconnect self.curr_liveness = 0 # reconnect will be triggered by hb timer elif msg_type == b'\x02': # request # remaining parts are the user message envelope, msg = split_address(msg) envelope.append(b'') envelope = [b'', self._proto_version, b'\x03'] + envelope # REPLY self.envelope = envelope self.on_request(msg) else: # invalid message # ignored pass return def on_request(self, msg): """Public method called when a request arrived. Must be overloaded! """ pass
class MDPWorker(object): """Class for the MDP worker side. Thin encapsulation of a zmq.XREQ socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. """ _proto_version = b'MDPW01' # TODO: integrate that into API HB_INTERVAL = 1000 # in milliseconds HB_LIVENESS = 3 # HBs to miss before connection counts as dead def __init__(self, context, endpoint, service): """Initialize the MDPWorker. context is the zmq context to create the socket from. service is a byte-string with the service name. """ self.context = context self.endpoint = endpoint self.service = service self.stream = None self._tmo = None self.need_handshake = True self.ticker = None self._delayed_cb = None self._create_stream() return def _create_stream(self): """Helper to create the socket and the stream. """ socket = self.context.socket(zmq.XREQ) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() self.ticker.start() return def _send_ready(self): """Helper method to prepare and send the workers READY message. """ ready_msg = [ b'', self._proto_version, chr(1), self.service ] self.stream.send_multipart(ready_msg) self.curr_liveness = self.HB_LIVENESS return def _tick(self): """Method called every HB_INTERVAL milliseconds. """ self.curr_liveness -= 1 ## print '%.3f tick - %d' % (time.time(), self.curr_liveness) self.send_hb() if self.curr_liveness >= 0: return print '%.3f lost connection' % time.time() # ouch, connection seems to be dead self.shutdown() # try to recreate it self._delayed_cb = DelayedCallback(self._create_stream, 5000) self._delayed_cb.start() return def send_hb(self): """Construct and send HB message to broker. """ msg = [ b'', self._proto_version, chr(4) ] self.stream.send_multipart(msg) return def shutdown(self): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. """ if self.ticker: self.ticker.stop() self.ticker = None if not self.stream: return self.stream.socket.close() self.stream.close() self.stream = None self.timed_out = False self.need_handshake = True self.connected = False return def reply(self, msg): """Send the given message. msg can either be a byte-string or a list of byte-strings. """ ## if self.need_handshake: ## raise ConnectionNotReadyError() # prepare full message to_send = self.envelope self.envelope = None if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self.stream.send_multipart(to_send) return def _on_message(self, msg): """Helper method called on message receive. msg is a list w/ the message parts """ # 1st part is empty msg.pop(0) # 2nd part is protocol version # TODO: version check proto = msg.pop(0) # 3nd part is message type msg_type = msg.pop(0) # XXX: hardcoded message types! # any message resets the liveness counter self.need_handshake = False self.curr_liveness = self.HB_LIVENESS if msg_type == '\x05': # disconnect print ' DISC' self.curr_liveness = 0 # reconnect will be triggered by hb timer elif msg_type == '\x02': # request # remaining parts are the user message envelope, msg = split_address(msg) envelope.append(b'') envelope = [ b'', self._proto_version, '\x03'] + envelope # REPLY self.envelope = envelope self.on_request(msg) else: # invalid message # ignored pass return def on_request(self, msg): """Public method called when a request arrived. Must be overloaded! """ pass
class ZmqSubscriber(object): def __init__(self, moduleName, centralHost=SUBSCRIBER_OPT_DEFAULTS['centralHost'], context=None, centralPublishEndpoint=SUBSCRIBER_OPT_DEFAULTS['centralPublishEndpoint'], replay=None): self.moduleName = moduleName self.centralHost = centralHost if context is None: context = zmq.Context.instance() self.context = context self.centralPublishEndpoint = parseEndpoint(centralPublishEndpoint, defaultPort=DEFAULT_CENTRAL_PUBLISH_PORT, centralHost=self.centralHost) self.replayPaths = replay if self.replayPaths is None: self.replayPaths = [] self.handlers = {} self.counter = 0 self.deserializer = serializers.get_deserializer('json') self.stream = None @classmethod def addOptions(cls, parser, defaultModuleName): if not parser.has_option('--centralHost'): parser.add_option('--centralHost', default=SUBSCRIBER_OPT_DEFAULTS['centralHost'], help='Host where central runs [%default]') if not parser.has_option('--moduleName'): parser.add_option('--moduleName', default=defaultModuleName, help='Name to use for this module [%default]') if not parser.has_option('--centralPublishEndpoint'): parser.add_option('--centralPublishEndpoint', default=SUBSCRIBER_OPT_DEFAULTS['centralPublishEndpoint'], help='Endpoint where central publishes messages [%default]') if not parser.has_option('--replay'): parser.add_option('--replay', action='append', help='Replay specified message log (can specify multiple times), or use - to read from stdin') @classmethod def getOptionValues(cls, opts): result = {} for key in SUBSCRIBER_OPT_DEFAULTS.iterkeys(): val = getattr(opts, key, None) if val is not None: result[key] = val return result def start(self): sock = self.context.socket(zmq.SUB) self.stream = ZMQStream(sock) # causes problems with multiple instances #self.stream.setsockopt(zmq.IDENTITY, self.moduleName) self.stream.connect(self.centralPublishEndpoint) logging.info('zmq.subscriber: connected to central at %s', self.centralPublishEndpoint) self.stream.on_recv(self.routeMessages) def routeMessages(self, messages): for msg in messages: self.routeMessage(msg) def routeMessage(self, msg): colonIndex = msg.find(':') topic = msg[:(colonIndex + 1)] body = msg[(colonIndex + 1):] handled = 0 for topicPrefix, registry in self.handlers.iteritems(): if topic.startswith(topicPrefix): for handler in registry.itervalues(): handler(topic[:-1], body) handled = 1 return handled def subscribeRaw(self, topicPrefix, handler): topicRegistry = self.handlers.setdefault(topicPrefix, {}) if not topicRegistry: logging.info('zmq.subscriber: subscribe %s', topicPrefix) self.stream.setsockopt(zmq.SUBSCRIBE, topicPrefix) handlerId = (topicPrefix, self.counter) topicRegistry[self.counter] = handler self.counter += 1 return handlerId def subscribeJson(self, topicPrefix, handler): def jsonHandler(topicPrefix, body): return handler(topicPrefix, convertToDotDictRecurse(json.loads(body))) return self.subscribeRaw(topicPrefix, jsonHandler) def subscribeDjango(self, topicPrefix, handler): def djangoHandler(topicPrefix, body): obj = json.loads(body) dataText = json.dumps([obj['data']]) modelInstance = list(self.deserializer(dataText))[0] return handler(topicPrefix, modelInstance.object) return self.subscribeRaw(topicPrefix, djangoHandler) def unsubscribe(self, handlerId): topicPrefix, index = handlerId topicRegistry = self.handlers[topicPrefix] del topicRegistry[index] if not topicRegistry: logging.info('zmq.subscriber: unsubscribe %s', topicPrefix) self.stream.setsockopt(zmq.UNSUBSCRIBE, topicPrefix) def connect(self, endpoint): self.stream.connect(endpoint) def replay(self): numReplayed = 0 numHandled = 0 for replayPath in self.replayPaths: print '=== replaying messages from %s' % replayPath if replayPath == '-': replayFile = sys.stdin else: replayFile = open(replayPath, 'rb') stream = LogParser(replayFile) for rec in stream: numReplayed += 1 numHandled += self.routeMessage(rec.msg) if numReplayed % 10000 == 0: print 'replayed %d messages, %d handled' % (numReplayed, numHandled)
class ZmqPublisher(object): def __init__(self, moduleName, context=None, centralSubscribeEndpoint=PUBLISHER_OPT_DEFAULTS[ 'centralSubscribeEndpoint'], publishEndpoint=PUBLISHER_OPT_DEFAULTS['publishEndpoint'], heartbeatPeriodMsecs=PUBLISHER_OPT_DEFAULTS[ 'heartbeatPeriodMsecs'], highWaterMark=PUBLISHER_OPT_DEFAULTS['highWaterMark']): self.moduleName = moduleName if context is None: context = zmq.Context.instance() self.context = context self.centralSubscribeEndpoint = parseEndpoint( centralSubscribeEndpoint, defaultPort=DEFAULT_CENTRAL_SUBSCRIBE_PORT) self.publishEndpoint = parseEndpoint(publishEndpoint, defaultPort='random') self.heartbeatPeriodMsecs = heartbeatPeriodMsecs self.highWaterMark = highWaterMark self.pubStream = None self.heartbeatTimer = None self.serializer = serializers.get_serializer('json')() @classmethod def addOptions(cls, parser, defaultModuleName): if not parser.has_option('--moduleName'): parser.add_option('--moduleName', default=defaultModuleName, help='Name to use for this module [%default]') if not parser.has_option('--centralSubscribeEndpoint'): parser.add_option( '--centralSubcribeEndpoint', default=PUBLISHER_OPT_DEFAULTS['centralSubscribeEndpoint'], help='Endpoint where central listens for messages [%default]') if not parser.has_option('--publishEndpoint'): parser.add_option( '--publishEndpoint', default=PUBLISHER_OPT_DEFAULTS['publishEndpoint'], help='Endpoint to publish messages on [%default]') if not parser.has_option('--heartbeatPeriodMsecs'): parser.add_option( '--heartbeatPeriodMsecs', default=PUBLISHER_OPT_DEFAULTS['heartbeatPeriodMsecs'], type='int', help='Period for sending heartbeats to central [%default]') if not parser.has_option('--highWaterMark'): parser.add_option( '--highWaterMark', default=PUBLISHER_OPT_DEFAULTS['highWaterMark'], type='int', help= 'High-water mark for publish socket (see 0MQ docs) [%default]') @classmethod def getOptionValues(cls, opts): result = {} for key in PUBLISHER_OPT_DEFAULTS.iterkeys(): val = getattr(opts, key, None) if val is not None: result[key] = val return result def heartbeat(self): logging.debug('ZmqPublisher: heartbeat') self.sendJson('central.heartbeat.%s' % self.moduleName, { 'host': getShortHostName(), 'pub': self.publishEndpoint }) def sendRaw(self, topic, body): self.pubStream.send('%s:%s' % (topic, body)) def sendJson(self, topic, obj): if isinstance(obj, dict): obj.setdefault('module', self.moduleName) obj.setdefault('timestamp', str(getTimestamp())) self.sendRaw(topic, json.dumps(obj)) def sendDjango(self, modelInstance, topic=None, topicSuffix=None): dataText = self.serializer.serialize([modelInstance]) data = json.loads(dataText)[0] if topic is None: topic = data['model'].encode('utf-8') if topicSuffix is not None: topic += topicSuffix self.sendJson(topic, {'data': data}) def start(self): pubSocket = self.context.socket(zmq.PUB) self.pubStream = ZMQStream(pubSocket) # self.pubStream.setsockopt(zmq.IDENTITY, self.moduleName) self.pubStream.setsockopt(zmq.HWM, self.highWaterMark) self.pubStream.connect(self.centralSubscribeEndpoint) if self.publishEndpoint.endswith(':random'): endpointWithoutPort = re.sub(r':random$', '', self.publishEndpoint) port = self.pubStream.bind_to_random_port(endpointWithoutPort) self.publishEndpoint = '%s:%d' % (endpointWithoutPort, port) else: self.pubStream.bind(self.publishEndpoint) self.heartbeatTimer = ioloop.PeriodicCallback( self.heartbeat, self.heartbeatPeriodMsecs) self.heartbeatTimer.start() self.heartbeat()
class ZmqCentral(object): def __init__(self, opts): self.opts = opts self.info = {} def announceConnect(self, moduleName, params): logging.info('module %s connected', moduleName) self.injectStream.send('central.connect.%s:%s' % (moduleName, json.dumps(params))) def announceDisconnect(self, moduleName): logging.info('module %s disconnected', moduleName) self.injectStream.send( 'central.disconnect.%s:%s' % (moduleName, json.dumps({'timestamp': str(getTimestamp())}))) def logMessage(self, msg, posixTime=None, attachmentDir='-'): mlog = self.messageLog mlog.write('@@@ %d %d %s ' % (getTimestamp(posixTime), len(msg), attachmentDir)) mlog.write(msg) mlog.write('\n') def logMessageWithAttachments0(self, msg): parsed = parseMessage(msg) posixTime = time.time() # construct attachment directory dt = datetime.datetime.utcfromtimestamp(posixTime) dateText = dt.strftime('%Y-%m-%d') timeText = dt.strftime('%H-%M-%S') + '.%06d' % dt.microsecond uniq = '%08x' % random.getrandbits(32) attachmentSuffix = os.path.join('attachments', dateText, timeText, parsed['topic'], uniq) attachmentPath = os.path.join(self.logDir, attachmentSuffix) os.makedirs(attachmentPath) # write attachments to attachment directory for attachment in parsed['attachments']: fullName = os.path.join(attachmentPath, attachment.get_filename()) open(fullName, 'wb').write(attachment.get_payload()) # log message with a pointer to the attachment directory self.logMessage(':'.join((parsed['topic'], parsed['json'])), posixTime, attachmentSuffix) def logMessageWithAttachments(self, msg): try: return self.logMessageWithAttachments0(msg) except: # pylint: disable=W0702 self.logException('logging message with attachments') def handleHeartbeat(self, params): moduleName = params['module'].encode('utf-8') now = getTimestamp() oldInfo = self.info.get(moduleName, None) if oldInfo: if oldInfo.get('pub', None) != params.get('pub', None): self.announceDisconnect(moduleName) self.announceConnect(moduleName, params) else: self.announceConnect(moduleName, params) self.info[moduleName] = params keepalive = params.get('keepalive', DEFAULT_KEEPALIVE_US) params['timeout'] = now + keepalive return 'ok' def handleInfo(self): return self.info def logException(self, whileClause): errClass, errObject, errTB = sys.exc_info()[:3] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) logging.warning(''.join(traceback.format_tb(errTB))) logging.warning(errText) logging.warning('[error while %s at time %s]', whileClause, getTimestamp()) def handleMessages(self, messages): for msg in messages: if hasAttachments(msg): self.logMessageWithAttachments(msg) else: self.logMessage(msg) if msg.startswith('central.heartbeat.'): try: _topic, body = msg.split(':', 1) self.handleHeartbeat(json.loads(body)) except: # pylint: disable=W0702 self.logException('handling heartbeat') def handleRpcCall(self, messages): for msg in messages: try: call = json.loads(msg) callId = call['id'] except: # pylint: disable=W0702 self.rpcStream.send( json.dumps({ 'result': None, 'error': 'malformed request' })) try: method = call['method'] _params = call['params'] if method == 'info': result = self.handleInfo() else: raise ValueError('unknown method %s' % method) self.rpcStream.send( json.dumps({ 'result': result, 'error': None, 'id': callId })) except: # pylint: disable=W0702 self.logException('handling rpc message') errClass, errObject = sys.exc_info()[:2] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) self.rpcStream.send( json.dumps({ 'result': None, 'error': errText, 'id': callId })) def handleDisconnectTimer(self): now = getTimestamp() disconnectModules = [] for moduleName, entry in self.info.iteritems(): timeout = entry.get('timeout', None) if timeout is not None and now > timeout: disconnectModules.append(moduleName) for moduleName in disconnectModules: self.announceDisconnect(moduleName) del self.info[moduleName] def readyLog(self, pathTemplate, timestamp): if '%s' in pathTemplate: timeText = timestamp.strftime('%Y-%m-%d-%H-%M-%S') logFile = pathTemplate % timeText else: logFile = pathTemplate if not os.path.exists(self.logDir): os.makedirs(self.logDir) logPath = os.path.join(self.logDir, logFile) if '%s' in pathTemplate: latestPath = os.path.join(self.logDir, pathTemplate % 'latest') if os.path.islink(latestPath): os.unlink(latestPath) os.symlink(logFile, latestPath) return logPath def start(self): # open log files now = datetime.datetime.utcnow() self.logDir = os.path.abspath(self.opts.logDir) self.messageLogPath = self.readyLog(self.opts.messageLog, now) self.messageLog = open(self.messageLogPath, 'a') self.consoleLogPath = self.readyLog(self.opts.consoleLog, now) rootLogger = logging.getLogger() rootLogger.setLevel(logging.DEBUG) fmt = logging.Formatter('%(asctime)s - %(levelname)-7s - %(message)s') fmt.converter = time.gmtime fh = logging.FileHandler(self.consoleLogPath) fh.setFormatter(fmt) fh.setLevel(logging.DEBUG) rootLogger.addHandler(fh) if self.opts.foreground: ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) ch.setFormatter(fmt) rootLogger.addHandler(ch) # daemonize if self.opts.foreground: logging.info('staying in foreground') else: logging.info('daemonizing') pid = os.fork() if pid != 0: os._exit(0) os.setsid() pid = os.fork() if pid != 0: os._exit(0) os.chdir('/') os.close(1) os.close(2) nullFd = os.open('/dev/null', os.O_RDWR) os.dup2(nullFd, 1) os.dup2(nullFd, 2) try: # set up zmq self.context = zmq.Context.instance() self.rpcStream = ZMQStream(self.context.socket(zmq.REP)) self.rpcStream.bind(self.opts.rpcEndpoint) self.rpcStream.on_recv(self.handleRpcCall) self.forwarder = ThreadDevice(zmq.FORWARDER, zmq.SUB, zmq.PUB) self.forwarder.setsockopt_in(zmq.IDENTITY, THIS_MODULE) self.forwarder.setsockopt_out(zmq.IDENTITY, THIS_MODULE) self.forwarder.setsockopt_in(zmq.SUBSCRIBE, '') self.forwarder.setsockopt_out(zmq.HWM, self.opts.highWaterMark) self.forwarder.bind_in(self.opts.subscribeEndpoint) self.forwarder.bind_in(INJECT_ENDPOINT) self.forwarder.bind_out(self.opts.publishEndpoint) self.forwarder.bind_out(MONITOR_ENDPOINT) for entry in self.opts.subscribeTo: try: moduleName, endpoint = entry.split('@') endpoint = parseEndpoint(endpoint) except ValueError: raise ValueError( '--subscribeTo argument "%s" is not in the format "<moduleName>@<endpoint>"' % entry) self.forwarder.connect_in(endpoint) self.info[moduleName] = {'module': moduleName, 'pub': endpoint} self.forwarder.start() time.sleep(0.1) # wait for forwarder to bind sockets self.monStream = ZMQStream(self.context.socket(zmq.SUB)) self.monStream.setsockopt(zmq.SUBSCRIBE, '') self.monStream.connect(MONITOR_ENDPOINT) self.monStream.on_recv(self.handleMessages) self.injectStream = ZMQStream(self.context.socket(zmq.PUB)) self.injectStream.connect(INJECT_ENDPOINT) self.disconnectTimer = ioloop.PeriodicCallback( self.handleDisconnectTimer, 5000) self.disconnectTimer.start() except: # pylint: disable=W0702 errClass, errObject, errTB = sys.exc_info()[:3] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) logging.error(''.join(traceback.format_tb(errTB))) logging.error(errText) logging.error('[error during startup -- exiting]') sys.exit(1) def shutdown(self): self.messageLog.flush()
class MDPWorker(object): """Class for the MDP worker side. Thin encapsulation of a zmq.DEALER socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. """ _proto_version = W_WORKER # TODO: integrate that into API HB_INTERVAL = 2500 # in milliseconds HB_LIVENESS = 5 # HBs to miss before connection counts as dead RECONNECT_PERIOD = 5000 def __init__(self, context, endpoint, hb_endpoint, service, endpoint_callback=None): """Initialize the MDPWorker. context is the zmq context to create the socket from. service is a byte-string with the service name. """ self.context = context self.endpoint = endpoint self.hb_endpoint = hb_endpoint self.service = service self.endpoint_callback = endpoint_callback self.stream = None self.hb_stream = None self.ticker = None self._delayed_reconnect = None self._unique_id = '' self._create_stream() def _create_stream(self): """Helper to create the socket and the stream. """ logging.debug('Worker creating stream') ioloop = IOLoop.instance() socket = self.context.socket(zmq.DEALER) self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) socket = self.context.socket(zmq.DEALER) self.hb_stream = ZMQStream(socket, ioloop) self.hb_stream.on_recv(self._on_message) self.hb_stream.socket.setsockopt(zmq.LINGER, 0) self.hb_stream.connect(self.hb_endpoint) self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() self.ticker.start() def _send_ready(self): """Helper method to prepare and send the workers READY message. """ ready_msg = [EMPTY_FRAME, self._proto_version, W_READY, self.service] self.stream.send_multipart(ready_msg) self.curr_liveness = self.HB_LIVENESS def _tick(self): """Method called every HB_INTERVAL milliseconds. """ self.curr_liveness -= 1 logging.debug('Worker HB tick, current liveness: %d' % self.curr_liveness) self.send_hb() if self.curr_liveness >= 0: return # # Ouch, connection seems to be dead # logging.debug('Worker lost connection') self.shutdown() # # try to recreate the connection # self._delayed_reconnect = DelayedCallback(self._recreate_stream, self.RECONNECT_PERIOD) self._delayed_reconnect.start() def _recreate_stream(self): logging.debug('Worker trying to recreate stream') if self.endpoint_callback is not None: # # Check, maybe the ip of the proxy changed. # try: self.endpoint, self.hb_endpoint = self.endpoint_callback() except: # # Probably some problem in accessing the server. # self._delayed_reconnect = DelayedCallback( self._recreate_stream, self.RECONNECT_PERIOD) self._delayed_reconnect.start() return self._create_stream() def send_hb(self): """Construct and send HB message to broker. """ msg = [EMPTY_FRAME, self._proto_version, W_HEARTBEAT, self._unique_id] self.hb_stream.send_multipart(msg) def shutdown(self): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. """ logging.debug('Shutdown of the worker') if self.ticker: logging.debug('Stopping the HB ticker') self.ticker.stop() self.ticker = None if not self.stream: return logging.debug('Closing the stream') self.stream.socket.close() self.stream.close() self.stream = None self.hb_stream.socket.close() self.hb_stream.close() self.hb_stream = None self.timed_out = False self.connected = False def reply(self, msg): """Send the given message. msg can either be a byte-string or a list of byte-strings. """ # # prepare full message # to_send = self.envelope self.envelope = None if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self.stream.send_multipart(to_send) def _on_message(self, msg): """Helper method called on message receive. msg is a list w/ the message parts """ logging.debug('Received message: {}'.format(msg)) # # 1st part is empty # msg.pop(0) # # 2nd part is protocol version # TODO: version check # proto = msg.pop(0) # # 3rd part is message type # msg_type = msg.pop(0) # # XXX: hardcoded message types! # any message resets the liveness counter # self.curr_liveness = self.HB_LIVENESS if msg_type == W_DISCONNECT: # # Disconnect. Reconnection will be triggered by hb timer # self.curr_liveness = 0 elif msg_type == W_READY: # # The message contains the unique id attached to the worker. # if len(msg) > 0: # # This above check is used for supporting older version of # the code. # self._unique_id = msg[0] elif msg_type == W_REQUEST: # # Request. Remaining parts are the user message # envelope, msg = split_address(msg) envelope.append(EMPTY_FRAME) envelope = [EMPTY_FRAME, self._proto_version, W_REPLY] + envelope self.envelope = envelope self.on_request(msg) else: # # invalid message # ignored # pass def on_request(self, msg): """Public method called when a request arrived. Must be overloaded! """ raise NotImplementedError( 'on_request must be implemented by the subclass.')
class ZmqSubscriber(object): def __init__(self, moduleName, context=None, centralPublishEndpoint=SUBSCRIBER_OPT_DEFAULTS['centralPublishEndpoint']): self.moduleName = moduleName if context is None: context = zmq.Context.instance() self.context = context self.centralPublishEndpoint = parseEndpoint(centralPublishEndpoint, defaultPort=DEFAULT_CENTRAL_PUBLISH_PORT) self.handlers = {} self.counter = 0 self.deserializer = serializers.get_deserializer('json') @classmethod def addOptions(cls, parser, defaultModuleName): if not parser.has_option('--moduleName'): parser.add_option('--moduleName', default=defaultModuleName, help='Name to use for this module [%default]') if not parser.has_option('--centralPublishEndpoint'): parser.add_option('--centralPublishEndpoint', default=SUBSCRIBER_OPT_DEFAULTS['centralPublishEndpoint'], help='Endpoint where central publishes messages [%default]') @classmethod def getOptionValues(cls, opts): result = {} for key in SUBSCRIBER_OPT_DEFAULTS.iterkeys(): val = getattr(opts, key, None) if val is not None: result[key] = val return result def start(self): sock = self.context.socket(zmq.SUB) self.stream = ZMQStream(sock) self.stream.setsockopt(zmq.IDENTITY, self.moduleName) self.stream.connect(self.centralPublishEndpoint) self.stream.on_recv(self.routeMessage) def routeMessage(self, messages): for msg in messages: colonIndex = msg.find(':') topic = msg[:(colonIndex + 1)] body = msg[(colonIndex + 1):] for topicPrefix, registry in self.handlers.iteritems(): if topic.startswith(topicPrefix): topicRegistry = registry break for handler in topicRegistry.itervalues(): handler(topic[:-1], body) def subscribeRaw(self, topic, handler): topicRegistry = self.handlers.setdefault(topic, {}) if not topicRegistry: self.stream.setsockopt(zmq.SUBSCRIBE, topic) handlerId = (topic, self.counter) topicRegistry[self.counter] = handler self.counter += 1 return handlerId def subscribeJson(self, topic, handler): def jsonHandler(topic, body): return handler(topic, json.loads(body)) return self.subscribeRaw(topic, jsonHandler) def subscribeDjango(self, topic, handler): def djangoHandler(topic, body): obj = json.loads(body) dataText = json.dumps([obj['data']]) modelInstance = self.deserializer(dataText)[0] return handler(topic, modelInstance) return self.subscribeRaw(topic, djangoHandler) def unsubscribe(self, handlerId): topic, index = handlerId topicRegistry = self.handlers[topic] del topicRegistry[index] if not topicRegistry: self.stream.setsockopt(zmq.UNSUBSCRIBE, topic) def connect(self, endpoint): self.stream.connect(endpoint)
class ZmqCentral(object): def __init__(self, opts): self.opts = opts self.info = {} def announceConnect(self, moduleName, params): logging.info('module %s connected', moduleName) self.injectStream.send('central.connect.%s:%s' % (moduleName, json.dumps(params))) def announceDisconnect(self, moduleName): logging.info('module %s disconnected', moduleName) self.injectStream.send('central.disconnect.%s:%s' % (moduleName, json.dumps({'timestamp': str(getTimestamp())}))) def logMessage(self, msg, posixTime=None, attachmentDir='-'): mlog = self.messageLog mlog.write('@@@ %d %d %s ' % (getTimestamp(posixTime), len(msg), attachmentDir)) mlog.write(msg) mlog.write('\n') def logMessageWithAttachments0(self, msg): parsed = parseMessage(msg) posixTime = time.time() # construct attachment directory dt = datetime.datetime.utcfromtimestamp(posixTime) dateText = dt.strftime('%Y-%m-%d') timeText = dt.strftime('%H-%M-%S') + '.%06d' % dt.microsecond uniq = '%08x' % random.getrandbits(32) attachmentSuffix = os.path.join('attachments', dateText, timeText, parsed['topic'], uniq) attachmentPath = os.path.join(self.logDir, attachmentSuffix) os.makedirs(attachmentPath) # write attachments to attachment directory for attachment in parsed['attachments']: fullName = os.path.join(attachmentPath, attachment.get_filename()) open(fullName, 'wb').write(attachment.get_payload()) # log message with a pointer to the attachment directory self.logMessage(':'.join((parsed['topic'], parsed['json'])), posixTime, attachmentSuffix) def logMessageWithAttachments(self, msg): try: return self.logMessageWithAttachments0(msg) except: # pylint: disable=W0702 self.logException('logging message with attachments') def handleHeartbeat(self, params): moduleName = params['module'].encode('utf-8') now = getTimestamp() oldInfo = self.info.get(moduleName, None) if oldInfo: if oldInfo.get('pub', None) != params.get('pub', None): self.announceDisconnect(moduleName) self.announceConnect(moduleName, params) else: self.announceConnect(moduleName, params) self.info[moduleName] = params keepalive = params.get('keepalive', DEFAULT_KEEPALIVE_US) params['timeout'] = now + keepalive return 'ok' def handleInfo(self): return self.info def logException(self, whileClause): errClass, errObject, errTB = sys.exc_info()[:3] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) logging.warning(''.join(traceback.format_tb(errTB))) logging.warning(errText) logging.warning('[error while %s at time %s]', whileClause, getTimestamp()) def handleMessages(self, messages): for msg in messages: if hasAttachments(msg): self.logMessageWithAttachments(msg) else: self.logMessage(msg) if msg.startswith('central.heartbeat.'): try: _topic, body = msg.split(':', 1) self.handleHeartbeat(json.loads(body)) except: # pylint: disable=W0702 self.logException('handling heartbeat') def handleRpcCall(self, messages): for msg in messages: try: call = json.loads(msg) callId = call['id'] except: # pylint: disable=W0702 self.rpcStream.send(json.dumps({'result': None, 'error': 'malformed request'})) try: method = call['method'] _params = call['params'] if method == 'info': result = self.handleInfo() else: raise ValueError('unknown method %s' % method) self.rpcStream.send(json.dumps({'result': result, 'error': None, 'id': callId})) except: # pylint: disable=W0702 self.logException('handling rpc message') errClass, errObject = sys.exc_info()[:2] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) self.rpcStream.send(json.dumps({'result': None, 'error': errText, 'id': callId})) def handleDisconnectTimer(self): now = getTimestamp() disconnectModules = [] for moduleName, entry in self.info.iteritems(): timeout = entry.get('timeout', None) if timeout is not None and now > timeout: disconnectModules.append(moduleName) for moduleName in disconnectModules: self.announceDisconnect(moduleName) del self.info[moduleName] def readyLog(self, pathTemplate, timestamp): if '%s' in pathTemplate: timeText = timestamp.strftime('%Y-%m-%d-%H-%M-%S') logFile = pathTemplate % timeText else: logFile = pathTemplate if not os.path.exists(self.logDir): os.makedirs(self.logDir) logPath = os.path.join(self.logDir, logFile) if '%s' in pathTemplate: latestPath = os.path.join(self.logDir, pathTemplate % 'latest') if os.path.islink(latestPath): os.unlink(latestPath) os.symlink(logFile, latestPath) return logPath def start(self): # open log files now = datetime.datetime.utcnow() self.logDir = os.path.abspath(self.opts.logDir) self.messageLogPath = self.readyLog(self.opts.messageLog, now) self.messageLog = open(self.messageLogPath, 'a') self.consoleLogPath = self.readyLog(self.opts.consoleLog, now) rootLogger = logging.getLogger() rootLogger.setLevel(logging.DEBUG) fmt = logging.Formatter('%(asctime)s - %(levelname)-7s - %(message)s') fmt.converter = time.gmtime fh = logging.FileHandler(self.consoleLogPath) fh.setFormatter(fmt) fh.setLevel(logging.DEBUG) rootLogger.addHandler(fh) if self.opts.foreground: ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) ch.setFormatter(fmt) rootLogger.addHandler(ch) # daemonize if self.opts.foreground: logging.info('staying in foreground') else: logging.info('daemonizing') pid = os.fork() if pid != 0: os._exit(0) os.setsid() pid = os.fork() if pid != 0: os._exit(0) os.chdir('/') os.close(1) os.close(2) nullFd = os.open('/dev/null', os.O_RDWR) os.dup2(nullFd, 1) os.dup2(nullFd, 2) try: # set up zmq self.context = zmq.Context.instance() self.rpcStream = ZMQStream(self.context.socket(zmq.REP)) self.rpcStream.bind(self.opts.rpcEndpoint) self.rpcStream.on_recv(self.handleRpcCall) self.forwarder = ThreadDevice(zmq.FORWARDER, zmq.SUB, zmq.PUB) self.forwarder.setsockopt_in(zmq.IDENTITY, THIS_MODULE) self.forwarder.setsockopt_out(zmq.IDENTITY, THIS_MODULE) self.forwarder.setsockopt_in(zmq.SUBSCRIBE, '') self.forwarder.setsockopt_out(zmq.HWM, self.opts.highWaterMark) self.forwarder.bind_in(self.opts.subscribeEndpoint) self.forwarder.bind_in(INJECT_ENDPOINT) self.forwarder.bind_out(self.opts.publishEndpoint) self.forwarder.bind_out(MONITOR_ENDPOINT) for entry in self.opts.subscribeTo: try: moduleName, endpoint = entry.split('@') endpoint = parseEndpoint(endpoint) except ValueError: raise ValueError('--subscribeTo argument "%s" is not in the format "<moduleName>@<endpoint>"' % entry) self.forwarder.connect_in(endpoint) self.info[moduleName] = {'module': moduleName, 'pub': endpoint} self.forwarder.start() time.sleep(0.1) # wait for forwarder to bind sockets self.monStream = ZMQStream(self.context.socket(zmq.SUB)) self.monStream.setsockopt(zmq.SUBSCRIBE, '') self.monStream.connect(MONITOR_ENDPOINT) self.monStream.on_recv(self.handleMessages) self.injectStream = ZMQStream(self.context.socket(zmq.PUB)) self.injectStream.connect(INJECT_ENDPOINT) self.disconnectTimer = ioloop.PeriodicCallback(self.handleDisconnectTimer, 5000) self.disconnectTimer.start() except: # pylint: disable=W0702 errClass, errObject, errTB = sys.exc_info()[:3] errText = '%s.%s: %s' % (errClass.__module__, errClass.__name__, str(errObject)) logging.error(''.join(traceback.format_tb(errTB))) logging.error(errText) logging.error('[error during startup -- exiting]') sys.exit(1) def shutdown(self): self.messageLog.flush()
def init_hub(self): """construct""" client_iface = "%s://%s:" % (self.client_transport, self.client_ip) + "%i" engine_iface = "%s://%s:" % (self.engine_transport, self.engine_ip) + "%i" ctx = self.context loop = self.loop # Registrar socket q = ZMQStream(ctx.socket(zmq.ROUTER), loop) q.bind(client_iface % self.regport) self.log.info("Hub listening on %s for registration.", client_iface % self.regport) if self.client_ip != self.engine_ip: q.bind(engine_iface % self.regport) self.log.info("Hub listening on %s for registration.", engine_iface % self.regport) ### Engine connections ### # heartbeat hpub = ctx.socket(zmq.PUB) hpub.bind(engine_iface % self.hb[0]) hrep = ctx.socket(zmq.ROUTER) hrep.bind(engine_iface % self.hb[1]) self.heartmonitor = HeartMonitor(loop=loop, config=self.config, log=self.log, pingstream=ZMQStream(hpub,loop), pongstream=ZMQStream(hrep,loop) ) ### Client connections ### # Notifier socket n = ZMQStream(ctx.socket(zmq.PUB), loop) n.bind(client_iface%self.notifier_port) ### build and launch the queues ### # monitor socket sub = ctx.socket(zmq.SUB) sub.setsockopt(zmq.SUBSCRIBE, b"") sub.bind(self.monitor_url) sub.bind('inproc://monitor') sub = ZMQStream(sub, loop) # connect the db self.log.info('Hub using DB backend: %r'%(self.db_class.split()[-1])) # cdir = self.config.Global.cluster_dir self.db = import_item(str(self.db_class))(session=self.session.session, config=self.config, log=self.log) time.sleep(.25) try: scheme = self.config.TaskScheduler.scheme_name except AttributeError: from .scheduler import TaskScheduler scheme = TaskScheduler.scheme_name.get_default_value() # build connection dicts self.engine_info = { 'control' : engine_iface%self.control[1], 'mux': engine_iface%self.mux[1], 'heartbeat': (engine_iface%self.hb[0], engine_iface%self.hb[1]), 'task' : engine_iface%self.task[1], 'iopub' : engine_iface%self.iopub[1], # 'monitor' : engine_iface%self.mon_port, } self.client_info = { 'control' : client_iface%self.control[0], 'mux': client_iface%self.mux[0], 'task' : (scheme, client_iface%self.task[0]), 'iopub' : client_iface%self.iopub[0], 'notification': client_iface%self.notifier_port } self.log.debug("Hub engine addrs: %s", self.engine_info) self.log.debug("Hub client addrs: %s", self.client_info) # resubmit stream r = ZMQStream(ctx.socket(zmq.DEALER), loop) url = util.disambiguate_url(self.client_info['task'][-1]) r.setsockopt(zmq.IDENTITY, self.session.bsession) r.connect(url) self.hub = Hub(loop=loop, session=self.session, monitor=sub, heartmonitor=self.heartmonitor, query=q, notifier=n, resubmit=r, db=self.db, engine_info=self.engine_info, client_info=self.client_info, log=self.log)
class UniWorker(object): """ Implementation of "simple" ZeroMQ Paranoid Pirate communication scheme. This class is the DEALER, and performs the "reply" in RPC calls. By design, only supports one remote client (ROUTER) in order to keep example simple. Supports a very basic RPC interface, using MessagePack for encoding/decoding. """ __metaclass__ = ABCMeta def __init__(self, endpoint, context=None): # type: (str, zmq.Context) -> None """ Initialize the worker. :param endpoint: ZeroMQ endpoint to connect to. :param context: ZeroMQ Context """ self._context = context or zmq.Context.instance() self._endpoint = endpoint self._stream = None # type: Optional[ZMQStream] self._tmo = None self._need_handshake = True self._ticker = None # type: Optional[PeriodicCallback] self._delayed_cb = None self._connected_event = Event() self._lock = Lock() self._create_stream() self._curr_liveness = HB_LIVENESS self._keep_running = True def _create_stream(self): # type: () -> None """ Helper function to create the ZMQ stream, configure callbacks. """ self.on_log_event("uniworker.connect", "Trying to connect to client") socket = self._context.socket(zmq.DEALER) self._stream = ZMQStream(socket, IOLoop()) self._stream.on_recv(self._on_message) self._stream.socket.setsockopt(zmq.LINGER, 0) self._stream.connect(self._endpoint) self._ticker = PeriodicCallback(self._tick, HB_INTERVAL) self._send_ready() self._ticker.start() def run(self): # type: () -> None """ Start the IOLoop, a blocking call to send/recv ZMQ messsages until the IOLoop is stopped. Note: The name of this function needs to stay the same so UniWorkerThread's run() is overridden with this function. """ if self._keep_running: self._stream.io_loop.start() def stop(self): # type: () -> None """ Stop the IOLoop. """ with self._lock: self._keep_running = False if self._stream is not None: self._stream.io_loop.stop() else: logger.warning("Can't stop worker-has shutdown() been called?") def shutdown(self): # type: () -> None """ Close the stream/socket. This should be called with the final flag when closing the connection for the last time. """ with self._lock: if self._ticker: self._ticker.stop() self._ticker = None if not self._stream: return self._stream.on_recv(None) self._send_disconnect() self._stream.close() self._stream = None self._need_handshake = True def wait_for_client(self, timeout): # type: (float) -> None """ Wait for the worker to establish a connection with the remote client. Will return immediately if already connected. Typically, the worker provides a service/responds to requests, so this is really only used for unit testing. :param timeout: Max time, in seconds, to wait for the connection to establish. """ event_status = self._connected_event.wait(timeout) if not event_status: raise LostRemoteError("No worker is connected.") def is_connected(self): # type: () -> bool """ Returns whether worker is connected to a client. :return: A boolean flag to indicate whether a connection to a client is established. """ return not self._need_handshake def send_reply(self, msg, partial=False, exception=False): # type: (Any, bool, bool) -> None """ Send a ZeroMQ message in reply to a client request. This should only be called out of the overridden do_work method. :param msg: The message to be sent out. :param partial: Flag indicating whether the response is a partial or final ZMQ message. """ msg = msgpack.Packer(default=XeroSerializer.encoder).pack(msg) if exception: to_send = [WORKER_EXCEPTION] elif partial: to_send = [WORKER_PARTIAL_REPLY] else: to_send = [WORKER_FINAL_REPLY] to_send.append(b'') if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self._stream.send_multipart(to_send, track=True, copy=False) def emit(self, msg): # type: (Any) -> None if not self.is_connected(): raise LostRemoteError("No client is connected.") msg = msgpack.Packer(default=XeroSerializer.encoder).pack(msg) to_send = [WORKER_EMIT] to_send.append(b'') if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self._stream.io_loop.add_callback( lambda x: self._stream.send_multipart(x, track=True, copy=False), to_send) def _tick(self): # type: () -> None """ Periodic callback to check connectivity to client. """ if self._curr_liveness >= 0: self._curr_liveness -= 1 if self._curr_liveness > 0: self._send_heartbeat() elif self._curr_liveness == 0: # Connection died, close on our side. self.on_log_event( "uniworker.tick", "Connection to uniclient timed out, disconnecting") self._connected_event.clear() else: self._send_ready() def _send_heartbeat(self): # type: () -> None """ Send a heartbeat message to the client. """ # Heartbeats should go out immediately, if a lot of messages to be emitted are queued up heartbeats should # still be sent out regularly. Therefore, send it out via the stream's socket, rather than the stream itself # See https://pyzmq.readthedocs.io/en/latest/eventloop.html#send self._stream.send_multipart([WORKER_HEARTBEAT]) def _send_disconnect(self): # type: () -> None """ Send a disconnect message to the client. """ # Send out via the socket, this message takes priority. self._stream.send_multipart([WORKER_DISCONNECT]) def _send_ready(self): # type: () -> None """ Send a ready message to the client. """ self.on_log_event("uniworker.ready", "Sending ready to client.") self._stream.send_multipart([WORKER_READY]) def _on_message(self, msg): # type: (List[bytes]) -> None """ Processes a received ZeroMQ message. :param msg: List of strings in the format: [ ZMQ Client ID, Header, StrMessagePart1, StrMessagePart2...] """ # 2nd part is protocol version protocol_version = msg.pop(0) if protocol_version != UNI_CLIENT_HEADER: # version check, ignore old versions logger.error( "Message doesn't start with {}".format(UNI_CLIENT_HEADER)) return # 3rd part is message type msg_type = msg.pop(0) # any message resets the liveness counter self._need_handshake = False self._connected_event.set() self._curr_liveness = HB_LIVENESS if msg_type == WORKER_DISCONNECT: # disconnect self._curr_liveness = 0 # reconnect will be triggered by hb timer elif msg_type == WORKER_REQUEST: # request # remaining parts are the user message self._on_request(msg) elif msg_type == WORKER_HEARTBEAT: # received hardbeat - timer handled above pass else: logger.error("Uniworker received unrecognized message") def _on_request(self, message): # type: (List[bytes]) -> None """ This gets called on incoming RPC messages, will break up the encoded message into something do_work() can process :param message: """ name = str(message[0], 'utf-8') args = msgpack.unpackb(message[1], object_hook=XeroSerializer.decoder, raw=False) kwargs = msgpack.unpackb(message[2], object_hook=XeroSerializer.decoder, raw=False) self.do_work(name, args, kwargs) def on_log_event(self, event, message): # type: (str, str) -> None """ Called on internal loggable events. Designed for override. :param event: The event type. :param message: Loggable message. """ logger.debug("{}: {}".format(event, message)) @abstractmethod def do_work(self, name, args, kwargs): # type: (str, List[Any], Dict[Any,Any]) -> None """ Override this method for worker-specific message handling. :param name: The 'name' of the function/rpc call. :param args: Function call arguments. :param kwargs: Function call key arguments. """ raise NotImplementedError()
class MQRep(object): """Class for the MDP worker side. Thin encapsulation of a zmq.DEALER socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. """ _proto_version = b'MDPW01' # TODO: integrate that into API HB_INTERVAL = 1000 # in milliseconds HB_LIVENESS = 3 # HBs to miss before connection counts as dead def __init__(self, context, service): """Initialize the MDPWorker. context is the zmq context to create the socket from. service is a byte-string with the service name. """ if DEBUG: print("MQRep > __init__") cfg = Loader('mq').load() config = dict(cfg[1]) if config['ip'].strip() == "*": config['ip'] = get_ip() self.endpoint = "tcp://{0}:{1}".format(config['ip'], config['req_rep_port']) self.context = context self.service = service self.stream = None self._tmo = None self.need_handshake = True self.ticker = None self._delayed_cb = None self._create_stream() ### patch fritz self._reconnect_in_progress = False ### end patch fritz return def _create_stream(self): """Helper to create the socket and the stream. """ if DEBUG: print("MQRep > _create_stream") socket = ZmqSocket(self.context, zmq.DEALER) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_mpd_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) if self.ticker != None: if DEBUG: print("MQRep > _create_stream - stop ticker") self.ticker.stop() self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() self.ticker.start() return def _send_ready(self): """Helper method to prepare and send the workers READY message. """ if DEBUG: print("MQREP > _send_ready") ready_msg = [ b'', self._proto_version, b'\x01', self.service ] self.stream.send_multipart(ready_msg) self.curr_liveness = self.HB_LIVENESS if DEBUG: print("MQREP > _send_ready > curr_liveness <= {0}".format(self.HB_LIVENESS)) return def _tick(self): """Method called every HB_INTERVAL milliseconds. """ if DEBUG: print("MQREP > _tick") self.curr_liveness -= 1 if DEBUG: print('MQREP > _tick - {0} tick = {1}'.format(time.time(), self.curr_liveness)) self.send_hb() if self.curr_liveness >= 0: return if DEBUG: print('MQREP > _tick - {0} lost connection'.format(time.time())) # ouch, connection seems to be dead self.shutdown() # try to recreate it self._delayed_cb = DelayedCallback(self._create_stream, self.HB_INTERVAL) self._delayed_cb.start() return def send_hb(self): """Construct and send HB message to broker. """ msg = [ b'', self._proto_version, b'\x04' ] self.stream.send_multipart(msg) return def shutdown(self): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. """ if self.ticker: self.ticker.stop() self.ticker = None if not self.stream: return self.stream.socket.close() self.stream.close() self.stream = None self.timed_out = False self.need_handshake = True self.connected = False return def reply(self, msg): """Send the given message. msg can either be a byte-string or a list of byte-strings. """ ## if self.need_handshake: ## raise ConnectionNotReadyError() # prepare full message to_send = self.envelope self.envelope = None if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self.stream.send_multipart(to_send) return def _on_mpd_message(self, msg): """Helper method called on message receive. msg is a list w/ the message parts """ if DEBUG: print("MQRep > _on_mpd_message : {0} - {1}".format(time.strftime("%H:%M:%S"), msg)) # 1st part is empty msg.pop(0) # 2nd part is protocol version # TODO: version check proto = msg.pop(0) # 3rd part is message type msg_type = msg.pop(0) # XXX: hardcoded message types! # any message resets the liveness counter self.need_handshake = False self.curr_liveness = self.HB_LIVENESS if DEBUG: print("MQREP > _on_mpd_message > curr_liveness <= {0}".format(self.HB_LIVENESS)) if msg_type == b'\x05': # disconnect if DEBUG: print("MQREP > _on_mpd_message > type x05 : disconnect") self.curr_liveness = 0 # reconnect will be triggered by hb timer elif msg_type == b'\x02': # request if DEBUG: print("MQREP > _on_mpd_message > type x02 : request") # remaining parts are the user message envelope, msg = split_address(msg) envelope.append(b'') envelope = [ b'', self._proto_version, b'\x03'] + envelope # REPLY self.envelope = envelope mes = MQMessage() mes.set(msg) #print("MQRep > before self.on_mdp_request") #print(self.on_mdp_request) #print(mes) try: self.on_mdp_request(mes) except: print("ERROR {0}".format(traceback.format_exc())) else: if DEBUG: print("MQREP > _on_mpd_message > type ??? : invalid or hbeat") # invalid message # ignored # if \x04, this is a hbeat message pass return def on_mdp_request(self, msg): """Public method called when a request arrived. Must be overloaded! """ pass
class ZmqPublisher(object): def __init__(self, moduleName, centralHost=PUBLISHER_OPT_DEFAULTS['centralHost'], context=None, centralSubscribeEndpoint=PUBLISHER_OPT_DEFAULTS['centralSubscribeEndpoint'], publishEndpoint=PUBLISHER_OPT_DEFAULTS['publishEndpoint'], heartbeatPeriodMsecs=PUBLISHER_OPT_DEFAULTS['heartbeatPeriodMsecs'], # highWaterMark=PUBLISHER_OPT_DEFAULTS['highWaterMark'] ): self.moduleName = moduleName self.centralHost = centralHost if context is None: context = zmq.Context.instance() self.context = context self.centralSubscribeEndpoint = parseEndpoint(centralSubscribeEndpoint, defaultPort=DEFAULT_CENTRAL_SUBSCRIBE_PORT, centralHost=self.centralHost) self.publishEndpoint = parseEndpoint(publishEndpoint, defaultPort='random') self.heartbeatPeriodMsecs = heartbeatPeriodMsecs #self.highWaterMark = highWaterMark self.pubStream = None self.heartbeatTimer = None self.serializer = serializers.get_serializer('json')() @classmethod def addOptions(cls, parser, defaultModuleName): if not parser.has_option('--centralHost'): parser.add_option('--centralHost', default=PUBLISHER_OPT_DEFAULTS['centralHost'], help='Host where central runs [%default]') if not parser.has_option('--moduleName'): parser.add_option('--moduleName', default=defaultModuleName, help='Name to use for this module [%default]') if not parser.has_option('--centralSubscribeEndpoint'): parser.add_option('--centralSubcribeEndpoint', default=PUBLISHER_OPT_DEFAULTS['centralSubscribeEndpoint'], help='Endpoint where central listens for messages [%default]') if not parser.has_option('--publishEndpoint'): parser.add_option('--publishEndpoint', default=PUBLISHER_OPT_DEFAULTS['publishEndpoint'], help='Endpoint to publish messages on [%default]') if not parser.has_option('--heartbeatPeriodMsecs'): parser.add_option('--heartbeatPeriodMsecs', default=PUBLISHER_OPT_DEFAULTS['heartbeatPeriodMsecs'], type='int', help='Period for sending heartbeats to central [%default]') #if not parser.has_option('--highWaterMark'): # parser.add_option('--highWaterMark', # default=PUBLISHER_OPT_DEFAULTS['highWaterMark'], # type='int', # help='High-water mark for publish socket (see 0MQ docs) [%default]') @classmethod def getOptionValues(cls, opts): result = {} for key in PUBLISHER_OPT_DEFAULTS.iterkeys(): val = getattr(opts, key, None) if val is not None: result[key] = val return result def heartbeat(self): logging.debug('ZmqPublisher: heartbeat') self.sendJson('central.heartbeat.%s' % self.moduleName, {'host': getShortHostName(), 'pub': self.publishEndpoint}) def sendRaw(self, topic, body): self.pubStream.send('%s:%s' % (topic, body)) self.pubStream.flush() def sendJson(self, topic, obj): if isinstance(obj, dict): obj.setdefault('module', self.moduleName) obj.setdefault('timestamp', str(getTimestamp())) self.sendRaw(topic, json.dumps(obj)) def sendDjango(self, modelInstance, topic=None, topicSuffix=None): dataText = self.serializer.serialize([modelInstance]) data = json.loads(dataText)[0] if topic is None: topic = data['model'].encode('utf-8') if topicSuffix is not None: topic += topicSuffix self.sendJson(topic, {'data': data}) def start(self): pubSocket = self.context.socket(zmq.PUB) self.pubStream = ZMQStream(pubSocket) # self.pubStream.setsockopt(zmq.IDENTITY, self.moduleName) # self.pubStream.setsockopt(zmq.HWM, self.highWaterMark) self.pubStream.connect(self.centralSubscribeEndpoint) logging.info('zmq.publisher: connected to central at %s', self.centralSubscribeEndpoint) if self.publishEndpoint.endswith(':random'): endpointWithoutPort = re.sub(r':random$', '', self.publishEndpoint) port = self.pubStream.bind_to_random_port(endpointWithoutPort) self.publishEndpoint = '%s:%d' % (endpointWithoutPort, port) else: self.pubStream.bind(self.publishEndpoint) self.heartbeatTimer = ioloop.PeriodicCallback(self.heartbeat, self.heartbeatPeriodMsecs) self.heartbeatTimer.start() self.heartbeat()
class ZmqSubscriber(object): def __init__(self, moduleName, centralHost=SUBSCRIBER_OPT_DEFAULTS['centralHost'], context=None, centralPublishEndpoint=SUBSCRIBER_OPT_DEFAULTS[ 'centralPublishEndpoint'], replay=None): self.moduleName = moduleName self.centralHost = centralHost if context is None: context = zmq.Context.instance() self.context = context self.centralPublishEndpoint = parseEndpoint( centralPublishEndpoint, defaultPort=DEFAULT_CENTRAL_PUBLISH_PORT, centralHost=self.centralHost) self.replayPaths = replay if self.replayPaths is None: self.replayPaths = [] self.handlers = {} self.counter = 0 self.deserializer = serializers.get_deserializer('json') self.stream = None @classmethod def addOptions(cls, parser, defaultModuleName): if not parser.has_option('--centralHost'): parser.add_option('--centralHost', default=SUBSCRIBER_OPT_DEFAULTS['centralHost'], help='Host where central runs [%default]') if not parser.has_option('--moduleName'): parser.add_option('--moduleName', default=defaultModuleName, help='Name to use for this module [%default]') if not parser.has_option('--centralPublishEndpoint'): parser.add_option( '--centralPublishEndpoint', default=SUBSCRIBER_OPT_DEFAULTS['centralPublishEndpoint'], help='Endpoint where central publishes messages [%default]') if not parser.has_option('--replay'): parser.add_option( '--replay', action='append', help= 'Replay specified message log (can specify multiple times), or use - to read from stdin' ) @classmethod def getOptionValues(cls, opts): result = {} for key in SUBSCRIBER_OPT_DEFAULTS.iterkeys(): val = getattr(opts, key, None) if val is not None: result[key] = val return result def start(self): sock = self.context.socket(zmq.SUB) self.stream = ZMQStream(sock) # causes problems with multiple instances #self.stream.setsockopt(zmq.IDENTITY, self.moduleName) self.stream.connect(self.centralPublishEndpoint) logging.info('zmq.subscriber: connected to central at %s', self.centralPublishEndpoint) self.stream.on_recv(self.routeMessages) def routeMessages(self, messages): for msg in messages: self.routeMessage(msg) def routeMessage(self, msg): colonIndex = msg.find(':') topic = msg[:(colonIndex + 1)] body = msg[(colonIndex + 1):] handled = 0 for topicPrefix, registry in self.handlers.iteritems(): if topic.startswith(topicPrefix): for handler in registry.itervalues(): handler(topic[:-1], body) handled = 1 return handled def subscribeRaw(self, topicPrefix, handler): topicRegistry = self.handlers.setdefault(topicPrefix, {}) if not topicRegistry: logging.info('zmq.subscriber: subscribe %s', topicPrefix) self.stream.setsockopt(zmq.SUBSCRIBE, topicPrefix) handlerId = (topicPrefix, self.counter) topicRegistry[self.counter] = handler self.counter += 1 return handlerId def subscribeJson(self, topicPrefix, handler): def jsonHandler(topicPrefix, body): return handler(topicPrefix, convertToDotDictRecurse(json.loads(body))) return self.subscribeRaw(topicPrefix, jsonHandler) def subscribeDjango(self, topicPrefix, handler): def djangoHandler(topicPrefix, body): obj = json.loads(body) dataText = json.dumps([obj['data']]) modelInstance = list(self.deserializer(dataText))[0] return handler(topicPrefix, modelInstance.object) return self.subscribeRaw(topicPrefix, djangoHandler) def unsubscribe(self, handlerId): topicPrefix, index = handlerId topicRegistry = self.handlers[topicPrefix] del topicRegistry[index] if not topicRegistry: logging.info('zmq.subscriber: unsubscribe %s', topicPrefix) self.stream.setsockopt(zmq.UNSUBSCRIBE, topicPrefix) def connect(self, endpoint): self.stream.connect(endpoint) def replay(self): numReplayed = 0 numHandled = 0 for replayPath in self.replayPaths: print '=== replaying messages from %s' % replayPath if replayPath == '-': replayFile = sys.stdin else: replayFile = open(replayPath, 'rb') stream = LogParser(replayFile) for rec in stream: numReplayed += 1 numHandled += self.routeMessage(rec.msg) if numReplayed % 10000 == 0: print 'replayed %d messages, %d handled' % (numReplayed, numHandled)
class Subscriber (object): def __init__(self, sub_uri=defaults.subscriber_sub_uri, patterns=None, callbacks=None, ): self.sub_uri = sub_uri if patterns: self.patterns = patterns else: self.patterns = [] if callbacks: self.callbacks = callbacks else: self.callbacks = [] self.setup_logging() self.setup_zmq() self.setup_sockets() self.setup_subscriptions() self.setup_events() def setup_logging(self): self.log = logging.getLogger('zmqevt.subscriber') def setup_zmq(self): self.context = zmq.Context() def setup_sockets(self): self.sub = ZMQStream(self.context.socket(zmq.SUB)) self.sub.connect(self.sub_uri) def setup_subscriptions(self): if self.patterns: for p in self.patterns: self.subscribe(p) def subscribe(self, pattern): self.log.debug('Subcribe to "%s".' % pattern) self.sub.setsockopt(zmq.SUBSCRIBE, pattern) def setup_events(self): self.sub.on_recv(self.on_recv) def register_callback(self, func, data=None): self.callbacks.append((func,data)) def unregister_callback(self, func): self.callbacks = [x for x in self.callbacks if x[0] is not func] def on_recv(self, msg): self.log.debug('Receive: %s' % (str(msg))) assert len(msg) == 2, 'Received invalid message.' # This regenerates the event to ensure that we don't # pass on invalid data. try: evt = event.Event.load(msg) except Exception, detail: self.log.error('Error processing message: %s' % detail) return self.log.debug('Event: %s' % (str(evt.dump()))) for func, data in self.callbacks: func(evt, data=data)
class Worker(object): """Class for the MDP worker side. Thin encapsulation of a zmq.DEALER socket. Provides a send method with optional timeout parameter. Will use a timeout to indicate a broker failure. """ max_forks = 10 ipc = 'ipc:///tmp/zmq-rpc-'+str(uuid4()) HB_INTERVAL = 1000 # in milliseconds HB_LIVENESS = 3 # HBs to miss before connection counts as dead def __init__(self, context, endpoint, service, multicasts=()): """Initialize the MDPWorker. :param context: is the zmq context to create the socket from :type context: zmq.Context :param service: service name - you can put hostname here :type service: str :param multicasts: list of groups to subscribe :type multicasts: list """ self.context = context self.endpoint = endpoint self.service = service.encode('utf-8') # convert to byte-string - required in python 3 self.multicasts = [m.encode('utf-8') for m in multicasts] # convert to byte-string self.stream = None self._tmo = None self.need_handshake = True self.ticker = None self._delayed_cb = None self._create_stream() self.forks = [] self.curr_liveness = self.HB_LIVENESS socket = self.context.socket(zmq.ROUTER) socket.bind(self.ipc) self.stream_w = ZMQStream(socket) self.stream_w.on_recv(self._on_fork_response) self.reply_socket = None return def _create_stream(self): """Helper to create the socket and the stream. """ self.on_log_event('broker.connect', 'Trying to connect do broker') socket = self.context.socket(zmq.DEALER) ioloop = IOLoop.instance() self.stream = ZMQStream(socket, ioloop) self.stream.on_recv(self._on_message) self.stream.socket.setsockopt(zmq.LINGER, 0) self.stream.connect(self.endpoint) self.ticker = PeriodicCallback(self._tick, self.HB_INTERVAL) self._send_ready() for m in self.multicasts: self._register_worker_to_multicast(m) self.ticker.start() return def _tick(self): """Method called every HB_INTERVAL milliseconds. """ self.curr_liveness -= 1 self.send_hb() if self.curr_liveness >= 0: return # ouch, connection seems to be dead self.on_log_event('broker.timeout', 'Connection to broker timeouted, disconnecting') self.shutdown(False) # try to recreate it self._delayed_cb = DelayedCallback(self._create_stream, 5000) self._delayed_cb.start() return def send_hb(self): """Construct and send HB message to broker. """ msg = [b'', MDP_WORKER_VERSION, b'\x05'] self.stream.send_multipart(msg) return def shutdown(self, final=True): """Method to deactivate the worker connection completely. Will delete the stream and the underlying socket. :param final: if shutdown is final and we want to close all sockets :type final: bool """ if self.ticker: self.ticker.stop() self.ticker = None if not self.stream: return self.stream.on_recv(None) self.disconnect() self.stream.socket.close() self.stream.close() self.stream = None self.need_handshake = True if final: self.stream_w.socket.close() self.stream_w.close() self.stream = None return def disconnect(self): """Helper method to send the workers DISCONNECT message. """ self.stream.socket.send_multipart([b'', MDP_WORKER_VERSION, b'\x06' ]) self.curr_liveness = self.HB_LIVENESS return def _send_ready(self): """Helper method to prepare and send the workers READY message. """ self.on_log_event('broker.ready', 'Sending ready to broker.') ready_msg = [b'', MDP_WORKER_VERSION, b'\x01', self.service] self.stream.send_multipart(ready_msg) self.curr_liveness = self.HB_LIVENESS return def _register_worker_to_multicast(self, name): """Helper method to register worker to multicast group :param name: group name :type name: str """ self.on_log_event('broker.register-group', 'Subscribing to group \'%s\'.' % name) reg_msg = [b'', MDP_WORKER_VERSION, b'\x07', name] self.stream.send_multipart(reg_msg) self.curr_liveness = self.HB_LIVENESS return def _on_message(self, msg): """Helper method called on message receive. :param msg: message parts :type msg: list """ # 1st part is empty msg.pop(0) # 2nd part is protocol version protocol_version = msg.pop(0) if protocol_version != MDP_WORKER_VERSION: # version check, ignore old versions return # 3rd part is message type msg_type = msg.pop(0) # any message resets the liveness counter self.need_handshake = False self.curr_liveness = self.HB_LIVENESS if msg_type == b'\x06': # disconnect self.curr_liveness = 0 # reconnect will be triggered by hb timer elif msg_type == b'\x02': # request # remaining parts are the user message addresses, msg = self.split_address(msg) self._on_request(addresses, msg) elif msg_type == b'\x05': # received hardbeat - timer handled above pass else: # invalid message ignored pass return def _on_fork_response(self, to_send): """Helper method to send message from forked worker. This message will be received by main worker process and resend to broker. :param to_send address and data to send :type to_send list """ self.stream.send_multipart(to_send) return def send_reply(self, addresses, msg, partial=False, exception=False): """Send reply from forked worker process. This method can be called only from do_work() method! This method will send messages to main worker listening on local socket in /tmp/zmq-rpc-... :param addresses: return address stack :type addresses: list of str :param msg: return value from called method :type msg: mixed :param partial: if the message is partial or final :type partial: bool :param exception: if the message is exception, msg format is: {'class':'c', 'message':'m', 'traceback':'t'} :type exception: bool """ if not self.reply_socket: context = zmq.Context() self.reply_socket = context.socket(zmq.DEALER) self.reply_socket.connect(self.ipc) msg = msgpack.Packer().pack(msg) if exception: to_send = [b'', MDP_WORKER_VERSION, b'\x08'] elif partial: to_send = [b'', MDP_WORKER_VERSION, b'\x03'] else: to_send = [b'', MDP_WORKER_VERSION, b'\x04'] to_send.extend(addresses) to_send.append(b'') if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) m = self.reply_socket.send_multipart(to_send, track=True, copy=False) m.wait() if not partial: self.reply_socket.close() self.reply_socket = None return def send_message(self, addresses, msg, partial=False, error=False): """Send response message from main worker process. Please do not call this method from do_work() :param addresses: return address stack :type addresses: list of str :param msg: return value from called method :type msg: mixed :param partial: if the message is partial or final :type partial: bool :param error: if the message is error :type error: bool """ to_send = [b'', MDP_WORKER_VERSION] if partial: to_send.append(b'\x03') elif error: to_send.append(b'\x09') else: to_send.append(b'\x04') to_send.extend(addresses) to_send.append(b'') if isinstance(msg, list): to_send.extend(msg) else: to_send.append(msg) self.stream.send_multipart(to_send) return def _on_request(self, addresses, message): """Helper method called on RPC message receive. """ # remove finished forks self._remove_finished_processes() # test max forks if len(self.forks) >= self.max_forks: self.send_message(addresses, b'max workers limit exceeded', error=True) self.on_max_forks(addresses, message) return name = message[0] args = msgpack.unpackb(message[1]) kwargs = msgpack.unpackb(message[2]) p = Process(target=self.do_work, args=(addresses, name, args, kwargs)) p.start() p._args = None # free memory self.forks.append(p) return def _remove_finished_processes(self): """Helper method dedicated to cleaning list of forked workers """ for f in [f for f in self.forks if not f.is_alive()]: self.forks.remove(f) return def split_address(self, msg): """Function to split return Id and message received by ROUTER socket. Returns 2-tuple with return Id and remaining message parts. Empty frames after the Id are stripped. """ ret_ids = [] for i, p in enumerate(msg): if p: ret_ids.append(p) else: break return ret_ids, msg[i + 1:] def on_log_event(self, event, message): """Override this method if you want to log events from broker :type event: str :param event: event type - used for filtering :type message: str :param message: log message :rtype: None """ pass def on_max_forks(self, addresses, message): """This method is called when max_forks limit is reached You can override this method. """ pass def do_work(self, addresses, name, args, kwargs): """Main method responsible for handling rpc calls, and sending response messages. Please override this method! :param addresses: return address stack :type addresses: list of str :param name: name of task :type name: str :param args: positional task arguments :type args: list :param kwargs: key-value task arguments :type kwargs: dict """ # this is example of simple response message self.send_reply(addresses, 'method not implemented') # and send message to main worker # you can also send partial message and exception - read 'send_reply' docs return