class InternalRouter(object):
def __init__(self, context):
socket = context.socket(zmq.ROUTER)
socket.setsockopt(zmq.LINGER, 0)
identity = 'router'
socket.identity = identity.encode('ascii')
session = Session()
router_internal = session.query(RoutingDefiniton.route) \
.filter_by(name="internal_router") \
.scalar()
session.close()
socket.bind(router_internal)
ioloop = IOLoop.instance()
self.internal_router = ZMQStream(socket, ioloop)
self.internal_router.on_recv(callback=self.router_recv)
return
def shutdown(self):
self.internal_router.on_recv(callback=None)
self.internal_router = None
return
def router_recv(self, msg):
return
class SoundServer(object):
def __init__(self):
self._context = zmq.Context()
self._socket = self._context.socket(zmq.REP)
self._socket.bind('tcp://127.0.0.1:{0}'.format(CONTROL_PORT))
self._stream = ZMQStream(self._socket)
self._stream.on_recv(self._handle_msg)
def start(self):
IOLoop.instance().start()
def _handle_msg(self, msg):
method = '_handler_{0}'.format(msg[0].decode("utf-8"))
try:
print("Trying method {0}".format(method))
getattr(self, method)()
except AttributeError:
sys.exit(1)
def _handler_hi(self):
self._socket.send_string('sup?')
def _handler_exit(self):
pass
def Run():
try:
print "Running web server on port: "+str(Config["Web server"]["Port"])
sessionsManager = SessionsManager()
socket = zmq.Context.instance().socket(zmq.SUB)
socket.bind("tcp://*:"+str(Config["Web server"]["IncomingZmqPort"]))
socket.setsockopt(zmq.SUBSCRIBE, "")
stream = ZMQStream(socket)
stream.on_recv(sessionsManager.ProcessZmqMessages)
settings = {
"debug" : False,
"template_path":os.path.join(os.path.dirname(__file__), "../Frontend"),
}
app = tornado.web.Application([
(r'/Static/(.*)', tornado.web.StaticFileHandler, {'path': os.path.join(os.path.dirname(__file__), "../Frontend/Static")}),
(r'/Tmp/(.*)', tornado.web.StaticFileHandler, {'path': Config["Web server"]["Temporary files folder"]}),
(r'/RCP', IndexHandler),
(r'/RCP/', IndexHandler),
(r'/StreamsTree/(.*)', StreamsTreeRequestHandler, dict(sessionsManager=sessionsManager)),
(r'/WebSockets/', WebSocketHandler, dict(sessionsManager=sessionsManager)),
], **settings)
periodic = ioloop.PeriodicCallback(CheckServerStopRequests, 500)
periodic.start()
app.listen(Config["Web server"]["Port"])
tornado.ioloop.IOLoop.instance().start()
except:
print traceback.format_exc()
class MonitorEvents(SockJSConnection):
def _zmq_msg(self, msg):
#logging.debug(msg)
try:
msg_obj = json.loads(msg[0])
logging.debug(msg_obj)
if self.monitor != 'All':
if 'mon_id' in msg_obj and msg_obj['mon_id'] == self.monitor:
self.send(msg_obj)
else:
self.send(msg_obj)
except Exception as ex:
logging.error(ex)
def on_open(self, info):
logging.debug("Monitor ticker open: "+self.monitor)
zmq_socket = zmq.Context.instance().socket(zmq.SUB)
zmq_socket.connect(zmq_local_endpoint)
zmq_socket.setsockopt(zmq.SUBSCRIBE, '')
self.stream = ZMQStream(zmq_socket)
self.stream.on_recv(self._zmq_msg)
def on_close(self):
logging.debug("Monitor ticker close: "+self.monitor)
self.stream.stop_on_recv()
class Client(object):
def __init__(self):
self.stream = None
self.result = None
def connect(self, port):
context = zmq.Context()
socket = context.socket(zmq.REQ)
socket.connect(port)
self.stream = ZMQStream(socket)
def __getattr__(self, item):
def wrapper(*args):
request = [item]
for param in args:
request.append(param)
return self._run(tuple(request))
return wrapper
def _run(self, request):
def on_response(message):
response = msgpack.unpackb(message[0], use_list=False)
if response[0] == 'OK':
self.result = response[1]
elif response[0] == 'ERR':
raise Exception(response[2])
ZMQIOLoop.instance().stop()
self.stream.send(msgpack.packb(request))
self.stream.on_recv(on_response)
ZMQIOLoop.instance().start()
return self.result
def disconnect(self):
self.stream.close()
class SwysSearchRequestHandler(BaseHandler):
def initialize(self):
socket = context.socket(zmq.REQ)
socket.connect(conf.SEARCH_WORKER_ZMQ_ENDPOINT)
self._zmq_stream = ZMQStream(socket)
self._zmq_stream.on_recv(self._recv_result, copy=True)
@tornado.web.asynchronous
def handle_request_async(self, *args, **kwargs):
files = self.request.files.get('image', [])
if len(files) == 0:
raise Exception("there is no file attached")
file = files[0]
temp_file = tempfile.NamedTemporaryFile('wb', delete=False)
temp_file.write(file.body)
self._zmq_stream.send_json({'filename': temp_file.name})
def _recv_result(self, msg):
result_str = "".join(( part.decode('utf-8') for part in msg ))
result = json.loads(result_str)['data']
return self.on_complete(result)
def start(self):
identifier = self.identifier
job = self.job
def execute_next(msg):
# We can't use zodb object from outside here because
# this code is executed in another thread (eventloop)
# We don't have site or interaction, so the job must be created
# before.
# we can't use push_event_callback_after_commit here because
# it will never commit in this thread (eventloop)
if identifier in callbacks:
callbacks[identifier].close()
del callbacks[identifier]
job.args = (msg, )
# wait 2s that the throw event transaction has committed
dc = DelayedCallback(job, 2000)
dc.start()
ctx = get_zmq_context()
s = ctx.socket(zmq.SUB)
s.setsockopt_string(zmq.SUBSCRIBE, u'')
s.connect(get_signal_socket_url())
stream = ZMQStream(s)
callbacks[identifier] = stream
stream.on_recv(execute_next)
def enable_depth_frames(self, kinect_id):
"""Enable streaming of depth frames. *kinect_id* is the id of the
device which should have streaming enabled.
:raises ValueError: if *kinect_id* does not correspond to a connected device
"""
try:
record = self._kinect_records[kinect_id]
except KeyError:
raise ValueError('Kinect id "{0}" does not correspond to a connected device'.format(
kinect_id))
# Create subscriber stream
socket = self._zmq_ctx.socket(zmq.SUB)
socket.connect(record.endpoints[EndpointType.depth])
socket.setsockopt_string(zmq.SUBSCRIBE, u'')
stream = ZMQStream(socket, self._io_loop)
record.streams[EndpointType.depth] = stream
# Fire signal on incoming depth frame
def on_recv(msg, kinect_id=kinect_id):
# TODO: decompress frame
self.on_depth_frame.send(self, kinect_id=kinect_id, depth_frame=msg)
# Wire up callback
stream.on_recv(on_recv)
class ZBus(object):
def __init__(self):
self._context = zmq.Context()
self._callback = {}
self._zstream = None
@staticmethod
def instance():
if not hasattr(ZBus, '_instance'):
ZBus._instance = ZBus()
return ZBus._instance
@staticmethod
def initialized():
return hasattr(ZBus, '_instance')
def connect(self, dist):
if self._zstream:
self._zstream.close()
self._zsock = self._context.socket(zmq.XREQ)
self._zsock.connect('tcp://{dist}'.format(dist=dist))
self._zstream = ZMQStream(self._zsock)
self._zstream.on_recv(self.on_recv)
def send(self, request, callback):
self._callback[request.seed_id] = callback
self._zstream.send_multipart(request.box())
def on_recv(self, frame):
response = ZResponse(frame)
callback = self._callback.pop(response.seed_id) if self._callback.get(response.seed_id) else None
if callback and callable(callback):
callback(response)
class UartzStream(object):
def __init__(self, addr, io_loop=None):
self.addr = addr
self.ioloop = io_loop
self.stream = None
self.reset_stream()
def reset_stream(self):
ctx = zmq.Context.instance()
sock = ctx.socket(zmq.SUB)
sock.connect(self.addr)
sock.setsockopt(zmq.SUBSCRIBE, b'')
self.stream = ZMQStream(sock, self.ioloop)
self.stream.on_recv(self._handle_msg)
def _handle_msg(self, msg):
assert len(msg) == 1
msg = msg[0]
chan_idx = msg.index(b":")
assert chan_idx > 0
self.handle_msg(UartzMsg(dev=msg[:chan_idx],
msg=msg[chan_idx + 1:]))
def main(pat):
fname = find_connection_file(pat)
with open(fname) as f:
cfg = json.load(f)
url = "%s://%s:%s" % (cfg.get('transport', 'tcp'), cfg['ip'], cfg['iopub_port'])
session = Session(key=cfg['key'])
ctx = zmq.Context.instance()
sub = ctx.socket(zmq.SUB)
sub.subscribe = b''
sub.connect(url)
# import IPython
# IPython.embed()
# return
stream = ZMQStream(sub)
stream.on_recv(lambda msg_list: log_msg(session, msg_list))
pc = PeriodicCallback(print_time, 5 * 60 * 1000)
pc.start()
IOLoop.instance().start()
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
class LRUQueue(object):
def __init__(self, backend_socket, frontend_socket, clients, workers):
self.avaliable_workers = 0
self.workers = []
self.worker_num = workers
self.client_num = clients
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
worker_addr, empty, client_addr = msg[:3]
assert self.avaliable_workers < self.worker_num
self.avaliable_workers += 1
self.workers.append(worker_addr)
assert empty == ""
if client_addr != "READY":
empty, reply = msg[3:]
assert empty == ""
self.frontend.send_multipart([client_addr, "", reply])
self.client_num -= 1
if 0 == self.client_num:
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.avaliable_workers == 1:
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
client_addr, empty, request = msg
assert empty == ""
self.avaliable_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart([worker_id, "", client_addr, "", request])
if self.avaliable_workers == 0:
self.frontend.stop_on_recv()
class Connection(object):
"""The base class for the connection between node and master
"""
def __init__(self, endpoint):
self.endpoint = endpoint
# init zeromq
self.context = zmq.Context()
self.socket = self.context.socket(zmq.REQ)
self.stream = ZMQStream(self.socket)
self.stream.on_recv(self.OnRecvMsg)
# get local endpoint
self.socket.bind("tcp://eth0:*")
self.local_endpoint = str(self.socket.getsockopt(zmq.LAST_ENDPOINT))
print "Local endpoint [%s]" % self.local_endpoint
# connect to target
self.socket.connect(endpoint)
print "Connected to [%s]" % endpoint
def SendMsg(self, msg):
self.socket.send(msg, copy=False)
print "Sending message [%s]" % msg
#msg_rsp = self.socket.recv( copy = False )
#print "Receiving message [%s]" % msg_rsp
def OnRecvMsg(self, msg):
#msg_rsp = self.socket.recv( copy = False )
#print "Receiving message [%s]" % msg_rsp
print "Receiving message ========== [%s]" % msg
def GetLocalEndpoint(self):
return self.local_endpoint
def GetEndpoint(self):
return self.endpoint
class SocketConnection(sockjs.tornado.SockJSConnection):
clients = set()
def on_open(self, request):
self.clients.add(self)
subscriber = context.socket(zmq.SUB)
subscriber.connect("tcp://localhost:%s" % str(ZMQ_PORT))
subscriber.setsockopt(zmq.SUBSCRIBE, '')
self.subscribe_stream = ZMQStream(subscriber)
self.subscribe_stream.on_recv(self.on_message_published)
def on_message(self, message):
logging.info(
'message received, publish it to %d clients' % len(self.clients)
)
publish_stream.send_unicode(message)
def on_message_published(self, message):
logging.info('client received new published message')
self.send(message)
def on_close(self):
self.clients.remove(self)
# Properly close ZMQ socket
self.subscribe_stream.close()
class Broker (object):
def __init__(self,
pub_uri=defaults.broker_pub_uri,
sub_uri=defaults.broker_sub_uri,
patterns=None,
):
self.pub_uri = pub_uri
self.sub_uri = sub_uri
if patterns:
self.patterns = patterns
else:
self.patterns = []
self.setup_logging()
self.setup_zmq()
self.setup_sockets()
self.setup_subscriptions()
self.setup_events()
def setup_logging(self):
self.log = logging.getLogger('zmqevt.broker')
def setup_zmq(self):
self.context = zmq.Context()
def setup_sockets(self):
self.sub = ZMQStream(self.context.socket(zmq.SUB))
self.sub.bind(self.sub_uri)
self.pub = ZMQStream(self.context.socket(zmq.PUB))
self.pub.bind(self.pub_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.publish)
def publish(self, 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())))
self.pub.send_multipart(evt.dump())
def connect(self):
if self.connected:
return
self.socket = ctx.socket(zmq.SUBSCRIBE)
self.socket.connect(self.socket_addr)
stream = ZMQStream(self.socket, self.io_loop)
stream.on_recv(self.on_message)
self.connected = True
class SubQueue(BaseQueue):
def __init__(self, connect, callback, ioloop=None):
super().__init__(zmq.PULL)
self.socket.connect(connect)
self.stream = ZMQStream(self.socket, ioloop)
self.stream.on_recv(callback)
self.stream.flush()
def call(self, message, callback):
context = zmq.Context()
socket = context.socket(zmq.DEALER)
socket.connect(self.endpoint)
stream = ZMQStream(socket)
stream.on_recv(callback)
socket.send(message)
class zmq_bonjour_bind_wrapper(object):
context = None
socket = None
stream = None
heartbeat_timer = None
method_callbacks = {}
def _hearbeat(self):
#print "Sending heartbeat"
self.stream.send_multipart(("HEARTBEAT", "1"))
def __init__(self, socket_type, service_name, service_port=None, service_type=None):
self.context = zmq.Context()
self.socket = self.context.socket(socket_type)
if not service_port:
service_port = self.socket.bind_to_random_port('tcp://*', min_port=49152, max_port=65535, max_tries=100)
else:
self.socket.bind("tcp://*:%d" % service_port)
print "Bound to port %d" % service_port
self.stream = ZMQStream(self.socket)
if not service_type:
service_type = socket_type_to_service(socket_type)
if socket_type == zmq.PUB:
# TODO: how to handle this with ROUTER/DEALER combinations...
self.heartbeat_timer = ioloop.PeriodicCallback(self._hearbeat, 1000)
self.heartbeat_timer.start()
if socket_type == zmq.ROUTER:
self.stream.on_recv(self._method_callback_wrapper)
bonjour_utilities.register_ioloop(ioloop.IOLoop.instance(), service_type, service_name, service_port)
def _method_callback_wrapper(self, datalist):
#print "_method_callback_wrapper called: %s" % repr(datalist)
if len(datalist) < 2:
return
client_id = datalist[0]
method = datalist[1]
args = datalist[2:]
#print "DEBUG: _method_callback_wrapper(%s, %s)" % (method, repr(args))
if not self.method_callbacks.has_key(method):
print "No such method: %s" % method
print "Methods: %s" % self.method_callbacks.keys()
return
for f in self.method_callbacks[method]:
resp = zmq_client_response(client_id, self.stream)
# TODO: make a wrapper object for sending responses and pass that instead of the client_id
#print "Calling f(resp, %s)" % repr(args)
f(resp, *args)
def register_method(self, name, callback):
if not self.method_callbacks.has_key(name):
self.method_callbacks[name] = []
self.method_callbacks[name].append(callback)
def run():
""" main method """
hostname = socket.gethostname()
context = zmq.Context()
# Configure the upstream socket
upstream = context.socket(zmq.SUB)
upstream_port = upstream.bind_to_random_port('tcp://*')
upstream_url = "tcp://%s:%d" % (hostname, upstream_port)
log.info("Upstream bound to %s" % upstream_url)
upstream.setsockopt_string(zmq.SUBSCRIBE, u"buildlight")
# Configure the downstream socket
downstream = context.socket(zmq.PUB)
downstream_port = downstream.bind_to_random_port('tcp://*')
downstream_url = "tcp://%s:%d" % (hostname, downstream_port)
log.info("Downstream bound to %s" % downstream_url)
ping_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM, socket.IPPROTO_UDP)
ping_socket.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
ping_socket.bind(('', 0))
loop = ioloop.IOLoop.instance()
def ping(serial_number=0):
ps = ping_socket
data = json.dumps({
'upstream':upstream_url,
'downstream':downstream_url,
'serial_number':serial_number
})
message = "%s\n%s" % (MAGIC_WORD, data)
ps.sendto(message.encode(), ('<broadcast>', 6969))
# Let's do this again some time?
callback = lambda : ping(serial_number=serial_number + 1)
loop.add_timeout(datetime.timedelta(seconds=PING_FREQUENCY), callback)
log.info("Ping: %s" % message)
loop.add_callback(ping)
def recieve(msg):
for m in msg:
downstream.send(m)
log.info('Proxied: %s' % (m))
zstr = ZMQStream(upstream)
zstr.on_recv(recieve)
try:
loop.start()
finally:
upstream.close()
downstream.close()
class MQAsyncSub():
def __init__(self, context, caller_id, category_filters):
cfg = Loader('mq').load()
self.cfg_mq = dict(cfg[1])
if self.cfg_mq['ip'].strip() == "*":
self.cfg_mq['ip'] = get_ip()
sub_addr = "tcp://{0}:{1}".format(self.cfg_mq['ip'], self.cfg_mq['sub_port'])
self.caller_id = caller_id
self.s_recv = context.socket(zmq.SUB)
self.s_recv.connect(sub_addr)
if len(category_filters) == 0:
self.s_recv.setsockopt_string(zmq.SUBSCRIBE, u'')
else:
for category_filter in category_filters:
self.s_recv.setsockopt_string(zmq.SUBSCRIBE, u"{0}".format(category_filter))
ioloop = IOLoop.instance()
self.stream_asyncmq = ZMQStream(self.s_recv, ioloop)
self.stream_asyncmq.on_recv(self._on_message)
def __del__(self):
# Not sure this is really mandatory
self.s_recv.close()
def _on_message(self, msg):
"""Received an event
will lookup the correct callback to use, and call it
:param: msg = the message received
"""
if len(msg) < 2:
# this sometimes happens, no idea why, probebly a bug in pyzmq
return
mid = msg[0].decode()
mid = mid.split('.')
if len(mid) < 3:
# this sometimes happens, no idea why, probebly a bug in pyzmq
return
# delete msg version
del mid[-1]
# delete timestamp
del mid[-1]
# build up the id again
mid = '.'.join(mid)
try:
jsons = json.loads(msg[1].decode())
self.on_message(mid, jsons)
except ValueError as e:
pass
def on_message(self, msg_id, content):
"""Public method called when a message arrived.
.. note:: Does nothing. Should be overloaded!
"""
pass
class RouterPubSubProxy:
'''
This is a proxy that has one front end socket, and two backend sockets. The
front end socket is a router that passes the messages to backend Pub. Pub
broadcasts them to all subscribers, which respond with results to backend
Sub. All communications on this proxy are done through IPC.
'''
def __init__(self,
front,
back_out,
back_in,
loop):
'''
Initializes the instance of RouterPubSubProxy.
@param front - channel name to be the routing stream
@param back_out - channel name of the publishing stream
@param back_in - channel name of result receiving stream
@param loop - IOLoop
'''
self._loop = loop
ctx = zmq.Context.instance()
# Create the front end stream
front_address = ZmqAddress(chan_name=front)
self._front_stream = ZMQStream(ctx.socket(zmq.ROUTER), io_loop=loop)
self._front_stream.setsockopt(zmq.ROUTER_MANDATORY, 1)
self._front_stream.bind(front_address.zmq_url())
# Create the back end streams
back_out_address = ZmqAddress(chan_name=back_out)
self._back_out_stream = ZMQStream(ctx.socket(zmq.PUB), io_loop=loop)
self._back_out_stream.bind(back_out_address.zmq_url())
back_in_address = ZmqAddress(chan_name=back_in)
self._back_in_stream = ZMQStream(ctx.socket(zmq.SUB), io_loop=loop)
self._back_in_stream.setsockopt(zmq.SUBSCRIBE, b'')
self._back_in_stream.bind(back_in_address.zmq_url())
def callback(from_name, to_name, zmq_stream, msgs):
log.info("Routing from {0} to {1} messages {2}"
.format(from_name, to_name, msgs))
zmq_stream.send_multipart(msgs)
zmq_stream.flush()
self._front_stream.on_recv(lambda msgs:
callback(front, back_out, self._back_out_stream, msgs))
self._back_in_stream.on_recv(lambda msgs:
callback(back_in, front, self._front_stream, msgs))
def start(self):
'''
Start this proxy.
'''
self._loop.start()
def recv(msg):
global pair_socket, pair_stream
print msg
port = int(msg[0])
print port
pair_socket = ctx.socket(zmq.PAIR)
pair_socket.connect("tcp://176.31.243.99:%d" % port)
pair_stream = ZMQStream(pair_socket, command_stream.io_loop)
pair_stream.on_recv(pair_recv)
send()
def start(self):
loop = self.loop
loop.add_handler(self.udp.handle.fileno(), self.handle_beacon, loop.READ)
stream = ZMQStream(self.pipe, loop)
stream.on_recv(self.control_message)
pc = PeriodicCallback(self.send_ping, PING_INTERVAL * 1000, loop)
pc.start()
pc = PeriodicCallback(self.reap_peers, PING_INTERVAL * 1000, loop)
pc.start()
loop.start()
def worker(vent_address,sink_address,sz,dtype):
import zmq
import theano
from zmq.eventloop import ioloop
ioloop.install()
from zmq.eventloop.zmqstream import ZMQStream
# Context
context = zmq.Context()
# Socket to receive messages on
receiver = context.socket(zmq.PULL)
receiver.connect(vent_address)
receiver_stream = ZMQStream(receiver)
# Socket to send messages to
sender = context.socket(zmq.PUSH)
sender.connect(sink_address)
def _worker(msg_list, sz=sz, dtype=dtype, sender=sender):
import theano
import numpy as np
msg = msg_list[0]
# if normalize_A does any inplace operation, we need to .copy() here:
new_A = np.frombuffer(buffer(msg), dtype=dtype).reshape(sz).copy()
new_A = l.model.normalize_A(new_A)
l.model.A.set_value(new_A.astype(theano.config.floatX))
x = l.get_databatch()
dA = l.model.gradient(x)['dA']
dA *= l.eta
param_max = np.max(np.abs(l.model.A.get_value()), axis=0)
update_max = np.max(np.abs(dA), axis=0)
update_max = np.max(update_max / param_max)
l._adapt_eta(update_max)
# no subset selection:
sender.send(dA,copy=False)
# subset selection:
#inds = np.argwhere(dA.sum(0) != 0.).ravel()
#subset_dA = dA[:, inds]
#sender.send_pyobj(dict(inds=inds, subset_dA=subset_dA))
receiver_stream.on_recv(_worker,copy=False)
iolooper = ioloop.IOLoop.instance()
iolooper.start()
return
class MQAsyncSub():
def __init__(self, context, caller_id, category_filters):
if ("domogik.common.configloader" in sys.modules):
cfg = Loader('mq').load()
self.cfg_mq = dict(cfg[1])
sub_addr = "tcp://{0}:{1}".format(self.cfg_mq['ip'], self.cfg_mq['sub_port'])
else:
ipaddr = Parameter.objects.get(key='mq-ip')
port = Parameter.objects.get(key='mq-sub_port')
sub_addr = "tcp://{0}:{1}".format(ipaddr.value, port.value)
self.caller_id = caller_id
self.s_recv = context.socket(zmq.SUB)
self.s_recv.connect(sub_addr)
if len(category_filters) == 0:
self.s_recv.setsockopt(zmq.SUBSCRIBE, '')
else:
for category_filter in category_filters:
self.s_recv.setsockopt(zmq.SUBSCRIBE, category_filter)
ioloop = IOLoop.instance()
self.stream = ZMQStream(self.s_recv, ioloop)
self.stream.on_recv(self._on_message)
def __del__(self):
# Not sure this is really mandatory
self.s_recv.close()
def _on_message(self, msg):
"""Received an event
will lookup the correct callback to use, and call it
:param: msg = the message received
"""
mid = msg[0]
mid = mid.split('.')
# delete msg version
del mid[-1]
# delete timestamp
del mid[-1]
# build up the id again
mid = '.'.join(mid)
try:
jsons = json.loads(msg[1])
self.on_message(mid, jsons)
except ValueError as e:
pass
def on_message(self, msg_id, content):
"""Public method called when a message arrived.
.. note:: Does nothing. Should be overloaded!
"""
pass
class Worker(object):
def __init__(self, robot,
data_in_sock='ipc:///tmp/robot-data-m2w.sock',
data_out_sock='ipc:///tmp/robot-data-w2m.sock',
msg_in_sock='ipc:///tmp/robot-msg-m2w.sock',
msg_out_sock='ipc:///tmp/robot-msg-w2m.sock',
io_loop=None):
self.identity = 'worker:%s:%s' % (socket.gethostname(), os.getpid())
context = zmq.Context()
self._io_loop = io_loop or IOLoop.instance()
self._in_socket = context.socket(zmq.PULL)
self._in_socket.connect(data_in_sock)
self._in_stream = ZMQStream(self._in_socket, io_loop)
self._out_socket = context.socket(zmq.PUB)
self._out_socket.connect(data_out_sock)
self._out_stream = ZMQStream(self._out_socket, io_loop)
self._running = False
self.robot = robot
self.robot.set_worker_identity(self.identity)
self.messenger = ClientMessenger(msg_in_sock, msg_out_sock,
context, io_loop)
def start(self):
logging.info('[%s] starting', self.identity)
self.messenger.start()
self.messenger.publish(CTRL_MSG_WORKER, self.identity,
CTRL_MSG_WORKER_ONLINE)
self._in_stream.on_recv(self._on_receive_request)
self._running = True
def stop(self):
self._running = False
self.messenger.stop()
def close(self):
self._in_stream.close()
self._in_socket.close()
self._out_stream.close()
self._out_socket.close()
self.messenger.close()
def _on_receive_request(self, zmq_msg):
msg = RequestMessage.deserialize(zmq_msg)
request = msg.request
logging.debug('[%s] receive request(%s)', self.identity, request.url)
self.robot.fetch(request)
class MainHandler(websocket.WebSocketHandler):
_first = True
@property
def ref(self):
return id(self)
def initialize(self):
print "WebSocket initialize"
self.push_socket = ctx.socket(zmq.PUSH)
self.sub_socket = ctx.socket(zmq.SUB)
self.push_socket.connect("ipc:///tmp/ws_push")
#self.sub_socket.connect("ipc:///tmp/ws_sub")
self.sub_socket.bind("ipc:///tmp/ws_sub")
self.sub_socket.setsockopt(zmq.SUBSCRIBE, "")
self.zmq_stream = ZMQStream(self.sub_socket)
self.zmq_stream.on_recv(self.zmq_msg_recv)
def open(self, *args, **kwargs):
print "WebSocket opened", args, kwargs
def on_message(self, message):
print "WebSocket on_message"
if self._first:
msg = {'message': message, 'id':self.ref, 'action':'connect'}
self._first = False
else:
msg = {'message': message, 'id':self.ref, 'action':'message'}
self.push_socket.send_pyobj(msg)
def on_close(self):
print "WebSocket closed"
msg = {'message': '', 'id': id(self), 'action': 'close'}
self.push_socket.send_pyobj(msg)
self.zmq_stream.close()
self.sub_socket.close()
self.push_socket.close()
def zmq_msg_recv(self, data):
print "zmq_msg_recv: %s" % repr(data)
for message in data:
message = pickle.loads(message)
#_id, _msg = message['id'], message['message']
print ' = ', repr(message)
#if _id != self.ref:
# continue
#self.write_message(_msg)
self.write_message(json.dumps(message, indent=4))
def setupZmqSubscriber():
ctx = zmq.Context()
s = ctx.socket(zmq.SUB)
s.connect('tcp://127.0.0.1:5000')
s.setsockopt(zmq.SUBSCRIBE, "")
stream = ZMQStream(s)
def echo(product):
print "receiving message: %s" % product
for socket in socket_connections:
socket.write_message(product[0])
stream.on_recv(echo)
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.pea.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def register_pollin(self):
use_uvloop()
import asyncio
asyncio.set_event_loop(asyncio.new_event_loop())
try:
import tornado.ioloop
self.io_loop = tornado.ioloop.IOLoop.current()
except (ModuleNotFoundError, ImportError):
self.logger.error(
'Since v0.3.6 Jina requires "tornado" as a base dependency, '
'we use its I/O event loop for non-blocking sockets. '
'Please try reinstall via "pip install -U jina" to include this dependency'
)
raise
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
self.in_sock.stop_on_recv()
def close(self):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`. """
if not self.is_closed:
# wait until the close signal is received
time.sleep(.01)
for s in self.opened_socks:
s.flush()
super().close()
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
except AttributeError as e:
self.logger.error(f'failed to stop. {e}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
self.in_sock.on_recv(self._in_sock_callback)
def start(self, callback: Callable[['jina_pb2.Message'], None]):
def _callback(msg, sock_type):
msg, num_bytes = _prepare_recv_msg(sock_type, msg,
self.args.check_version)
self.bytes_recv += num_bytes
self.msg_recv += 1
msg = callback(msg)
if msg:
self.send_message(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class CloneServer(object):
# Our server is defined by these properties
ctx = None # Context wrapper
kvmap = None # Key-value store
loop = None # IOLoop reactor
port = None # Main port we're working on
sequence = 0 # How many updates we're at
snapshot = None # Handle snapshot requests
publisher = None # Publish updates to clients
collector = None # Collect updates from clients
def __init__(self, port=5556):
self.port = port
self.ctx = zmq.Context()
self.kvmap = {}
self.loop = IOLoop.instance()
# Set up our clone server sockets
self.snapshot = self.ctx.socket(zmq.ROUTER)
self.publisher = self.ctx.socket(zmq.PUB)
self.collector = self.ctx.socket(zmq.PULL)
self.snapshot.bind("tcp://*:%d" % self.port)
self.publisher.bind("tcp://*:%d" % (self.port + 1))
self.collector.bind("tcp://*:%d" % (self.port + 2))
# Wrap sockets in ZMQStreams for IOLoop handlers
self.snapshot = ZMQStream(self.snapshot)
self.publisher = ZMQStream(self.publisher)
self.collector = ZMQStream(self.collector)
# Register our handlers with reactor
self.snapshot.on_recv(self.handle_snapshot)
self.collector.on_recv(self.handle_collect)
self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)
# basic log formatting:
logging.basicConfig(format="%(asctime)s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO)
def start(self):
# Run reactor until process interrupted
self.flush_callback.start()
try:
self.loop.start()
except KeyboardInterrupt:
pass
def handle_snapshot(self, msg):
"""snapshot requests"""
if len(msg) != 3 or msg[1] != "ICANHAZ?":
print "E: bad request, aborting"
dump(msg)
self.loop.stop()
return
identity, request, subtree = msg
if subtree:
# Send state snapshot to client
route = Route(self.snapshot, identity, subtree)
# For each entry in kvmap, send kvmsg to client
for k, v in self.kvmap.items():
send_single(k, v, route)
# Now send END message with sequence number
logging.info("I: Sending state shapshot=%d" % self.sequence)
self.snapshot.send(identity, zmq.SNDMORE)
kvmsg = KVMsg(self.sequence)
kvmsg.key = "KTHXBAI"
kvmsg.body = subtree
kvmsg.send(self.snapshot)
def handle_collect(self, msg):
"""Collect updates from clients"""
kvmsg = KVMsg.from_msg(msg)
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
ttl = kvmsg.get('ttl')
if ttl is not None:
kvmsg['ttl'] = time.time() + ttl
kvmsg.store(self.kvmap)
logging.info("I: publishing update=%d", self.sequence)
def flush_ttl(self):
"""Purge ephemeral values that have expired"""
for key, kvmsg in self.kvmap.items():
self.flush_single(kvmsg)
def flush_single(self, kvmsg):
"""If key-value pair has expired, delete it and publish the fact
to listening clients."""
if kvmsg.get('ttl', 0) <= time.time():
kvmsg.body = ""
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
del self.kvmap[kvmsg.key]
logging.info("I: publishing delete=%d", self.sequence)
class TagHandler(object):
def __init__(self, connection):
super(TagHandler, self).__init__()
self.connection = connection
self.send_element = connection.send_element
self.send_string = connection.send_string
self.jid = None
self.hostname = None
self.authenticated = False
self.session_active = False
self.publisher = get_publisher()
self.pull_url = None
self.pull_socket = None
def close(self):
"""Is called when the client connection is closed to do cleanup work"""
# unregister from forwarder
if self.pull_socket is not None:
reg_msg = ZMQForwarder_message('UNREGISTER')
reg_msg.attributes = self.pull_url
self.publisher.send_pyobj(reg_msg)
self.processed_stream.close()
self.pull_socket.close()
self.pull_socket = None
def contenthandler(self, tree):
"""Handles an incomming content tree"""
# set/replace the from attribute in stanzas as required
# by RFC 6120 Section 8.1.2.1
if self.authenticated:
tree.set("from", str(self.jid))
try:
if tree.tag == "auth":
if self.authenticated:
raise NotAllowedError
self.authenticate(tree)
elif tree.tag == "iq":
if not self.authenticated:
raise NotAuthorizedError
if self.jid.resource is None:
self.set_resource(tree)
else:
if not self.session_active:
first_element = tree[0]
if first_element.tag == "session" and \
first_element.get("xmlns") == SESSION_NS:
response_element = ET.Element("iq")
response_element.set("type", "result")
response_element.set("id", tree.get("id"))
session_element = ET.SubElement(
response_element, "session")
session_element.set("xmlns", SESSION_NS)
self.send_element(response_element)
log.debug("Sent empty session element")
self.processed_stream.stop_on_recv()
self.processed_stream.on_recv(
self.send_list, False)
else:
self.publish_stanza(tree)
else:
self.publish_stanza(tree)
elif tree.tag in ["message", "presence"]:
if not self.authenticated:
raise NotAuthorizedError
self.publish_stanza(tree)
except (StreamError, e):
self.send_string(unicode(e))
self.connection.stop_connection()
def publish_stanza(self, tree):
log.debug("Publishing Stanza %s" % (ET.tostring(tree)))
self.publisher.send(cPickle.dumps(tree))
def masked_send_list(self, msgs):
"""Unmark waiting for a session element if we received another stanza response"""
self.session_active = True
self.processed_stream.stop_on_recv()
self.processed_stream.on_recv(self.send_list, False)
self.send_list(msgs)
def send_list(self, msgs):
try:
for msg in msgs:
tmp = cPickle.loads(msg.bytes)
if tmp.get("to") == str(
self.jid) or tmp.get("to") == self.jid.bare:
self.send_element(tmp)
except IOError:
self.connection.stop_connection()
def set_resource(self, tree):
"""Set a resource on our JID"""
bind_element = tree[0]
if tree.get("type") != "set" or \
bind_element.tag != 'bind' or\
bind_element.get("xmlns") != BIND_NS:
raise NotAuthorizedError
resource_element = bind_element.find("resource")
if resource_element is None:
# No prefered resource was set, generate one
self.jid.resource = uuid.uuid4().hex
else:
self.jid.resource = resource_element.text
if not self.jid.validate():
raise BadRequestError
# Check if given resource is already in use
known_jids = get_known_jids()
try:
known_jids.append(self.jid)
except ValueError:
raise ConflictError
log.info("Bound connection as %s" % str(self.jid))
# Connect to forwarder to receive stanzas sent back to client
log.debug('Registering Client at forwarder..')
self.pull_socket = zmq.Context().socket(zmq.PULL)
self.processed_stream = ZMQStream(self.pull_socket,
self.connection.stream.io_loop)
self.processed_stream.on_recv(self.masked_send_list, False)
port = self.pull_socket.bind_to_random_port('tcp://127.0.0.1')
self.pull_url = 'tcp://127.0.0.1:' + str(port)
reg_msg = ZMQForwarder_message('REGISTER')
reg_msg.attributes = (config.get('ipc',
'password'), self.pull_url, self.jid)
self.publisher.send_pyobj(reg_msg)
# Send registered resource back to client
response_element = ET.Element("iq")
response_element.set("type", "result")
response_element.set("id", tree.get("id"))
bind_element = ET.SubElement(response_element, "bind")
bind_element.set("xmlns", BIND_NS)
jid_element = ET.SubElement(bind_element, "jid")
jid_element.text = str(self.jid)
self.send_element(response_element)
def authenticate(self, tree):
"""Authenticates user for session"""
# Currently RFC specifies only SASL as supported way of auth'ing
handler = SASLAuthHandler()
if tree.get('xmlns') != handler.namespace:
raise MalformedRequestError
handler.process(tree)
self.connection.parser.reset()
self.jid = JID("@".join([handler.authenticated_user, self.hostname]))
self.authenticated = True
response_element = ET.Element("success")
response_element.set("xmlns", handler.namespace)
self.send_element(response_element)
def add_auth_options(self, feature_element):
"""Add supported auth mechanisms to feature element"""
handler = SASLAuthHandler()
mechtype_element = ET.SubElement(feature_element, "mechanisms")
mechtype_element.set("xmlns", handler.namespace)
for mech in handler.supported_mechs:
mech_element = ET.SubElement(mechtype_element, 'mechanism')
mech_element.text = mech
def add_server_features(self, feature_element):
bind = ET.SubElement(feature_element, "bind")
bind.set("xmlns", "urn:ietf:params:xml:ns:xmpp-bind")
# Session establishment is deprecated in RFC6121 but Appendix E
# suggests to still advertise it as feature for compatibility.
session = ET.SubElement(feature_element, "session")
session.set("xmlns", "urn:ietf:params:xml:ns:xmpp-session")
def streamhandler(self, attrs):
"""Handles a stream start"""
if attrs == {}:
# </stream:stream> received
self.connection.stop_connection()
else:
# check if we are responsible for this stream
self.hostname = attrs.getValue("to")
if self.hostname not in config.getlist("listeners", "domains"):
raise HostUnknownError
# Stream restart
stream = ET.Element("stream:stream")
stream.set("xmlns", attrs.getValue("xmlns"))
stream.set("from", self.hostname)
stream.set("id", uuid.uuid4().hex)
stream.set("xml:lang", "en")
stream.set("xmlns:stream", "http://etherx.jabber.org/streams")
# only include version in response if client sent its max supported
# version (RFC6120 Section 4.7.5)
try:
if attrs.getValue("version") != "1.0":
raise UnsupportedVersionError
stream.set("version", "1.0")
except KeyError:
pass
try:
from_jid = JID(attrs.getValue("from"))
stream.set("to", unicode(from_jid))
except ValueError:
raise InvalidFromError
except KeyError:
pass
start_stream = """<?xml version="1.0"?>""" + ET.tostring(stream)
# Element has subitems but are added later to the stream,
# so don't mark it a single element
self.send_string(start_stream.replace("/>", ">"))
# Make a list of supported features
features = ET.Element("stream:features")
if not self.authenticated:
self.add_auth_options(features)
else:
self.add_server_features(features)
self.send_element(features)
class Emperor(object):
"""Обертка вокруг uwsgi-emperor, отвечающая за его запуск и управление вассалами."""
def __init__(self, root_dir):
"""Инициализирует uwsgi-emperor.
:param root_dir: Полный путь к корневой директории uwsgi-emperor
:type root_dir: str
"""
self.__root_dir__ = root_dir
if not os.path.exists(self.vassal_dir):
log_message("Vassal directory does not exist, creating one",
component="Emperor")
os.mkdir(self.vassal_dir)
emperor_pid = 0
if os.path.exists(self.pidfile):
with open(self.pidfile) as pid_file:
try:
emperor_pid = int(pid_file.read())
psutil.Process(emperor_pid)
log_message("Found running emperor server",
component="Emperor")
except (ValueError, psutil.NoSuchProcess):
os.remove(self.pidfile)
if not emperor_pid:
emperor = subprocess.Popen([
self.uwsgi_binary, "--plugins-dir", self.binary_dir,
"--emperor", self.vassal_dir, "--pidfile", self.pidfile,
"--logger", "zeromq:tcp://127.0.0.1:5123", "--daemonize",
"/dev/null", "--emperor-stats", "127.0.0.1:1777",
"--emperor-required-heartbeat", "40", "--emperor-throttle",
"10000", "--vassal-set", "plugins-dir={}".format(
self.binary_dir)
],
bufsize=1,
close_fds=True)
code = emperor.wait()
assert code == 0, "Error starting emperor server"
log_message("Started emperor server", component="Emperor")
self.vassals = {}
ctx = zmq.Context()
s = ctx.socket(zmq.PULL)
s.bind('tcp://127.0.0.1:5123')
self.stream = ZMQStream(s)
self.stream.on_recv(self.log_message)
@property
def root_dir(self):
"""Корневая директория uwsgi-emperor.
:returns: Полный путь к корневой директории
:rtype: str
"""
return self.__root_dir__
@property
def binary_dir(self):
"""Директория с исполняемыми файлами и плагинами uwsgi-emperor.
:returns: Полный путь к директории с исполняемыми файлами
:rtype: str
"""
return os.path.join(self.root_dir, "bin")
@property
def uwsgi_binary(self):
"""Основной исполняемый файл uwsgi.
:returns: Полный путь к исполняемому файлу uwsgi
:rtype: str
"""
return os.path.join(self.binary_dir, "uwsgi")
@property
def vassal_dir(self):
"""Директория с вассалами uwsgi.
:returns: Полный путь к директории вассалов
:rtype: str
"""
return os.path.join(self.root_dir, "vassals")
@property
def pidfile(self):
"""Pid-файл uwsgi-emperor.
:returns: Полный путь к pid-файлу
:rtype: str
"""
return os.path.join(self.root_dir, "emperor.pid")
@property
def vassal_names(self):
"""Возвращает спиоск имен вассалов, активных в данный момент. Имена передаются без расширения.
:returns: Список имен активных вассалов
:rtype: list
"""
raw_names = os.listdir(self.vassal_dir)
return [name[:-4] for name in raw_names]
@coroutine
def call_vassal_rpc(self, vassal, *args):
"""Вызывает rpc-функцию вассала.
:param vassal: Имя вассала
:type vassal: str
:returns: Результат выполнения функции
:rtype: dict
"""
stats = self.stats(vassal)
try:
assert "pid" in stats
process = psutil.Process(stats["pid"])
host, port = process.connections()[0].laddr
client = yield TCPClient().connect(host, port)
yield client.write(
pack('<BHB', 173, sum(2 + len(arg) for arg in args), 0))
for arg in args:
yield client.write(pack('<H', len(arg)) + arg)
data = yield client.read_bytes(4)
modifier1, datasize, modifier2 = unpack("<BHB", data)
data = yield client.read_bytes(datasize)
raise Return({"result": "success", "data": data})
except (AssertionError, psutil.NoSuchProcess):
raise Return({"result": "failure", "message": "Not running"})
except StreamClosedError:
raise Return({"result": "failure", "message": "Call failure"})
def stop(self):
"""Останавливает uwsgi-emperor и очищает директорию вассалов."""
log_message("Stopping uwsgi emperor", component="Emperor")
subprocess.call([self.uwsgi_binary, "--stop", self.pidfile])
os.remove(self.pidfile)
for name in os.listdir(self.vassal_dir):
os.remove(os.path.join(self.vassal_dir, name))
def start_vassal(self, vassal):
"""Запускает указанного вассала.
:param vassal: Запускаемый вассал
:type vassal: Vassal
"""
cfg_path = os.path.join(self.vassal_dir, "{}.ini".format(vassal.id))
self.vassals[str(vassal.id)] = vassal
if os.path.exists(cfg_path):
with open(cfg_path, "r") as cfg:
data = cfg.read()
if data == vassal.get_config():
return
log_message(
"Leaf {} have stale configuration, will restart".format(
vassal.id))
with open(cfg_path, "w") as cfg:
cfg.write(vassal.get_config())
def stop_vassal(self, vassal):
"""Останавливает указанного вассала.
:param vassal: Останавливаемый вассал
:type vassal: Vassal
"""
cfg_path = os.path.join(self.vassal_dir, "{}.ini".format(vassal.id))
if str(vassal.id) in self.vassals:
del self.vassals[str(vassal.id)]
if os.path.exists(cfg_path):
os.remove(cfg_path)
def soft_restart_vassal(self, vassal):
"""Выполняет плавный перезапуск вассала.
:param vassal: Перезапускаемый вассал
:type vassal: Vassal
"""
cfg_path = os.path.join(self.vassal_dir, "{}.ini".format(vassal.id))
if os.path.exists(cfg_path):
os.utime(cfg_path, None)
def stats(self, vassal):
"""Возвращает статистику по указанному вассалу.
:param vassal: Имя вассала
:type vassal: str
:returns: Статистика по вассалу
:rtype: dict
"""
for l in self.__stats__()["vassals"]:
if l["id"] == "{}.ini".format(vassal):
return l
return {}
def __stats__(self):
"""Возвращает внутреннюю статистику uwsgi-emperor.
:returns: Словарь со статистикой
:rtype: dict
"""
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.connect(("127.0.0.1", 1777))
data = ""
while True:
new_data = s.recv(4096)
if len(new_data) < 1:
break
data += new_data.decode('utf8')
return json.loads(data)
@coroutine
def log_message(self, message):
"""Обрабатывает входящее сообщение uwsgi-emperor.
:param message: Входящее сообщение ZeroMQ
:type message: list
"""
for m in (_.strip() for _ in message if _.strip()):
data = logparse_emperor(m)
if data.get("log_type") == "emperor_vassal_ready":
vassal_id = data.get("vassal")
if vassal_id in self.vassals:
self.vassals[vassal_id].status = "Running"
elif data.get("log_type") == "emperor_vassal_removed":
vassal_id = data.get("vassal")
if vassal_id in self.vassals:
if self.vassals[vassal_id].status in ("Started", "Failed"):
self.vassals[vassal_id].status = "Failed"
else:
self.vassals[vassal_id].status = "Stopped"
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 MainHandler(websocket.WebSocketHandler):
_first = True
@property
def ref(self):
return id(self)
def initialize(self):
print 'Initializing tornado websocket'
self.context = zmq.Context()
self.sub_socket = self.context.socket(zmq.SUB)
self.sub_socket.connect(PUB_ADDRESS)
self.zmq_stream = ZMQStream(self.sub_socket)
self.zmq_stream.on_recv(self.recv_func)
def check_origin(self, origin):
return True
def open(self, *args, **kwargs):
self.write_message("Open_Success")
self.zmq_subscribe(args[0])
def on_message(self, message):
if self._first:
msg = {'message': message, 'id': self.ref, 'action': 'connect'}
print 'in if part - tornado server'
self._first = False
else:
msg = {'message': message, 'id': self.ref, 'action': 'message'}
print 'in else part - tornado server'
print msg
self.write_message(msg)
def on_close(self):
print("WebSocket closed")
self.sub_socket.setsockopt(zmq.LINGER, 0)
self.sub_socket.close()
self.zmq_stream.close()
self.context.term()
def recv_func(self, message):
#self.process_data(message)
try:
message = message[0]
topic, messagedata = message.split(ZMQ_SEPARATOR)
messagedata = json.loads(messagedata)
# self.vip.pubsub.publish('pubsub',topic,message=messagedata)
print topic, messagedata
self.write_message({'message': messagedata, 'topic': topic})
except Exception as er:
print er
pass
def process_data(self, data):
print data
zmessage = {'topic': '', 'headers': {}, 'message': ''}
for item in data:
if '/agent/ui' in item or '/app/ui' in item:
zmessage['topic'] = item
elif 'Date' in str(item):
mesg = json.loads(item)
zmessage['headers'] = mesg
else:
if '[' in item:
item = eval(item)
print type(item)
item = item[0]
if item[0] == '{':
item = json.loads(item)
zmessage['message'] = item
else:
zmessage['message'] = item
self.write_message(zmessage)
def zmq_subscribe(self, agent_id):
#self.sub_socket.setsockopt(zmq.SUBSCRIBE, "")
topicfilter = "to/ui/"
self.sub_socket.setsockopt_string(zmq.SUBSCRIBE,
"to/ui/from/" + agent_id)
class BinaryStar(object):
ctx = None # Our private context
loop = None # Reactor loop
statepub = None # State publisher
statesub = None # State subscriber
state = None # Current state
event = None # Current event
peer_expiry = 0 # When peer is considered 'dead'
voter_callback = None # Voting socket handler
master_callback = None # Call when become master
slave_callback = None # Call when become slave
heartbeat = None # PeriodicCallback for
def __init__(self, primary, local, remote):
# initialize the Binary Star
self.ctx = zmq.Context()
self.loop = IOLoop.instance()
self.state = STATE_PRIMARY if primary else STATE_BACKUP
# Create publisher for state going to peer
self.statepub = self.ctx.socket(zmq.PUB)
self.statepub.bind(local)
# Create subscriber for state coming from peer
self.statesub = self.ctx.socket(zmq.SUB)
self.statesub.setsockopt(zmq.SUBSCRIBE, '')
self.statesub.connect(remote)
# wrap statesub in ZMQStream for event triggers
self.statesub = ZMQStream(self.statesub, self.loop)
# setup basic reactor events
self.heartbeat = PeriodicCallback(self.send_state, HEARTBEAT,
self.loop)
self.statesub.on_recv(self.recv_state)
# setup log formmater
def update_peer_expiry(self):
"""Update peer expiry time to be 2 heartbeats from now."""
self.peer_expiry = time.time() + 2e-3 * HEARTBEAT
def start(self):
self.update_peer_expiry()
self.heartbeat.start()
return self.loop.start()
def execute_fsm(self):
"""Binary Star finite state machine (applies event to state)
returns True if connections should be accepted, False otherwise.
"""
accept = True
if (self.state == STATE_PRIMARY):
# Primary server is waiting for peer to connect
# Accepts CLIENT_REQUEST events in this state
if (self.event == PEER_BACKUP):
print("I: connected to backup (slave), ready as master")
self.state = STATE_ACTIVE
if (self.master_callback):
self.loop.add_callback(self.master_callback)
elif (self.event == PEER_ACTIVE):
print("I: connected to backup (master), ready as slave")
self.state = STATE_PASSIVE
if (self.slave_callback):
self.loop.add_callback(self.slave_callback)
elif (self.event == CLIENT_REQUEST):
if (time.time() >= self.peer_expiry):
print("I: request from client, ready as master")
self.state = STATE_ACTIVE
if (self.master_callback):
self.loop.add_callback(self.master_callback)
else:
# don't respond to clients yet - we don't know if
# the backup is currently Active as a result of
# a successful failover
accept = False
elif (self.state == STATE_BACKUP):
# Backup server is waiting for peer to connect
# Rejects CLIENT_REQUEST events in this state
if (self.event == PEER_ACTIVE):
print("I: connected to primary (master), ready as slave")
self.state = STATE_PASSIVE
if (self.slave_callback):
self.loop.add_callback(self.slave_callback)
elif (self.event == CLIENT_REQUEST):
accept = False
elif (self.state == STATE_ACTIVE):
# Server is active
# Accepts CLIENT_REQUEST events in this state
# The only way out of ACTIVE is death
if (self.event == PEER_ACTIVE):
# Two masters would mean split-brain
print("E: fatal error - dual masters, aborting")
raise FSMError("Dual Masters")
elif (self.state == STATE_PASSIVE):
# Server is passive
# CLIENT_REQUEST events can trigger failover if peer looks dead
if (self.event == PEER_PRIMARY):
# Peer is restarting - become active, peer will go passive
print("I: primary (slave) is restarting, ready as master")
self.state = STATE_ACTIVE
elif (self.event == PEER_BACKUP):
# Peer is restarting - become active, peer will go passive
print("I: backup (slave) is restarting, ready as master")
self.state = STATE_ACTIVE
elif (self.event == PEER_PASSIVE):
# Two passives would mean cluster would be non-responsive
print("E: fatal error - dual slaves, aborting")
raise FSMError("Dual slaves")
elif (self.event == CLIENT_REQUEST):
# Peer becomes master if timeout has passed
# It's the client request that triggers the failover
assert (self.peer_expiry > 0)
if (time.time() >= self.peer_expiry):
# If peer is dead, switch to the active state
print("I: failover successful, ready as master")
self.state = STATE_ACTIVE
else:
# If peer is alive, reject connections
accept = False
# Call state change handler if necessary
if (self.state == STATE_ACTIVE and self.master_callback):
self.loop.add_callback(self.master_callback)
return accept
# ---------------------------------------------------------------------
# Reactor event handlers...
def send_state(self):
"""Publish our state to peer"""
self.statepub.send("%d" % self.state)
def recv_state(self, msg):
"""Receive state from peer, execute finite state machine"""
state = msg[0]
if state:
self.event = int(state)
self.update_peer_expiry()
self.execute_fsm()
def voter_ready(self, msg):
"""Application wants to speak to us, see if it's possible"""
# If server can accept input now, call appl handler
self.event = CLIENT_REQUEST
if self.execute_fsm():
print "CLIENT REQUEST"
self.voter_callback(self.voter_socket, msg)
else:
# Message will be ignored
pass
# -------------------------------------------------------------------------
#
def register_voter(self, endpoint, type, handler):
"""Create socket, bind to local endpoint, and register as reader for
voting. The socket will only be available if the Binary Star state
machine allows it. Input on the socket will act as a "vote" in the
Binary Star scheme. We require exactly one voter per bstar instance.
handler will always be called with two arguments: (socket,msg)
where socket is the one we are creating here, and msg is the message
that triggered the POLLIN event.
"""
assert self.voter_callback is None
socket = self.ctx.socket(type)
socket.bind(endpoint)
self.voter_socket = socket
self.voter_callback = handler
stream = ZMQStream(socket, self.loop)
stream.on_recv(self.voter_ready)
class IOPubThread(object):
"""An object for sending IOPub messages in a background thread
prevents a blocking main thread
IOPubThread(pub_socket).background_socket is a Socket-API-providing object
whose IO is always run in a thread.
"""
def __init__(self, socket, pipe=False):
self.socket = socket
self.background_socket = BackgroundSocket(self)
self._master_pid = os.getpid()
self._pipe_flag = pipe
self.io_loop = IOLoop()
if pipe:
self._setup_pipe_in()
self.thread = threading.Thread(target=self._thread_main)
self.thread.daemon = True
def _thread_main(self):
"""The inner loop that's actually run in a thread"""
self.io_loop.start()
self.io_loop.close()
def _setup_pipe_in(self):
"""setup listening pipe for subprocesses"""
ctx = self.socket.context
# use UUID to authenticate pipe messages
self._pipe_uuid = uuid.uuid4().bytes
pipe_in = ctx.socket(zmq.PULL)
pipe_in.linger = 0
try:
self._pipe_port = pipe_in.bind_to_random_port("tcp://127.0.0.1")
except zmq.ZMQError as e:
warnings.warn("Couldn't bind IOPub Pipe to 127.0.0.1: %s" % e +
"\nsubprocess output will be unavailable.")
self._pipe_flag = False
pipe_in.close()
return
self._pipe_in = ZMQStream(pipe_in, self.io_loop)
self._pipe_in.on_recv(self._handle_pipe_msg)
def _handle_pipe_msg(self, msg):
"""handle a pipe message from a subprocess"""
if not self._pipe_flag or not self._is_master_process():
return
if msg[0] != self._pipe_uuid:
print("Bad pipe message: %s", msg, file=sys.__stderr__)
return
self.send_multipart(msg[1:])
def _setup_pipe_out(self):
# must be new context after fork
ctx = zmq.Context()
pipe_out = ctx.socket(zmq.PUSH)
pipe_out.linger = 3000 # 3s timeout for pipe_out sends before discarding the message
pipe_out.connect("tcp://127.0.0.1:%i" % self._pipe_port)
return ctx, pipe_out
def _is_master_process(self):
return os.getpid() == self._master_pid
def _check_mp_mode(self):
"""check for forks, and switch to zmq pipeline if necessary"""
if not self._pipe_flag or self._is_master_process():
return MASTER
else:
return CHILD
def start(self):
"""Start the IOPub thread"""
self.thread.start()
# make sure we don't prevent process exit
# I'm not sure why setting daemon=True above isn't enough, but it doesn't appear to be.
atexit.register(self.stop)
def stop(self):
"""Stop the IOPub thread"""
if not self.thread.is_alive():
return
self.io_loop.add_callback(self.io_loop.stop)
self.thread.join()
def close(self):
self.socket.close()
self.socket = None
@property
def closed(self):
return self.socket is None
def send_multipart(self, *args, **kwargs):
"""send_multipart schedules actual zmq send in my thread.
If my thread isn't running (e.g. forked process), send immediately.
"""
if self.thread.is_alive():
self.io_loop.add_callback(
lambda: self._really_send(*args, **kwargs))
else:
self._really_send(*args, **kwargs)
def _really_send(self, msg, *args, **kwargs):
"""The callback that actually sends messages"""
mp_mode = self._check_mp_mode()
if mp_mode != CHILD:
# we are master, do a regular send
self.socket.send_multipart(msg, *args, **kwargs)
else:
# we are a child, pipe to master
# new context/socket for every pipe-out
# since forks don't teardown politely, use ctx.term to ensure send has completed
ctx, pipe_out = self._setup_pipe_out()
pipe_out.send_multipart([self._pipe_uuid] + msg, *args, **kwargs)
pipe_out.close()
ctx.term()
class Test_MDPWorker(unittest.TestCase):
endpoint = b'tcp://127.0.0.1:7777'
service = b'test'
def setUp(self):
print 'set up'
sys.stdout.flush()
self.context = zmq.Context()
self.broker = None
self._msgs = []
return
def tearDown(self):
print 'tear down'
sys.stdout.flush()
if self.broker:
self._stop_broker()
self.broker = None
## self.context.term()
self.context = None
return
def _on_msg(self, msg):
if _do_print:
print 'broker received:',
pprint(msg)
self.target = msg.pop(0)
if msg[1] == chr(1): # ready
print 'READY'
self.target = msg[0]
return
if msg[1] == chr(4): # ready
print 'HB'
return
if msg[1] == chr(3): # reply
IOLoop.instance().stop()
return
return
def _start_broker(self, do_reply=False):
"""Helper activating a fake broker in the ioloop.
"""
socket = self.context.socket(zmq.XREP)
self.broker = ZMQStream(socket)
self.broker.socket.setsockopt(zmq.LINGER, 0)
self.broker.bind(self.endpoint)
self.broker.on_recv(self._on_msg)
self.broker.do_reply = do_reply
self.broker.ticker = PeriodicCallback(self._tick, MyWorker.HB_INTERVAL)
self.broker.ticker.start()
self.target = None
return
def _stop_broker(self):
if self.broker:
self.broker.ticker.stop()
self.broker.ticker = None
self.broker.socket.close()
self.broker.close()
self.broker = None
return
def _tick(self):
if self.broker and self.target:
msg = [self.target, b'MPDW01', chr(4)]
self.broker.send_multipart(msg)
return
def send_req(self):
data = ['AA', 'bb']
msg = [self.target, b'MPDW01', chr(2), self.target, b''] + data
print 'borker sending:',
pprint(msg)
self.broker.send_multipart(msg)
return
# tests follow
def test_01_simple_01(self):
"""Test MDPWorker simple req/reply.
"""
self._start_broker()
time.sleep(0.2)
worker = MyWorker(self.context, self.endpoint, self.service)
sender = DelayedCallback(self.send_req, 1000)
sender.start()
IOLoop.instance().start()
worker.shutdown()
self._stop_broker()
return
class MarketDataSubscriber(object):
"""" MarketDataSubscriber. """
@classmethod
def instance(cls):
if not hasattr(cls, "_instance"):
cls._instance = cls()
return cls._instance
def __init__(self, symbol="ALL", application=None):
self.symbol = str(symbol)
self.buy_side = []
self.sell_side = []
self.volume_dict = {}
self.inst_status = InstrumentStatusHelper(symbol)
self.md_pub_socket = None
self.md_pub_socket_stream = None
self.is_ready = False
self.process_later = []
self.application = application
self.db_session = application.db_session
def subscribe(self, zmq_context, trade_pub_connection_string,
trade_client):
"""" subscribe. """
self.md_pub_socket = zmq_context.socket(zmq.SUB)
self.md_pub_socket.connect(trade_pub_connection_string)
self.md_pub_socket.setsockopt(zmq.SUBSCRIBE,
"^MD_FULL_REFRESH_" + self.symbol + '$')
self.md_pub_socket.setsockopt(zmq.SUBSCRIBE,
"^MD_TRADE_" + self.symbol + '$')
self.md_pub_socket.setsockopt(zmq.SUBSCRIBE,
"^MD_INCREMENTAL_" + self.symbol + ".0$")
self.md_pub_socket.setsockopt(zmq.SUBSCRIBE,
"^MD_INCREMENTAL_" + self.symbol + ".1$")
self.md_pub_socket_stream = ZMQStream(self.md_pub_socket)
self.md_pub_socket_stream.on_recv(self.on_md_publish)
md_subscription_msg = {
'MsgType': 'V',
'MDReqID': '0', # not important.
'SubscriptionRequestType': '0',
'MarketDepth': 0,
'TradeDate': time.strftime("%Y%m%d", time.localtime()),
'MDUpdateType': '0',
'MDEntryTypes': ['0', '1', '2'],
'Instruments': [self.symbol]
}
self.application.log('DEBUG', 'MARKET_DATA_SUBSCRIBER', 'SUBSCRIBE')
return trade_client.sendJSON(md_subscription_msg)
def ready(self):
self.is_ready = True
for trade in self.process_later:
self.on_trade(trade)
self.process_later = []
@staticmethod
def get(symbol, application=None):
"""" get. """
global MDSUBSCRIBEDICT
if symbol not in MDSUBSCRIBEDICT:
MDSUBSCRIBEDICT[symbol] = MarketDataSubscriber(symbol, application)
return MDSUBSCRIBEDICT[symbol]
def get_last_trades(self):
"""" get_last_trades. """
return Trade.get_last_trades(self.db_session)
def get_trades(self, symbol, since):
"""" get_trades. """
return Trade.get_trades(self.db_session, symbol, since)
def on_md_publish(self, publish_msg):
"""" on_md_publish. """
start = datetime.datetime.now()
topic = publish_msg[0]
raw_message = publish_msg[1]
msg = JsonMessage(raw_message)
if msg.type == 'W': # Full Refresh
self.on_md_full_refresh(msg)
elif msg.type == 'X': # Incremental
self.on_md_incremental(msg)
finish = datetime.datetime.now()
self.application.log(
"DEBUG", "PERF",
str([(finish - start).total_seconds(),
"MarketDataSubscriber.on_md_publish", "1",
[topic, raw_message]]))
def on_md_full_refresh(self, msg):
"""" on_md_full_refresh. """
# TODO: Check if our current order book is sync with the full refresh
if msg.get('MarketDepth') != 1: # Has Market Depth
self.on_book_clear()
self.on_trade_clear()
group = msg.get('MDFullGrp')
for entry in group:
entry_type = entry.get('MDEntryType')
if entry_type == '0' or entry_type == '1':
self.on_book_new_order(entry)
elif entry_type == '2':
self.on_trade(entry)
def on_md_incremental(self, msg):
"""" on_md_incremental. """
if msg.get('MDBkTyp') == '3': # Order Depth
group = msg.get('MDIncGrp')
for entry in group:
entry_type = entry.get('MDEntryType')
signal_order_depth_entry(self.symbol + '.3.' + entry_type,
entry)
if entry_type == '0' or entry_type == '1':
update_action = entry.get('MDUpdateAction')
if update_action == '0':
self.on_book_new_order(entry)
elif update_action == '1':
self.on_book_update_order(entry)
elif update_action == '2':
self.on_book_delete_order(entry)
elif update_action == '3':
self.on_book_delete_orders_thru(entry)
elif entry_type == '2':
self.on_trade(entry)
signal_publish_md_order_depth_incremental(self.symbol + '.3',
{"MsgType": "X"})
def on_book_clear(self):
"""" on_book_clear. """
self.buy_side = []
self.sell_side = []
def on_trade_clear(self):
"""" on_trade_clear. """
self.volume_dict = {}
def on_book_delete_orders_thru(self, msg):
"""" on_book_delete_orders_thru. """
index = msg.get('MDEntryPositionNo')
side = msg.get('MDEntryType')
if side == '0':
self.buy_side = self.buy_side[index:]
if self.buy_side:
self.inst_status.set_best_bid(self.buy_side[0]['price'])
else:
self.inst_status.set_best_bid(None)
elif side == '1':
self.sell_side = self.sell_side[index:]
if self.sell_side:
self.inst_status.set_best_ask(self.sell_side[0]['price'])
else:
self.inst_status.set_best_ask(None)
def on_book_delete_order(self, msg):
"""" on_book_delete_order. """
index = msg.get('MDEntryPositionNo') - 1
side = msg.get('MDEntryType')
if side == '0':
del self.buy_side[index]
if index == 0:
if self.buy_side:
self.inst_status.set_best_bid(self.buy_side[0]['price'])
else:
self.inst_status.set_best_bid(None)
elif side == '1':
del self.sell_side[index]
if index == 0:
if self.sell_side:
self.inst_status.set_best_ask(self.sell_side[0]['price'])
else:
self.inst_status.set_best_ask(None)
def on_book_new_order(self, msg):
"""" on_book_new_order. """
index = msg.get('MDEntryPositionNo') - 1
order = {
'price': msg.get('MDEntryPx'),
'qty': msg.get('MDEntrySize'),
'username': msg.get('Username'),
'user_id': msg.get('UserID'),
'broker': msg.get('Broker'),
'order_id': msg.get('OrderID'),
'side': msg.get('MDEntryType'),
'order_time': msg.get('MDEntryTime'),
'order_date': msg.get('MDEntryDate')
}
if msg.get('MDEntryType') == '0': # buy
self.buy_side.insert(index, order)
if index == 0:
self.inst_status.set_best_bid(msg.get('MDEntryPx'))
elif msg.get('MDEntryType') == '1': # sell
self.sell_side.insert(index, order)
if index == 0:
self.inst_status.set_best_ask(msg.get('MDEntryPx'))
def on_book_update_order(self, msg):
"""" on_book_new_order. """
index = msg.get('MDEntryPositionNo') - 1
order = {
'price': msg.get('MDEntryPx'),
'qty': msg.get('MDEntrySize'),
'username': msg.get('Username'),
'user_id': msg.get('UserID'),
'broker': msg.get('Broker'),
'order_id': msg.get('OrderID'),
'side': msg.get('MDEntryType'),
'order_time': msg.get('MDEntryTime'),
'order_date': msg.get('MDEntryDate')
}
if msg.get('MDEntryType') == '0': # sell
self.buy_side[index] = order
if index == 0:
self.inst_status.set_best_bid(msg.get('MDEntryPx'))
elif msg.get('MDEntryType') == '1': # sell
self.sell_side[index] = order
if index == 0:
self.inst_status.set_best_ask(msg.get('MDEntryPx'))
def on_trade(self, msg):
if not self.is_ready:
self.process_later.append(msg)
return
trade = {
"price": msg.get('MDEntryPx'),
"symbol": msg.get('Symbol'),
"size": msg.get('MDEntrySize'),
"trade_date": msg.get('MDEntryDate'),
"trade_time": msg.get('MDEntryTime'),
"order_id": msg.get('OrderID'),
"side": msg.get('Side'),
"counter_order_id": msg.get('SecondaryOrderID'),
"id": msg.get('TradeID'),
"buyer_id": msg.get('MDEntryBuyerID'),
"seller_id": msg.get('MDEntrySellerID'),
"buyer_username": msg.get('MDEntryBuyer'),
"seller_username": msg.get('MDEntrySeller'),
}
Trade.create(self.db_session, trade)
# BTC BRL
price_currency = self.symbol[3:]
size_currency = self.symbol[:3]
if price_currency not in self.volume_dict:
self.volume_dict[price_currency] = 0
if size_currency not in self.volume_dict:
self.volume_dict[size_currency] = 0
volume_price = int(
msg.get('MDEntryPx') * msg.get('MDEntrySize') / 1.e8)
volume_size = msg.get('MDEntrySize')
self.volume_dict[price_currency] += volume_price
self.volume_dict[size_currency] += volume_size
self.volume_dict['MDEntryType'] = '4'
signal_publish_md_status('MD_STATUS', self.volume_dict)
self.inst_status.push_trade(trade)
class MDPBroker(object):
"""The MDP broker class.
The broker routes messages from clients to appropriate workers based on the
requested service.
This base class defines the overall functionality and the API. Subclasses are
meant to implement additional features (like logging).
The broker uses ZMQ ROUTER sockets to deal with clients and workers. These sockets
are wrapped in pyzmq streams to fit well into IOLoop.
.. note::
The workers will *always* be served by the `main_ep` endpoint.
In a two-endpoint setup clients will be handled via the `opt_ep`
endpoint.
:param context: the context to use for socket creation.
:type context: zmq.Context
:param main_ep: the primary endpoint for workers.
:type main_ep: str
:param client_ep: the clients endpoint
:type client_ep: str
:param hb_ep: the heart beat endpoint for workers.
:type hb_ep: str
:param service_q: the class to be used for the service worker-queue.
:type service_q: class
"""
CLIENT_PROTO = C_CLIENT #: Client protocol identifier
WORKER_PROTO = W_WORKER #: Worker protocol identifier
def __init__(self, context, main_ep, client_ep, hb_ep, service_q=None):
"""Init MDPBroker instance.
"""
if service_q is None:
self.service_q = ServiceQueue
else:
self.service_q = service_q
#
# Setup the zmq sockets.
#
socket = context.socket(zmq.ROUTER)
socket.bind(main_ep)
self.main_stream = ZMQStream(socket)
self.main_stream.on_recv(self.on_message)
socket = context.socket(zmq.ROUTER)
socket.bind(client_ep)
self.client_stream = ZMQStream(socket)
self.client_stream.on_recv(self.on_message)
socket = context.socket(zmq.ROUTER)
socket.bind(hb_ep)
self.hb_stream = ZMQStream(socket)
self.hb_stream.on_recv(self.on_message)
self._workers = {}
#
# services contain the service queue and the request queue
#
self._services = {}
#
# Mapping of worker commands and callbacks.
#
self._worker_cmds = {
W_READY: self.on_ready,
W_REPLY: self.on_reply,
W_HEARTBEAT: self.on_heartbeat,
W_DISCONNECT: self.on_disconnect,
}
#
# 'Cleanup' timer for workers without heartbeat.
#
self.hb_check_timer = PeriodicCallback(self.on_timer, HB_INTERVAL)
self.hb_check_timer.start()
def register_worker(self, wid, service):
"""Register the worker id and add it to the given service.
Does nothing if worker is already known.
:param wid: the worker id.
:type wid: str
:param service: the service name.
:type service: str
:rtype: None
"""
if wid in self._workers:
logging.info('Worker {} already registered'.format(service))
return
logging.info('Registering new worker {}'.format(service))
self._workers[wid] = WorkerRep(self.WORKER_PROTO, wid, service, self.main_stream)
if service in self._services:
wq, wr = self._services[service]
wq.put(wid)
else:
q = self.service_q()
q.put(wid)
self._services[service] = (q, [])
def unregister_worker(self, wid):
"""Unregister the worker with the given id.
If the worker id is not registered, nothing happens.
Will stop all timers for the worker.
:param wid: the worker id.
:type wid: str
:rtype: None
"""
try:
wrep = self._workers[wid]
except KeyError:
#
# Not registered, ignore
#
return
logging.info('Unregistering worker {}'.format(wrep.service))
wrep.shutdown()
service = wrep.service
if service in self._services:
wq, wr = self._services[service]
wq.remove(wid)
del self._workers[wid]
def disconnect(self, wid):
"""Send disconnect command and unregister worker.
If the worker id is not registered, nothing happens.
:param wid: the worker id.
:type wid: str
:rtype: None
"""
try:
wrep = self._workers[wid]
except KeyError:
#
# Not registered, ignore
#
return
logging.info('Disconnecting worker {}'.format(wrep.service))
to_send = [wid, self.WORKER_PROTO, W_DISCONNECT]
self.main_stream.send_multipart(to_send)
self.unregister_worker(wid)
def client_response(self, rp, service, msg):
"""Package and send reply to client.
:param rp: return address stack
:type rp: list of str
:param service: name of service
:type service: str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
if service == MMI_SERVICE:
logging.debug('Send reply to client from worker {}'.format(service))
else:
logging.info('Send reply to client from worker {}'.format(service))
to_send = rp[:]
to_send.extend([EMPTY_FRAME, self.CLIENT_PROTO, service])
to_send.extend(msg)
self.client_stream.send_multipart(to_send)
def shutdown(self):
"""Shutdown broker.
Will unregister all workers, stop all timers and ignore all further
messages.
.. warning:: The instance MUST not be used after :func:`shutdown` has been called.
:rtype: None
"""
logging.debug('Shutting down')
self.main_stream.on_recv(None)
self.main_stream.socket.setsockopt(zmq.LINGER, 0)
self.main_stream.socket.close()
self.main_stream.close()
self.main_stream = None
self.client_stream.on_recv(None)
self.client_stream.socket.setsockopt(zmq.LINGER, 0)
self.client_stream.socket.close()
self.client_stream.close()
self.client_stream = None
self.hb_stream.on_recv(None)
self.hb_stream.socket.setsockopt(zmq.LINGER, 0)
self.hb_stream.socket.close()
self.hb_stream.close()
self.hb_stream = None
self._workers = {}
self._services = {}
def on_timer(self):
"""Method called on timer expiry.
Checks which workers are dead and unregisters them.
:rtype: None
"""
#
# Remove 'dead' (not responding to heartbeats) workers.
#
for wrep in self._workers.values():
if not wrep.is_alive():
self.unregister_worker(wrep.id)
def on_ready(self, rp, msg):
"""Process worker READY command.
Registers the worker for a service.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
ret_id = rp[0]
logging.debug('Worker sent ready msg: {} ,{}'.format(rp, msg))
self.register_worker(ret_id, msg[0])
def on_reply(self, rp, msg):
"""Process worker REPLY command.
Route the `msg` to the client given by the address(es) in front of `msg`.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
ret_id = rp[0]
wrep = self._workers.get(ret_id)
if not wrep:
#
# worker not found, ignore message
#
logging.error(
"Worker with return id {} not found. Ignore message.".format(
ret_id))
return
service = wrep.service
logging.info("Worker {} sent reply.".format(service))
try:
wq, wr = self._services[service]
#
# Send response to client
#
cp, msg = split_address(msg)
self.client_response(cp, service, msg)
#
# make worker available again
#
wq.put(wrep.id)
if wr:
logging.info("Sending queued message to worker {}".format(service))
proto, rp, msg = wr.pop(0)
self.on_client(proto, rp, msg)
except KeyError:
#
# unknown service
#
self.disconnect(ret_id)
def on_heartbeat(self, rp, msg):
"""Process worker HEARTBEAT command.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
#
# Note:
# The modified heartbeat of the worker is sent over a separate socket
# stream (self.hb_stream). Therefore the ret_id is wrong. Instead the
# worker sends its id in the message.
#
if len(msg) > 0:
ret_id = msg[0]
else:
ret_id = rp[0]
try:
worker = self._workers[ret_id]
if worker.is_alive():
worker.on_heartbeat()
except KeyError:
#
# Ignore HB for unknown worker
#
pass
def on_disconnect(self, rp, msg):
"""Process worker DISCONNECT command.
Unregisters the worker who sent this message.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
wid = rp[0]
self.unregister_worker(wid)
def on_mmi(self, rp, service, msg):
"""Process MMI request.
mmi.service is used for querying if a specific service is available.
mmi.services is used for querying the list of services available.
:param rp: return address stack
:type rp: list of str
:param service: the protocol id sent
:type service: str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
if service == MMI_SERVICE:
s = msg[0]
ret = [UNKNOWN_SERVICE]
for wr in self._workers.values():
if s == wr.service:
ret = [KNOWN_SERVICE]
break
elif service == MMI_SERVICES:
#
# Return list of services
#
ret = [wr.service for wr in self._workers.values()]
elif service == MMI_TUNNELS:
#
# Read the tunnel files, and send back the network info.
#
tunnel_paths = glob.glob(os.path.expanduser("~/tunnel_port_*.txt"))
tunnels_data = {}
for path in tunnel_paths:
filename = os.path.split(path)[-1]
service_name = filename[-7:-4]
with open(path, 'r') as f:
tunnels_data[service_name] = json.load(f)
ret = [cPickle.dumps(tunnels_data)]
else:
#
# Unknown command.
#
ret = [UNKNOWN_COMMAND]
self.client_response(rp, service, ret)
def on_client(self, proto, rp, msg):
"""Method called on client message.
Frame 0 of msg is the requested service.
The remaining frames are the request to forward to the worker.
.. note::
If the service is unknown to the broker the message is
ignored.
.. note::
If currently no worker is available for a known service,
the message is queued for later delivery.
If a worker is available for the requested service, the
message is repackaged and sent to the worker. The worker in
question is removed from the pool of available workers.
If the service name starts with `mmi.`, the message is passed to
the internal MMI_ handler.
.. _MMI: http://rfc.zeromq.org/spec:8
:param proto: the protocol id sent
:type proto: str
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
service = msg.pop(0)
if service.startswith(b'mmi.'):
logging.debug("Got MMI message from client.")
self.on_mmi(rp, service, msg)
return
logging.info("Client sends message (possibly queued) to worker {}".format(service))
try:
wq, wr = self._services[service]
wid = wq.get()
if not wid:
#
# No worker ready. Queue message
#
logging.info("Worker {} missing. Queuing message.".format(service))
msg.insert(0, service)
wr.append((proto, rp, msg))
return
wrep = self._workers[wid]
to_send = [wrep.id, EMPTY_FRAME, self.WORKER_PROTO, W_REQUEST]
to_send.extend(rp)
to_send.append(EMPTY_FRAME)
to_send.extend(msg)
self.main_stream.send_multipart(to_send)
except KeyError:
#
# Unknwon service. Ignore request
#
logging.info('broker has no service "{}"'.format(service))
def on_worker(self, proto, rp, msg):
"""Method called on worker message.
Frame 0 of msg is the command id.
The remaining frames depend on the command.
This method determines the command sent by the worker and
calls the appropriate method. If the command is unknown the
message is ignored and a DISCONNECT is sent.
:param proto: the protocol id sent
:type proto: str
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
cmd = msg.pop(0)
if cmd in self._worker_cmds:
fnc = self._worker_cmds[cmd]
fnc(rp, msg)
else:
#
# Ignore unknown command. Disconnect worker.
#
logging.error("Unknown worker command: {}".format(cmd))
self.disconnect(rp[0])
def on_message(self, msg):
"""Processes given message.
Decides what kind of message it is -- client or worker -- and
calls the appropriate method. If unknown, the message is
ignored.
:param msg: message parts
:type msg: list of str
:rtype: None
"""
rp, msg = split_address(msg)
try:
#
# Dispatch on first frame after path
#
t = msg.pop(0)
if t.startswith(b'MDPW'):
logging.debug('Recieved message from worker {}'.format(rp))
self.on_worker(t, rp, msg)
elif t.startswith(b'MDPC'):
logging.debug('Recieved message from client {}'.format(rp))
self.on_client(t, rp, msg)
else:
logging.error('Broker unknown Protocol: "{}"'.format(t))
except:
logging.error(
"An error occured while trying to process message: rp: {}, msg: {}\n{}".format(
rp, msg, traceback.format_exc()
)
)
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)
result_dict = {}
def update_result(msg_list):
"""
for collecting the result
msg_list - the list containing the result.
PS: loop is necessary, as it may get more than one result.
"""
for msg in msg_list:
m = json.loads(msg)
result_dict[m['uuid']] = m
# registering the callback
zstream_pull.on_recv(update_result)
class MDHandler(tornado.web.RequestHandler):
reult = None
@asynchronous
@gen.engine
def post(self):
# key for this task
self.result_key = str(uuid())
self.result = None
# pack the task
src = self.get_argument('src', None)
msg = {
class MDPClient(object):
"""Class for the MDP client side.
Thin asynchronous encapsulation of a zmq.REQ socket.
Provides a :func:`request` method with optional timeout.
Objects of this class are ment to be integrated into the
asynchronous IOLoop of pyzmq.
:param context: the ZeroMQ context to create the socket in.
:type context: zmq.Context
:param endpoint: the enpoint to connect to.
:type endpoint: str
:param service: the service the client should use
:type service: str
"""
_proto_version = C_CLIENT
def __init__(self, context, endpoint):
"""Initialize the MDPClient.
"""
self.context = context
self.endpoint = endpoint
def start(self):
"""
Initialize the zmq sockets on a ioloop stream.
The separation of this part from the init is useful if
we start the client on a separate thread with a new ioloop
(for example to enable use in an ipython notebook)
"""
socket = self.context.socket(zmq.DEALER)
ioloop = IOLoop.instance()
self.stream = ZMQStream(socket, ioloop)
self.stream.on_recv(self._on_message)
self._proto_prefix = [EMPTY_FRAME, self._proto_version]
self._delayed_timeout = None
self.timed_out = False
socket.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.setsockopt(zmq.LINGER, 0)
self.stream.socket.close()
self.stream.close()
self.stream = None
def request(self, service, msg, msg_extra=STANDARD, timeout=None):
"""Send the given message.
:param msg: message parts to send.
:type msg: list of str
:param msg_extra: Extra message flags (e.g. STANDARD or BROADCAST)
:type msg_extra: int
:param timeout: time to wait in milliseconds.
:type timeout: int
:rtype None:
"""
if type(msg) in (bytes, unicode):
msg = [msg]
#
# prepare full message
#
to_send = [chr(msg_extra)] + self._proto_prefix[:]
to_send.extend([service])
to_send.extend(msg)
self.stream.send_multipart(to_send)
if timeout:
self._start_timeout(timeout)
def _on_timeout(self):
"""Helper called after timeout.
"""
self.timed_out = True
self._delayed_timeout = None
self.on_timeout()
def _start_timeout(self, timeout):
"""Helper for starting the timeout.
:param timeout: the time to wait in milliseconds.
:type timeout: int
"""
self._delayed_timeout = DelayedCallback(self._on_timeout, timeout)
self._delayed_timeout.start()
def _on_message(self, msg):
"""Helper method called on message receive.
:param msg: list of message parts.
:type msg: list of str
"""
if self._delayed_timeout:
#
# disable timout
#
self._delayed_timeout.stop()
self._delayed_timeout = None
self.on_message(msg)
def on_message(self, msg):
"""Public method called when a message arrived.
.. note:: Does nothing. Should be overloaded!
"""
raise NotImplementedError('on_message must be implemented by the subclass.')
def on_timeout(self):
"""Public method called when a timeout occured.
.. note:: Does nothing. Should be overloaded!
"""
raise NotImplementedError('on_timeout must be implemented by the subclass.')
class Master(object):
def __init__(self,
frontier,
data_in_sock='ipc:///tmp/robot-data-w2m.sock',
data_out_sock='ipc:///tmp/robot-data-m2w.sock',
msg_in_sock='ipc:///tmp/robot-msg-w2m.sock',
msg_out_sock='ipc:///tmp/robot-msg-m2w.sock',
io_loop=None):
self.identity = 'master:%s:%s' % (socket.gethostname(), os.getpid())
context = zmq.Context()
self._io_loop = io_loop or IOLoop.instance()
self._in_socket = context.socket(zmq.SUB)
self._in_socket.setsockopt(zmq.SUBSCRIBE, '')
self._in_socket.bind(data_in_sock)
self._in_stream = ZMQStream(self._in_socket, io_loop)
self._out_socket = context.socket(zmq.PUSH)
self._out_socket.bind(data_out_sock)
self._out_stream = ZMQStream(self._out_socket, io_loop)
self._online_workers = set()
self._running = False
self._updater = PeriodicCallback(self._send_next, 100, io_loop=io_loop)
self._reloader = PeriodicCallback(self.reload, 1000, io_loop=io_loop)
self.frontier = frontier
self.messenger = ServerMessenger(msg_in_sock, msg_out_sock, context,
io_loop)
def start(self):
logging.info('[%s] starting', self.identity)
self.messenger.add_callback(CTRL_MSG_WORKER, self._on_worker_msg)
self.messenger.start()
self._in_stream.on_recv(self._on_receive_processed)
self._updater.start()
self._reloader.start()
self._running = True
def stop(self):
self._running = False
self._reloader.stop()
self._updater.stop()
self.messenger.stop()
# self.messenger.publish(CTRL_MSG_WORKER, self.identity,
# CTRL_MSG_WORKER_QUIT)
def close(self):
self._in_stream.close()
self._in_socket.close()
self._out_stream.close()
self._out_socket.close()
self.messenger.close()
def reload(self):
pass
def _on_worker_msg(self, msg):
if msg.data == CTRL_MSG_WORKER_ONLINE:
self._online_workers.add(msg.identity)
logging.info('[%s] append [%s]', self.identity, msg.identity)
self._send_next()
# if msg.data == CTRL_MSG_WORKER_QUIT_ACK:
# if msg.identity in self._online_workers:
# self._online_workers.remove(msg.identity)
def _send_next(self):
if not self._running:
return
worker_num = len(self._online_workers)
if self._running and worker_num > 0:
while self._out_stream._send_queue.qsize() < worker_num * 4:
request = self.frontier.get_next_request()
if not request:
break
msg = RequestMessage(self.identity, request)
self._out_stream.send_multipart(msg.serialize())
logging.debug('[%s] send request(%s)', self.identity,
request.url)
self.frontier.reload_request(request)
def _on_receive_processed(self, zmq_msg):
msg = ResponseMessage.deserialize(zmq_msg)
request = msg.response.request
logging.debug('[%s] receive response(%s)', self.identity, request.url)
self._send_next()
class ZMQ(Client):
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback):
# do our best to clean up potentially leaky FDs
if hasattr(self, 'stream') and not self.stream.closed():
self.stream.close()
if hasattr(self, 'loop'):
try:
self.loop.close()
except KeyError:
pass
self.socket.close()
self.context.term()
def __init__(self, remote, token, **kwargs):
super(ZMQ, self).__init__(remote, token)
self.context = zmq.Context()
self.socket = self.context.socket(zmq.REQ)
self.socket.RCVTIMEO = RCVTIMEO
self.socket.SNDTIMEO = SNDTIMEO
self.socket.setsockopt(zmq.LINGER, LINGER)
self.nowait = kwargs.get('nowait', False)
self.autoclose = kwargs.get('autoclose', True)
if self.nowait:
self.socket = self.context.socket(zmq.DEALER)
self.autoclose = kwargs.get('autoclose', True)
def _handle_message_fireball(self, m):
logger.debug('message received')
m = json.loads(m[2].decode('utf-8'))
self.response.append(m)
self.num_responses -= 1
logger.debug('num responses remaining: %i' % self.num_responses)
if self.num_responses == 0:
logger.debug('finishing up...')
self.loop.stop()
def _fireball_timeout(self):
logger.info('fireball timeout')
self.loop.stop()
def _send_fireball(self, mtype, data, f_size):
if len(data) < 3:
logger.error('no data to send')
return []
self.loop = IOLoop().instance()
self.socket.close()
self.socket = self.context.socket(zmq.DEALER)
self.socket.connect(self.remote)
self.stream = ZMQStream(self.socket)
self.stream.on_recv(self._handle_message_fireball)
self.stream.io_loop.call_later(SNDTIMEO, self._fireball_timeout)
self.response = []
if PYVERSION == 3:
if isinstance(data, bytes):
data = data.decode('utf-8')
data = json.loads(data)
if not isinstance(data, list):
data = [data]
if (len(data) % f_size) == 0:
self.num_responses = int((len(data) / f_size))
else:
self.num_responses = int((len(data) / f_size)) + 1
logger.debug('responses expected: %i' % self.num_responses)
batch = []
for d in data:
batch.append(d)
if len(batch) == f_size:
Msg(mtype=Msg.INDICATORS_CREATE, token=self.token, data=batch).send(self.socket)
batch = []
if len(batch):
Msg(mtype=Msg.INDICATORS_CREATE, token=self.token, data=batch).send(self.socket)
logger.debug("starting loop to receive")
self.loop.start()
# clean up FDs
self.loop.close()
self.stream.close()
self.socket.close()
return self.response
def _recv(self, decode=True, close=True):
mtype, data = Msg().recv(self.socket)
if close:
self.socket.close()
if not decode:
return data
data = json.loads(data)
if data.get('message') == 'unauthorized':
raise AuthError()
if data.get('message') == 'busy':
raise CIFBusy()
if data.get('message') == 'invalid search':
raise InvalidSearch()
if data.get('status') != 'success':
raise RuntimeError(data.get('message'))
if data.get('data') is None:
raise RuntimeError('invalid response')
if isinstance(data.get('data'), bool):
return data['data']
# is this a straight up elasticsearch string?
if data['data'] == '{}':
return []
if isinstance(data['data'], basestring) and data['data'].startswith('{"hits":{"hits":[{"_source":'):
data['data'] = json.loads(data['data'])
data['data'] = [r['_source'] for r in data['data']['hits']['hits']]
try:
data['data'] = zlib.decompress(data['data'])
except (zlib.error, TypeError):
pass
return data.get('data')
def _send(self, mtype, data='[]', nowait=False, decode=True):
self.socket.connect(self.remote)
if isinstance(data, str):
data = data.encode('utf-8')
Msg(mtype=mtype, token=self.token, data=data).send(self.socket)
if self.nowait or nowait:
if self.autoclose:
self.socket.close()
return
rv = self._recv(decode=decode)
return rv
def ping(self):
try:
return self._send(Msg.PING)
except zmq.error.Again:
raise TimeoutError
def ping_write(self):
try:
return self._send(Msg.PING_WRITE)
except zmq.error.Again:
raise TimeoutError
def indicators_search(self, filters, decode=True):
return self._send(Msg.INDICATORS_SEARCH, json.dumps(filters), decode=decode)
def graph_search(self, filters, decode=True):
return self._send(Msg.GRAPH_SEARCH, json.dumps(filters), decode=decode)
def stats_search(self, filters, decode=True):
return self._send(Msg.STATS_SEARCH, json.dumps(filters), decode=decode)
def indicators_create(self, data, nowait=False, fireball=False, f_size=FIREBALL_SIZE):
if isinstance(data, dict):
data = self._kv_to_indicator(data)
if isinstance(data, Indicator):
data = str(data)
if fireball:
return self._send_fireball(Msg.INDICATORS_CREATE, data, f_size)
return self._send(Msg.INDICATORS_CREATE, data, nowait=nowait)
def indicators_delete(self, data):
if isinstance(data, dict):
data = self._kv_to_indicator(data)
if isinstance(data, Indicator):
data = str(data)
return self._send(Msg.INDICATORS_DELETE, data)
def tokens_search(self, filters={}):
return self._send(Msg.TOKENS_SEARCH, json.dumps(filters))
def tokens_create(self, data):
return self._send(Msg.TOKENS_CREATE, data)
def tokens_delete(self, data):
return self._send(Msg.TOKENS_DELETE, data)
def tokens_edit(self, data):
return self._send(Msg.TOKENS_EDIT, data)
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 CloneServer(object):
# Our server is defined by these properties
ctx = None # Context wrapper
kvmap = None # Key-value store
bstar = None # Binary Star
sequence = 0 # How many updates so far
port = None # Main port we're working on
peer = None # Main port of our peer
publisher = None # Publish updates and hugz
collector = None # Collect updates from clients
subscriber = None # Get updates from peer
pending = None # Pending updates from client
primary = False # True if we're primary
master = False # True if we're master
slave = False # True if we're slave
def __init__(self, primary=True, ports=(5556, 5566)):
self.primary = primary
if primary:
self.port, self.peer = ports
frontend = "tcp://*:5003"
backend = "tcp://localhost:5004"
self.kvmap = {}
else:
self.peer, self.port = ports
frontend = "tcp://*:5004"
backend = "tcp://localhost:5003"
self.ctx = zmq.Context.instance()
self.pending = []
self.bstar = BinaryStar(primary, frontend, backend)
self.bstar.register_voter("tcp://*:%i" % self.port, zmq.ROUTER,
self.handle_snapshot)
# Set up our clone server sockets
self.publisher = self.ctx.socket(zmq.PUB)
self.collector = self.ctx.socket(zmq.SUB)
self.collector.setsockopt(zmq.SUBSCRIBE, b'')
self.publisher.bind("tcp://*:%d" % (self.port + 1))
self.collector.bind("tcp://*:%d" % (self.port + 2))
# Set up our own clone client interface to peer
self.subscriber = self.ctx.socket(zmq.SUB)
self.subscriber.setsockopt(zmq.SUBSCRIBE, b'')
self.subscriber.connect("tcp://localhost:%d" % (self.peer + 1))
# Register state change handlers
self.bstar.master_callback = self.become_master
self.bstar.slave_callback = self.become_slave
# Wrap sockets in ZMQStreams for IOLoop handlers
self.publisher = ZMQStream(self.publisher)
self.subscriber = ZMQStream(self.subscriber)
self.collector = ZMQStream(self.collector)
# Register our handlers with reactor
self.collector.on_recv(self.handle_collect)
self.flush_callback = PeriodicCallback(self.flush_ttl, 1000)
self.hugz_callback = PeriodicCallback(self.send_hugz, 1000)
# basic log formatting:
logging.basicConfig(format="%(asctime)s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO)
def start(self):
# start periodic callbacks
self.flush_callback.start()
self.hugz_callback.start()
# Run bstar reactor until process interrupted
try:
self.bstar.start()
except KeyboardInterrupt:
pass
def handle_snapshot(self, socket, msg):
"""snapshot requests"""
if msg[1] != "ICANHAZ?" or len(msg) != 3:
logging.error("E: bad request, aborting")
dump(msg)
self.bstar.loop.stop()
return
identity, request = msg[:2]
if len(msg) >= 3:
subtree = msg[2]
# Send state snapshot to client
route = Route(socket, identity, subtree)
# For each entry in kvmap, send kvmsg to client
for k, v in self.kvmap.items():
send_single(k, v, route)
# Now send END message with sequence number
logging.info("I: Sending state shapshot=%d" % self.sequence)
socket.send(identity, zmq.SNDMORE)
kvmsg = KVMsg(self.sequence)
kvmsg.key = "KTHXBAI"
kvmsg.body = subtree
kvmsg.send(socket)
def handle_collect(self, msg):
"""Collect updates from clients
If we're master, we apply these to the kvmap
If we're slave, or unsure, we queue them on our pending list
"""
kvmsg = KVMsg.from_msg(msg)
if self.master:
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
ttl = float(kvmsg.get('ttl', 0))
if ttl:
kvmsg['ttl'] = time.time() + ttl
kvmsg.store(self.kvmap)
logging.info("I: publishing update=%d", self.sequence)
else:
# If we already got message from master, drop it, else
# hold on pending list
if not self.was_pending(kvmsg):
self.pending.append(kvmsg)
def was_pending(self, kvmsg):
"""If message was already on pending list, remove and return True.
Else return False.
"""
found = False
for idx, held in enumerate(self.pending):
if held.uuid == kvmsg.uuid:
found = True
break
if found:
self.pending.pop(idx)
return found
def flush_ttl(self):
"""Purge ephemeral values that have expired"""
if self.kvmap:
for key, kvmsg in self.kvmap.items():
self.flush_single(kvmsg)
def flush_single(self, kvmsg):
"""If key-value pair has expired, delete it and publish the fact
to listening clients."""
ttl = float(kvmsg.get('ttl', 0))
if ttl and ttl <= time.time():
kvmsg.body = ""
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.send(self.publisher)
del self.kvmap[kvmsg.key]
logging.info("I: publishing delete=%d", self.sequence)
def send_hugz(self):
"""Send hugz to anyone listening on the publisher socket"""
kvmsg = KVMsg(self.sequence)
kvmsg.key = "HUGZ"
kvmsg.body = ""
kvmsg.send(self.publisher)
# ---------------------------------------------------------------------
# State change handlers
def become_master(self):
"""We're becoming master
The backup server applies its pending list to its own hash table,
and then starts to process state snapshot requests.
"""
self.master = True
self.slave = False
# stop receiving subscriber updates while we are master
self.subscriber.stop_on_recv()
# Apply pending list to own kvmap
while self.pending:
kvmsg = self.pending.pop(0)
self.sequence += 1
kvmsg.sequence = self.sequence
kvmsg.store(self.kvmap)
logging.info("I: publishing pending=%d", self.sequence)
def become_slave(self):
"""We're becoming slave"""
# clear kvmap
self.kvmap = None
self.master = False
self.slave = True
self.subscriber.on_recv(self.handle_subscriber)
def handle_subscriber(self, msg):
"""Collect updates from peer (master)
We're always slave when we get these updates
"""
if self.master:
logging.warn("received subscriber message, but we are master %s",
msg)
return
# Get state snapshot if necessary
if self.kvmap is None:
self.kvmap = {}
snapshot = self.ctx.socket(zmq.DEALER)
snapshot.linger = 0
snapshot.connect("tcp://localhost:%i" % self.peer)
logging.info("I: asking for snapshot from: tcp://localhost:%d",
self.peer)
snapshot.send_multipart(["ICANHAZ?", ''])
while True:
try:
kvmsg = KVMsg.recv(snapshot)
except KeyboardInterrupt:
# Interrupted
self.bstar.loop.stop()
return
if kvmsg.key == "KTHXBAI":
self.sequence = kvmsg.sequence
break # Done
kvmsg.store(self.kvmap)
logging.info("I: received snapshot=%d", self.sequence)
# Find and remove update off pending list
kvmsg = KVMsg.from_msg(msg)
# update float ttl -> timestamp
ttl = float(kvmsg.get('ttl', 0))
if ttl:
kvmsg['ttl'] = time.time() + ttl
if kvmsg.key != "HUGZ":
if not self.was_pending(kvmsg):
# If master update came before client update, flip it
# around, store master update (with sequence) on pending
# list and use to clear client update when it comes later
self.pending.append(kvmsg)
# If update is more recent than our kvmap, apply it
if (kvmsg.sequence > self.sequence):
self.sequence = kvmsg.sequence
kvmsg.store(self.kvmap)
logging.info("I: received update=%d", self.sequence)
class MDPClient(object):
"""Class for the MDP client side.
Thin asynchronous encapsulation of a zmq.REQ socket.
Provides a :func:`request` method with optional timeout.
Objects of this class are ment to be integrated into the
asynchronous IOLoop of pyzmq.
:param context: the ZeroMQ context to create the socket in.
:type context: zmq.Context
:param endpoint: the enpoint to connect to.
:type endpoint: str
:param service: the service the client should use
:type service: str
"""
_proto_version = b'MDPC01'
def __init__(self, context, endpoint, service):
"""Initialize the MDPClient.
"""
socket = context.socket(zmq.REQ)
ioloop = IOLoop.instance()
self.service = service
self.endpoint = endpoint
self.stream = ZMQStream(socket, ioloop)
self.stream.on_recv(self._on_message)
self.can_send = True
self._proto_prefix = [PROTO_VERSION, service]
self._tmo = None
self.timed_out = False
socket.connect(endpoint)
return
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.setsockopt(zmq.LINGER, 0)
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()
# prepare full message
to_send = self._proto_prefix[:]
to_send.extend(msg)
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 occured.
.. note:: Does nothing. Should be overloaded!
"""
pass
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.runtime.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def __init__(
self,
*args,
**kwargs,
):
super().__init__(**kwargs)
def _register_pollin(self):
"""Register :attr:`in_sock`, :attr:`ctrl_sock` and :attr:`out_sock` in poller."""
with ImportExtensions(required=True):
import tornado.ioloop
get_or_reuse_loop()
self.io_loop = tornado.ioloop.IOLoop.current()
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
if self.in_connect_sock is not None:
self.in_connect_sock = ZMQStream(self.in_connect_sock, self.io_loop)
self.in_sock.stop_on_recv()
def _get_dynamic_out_socket(self, target_pod):
return super()._get_dynamic_out_socket(target_pod, as_streaming=True)
def close(self, flush: bool = True, *args, **kwargs):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`.
.. note::
This method is idempotent.
:param flush: flag indicating if `sockets` need to be flushed before close is done
:param args: Extra positional arguments
:param kwargs: Extra key-value arguments
"""
# if Address already in use `self.in_sock_type` is not set
if (
not self.is_closed
and hasattr(self, 'in_sock_type')
and self.in_sock_type == zmq.DEALER
):
try:
if self._active:
self._send_cancel_to_router(raise_exception=True)
except zmq.error.ZMQError:
self.logger.debug(
f'The dealer {self.name} can not unsubscribe from the router. '
f'In case the router is down this is expected.'
)
self._active = (
False # Important to avoid sending idle back while flushing in socket
)
if not self.is_closed:
# wait until the close signal is received
time.sleep(0.01)
if flush:
for s in self.opened_socks:
events = s.flush()
self.logger.debug(f'Handled #{events} during flush of socket')
super().close()
if hasattr(self, 'io_loop'):
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.in_connect_sock, '_handle_events'):
self.in_connect_sock._handle_events = (
lambda *args, **kwargs: None
)
except AttributeError as e:
self.logger.error(f'failed to stop. {e!r}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
self.is_polling_paused = True
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
if self.is_polling_paused:
self.in_sock.on_recv(self._in_sock_callback)
self.is_polling_paused = False
def start(self, callback: Callable[['Message'], None]):
"""
Open all sockets and start the ZMQ context associated to this `Zmqlet`.
:param callback: callback function to receive the protobuf message
"""
def _callback(msg, sock_type):
msg = _parse_from_frames(sock_type, msg)
self.bytes_recv += msg.size
self.msg_recv += 1
callback(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER and not self.args.dynamic_routing_out:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
if self.in_connect_sock is not None:
self.in_connect_sock.on_recv(
lambda x: _callback(x, SocketType.ROUTER_CONNECT)
)
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class IOPubThread(object):
"""An object for sending IOPub messages in a background thread
Prevents a blocking main thread from delaying output from threads.
IOPubThread(pub_socket).background_socket is a Socket-API-providing object
whose IO is always run in a thread.
"""
def __init__(self, socket, pipe=False):
"""Create IOPub thread
Parameters
----------
socket: zmq.PUB Socket
the socket on which messages will be sent.
pipe: bool
Whether this process should listen for IOPub messages
piped from subprocesses.
"""
self.socket = socket
self.background_socket = BackgroundSocket(self)
self._master_pid = os.getpid()
self._pipe_flag = pipe
self.io_loop = IOLoop(make_current=False)
if pipe:
self._setup_pipe_in()
self._local = threading.local()
self._events = deque()
self._setup_event_pipe()
self.thread = threading.Thread(target=self._thread_main)
self.thread.daemon = True
def _thread_main(self):
"""The inner loop that's actually run in a thread"""
self.io_loop.make_current()
self.io_loop.start()
self.io_loop.close(all_fds=True)
def _setup_event_pipe(self):
"""Create the PULL socket listening for events that should fire in this thread."""
ctx = self.socket.context
pipe_in = ctx.socket(zmq.PULL)
pipe_in.linger = 0
_uuid = b2a_hex(os.urandom(16)).decode('ascii')
iface = self._event_interface = 'inproc://%s' % _uuid
pipe_in.bind(iface)
self._event_puller = ZMQStream(pipe_in, self.io_loop)
self._event_puller.on_recv(self._handle_event)
@property
def _event_pipe(self):
"""thread-local event pipe for signaling events that should be processed in the thread"""
try:
event_pipe = self._local.event_pipe
except AttributeError:
# new thread, new event pipe
ctx = self.socket.context
event_pipe = ctx.socket(zmq.PUSH)
event_pipe.linger = 0
event_pipe.connect(self._event_interface)
self._local.event_pipe = event_pipe
return event_pipe
def _handle_event(self, msg):
"""Handle an event on the event pipe
Content of the message is ignored.
Whenever *an* event arrives on the event stream,
*all* waiting events are processed in order.
"""
# freeze event count so new writes don't extend the queue
# while we are processing
n_events = len(self._events)
for i in range(n_events):
event_f = self._events.popleft()
event_f()
def _setup_pipe_in(self):
"""setup listening pipe for IOPub from forked subprocesses"""
ctx = self.socket.context
# use UUID to authenticate pipe messages
self._pipe_uuid = os.urandom(16)
pipe_in = ctx.socket(zmq.PULL)
pipe_in.linger = 0
try:
self._pipe_port = pipe_in.bind_to_random_port("tcp://127.0.0.1")
except zmq.ZMQError as e:
warnings.warn("Couldn't bind IOPub Pipe to 127.0.0.1: %s" % e +
"\nsubprocess output will be unavailable.")
self._pipe_flag = False
pipe_in.close()
return
self._pipe_in = ZMQStream(pipe_in, self.io_loop)
self._pipe_in.on_recv(self._handle_pipe_msg)
def _handle_pipe_msg(self, msg):
"""handle a pipe message from a subprocess"""
if not self._pipe_flag or not self._is_master_process():
return
if msg[0] != self._pipe_uuid:
print("Bad pipe message: %s", msg, file=sys.__stderr__)
return
self.send_multipart(msg[1:])
def _setup_pipe_out(self):
# must be new context after fork
ctx = zmq.Context()
pipe_out = ctx.socket(zmq.PUSH)
pipe_out.linger = 3000 # 3s timeout for pipe_out sends before discarding the message
pipe_out.connect("tcp://127.0.0.1:%i" % self._pipe_port)
return ctx, pipe_out
def _is_master_process(self):
return os.getpid() == self._master_pid
def _check_mp_mode(self):
"""check for forks, and switch to zmq pipeline if necessary"""
if not self._pipe_flag or self._is_master_process():
return MASTER
else:
return CHILD
def start(self):
"""Start the IOPub thread"""
self.thread.start()
# make sure we don't prevent process exit
# I'm not sure why setting daemon=True above isn't enough, but it doesn't appear to be.
atexit.register(self.stop)
def stop(self):
"""Stop the IOPub thread"""
if not self.thread.is_alive():
return
self.io_loop.add_callback(self.io_loop.stop)
self.thread.join()
if hasattr(self._local, 'event_pipe'):
self._local.event_pipe.close()
def close(self):
if self.closed:
return
self.socket.close()
self.socket = None
@property
def closed(self):
return self.socket is None
def schedule(self, f):
"""Schedule a function to be called in our IO thread.
If the thread is not running, call immediately.
"""
if self.thread.is_alive():
self._events.append(f)
# wake event thread (message content is ignored)
self._event_pipe.send(b'')
else:
f()
def send_multipart(self, *args, **kwargs):
"""send_multipart schedules actual zmq send in my thread.
If my thread isn't running (e.g. forked process), send immediately.
"""
self.schedule(lambda: self._really_send(*args, **kwargs))
def _really_send(self, msg, *args, **kwargs):
"""The callback that actually sends messages"""
mp_mode = self._check_mp_mode()
if mp_mode != CHILD:
# we are master, do a regular send
self.socket.send_multipart(msg, *args, **kwargs)
else:
# we are a child, pipe to master
# new context/socket for every pipe-out
# since forks don't teardown politely, use ctx.term to ensure send has completed
ctx, pipe_out = self._setup_pipe_out()
pipe_out.send_multipart([self._pipe_uuid] + msg, *args, **kwargs)
pipe_out.close()
ctx.term()
class Messager(object):
def __init__(self, callback):
self.callback = callback
self.stamps = {}
self.topics = []
self._sock = ctx.socket(zmq.SUB)
self._sock.connect('inproc:///tmp/monster_{}'.format(pid))
self._stream = ZMQStream(self._sock)
self._stream.on_recv(self.recv)
def destroy(self):
self.callback = None
self.topics = None
self._stream.close()
def subscribe(self, raw_topic, stamp):
'''订阅
订阅的主题支持多级主题,用冒号分隔,例如
WEATHER:CHINA:HANGZHOU
订阅父主题将会收到所有该类主题消息,比如订阅了WEATHER:CHINA,将收到所有中国城市的天气
但由于zmq的订阅规则中并不支持多级主题,于是需要自己在内容中维护多级主题关系,将顶级主题送到zmq中
'''
topic = raw_topic.split(':')
if not topic:
return
if [x for x in self.topics if x == topic]:
return
self.stamps.update({raw_topic: stamp})
self._sock.setsockopt(zmq.SUBSCRIBE, str(topic[0]))
#TODO 目前的多级主题(topics)的内存结构为列表,消息送达之后需整体遍历列表以判断用户是否订阅该主题,可以考虑使用trie树来优化
self.topics.append(topic)
def unsubscribe(self, raw_topic):
'''取消订阅
比订阅操作更复杂的是退订的时候必须检查顶级主题下的所以子主题都已经退订,如果有一个子主题还在就不能
完全从zmq当中退订该顶级主题
'''
topic = raw_topic.split(':')
if not topic:
return
try:
self.topics.remove(raw_topic.split(':'))
self.stamps.pop(raw_topic)
for tt in self.topics:
if tt[0] == topic[0]:
break
else:
self._sock.setsockopt(zmq.UNSUBSCRIBE, str(topic[0]))
except:
pass
def recv(self, frame):
_, data = frame
try:
dd = json.loads(data)
top = dd.get('topic')
if not top:
return
top = top.split(':')
for topic in self.topics:
if topic == top[:len(topic)]:
raw_topic = ':'.join(topic)
stamp = self.stamps.get(raw_topic)
dd['stamp'] = stamp
self.callback(json.dumps(dd))
except:
pass
class Server(object):
ctx = None
loop = None
stats = None
spiders_in = None
spiders_out = None
sw_in = None
sw_out = None
db_in = None
db_out = None
def __init__(self, hostname, base_port):
self.ctx = zmq.Context()
self.loop = IOLoop.instance()
self.stats = {
'started': time(),
'spiders_out_recvd': 0,
'spiders_in_recvd': 0,
'db_in_recvd': 0,
'db_out_recvd': 0,
'sw_in_recvd': 0,
'sw_out_recvd': 0
}
socket_config = SocketConfig(hostname, base_port)
spiders_in_s = self.ctx.socket(zmq.XPUB)
spiders_out_s = self.ctx.socket(zmq.XSUB)
sw_in_s = self.ctx.socket(zmq.XPUB)
sw_out_s = self.ctx.socket(zmq.XSUB)
db_in_s = self.ctx.socket(zmq.XPUB)
db_out_s = self.ctx.socket(zmq.XSUB)
spiders_in_s.bind(socket_config.spiders_in())
spiders_out_s.bind(socket_config.spiders_out())
sw_in_s.bind(socket_config.sw_in())
sw_out_s.bind(socket_config.sw_out())
db_in_s.bind(socket_config.db_in())
db_out_s.bind(socket_config.db_out())
self.spiders_in = ZMQStream(spiders_in_s)
self.spiders_out = ZMQStream(spiders_out_s)
self.sw_in = ZMQStream(sw_in_s)
self.sw_out = ZMQStream(sw_out_s)
self.db_in = ZMQStream(db_in_s)
self.db_out = ZMQStream(db_out_s)
self.spiders_out.on_recv(self.handle_spiders_out_recv)
self.sw_out.on_recv(self.handle_sw_out_recv)
self.db_out.on_recv(self.handle_db_out_recv)
self.sw_in.on_recv(self.handle_sw_in_recv)
self.db_in.on_recv(self.handle_db_in_recv)
self.spiders_in.on_recv(self.handle_spiders_in_recv)
logging.basicConfig(format="%(asctime)s %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO)
self.logger = logging.getLogger(
"distributed_frontera.messagebus.zeromq.broker.Server")
def start(self):
self.logger.info("Distributed Frontera ZeroMQ broker is started.")
self.log_stats()
try:
self.loop.start()
except KeyboardInterrupt:
pass
def log_stats(self):
self.logger.info(self.stats)
self.loop.add_timeout(timedelta(seconds=10), self.log_stats)
def handle_spiders_out_recv(self, msg):
self.sw_in.send_multipart(msg)
self.db_in.send_multipart(msg)
self.stats['spiders_out_recvd'] += 1
def handle_sw_out_recv(self, msg):
self.db_in.send_multipart(msg)
self.stats['sw_out_recvd'] += 1
def handle_db_out_recv(self, msg):
self.spiders_in.send_multipart(msg)
self.stats['db_out_recvd'] += 1
def handle_db_in_recv(self, msg):
self.stats['db_in_recvd'] += 1
if msg[0][0] in ['\x01', '\x00']:
action, identity, partition_id = self.decode_subscription(msg[0])
if identity == 'sl':
self.spiders_out.send_multipart(msg)
return
if identity == 'us':
self.sw_out.send_multipart(msg)
return
raise AttributeError('Unknown identity in channel subscription.')
def handle_sw_in_recv(self, msg):
if msg[0][0] in ['\x01', '\x00']:
self.spiders_out.send_multipart(msg)
self.stats['sw_in_recvd'] += 1
def handle_spiders_in_recv(self, msg):
if msg[0][0] in ['\x01', '\x00']:
self.db_out.send_multipart(msg)
self.stats['spiders_in_recvd'] += 1
def decode_subscription(self, msg):
"""
:param msg:
:return: tuple of action, identity, partition_id
where
action is 1 - subscription, 0 - unsubscription,
identity - 2 characters,
partition_id - 8 bit unsigned integer (None if absent)
"""
if len(msg) == 4:
return unpack(">B2sB", msg)
elif len(msg) == 3:
action, identity = unpack(">B2s", msg)
return action, identity, None
raise ValueError("Can't decode subscription correctly.")
class TestMNWorker(TestCase):
endpoint = b'tcp://127.0.0.1:5555'
service = b'test'
def setUp(self):
if _do_print:
print('Setting up...')
sys.stdout.flush()
self.context = zmq.Context()
self.broker = None
self._msgs = []
return
def tearDown(self):
if _do_print:
print('Tearing down...')
sys.stdout.flush()
if self.broker:
self._stop_broker()
self.broker = None
self.context = None
return
def _on_msg(self, msg):
if _do_print:
print('Broker received:', msg)
self.target = msg.pop(0)
marker_frame = msg.pop(0)
if msg[1] == b'\x01': # ready
if _do_print:
print('READY received')
return
if msg[1] == b'\x04': # ready
if _do_print:
print('HB received')
return
if msg[1] == b'\x03': # reply
IOLoop.instance().stop()
return
return
def _start_broker(self, do_reply=False):
"""Helper activating a fake broker in the ioloop.
"""
socket = self.context.socket(zmq.ROUTER)
self.broker = ZMQStream(socket)
self.broker.socket.setsockopt(zmq.LINGER, 0)
self.broker.bind(self.endpoint)
self.broker.on_recv(self._on_msg)
self.broker.do_reply = do_reply
self.broker.ticker = PeriodicCallback(self._tick,
WorkerRunner.HB_INTERVAL)
self.broker.ticker.start()
self.target = None
if _do_print:
print("Broker started")
return
def _stop_broker(self):
if self.broker:
self.broker.ticker.stop()
self.broker.ticker = None
self.broker.socket.close()
self.broker.close()
self.broker = None
if _do_print:
print("Broker stopped")
return
def _tick(self):
if self.broker and self.target:
msg = [self.target, b'', b'MNPW01', b'\x04']
self.broker.send_multipart(msg)
if _do_print:
print("Tick sent:", msg)
return
def send_req(self):
data = [b'AA', b'bb']
msg = [self.target, b'', b'MNPW01', b'\x02', self.target, b''] + data
self.broker.send_multipart(msg)
if _do_print:
print('broker sent:', msg)
return
@staticmethod
def stop_loop():
IOLoop.instance().stop()
return
# Tests follow
def test_simple_worker(self):
"""Test MNWorker simple req/reply.
"""
self._start_broker()
time.sleep(0.2)
worker = WorkerRunner(self.context, self.endpoint, self.service)
sender = DelayedCallback(self.send_req, 500)
stopper = DelayedCallback(self.stop_loop, 2500)
sender.start()
stopper.start()
IOLoop.instance().start()
worker.shutdown()
self._stop_broker()
return
class MDPBroker(object):
"""The MDP broker class.
The broker routes messages from clients to appropriate workers based on the
requested service.
This base class defines the overall functionality and the API. Subclasses are
ment to implement additional features (like logging).
The broker uses ØMQ ROUTER sockets to deal witch clients and workers. These sockets
are wrapped in pyzmq streams to fit well into IOLoop.
.. note::
The workers will *always* be served by the `main_ep` endpoint.
In a two-endpoint setup clients will be handled via the `opt_ep`
endpoint.
:param context: the context to use for socket creation.
:type context: zmq.Context
:param main_ep: the primary endpoint for workers and clients.
:type main_ep: str
:param opt_ep: is an optional 2nd endpoint.
:type opt_ep: str
:param worker_q: the class to be used for the worker-queue.
:type worker_q: class
"""
CLIENT_PROTO = b'MDPC01' #: Client protocol identifier
WORKER_PROTO = b'MDPW01' #: Worker protocol identifier
def __init__(self, context, main_ep, opt_ep=None):
"""Init MDPBroker instance.
"""
l = logger.Logger('mq_broker')
self.log = l.get_logger()
self.log.info("MDP broker startup...")
socket = ZmqSocket(context, zmq.ROUTER)
socket.bind(main_ep)
self.main_stream = ZMQStream(socket)
self.main_stream.on_recv(self.on_message)
if opt_ep:
socket = ZmqSocket(context, zmq.ROUTER)
socket.bind(opt_ep)
self.client_stream = ZMQStream(socket)
self.client_stream.on_recv(self.on_message)
else:
self.client_stream = self.main_stream
self.log.debug("Socket created...")
self._workers = {}
# services contain the worker queue and the request queue
self._services = {}
self._worker_cmds = {
b'\x01': self.on_ready,
b'\x03': self.on_reply,
b'\x04': self.on_heartbeat,
b'\x05': self.on_disconnect,
}
self.log.debug("Launch the timer...")
self.hb_check_timer = PeriodicCallback(self.on_timer, HB_INTERVAL)
self.hb_check_timer.start()
self.log.info("MDP broker started")
return
def register_worker(self, wid, service):
"""Register the worker id and add it to the given service.
Does nothing if worker is already known.
:param wid: the worker id.
:type wid: str
:param service: the service name.
:type service: str
:rtype: None
"""
self.log.debug(
"Try to register a worker : wid={0}, service={1}".format(
wid, service))
try:
if wid in self._workers:
self.log.debug("Worker %s already registered" % service)
return
self._workers[wid] = WorkerRep(self.WORKER_PROTO, wid, service,
self.main_stream)
if service in self._services:
wq, wr = self._services[service]
wq.put(wid)
else:
q = ServiceQueue()
q.put(wid)
self._services[service] = (q, [])
self.log.info("Registered worker : wid={0}, service={1}".format(
wid, service))
except:
self.log.error(
"Error while registering a worker : wid={0}, service={1}, trace={2}"
.format(wid, service, traceback.format_exc()))
return
def unregister_worker(self, wid):
"""Unregister the worker with the given id.
If the worker id is not registered, nothing happens.
Will stop all timers for the worker.
:param wid: the worker id.
:type wid: str
:rtype: None
"""
self.log.debug("Try to unregister a worker : wid={0}".format(wid))
try:
try:
wrep = self._workers[wid]
except KeyError:
# not registered, ignore
self.log.warning(
"The worker wid={0} is not registered, ignoring the unregister request"
.format(wid))
return
wrep.shutdown()
service = wrep.service
if service in self._services:
wq, wr = self._services[service]
wq.remove(wid)
del self._workers[wid]
self.log.info("Unregistered worker : wid={0}".format(wid))
except:
self.log.error(
"Error while unregistering a worker : wid={0}, trace={1}".
format(wid, traceback.format_exc()))
return
def disconnect(self, wid):
"""Send disconnect command and unregister worker.
If the worker id is not registered, nothing happens.
:param wid: the worker id.
:type wid: str
:rtype: None
"""
self.log.debug("Try to disconnect a worker : wid={0}".format(wid))
try:
try:
wrep = self._workers[wid]
except KeyError:
# not registered, ignore
self.log.warning(
"The worker wid={0} service={1} is not registered, ignoring the disconnect request"
.format(wid, wrep.service))
return
to_send = [wid, self.WORKER_PROTO, b'\x05']
self.main_stream.send_multipart(to_send)
self.log.info(
"Request to unregister a worker : wid={0} service={1}".format(
wid, wrep.service))
except:
self.log.error(
"Error while disconnecting a worker : wid={0}, trace={1}".
format(wid, traceback.format_exc()))
self.unregister_worker(wid)
return
def client_response(self, rp, service, msg):
"""Package and send reply to client.
:param rp: return address stack
:type rp: list of str
:param service: name of service
:type service: str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
to_send = rp[:]
to_send.extend([b'', self.CLIENT_PROTO, service])
to_send.extend(msg)
self.client_stream.send_multipart(to_send)
return
def shutdown(self):
"""Shutdown broker.
Will unregister all workers, stop all timers and ignore all further
messages.
.. warning:: The instance MUST not be used after :func:`shutdown` has been called.
:rtype: None
"""
self.log.info("Shutdown starting...")
try:
self.log.debug("Closing the socket...")
if self.client_stream == self.main_stream:
self.client_stream = None
self.main_stream.on_recv(None)
self.main_stream.socket.setsockopt(zmq.LINGER, 0)
self.main_stream.socket.close()
self.main_stream.close()
self.main_stream = None
if self.client_stream:
self.client_stream.on_recv(None)
self.client_stream.socket.setsockopt(zmq.LINGER, 0)
self.client_stream.socket.close()
self.client_stream.close()
self.client_stream = None
self.log.debug("Clean workers and services...")
self._workers = {}
self._services = {}
except:
self.log.error("Error during shutdown : trace={0}".format(
traceback.format_exc()))
return
def on_timer(self):
"""Method called on timer expiry.
Checks which workers are dead and unregisters them.
:rtype: None
"""
self.log.debug("Check for dead workers...")
for wrep in list(self._workers.values()):
if not wrep.is_alive():
self.log.info(
"A worker seems to be dead : wid={0} service={1}".format(
wrep.id, wrep.service))
self.unregister_worker(wrep.id)
return
def on_ready(self, rp, msg):
"""Process worker READY command.
Registers the worker for a service.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
ret_id = rp[0]
self.register_worker(ret_id, msg[0])
return
def on_reply(self, rp, msg):
"""Process worker REPLY command.
Route the `msg` to the client given by the address(es) in front of `msg`.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
ret_id = rp[0]
# make worker available again
try:
wrep = self._workers[ret_id]
service = wrep.service
wq, wr = self._services[service]
cp, msg = split_address(msg)
self.client_response(cp, service, msg)
wq.put(wrep.id)
if wr:
proto, rp, msg = wr.pop(0)
self.on_client(proto, rp, msg)
except KeyError:
# unknown service
self.disconnect(ret_id)
return
def on_heartbeat(self, rp, msg):
"""Process worker HEARTBEAT command.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
ret_id = rp[0]
try:
worker = self._workers[ret_id]
if worker.is_alive():
worker.on_heartbeat()
except KeyError:
# ignore HB for unknown worker
pass
return
def on_disconnect(self, rp, msg):
"""Process worker DISCONNECT command.
Unregisters the worker who sent this message.
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
wid = rp[0]
self.log.info("A worker disconnects itself : wid={0}".format(wid))
self.unregister_worker(wid)
return
def on_mmi(self, rp, service, msg):
"""Process MMI request.
For now only mmi.service is handled.
:param rp: return address stack
:type rp: list of str
:param service: the protocol id sent
:type service: str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
if service == b'mmi.service':
s = msg[0]
ret = b'404'
for wr in list(self._workers.values()):
if s == wr.service:
ret = b'200'
break
self.client_response(rp, service, [ret])
elif service == b'mmi.services':
ret = []
for wr in list(self._workers.values()):
ret.append(wr.service)
self.client_response(rp, service, [b', '.join(ret)])
else:
self.client_response(rp, service, [b'501'])
return
def on_client(self, proto, rp, msg):
"""Method called on client message.
Frame 0 of msg is the requested service.
The remaining frames are the request to forward to the worker.
.. note::
If the service is unknown to the broker the message is
ignored.
.. note::
If currently no worker is available for a known service,
the message is queued for later delivery.
If a worker is available for the requested service, the
message is repackaged and sent to the worker. The worker in
question is removed from the pool of available workers.
If the service name starts with `mmi.`, the message is passed to
the internal MMI_ handler.
.. _MMI: http://rfc.zeromq.org/spec:8
:param proto: the protocol id sent
:type proto: str
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
service = msg.pop(0)
if service.startswith(b'mmi.'):
self.on_mmi(rp, service, msg)
return
try:
wq, wr = self._services[service]
wid = wq.get()
if not wid:
# no worker ready
# queue message
msg.insert(0, service)
wr.append((proto, rp, msg))
return
wrep = self._workers[wid]
to_send = [wrep.id, b'', self.WORKER_PROTO, b'\x02']
to_send.extend(rp)
to_send.append(b'')
to_send.extend(msg)
self.main_stream.send_multipart(to_send)
except KeyError:
# unknwon service
# ignore request
msg = "broker has no service {0}".format(service)
print(msg)
self.log.warning(msg)
return
def on_worker(self, proto, rp, msg):
"""Method called on worker message.
Frame 0 of msg is the command id.
The remaining frames depend on the command.
This method determines the command sent by the worker and
calls the appropriate method. If the command is unknown the
message is ignored and a DISCONNECT is sent.
:param proto: the protocol id sent
:type proto: str
:param rp: return address stack
:type rp: list of str
:param msg: message parts
:type msg: list of str
:rtype: None
"""
cmd = msg.pop(0)
if cmd in self._worker_cmds:
fnc = self._worker_cmds[cmd]
fnc(rp, msg)
else:
# ignore unknown command
# DISCONNECT worker
self.log.warning(
"Unknown command from worker (it will be disconnect) : wid={0}, cmd={1}"
.format(rp[0], cmd))
self.disconnect(rp[0])
return
def on_message(self, msg):
"""Processes given message.
Decides what kind of message it is -- client or worker -- and
calls the appropriate method. If unknown, the message is
ignored.
:param msg: message parts
:type msg: list of str
:rtype: None
"""
self.log.debug("Message received: {0}".format(msg))
rp, msg = split_address(msg)
# dispatch on first frame after path
t = msg.pop(0)
if t.startswith(b'MDPW'):
self.on_worker(t, rp, msg)
elif t.startswith(b'MDPC'):
self.on_client(t, rp, msg)
else:
self.log.warning("Broker unknown Protocol: {0}".format(t))
return
class LRUQueue(object):
"""LRUQueue class using ZMQStream/IOLoop for event dispatching"""
def __init__(self, backend_socket, frontend_socket):
self.available_workers = 0
self.workers = []
self.client_nbr = NBR_CLIENTS
self.backend = ZMQStream(backend_socket)
self.frontend = ZMQStream(frontend_socket)
self.backend.on_recv(self.handle_backend)
self.loop = IOLoop.instance()
def handle_backend(self, msg):
# Queue worker address for LRU routing
worker_addr, empty, client_addr = msg[:3]
assert self.available_workers < NBR_WORKERS
# add worker back to the list of workers
self.available_workers += 1
self.workers.append(worker_addr)
# Second frame is empty
assert empty == b""
# Third frame is READY or else a client reply address
# If client reply, send rest back to frontend
if client_addr != b"READY":
empty, reply = msg[3:]
# Following frame is empty
assert empty == b""
self.frontend.send_multipart([client_addr, b'', reply])
self.client_nbr -= 1
if self.client_nbr == 0:
# Exit after N messages
self.loop.add_timeout(time.time() + 1, self.loop.stop)
if self.available_workers == 1:
# on first recv, start accepting frontend messages
self.frontend.on_recv(self.handle_frontend)
def handle_frontend(self, msg):
# Now get next client request, route to LRU worker
# Client request is [address][empty][request]
client_addr, empty, request = msg
assert empty == b""
# Dequeue and drop the next worker address
self.available_workers -= 1
worker_id = self.workers.pop()
self.backend.send_multipart(
[worker_id, b'', client_addr, b'', request])
if self.available_workers == 0:
# stop receiving until workers become available again
self.frontend.stop_on_recv()
class ZmqStreamlet(Zmqlet):
"""A :class:`ZmqStreamlet` object can send/receive data to/from ZeroMQ stream and invoke callback function. It
has three sockets for input, output and control.
.. warning::
Starting from v0.3.6, :class:`ZmqStreamlet` replaces :class:`Zmqlet` as one of the key components in :class:`jina.peapods.runtime.BasePea`.
It requires :mod:`tornado` and :mod:`uvloop` to be installed.
"""
def register_pollin(self):
"""Register :attr:`in_sock`, :attr:`ctrl_sock` and :attr:`out_sock` in poller."""
with ImportExtensions(required=True):
import tornado.ioloop
get_or_reuse_loop()
self.io_loop = tornado.ioloop.IOLoop.current()
self.in_sock = ZMQStream(self.in_sock, self.io_loop)
self.out_sock = ZMQStream(self.out_sock, self.io_loop)
self.ctrl_sock = ZMQStream(self.ctrl_sock, self.io_loop)
self.in_sock.stop_on_recv()
def close(self):
"""Close all sockets and shutdown the ZMQ context associated to this `Zmqlet`.
.. note::
This method is idempotent.
"""
if not self.is_closed:
# wait until the close signal is received
time.sleep(0.01)
for s in self.opened_socks:
s.flush()
super().close()
if hasattr(self, 'io_loop'):
try:
self.io_loop.stop()
# Replace handle events function, to skip
# None event after sockets are closed.
if hasattr(self.in_sock, '_handle_events'):
self.in_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.out_sock, '_handle_events'):
self.out_sock._handle_events = lambda *args, **kwargs: None
if hasattr(self.ctrl_sock, '_handle_events'):
self.ctrl_sock._handle_events = lambda *args, **kwargs: None
except AttributeError as e:
self.logger.error(f'failed to stop. {e!r}')
def pause_pollin(self):
"""Remove :attr:`in_sock` from the poller """
self.in_sock.stop_on_recv()
def resume_pollin(self):
"""Put :attr:`in_sock` back to the poller """
self.in_sock.on_recv(self._in_sock_callback)
def start(self, callback: Callable[['Message'], 'Message']):
"""
Open all sockets and start the ZMQ context associated to this `Zmqlet`.
:param callback: callback function to receive the protobuf message
"""
def _callback(msg, sock_type):
msg = _parse_from_frames(sock_type, msg)
self.bytes_recv += msg.size
self.msg_recv += 1
msg = callback(msg)
if msg:
self.send_message(msg)
self._in_sock_callback = lambda x: _callback(x, self.in_sock_type)
self.in_sock.on_recv(self._in_sock_callback)
self.ctrl_sock.on_recv(lambda x: _callback(x, self.ctrl_sock_type))
if self.out_sock_type == zmq.ROUTER:
self.out_sock.on_recv(lambda x: _callback(x, self.out_sock_type))
self.io_loop.start()
self.io_loop.clear_current()
self.io_loop.close(all_fds=True)
class ControlledProcess(mp.Process):
"""
Process subclass with control port.
Parameters
----------
port_ctrl : int
Network port for receiving control messages.
quit_sig : int
OS signal to use when quitting the proess.
id : str
Unique object identifier. Used for communication and logging.
See Also
--------
multiprocessing.Process
"""
def __init__(self, port_ctrl, id, *args, **kwargs):
# Unique object identifier:
self.id = id
# Logging:
self.logger = twiggy.log.name(self.id)
# Control port:
self.port_ctrl = port_ctrl
# Flag to use when stopping the process:
self.running = False
super(ControlledProcess, self).__init__(*args, **kwargs)
def _ctrl_handler(self, msg):
"""
Control port handler.
"""
self.logger.info('recv: %s' % str(msg))
if msg[0] == 'quit':
try:
self.stream_ctrl.flush()
self.stream_ctrl.stop_on_recv()
self.ioloop_ctrl.stop()
except IOError:
self.logger.info('streams already closed')
except Exception as e:
self.logger.info('other error occurred: ' + e.message)
self.running = False
def _init_ctrl_handler(self):
"""
Initialize control port handler.
"""
# Set the linger period to prevent hanging on unsent messages
# when shutting down:
self.logger.info('initializing ctrl handler')
self.sock_ctrl = self.zmq_ctx.socket(zmq.DEALER)
self.sock_ctrl.setsockopt(zmq.IDENTITY, self.id)
self.sock_ctrl.setsockopt(zmq.LINGER, LINGER_TIME)
self.sock_ctrl.connect('tcp://localhost:%i' % self.port_ctrl)
self.stream_ctrl = ZMQStream(self.sock_ctrl, self.ioloop_ctrl)
self.stream_ctrl.on_recv(self._ctrl_handler)
def _init_net(self, event_thread=True):
"""
Initialize network connection.
Parameters
----------
event_thread : bool
If True, start the control event loop in a new thread.
"""
# Set up zmq context and event loop:
self.zmq_ctx = zmq.Context()
self.ioloop_ctrl = IOLoop.instance()
# Set up event loop handlers:
self._init_ctrl_handler()
# Start event loop:
if event_thread:
th.Thread(target=self.ioloop_ctrl.start).start()
else:
self.ioloop_ctrl.start()
def run(self):
"""
Body of process.
"""
self._init_net()
self.running = True
while True:
self.logger.info('idling')
if not self.running:
self.logger.info('stopping run loop')
break
self.logger.info('done')
class ZMQPubSub(PubSubBase):
def __init__(self,
device_ip='127.0.0.1',
fport=5559,
bport=5560,
*args,
**kwargs):
self.channels = list()
self.device_ip = '127.0.0.1'
self.fport = fport
self.bport = bport
super(ZMQPubSub, self).__init__(*args, **kwargs)
##
## pubsub api
##
def connect(self):
self.context = zmq.Context()
self.socket = self.context.socket(zmq.SUB)
self.socket.connect('tcp://%s:%s' % (self.device_ip, self.bport))
self.stream = ZMQStream(self.socket)
self.stream.on_recv(self.on_streaming_data)
self.connected()
def disconnect(self):
self.disconnected()
def subscribe(self, channel_id):
self.socket.setsockopt(zmq.SUBSCRIBE, str(channel_id))
self.channels.append(channel_id)
self.subscribed(channel_id)
def unsubscribe(self, channel_id=None):
channels = [channel_id] if channel_id else self.channels
for channel_id in channels:
self.socket.setsockopt(zmq.UNSUBSCRIBE, str(channel_id))
self.unsubscribed(channel_id)
self.channels.remove(channel_id)
@staticmethod
def publish(channel_id, message, device_ip='127.0.0.1', fport=5559):
context = zmq.Context()
socket = context.socket(zmq.PUSH)
socket.connect('tcp://%s:%s' % (device_ip, fport))
socket.send_unicode('%s %s' % (channel_id, message))
##
## other methods
##
def on_streaming_data(self, data):
for l in data:
reply = l.split(' ', 1)
self.on_message(reply[0], reply[1])
@staticmethod
def start_service(fport=5559, bport=5560):
try:
context = zmq.Context(1)
frontend = context.socket(zmq.PULL)
frontend.bind('tcp://*:%s' % fport)
backend = context.socket(zmq.PUB)
backend.bind('tcp://*:%s' % bport)
logger.info('starting zmq device')
zmq.device(zmq.FORWARDER, frontend, backend)
except KeyboardInterrupt:
pass
except Exception as e:
logger.exception(e)
finally:
frontend.close()
backend.close()
context.term()
def setup_zmq(self, url):
zmq_socket = self.context.socket(zmq.REP)
zmq_socket.bind(url)
zmq_stream = ZMQStream(zmq_socket)
zmq_stream.on_recv(self.handle_zmq)
return zmq_socket, zmq_stream, url
