def recv(self):
"""Returns the reply message or None if there was no reply"""
try:
items = self.poller.poll(self.timeout)
except KeyboardInterrupt:
return # interrupted
if items:
msg = self.client.recv_multipart()
if self.verbose:
logging.info("I: received reply: ")
dump(msg)
# Don't try to handle errors, just assert noisily
assert len(msg) >= 4
empty = msg.pop(0)
assert empty == ''
header = msg.pop(0)
assert C_CLIENT == header
# pop service part of message off
msg.pop(0)
return msg
else:
logging.warn("W: permanent error, abandoning request")
def recv(self):
"""Returns the reply message or None if there was no reply"""
try:
items = self.poller.poll(self.timeout)
except KeyboardInterrupt:
return # interrupted
if items:
msg = self.client.recv_multipart()
if self.verbose:
logging.info("I: received reply: ")
dump(msg)
# Don't try to handle errors, just assert noisily
assert len(msg) >= 4
empty = msg.pop(0)
assert empty == ''
header = msg.pop(0)
assert C_CLIENT == header
# pop service part of message off
msg.pop(0)
return msg
else:
logging.warn("W: permanent error, abandoning request")
def mediate(self):
"""Main broker work happens here"""
while True:
try:
items = self.poller.poll(self.HEARTBEAT_INTERVAL)
except KeyboardInterrupt:
break # Interrupt
if items:
msg = self.socket.recv_multipart()
if self.verbose:
log.info("I: received message: ")
dump(msg)
sender = msg.pop(0)
empty = msg.pop(0)
assert empty == ''
header = msg.pop(0)
if (C_CLIENT == header):
self.process_client(sender, msg)
elif (W_WORKER == header):
self.process_worker(sender, msg)
else:
log.error("E: invalid message: ")
dump(msg)
self.purge_workers()
self.send_heartbeats()
def send(self, service, request):
"""Send request to broker and get reply by hook or crook.
Takes ownership of request message and destroys it when sent.
Returns the reply message or None if there was no reply.
"""
if not isinstance(request, list):
request = [request]
request = [C_CLIENT, service] + request
if self.verbose:
logging.warn("I: send request to '%s' service: ", service)
dump(request)
reply = None
retries = self.retries
while retries > 0:
self.client.send_multipart(request)
try:
items = self.poller.poll(self.timeout)
except KeyboardInterrupt:
break # interrupted
if items:
msg = self.client.recv_multipart()
if self.verbose:
logging.info("I: received reply")
dump(msg)
# Don't try to handle errors, just assert noisily
assert len(msg) >= 3
# make sure that we have the correct header
header = msg.pop(0)
assert C_CLIENT == header
# make sure we got the correct response
reply_service = msg.pop(0)
assert service == reply_service
reply = msg
break
else:
if retries:
logging.warn("W: no reply, reconnecting...")
self.reconnect_to_broker()
else:
logging.warn("W: permanent error, abandoning")
break
retries -= 1
return reply
def send(self, service, request):
"""Send request to broker"""
if not isinstance(request, list):
request = [request]
# Prefix request qwith protocol frames
# Frame 0: empty (REQ emulation)
# Frame 1: "MDPCxy" (six bytes, MDP/Client x.y)
# Frame 2: Service name (printable string)
request = ['', C_CLIENT, service] + request
if self.verbose:
logging.warn("I: send request to '%s' service: ", service)
dump(request)
self.client.send_multipart(request)
def send(self, service, request):
"""Send request to broker"""
if not isinstance(request, list):
request = [request]
# Prefix request qwith protocol frames
# Frame 0: empty (REQ emulation)
# Frame 1: "MDPCxy" (six bytes, MDP/Client x.y)
# Frame 2: Service name (printable string)
request = ['', C_CLIENT, service] + request
if self.verbose:
logging.warn("I: send request to '%s' service: ", service)
dump(request)
self.client.send_multipart(request)
def process_worker(self, sender, msg):
"""Process message sent to us by a worker."""
assert len(msg) >= 1 # At least, command
command = msg.pop(0)
worker_ready = hexlify(sender) in self.workers
worker = self.require_worker(sender)
if (W_READY == command):
assert len(msg) >= 1 # At least, a service name
service = msg.pop(0)
# Not first command in session or Reserved service name
if (worker_ready or service.startswith(
self.INTERNAL_SERVICE_PREFIX
)):
self.delete_worker(worker, True)
else:
# Attach worker to service and mark as idle
worker.service = self.require_service(service)
self.worker_waiting(worker)
elif (W_REPLY == command):
if (worker_ready):
# Remove and save client return envelope and insert the
# protocol header and service name, then rewrap envelope
client = msg.pop(0)
empty = msg.pop(0)
msg = [client, empty, C_CLIENT, worker.service.name] + msg
self.socket.send_multipart(msg)
self.worker_waiting(worker)
else:
self.delete_worker(worker, True)
elif (W_HEARTBEAT == command):
if (worker_ready):
worker.expiry = time.time() + 1e-3 * self.HEARTBEAT_EXPIRY
else:
self.delete_worker(worker, True)
elif (W_DISCONNECT == command):
self.delete_worker(worker, False)
else:
log.error("E: invalid message: ")
dump(msg)
def send_to_broker(self, command, option=None, msg=None):
"""Send message to broker.
If no msg is provided, creates one internally
"""
if msg is None:
msg = []
elif not isinstance(msg, list):
msg = [msg]
if option:
msg = [option] + msg
msg = [b'', W_WORKER, command] + msg
if self.verbose:
logging.info("I: sending %s to broker", command)
dump(msg)
self.worker.send_multipart(msg)
def send_to_worker(self, worker, command, option, msg=None):
"""Send message to worker.
If message is provided, sends that message
"""
if msg is None:
msg = []
elif not isinstance(msg, list):
msg = [msg]
# Stack routing and protocol envelopes to start of message
# and routing envelope
if option is not None:
msg = [option] + msg
msg = [worker.address, b'', W_WORKER, command] + msg
if self.verbose:
log.info("I: sending %r to worker: ", command)
dump(msg)
self.socket.send_multipart(msg)
def test_dump_non_ascii(self):
"""Test dump of non ascii chars"""
msg = ['ã']
with support.captured_stdout() as s:
dump(msg)
assert "%s\n%s\n%s\n" % (BUFFER, "[001]", msg[0]) == s.getvalue()
def test_dump_multipart(self):
"""Test dump function"""
msg = ['hello', 'world']
with support.captured_stdout() as s:
dump(msg)
assert "%s\n%s\n" % (BUFFER, "[005]\nhello\n[005]\nworld") == s.getvalue()
def test_dump(self):
"""Test dump function"""
msg = ['hello']
with support.captured_stdout() as s:
dump(msg)
assert "%s\n%s\n" % (BUFFER, "[005]\nhello") == s.getvalue()
def recv(self, reply=None):
"""Send reply, if any, to broker and wait for next request"""
# Format and send the reply if we were provided one
assert reply is not None or not self.expect_reply
if reply is not None:
assert self.reply_to is not None
reply = [self.reply_to, ''] + reply
self.send_to_broker(W_REPLY, msg=reply)
self.expect_reply = True
while True:
# Poll socket for a reply, with timeout
try:
items = self.poller.poll(self.timeout)
except KeyboardInterrupt:
break # Interrupt
if items:
msg = self.worker.recv_multipart()
if self.verbose:
logging.info("I: received message from broker: ")
dump(msg)
self.liveness = self.HEARTBEAT_LIVENESS
# Don't try to handle errors, just assert noisily
assert len(msg) >= 3
empty = msg.pop(0)
assert empty == ''
header = msg.pop(0)
assert header == W_WORKER
command = msg.pop(0)
if command == W_REQUEST:
# We should pop and save as many addresses as there are
# Up to a null part, but for now, just save one...
self.reply_to = msg.pop(0)
# pop empty
assert msg.pop(0) == ''
return msg # We have a request to process
elif command == W_HEARTBEAT:
# Do nothing for heartbeats
pass
elif command == W_DISCONNECT:
self.reconnect_to_broker()
else:
logging.error("E: invalid input message: ")
dump(msg)
else:
self.liveness -= 1
if self.liveness == 0:
if self.verbose:
logging.warn(
"W: disconnected from broker - retrying..."
)
try:
time.sleep(1e-3 * self.reconnect)
except KeyboardInterrupt:
break
self.reconnect_to_broker()
# Send HEARTBEAT if it's time
if time.time() > self.heartbeat_at:
self.send_to_broker(W_HEARTBEAT)
self.heartbeat_at = time.time() + 1e-3 * self.heartbeat
logging.warn("W: interrupt received, killing worker...")
return None
