def connect(self):
"""
Connect To the 0mq REPL socket and attempt a handshake to
ensure we're properly connected.
"""
# Socket to send handshake signals
self.handshake_socket = None
try:
self.handshake_socket = self.context.socket(zmq.REQ)
self.handshake_socket.connect(self.handshake_bind)
self.handshake_socket.setsockopt(zmq.LINGER, 0)
poll = zmq.Poller()
poll.register(self.handshake_socket, zmq.POLLIN)
self.handshake_socket.send("")
socks = dict(poll.poll(self.handshake_timeout))
if socks.get(self.handshake_socket) == zmq.POLLIN:
self.handshake_socket.recv()
return self.set_connected(True)
else:
return self.set_connected(False)
finally:
# Shutdown the handshake
if self.handshake_socket != None:
self.handshake_socket.close()
def _msg(self, req, rep):
self.pinger.disable()
try:
with self.locker:
poller = zmq.Poller()
poller.register(self.master, zmq.POLLIN)
# ping the master we are online, with an ID
try:
data = serialize(req, self.identity)
self.master.send(data, zmq.NOBLOCK)
except zmq.ZMQError, e:
raise RegisterError(str(e))
try:
events = dict(poller.poll(self.timeout))
except zmq.ZMQError, e:
raise RegisterError(str(e))
if events == {}:
raise RegisterError("Nothing came back")
else:
for socket in events:
res = socket.recv()
if res != rep:
raise RegisterError(res)
def _execute(self, job_id, job_data, timeout=1.):
worker = None
timeout *= 1000. # timeout is in ms
data = serialize("JOB", str(job_id), job_data)
try:
with self.workers.worker() as worker:
try:
worker.send(data, zmq.NOBLOCK)
except zmq.ZMQError, e:
raise ExecutionError(str(e))
poller = zmq.Poller()
poller.register(worker, zmq.POLLIN)
try:
events = dict(poller.poll(timeout))
except zmq.ZMQError, e:
raise ExecutionError(str(e))
if events == {}:
raise TimeoutError()
for socket in events:
try:
msg = unserialize(socket.recv())
except zmq.ZMQError, e:
raise ExecutionError(str(e))
if msg[0] == 'JOBRES':
# we got a result
return msg[-1]
else:
raise NotImplementedError(str(msg))
def kill_workers(self, timeout=5):
"""
Send a suicide message to all workers, with some kind of timeout.
"""
logging.info('Killing workers, taking up to %d seconds.', int(timeout))
poller = zmq.Poller()
poller.register(self.results_pull, zmq.POLLIN)
while True:
# Seems to get stuck gevent-blocking in the work_push.send() after
# all the workers have died. Also, gevent.Timeout() doesn't seem
# to work here?!
signal.alarm(int(timeout))
self.work_push.send(msgpack.dumps([{'type': 'PING'}]))
socks = dict(poller.poll(timeout * 1500))
if self.results_pull in socks \
and socks[self.results_pull] == zmq.POLLIN:
result_packed = self.results_pull.recv()
result = msgpack.loads(result_packed, use_list=False)
logging.info('Heard from worker id=%d; sending SUICIDE',
result[0]['worker_id'])
self.work_push.send(msgpack.dumps([{'type': 'SUICIDE'}]))
gevent.sleep(0.1)
else:
break
signal.alarm(0)
def recv_timeout(socket, timeout): poll = zmq.Poller() poll.register(socket, zmq.POLLIN) socks = dict(poll.poll(timeout=timeout)) if socks.get(socket, None) == zmq.POLLIN: return socket.recv_multipart() else: return None
def _callback_worker(self, wrk_num, worker_func):
log.debug('Initializing worker number %d', wrk_num)
context = zmq.Context()
# channel <- run_completion_task
task_receiver = context.socket(zmq.PULL)
task_receiver.connect(ZMQ_TASK_LAUNCH_CH)
# channel -> result_manager
result_sender = context.socket(zmq.PUSH)
result_sender.connect(ZMQ_TASK_RESULT_CH)
# channel <- result_manager (CONTROL)
control_receiver = context.socket(zmq.SUB)
control_receiver.connect(ZMQ_WORKER_CONTROL_CH)
control_receiver.setsockopt(zmq.SUBSCRIBE, "")
poller = zmq.Poller()
poller.register(task_receiver, zmq.POLLIN)
poller.register(control_receiver, zmq.POLLIN)
# tell callback handler that worker has initialized
syncservice = context.socket(zmq.REQ)
syncservice.connect(ZMQ_WORKER_SYNC_CH)
syncservice.send(ZMQ_CONTROL_SYNC)
syncservice.recv()
# loop and accept messages from both task and control channels
while True:
socks = dict(poller.poll())
# handle callback launch requests
if socks.get(task_receiver) == zmq.POLLIN:
worker_params = task_receiver.recv_json()
#log.debug('Launching post-computation for file: %s',
# out_dict['clean_fn'])
if DEBUG_CB:
launched_tasks = worker_params['kwargs']['launched_tasks']
log.debug('Launching task %d', launched_tasks)
del worker_params['kwargs']['launched_tasks']
else:
log.debug('Launching task')
worker_func(*worker_params['args'], **worker_params['kwargs'])
if DEBUG_CB:
log.debug('Completed task %d', launched_tasks)
else:
log.debug('Completed task')
result_sender.send(ZMQ_RESULT_DONE)
# handle control commands from manager
if socks.get(control_receiver) == zmq.POLLIN:
control_message = control_receiver.recv()
if control_message == ZMQ_CONTROL_DONE:
log.debug('Terminating worker number %d', wrk_num)
break
def __init__(self,
address,
timeout=2500,
retries=3,
log_level=DEFAULT_LOG_LEVEL):
self.address = address
self.timeout = timeout
self.retries = retries
self.context = zmq.Context(1)
self.poller = zmq.Poller()
self.client = None
self.logger = get_logger(self, log_level)
def run(self):
self.alive = True
# channel to communicate with the workers
logger.debug('Starting [workermgr]')
client = self.context.socket(zmq.REP)
client.identity = 'master'
client.bind(self.endpoint)
poller = zmq.Poller()
poller.register(client, zmq.POLLIN)
poll_timeout = 1000
while self.alive:
try:
events = dict(poller.poll(poll_timeout))
except zmq.ZMQError, e:
logger.debug("The poll failed")
logger.debug(str(e))
break
for socket in events:
msg = unserialize(socket.recv())
if len(msg) < 2:
# XXX log
socket.send('ERROR')
if msg[-2] == 'PING':
logger.debug("[workermgr] Got a PING")
if msg[-1] not in self.workers:
name = msg[-1]
logger.debug("Registered " + name)
# keep track of that worker
work = self.context.socket(zmq.REQ)
work.connect(name)
work.identity = name
self.workers.add(work)
# in any case we pong back
logger.debug("[workermgr] sent a PONG")
socket.send('PONG')
elif msg[-2] == 'REMOVE':
if msg[-1] in self.workers:
logger.debug("[workermgr] Removing` " + msg[-1])
self.workers.delete(msg[-1])
socket.send('REMOVED')
else:
logger.debug('Error')
socket.send('ERROR')
time.sleep(.1)
def __init__(self, identity, socket, locker, fail_callable,
duration=5., max_fails=10.):
threading.Thread.__init__(self)
self.duration = duration
self.identity = identity
logger.debug('starting pinger from %s' % self.identity)
self.socket = socket
self.locker = locker
self.running = False
self.fail_callable = fail_callable
self.max_fails = max_fails
self.poller = zmq.Poller()
self.poller.register(self.socket, zmq.POLLIN)
self.disabled = False
self.unresponsive = False
def __init__(self, endpoint, identity, target, timeout=1.):
self.identity = identity
self.ctx = zmq.Context()
self.timeout = timeout * 1000
self.master = self.ctx.socket(zmq.REQ)
self.master.connect(endpoint)
self.work = self.ctx.socket(zmq.REP)
self.work.bind(identity)
self.registered = self.running = False
# setting a poller
self.poller = zmq.Poller()
self.poller.register(self.work, zmq.POLLIN)
self.locker = threading.RLock()
self.pinger = Pinger(self.identity, self.master, self.locker,
self.failed)
self.target = target
def result_manager(num_tasks):
"""Function used internally by CallbackHandler to collect results of tasks"""
context = zmq.Context()
# channel <- completion_worker
result_receiver = context.socket(zmq.PULL)
result_receiver.bind(ZMQ_TASK_RESULT_CH)
# channel -> completion_worker (CONTROL)
control_sender = context.socket(zmq.PUB)
control_sender.bind(ZMQ_WORKER_CONTROL_CH)
# channel <- skip_completion_task
tcdec_receiver = context.socket(zmq.SUB)
tcdec_receiver.connect(ZMQ_TASK_COUNT_DEC_CH)
tcdec_receiver.setsockopt(zmq.SUBSCRIBE, "")
poller = zmq.Poller()
poller.register(result_receiver, zmq.POLLIN)
poller.register(tcdec_receiver, zmq.POLLIN)
while True:
socks = dict(poller.poll())
# handle results from workers
if socks.get(result_receiver) == zmq.POLLIN:
result_message = result_receiver.recv()
if result_message == ZMQ_RESULT_DONE:
num_tasks -= 1
log.debug('Completed post-computation, remaining tasks: %d',
num_tasks)
# handle counter decrementing
if socks.get(tcdec_receiver) == zmq.POLLIN:
result_message = tcdec_receiver.recv()
if result_message == ZMQ_RESULT_SKIPPING:
num_tasks -= 1
log.debug('Skipped post-computation, remaining tasks: %d',
num_tasks)
# break from loop when all tasks completed
if num_tasks == 0:
log.debug('All tasks done - terminating workers')
break
# terminate all workers when all tasks completed
control_sender.send(ZMQ_CONTROL_DONE)
def collect(self):
self.register()
poller = zmq.Poller()
poller.register(self.pull_socket, zmq.POLLIN)
poller.register(self.control_socket, zmq.POLLIN)
# multiplex the pull and control ports
pull_socket = self.pull_socket
control_socket = self.control_socket
while True:
# Wait until the server pushes bytecode to use to
# listening for data, ( this is a race condition
# otherwise )
if self.have_bytecode:
try:
socks = dict(poller.poll())
except zmq.ZMQError:
# Die gracefully if the user sends a SIGQUIT
self._kill()
break
if pull_socket in socks and socks[pull_socket] == zmq.POLLIN:
msg = self.pull_socket.recv()
if msg:
command, data = msg.split(self.delim)
self.process_command(command, data)
if control_socket in socks and socks[
control_socket] == zmq.POLLIN:
msg = self.control_socket.recv()
if msg:
worker, command, data = msg.split(self.delim)
self.process_command(command, data)
else:
msg = self.control_socket.recv()
if msg:
worker, command, data = msg.split(self.delim)
self.process_command(command, data)
def __init__(self, reqrep_socket, timeout=1000.0):
self.reqrep_socket = reqrep_socket
self.poller = zmq.Poller()
self.poller.register(self.reqrep_socket, zmq.POLLIN)
self.kill = False
self.timeout = timeout
def __init__(self, socket, ident, timeout=1000.0):
self.socket = socket
self.ident = ident
self.timeout = timeout
self.poller = zmq.Poller()
self.poller.register(self.socket, zmq.POLLIN)
def eventhub_client():
if not request.environ.get('wsgi.websocket'):
return ""
websocket = request.environ['wsgi.websocket']
push_socket = zmq_context.socket(zmq.PUSH)
push_socket.connect(current_app.config["PUSH_ADDRESS"])
subscribe_socket = zmq_context.socket(zmq.SUB)
client_id = str(uuid.uuid4())
subscribe_socket.connect(current_app.config["SUBSCRIBE_ADDRESS"])
# add yourself to your current pool of open clients
add_to_user_clients(g.user, client_id)
# listen for messages that happen on channels the user is subscribed to
for channel in g.user["channels"]:
channel_id = zmq_channel_key(channel)
subscribe_socket.setsockopt(zmq.SUBSCRIBE, channel_id)
# if redis was cleared, we'll need to resend the join channel event to populate the user status of open
# clients
channel_status = get_user_channel_status(g.user, channel)
if channel_status is None:
send_join_channel(channel, g.user, push_socket)
set_user_channel_status(g.user, channel, "active")
send_user_status_update(g.user, channel, push_socket, "active")
# subscribe to events the user triggered that could affect the user's other open clients
subscribe_socket.setsockopt(zmq.SUBSCRIBE, str(g.user["email"]))
poller = zmq.Poller()
poller.register(subscribe_socket, zmq.POLLIN)
poller.register(websocket.socket, zmq.POLLIN)
try:
message = None
while True:
events = dict(poller.poll())
# Server -> Client
if subscribe_socket in events:
message = subscribe_socket.recv()
# the message is prepended by the channel_id (for PUB/SUB reasons)
channel_id, packed = message.split(" ", 1)
g.msg_unpacker.feed(packed)
unpacked = g.msg_unpacker.unpack()
action = unpacked["action"]
if action in ["publish_message", "join_channel", "leave_channel", "user_active", "user_offline"]:
websocket.send(json.dumps(unpacked))
elif action in ["self_join_channel", "self_leave_channel", "self_reorder_channels"]:
event_type = action.split("_")[1]
handle_self_channel_event(client_id, websocket, subscribe_socket, unpacked["data"], event_type)
# Client -> Server
if websocket.socket.fileno() in events:
socket_data = websocket.receive()
if socket_data is None:
break
socket_data = json.loads(socket_data)
action = socket_data["action"]
data = socket_data["data"]
if action == "switch_channel":
handle_switch_channel(data["channel"])
elif action == "publish_message":
handle_publish_message(data, push_socket)
elif action == "preview_message":
handle_preview_message(data, websocket)
elif action == "join_channel":
handle_join_channel(data["channel"], subscribe_socket, push_socket, client_id)
elif action == "leave_channel":
handle_leave_channel(data["channel"], subscribe_socket, push_socket, client_id)
elif action == "reorder_channels":
handle_reorder_channels(data["channels"], push_socket, client_id)
except geventwebsocket.WebSocketError, e:
print "{0} {1}".format(e.__class__.__name__, e)
for i in xrange(10):
sender.send('test %d' % i)
gevent.sleep(1)
# create zmq context, and bind to pull sockets
context = zmq.Context()
receiver1 = context.socket(zmq.PULL)
receiver1.bind('inproc://polltest1')
receiver2 = context.socket(zmq.PULL)
receiver2.bind('inproc://polltest2')
gevent.spawn(sender)
# Create poller and register both reciever sockets
poller = zmq.Poller()
poller.register(receiver1, zmq.POLLIN)
poller.register(receiver2, zmq.POLLIN)
# Read 10 messages from both reciever sockets
msgcnt = 0
while msgcnt < 10:
socks = dict(poller.poll())
if receiver1 in socks and socks[receiver1] == zmq.POLLIN:
print "Message from receiver1: %s" % receiver1.recv()
msgcnt += 1
if receiver2 in socks and socks[receiver2] == zmq.POLLIN:
print "Message from receiver2: %s" % receiver2.recv()
msgcnt += 1
