def send_and_recv(self,
method,
endpoint,
parameters=None,
data=None,
decode_response=True):
logger.debug('Request: {} {}'.format(method, endpoint))
message = {
'method': method,
'endpoint': endpoint,
}
if parameters:
message['parameters'] = parameters
if data:
message['data'] = data
self._socket.send(msgpack.packb(message))
zmq_response = get_data_from_zmq(self._socket, self._timeout)
zmq_response = msgpack.unpackb(zmq_response)
if decode_response:
response = Response(decode(zmq_response[b'response']),
decode(zmq_response[b'status']))
else:
response = Response(zmq_response[b'response'],
zmq_response[b'status'])
logger.debug('Response: {} - {}'.format(str(response.status_code),
str(response.response)))
return response
def __init__(self, source, content):
self._id = uuid.uuid4().hex # str
self._source = copy.deepcopy(source) # list of bytes
self._content = msgpack.unpackb(content) # dictionary of bytes
self._method = decode(self._content[b'method']) # str
self._endpoint = decode(self._content[b'endpoint']) # str
logger.debug('Creating or loading ClientRequest {}@{}: {}/{}'.format(
self._id, self._endpoint, self._source, self._content))
def __init__(self, request):
self._id = decode(request[b'id']) # str
self._method = decode(request[b'method']) # str
self._endpoint = decode(request[b'endpoint']) # str
if b'parameters' in request:
self._parameters = decode(request[b'parameters']) # dict of str
else:
self._parameters = {}
if b'data' in request:
self._data = request[b'data'] # dict of bytes
else:
self._data = {}
def run(self):
self._connect()
poller = zmq.Poller()
poller.register(self._socket_control, zmq.POLLIN)
logger.info('Starting with endpoints {}'.format(ctx_request.endpoints))
state_manager = BrokerStateManager(
seconds_before_contact_check=SECONDS_BEFORE_CONTACT_CHECK,
seconds_before_disconnect=SECONDS_BEFORE_DISCONNECT)
loop = True
init_connection = True
poller_timeout = max([
int(
min([SECONDS_BEFORE_CONTACT_CHECK, SECONDS_BEFORE_DISCONNECT])
/ 10.0 * 1000), 500
])
while loop:
if init_connection:
self._socket_control.send_multipart([
b'',
msgpack.packb({
'endpoints': ctx_request.endpoints,
'w': True
})
])
init_connection = False
sockets = dict(poller.poll(poller_timeout))
# if sockets:
# logger.debug('Received: {}'.format(sockets))
if self._socket_control in sockets and sockets[
self._socket_control] == zmq.POLLIN:
# get message content
content = extract_content_from_message(
self._socket_control.recv_multipart())
content = msgpack.unpackb(content)
state_manager.contact_from_broker()
# process control messages
if b'control' in content:
control = decode(content[b'control'])
if control == 'ping':
# logger.debug('Received ping from broker')
self._socket_control.send_multipart(
[b'', msgpack.packb({'pong': True})])
elif control == 'pong':
# logger.debug('Received pong from broker')
pass
elif control == 'kick':
init_connection = True
# process client request messages
else:
message = Request(content)
self._socket_control.send_multipart(
[b'', msgpack.packb({'r': message.id})])
try:
logger.info('Received: {} {}'.format(
message.method, message.endpoint))
service = ctx_request.get_service_by_endpoint(
message.endpoint, message.method)
response = service(message)
if isinstance(response, tuple):
response, status = response
else:
status = codes.OK
except RequestError as e:
response = str(e)
status = e.status_code
except:
logger.error(sys.exc_info()[0])
logger.error(sys.exc_info()[1])
logger.error(traceback.extract_tb(sys.exc_info()[2]))
response = {}
status = codes.SERVER_ERROR
self._socket_response.send(
msgpack.packb({
'id': message.id,
'status': status,
'response': response
}))
self._socket_control.send_multipart(
[b'', msgpack.packb({'w': True})])
self._socket_response.recv()
if state_manager.ping_broker():
# logger.debug('Pinging broker')
self._socket_control.send_multipart(
[b'', msgpack.packb({'ping': True})])
if state_manager.disconnect_broker():
logger.info('Disconnecting from broker')
self._socket_control.send_multipart(
[b'', msgpack.packb({'disconnect': True})])
raise RuntimeError
def __init__(self, publication):
publication = msgpack.unpackb(publication)
self._id = decode(publication[b'id']) # str
self._topic = decode(publication[b'topic']) # str
self._data = publication[b'data'] # dict of bytes
def run(self):
poller = zmq.Poller()
poller.register(self._client_socket, zmq.POLLIN)
poller.register(self._worker_control_socket, zmq.POLLIN)
poller.register(self._worker_response_socket, zmq.POLLIN)
request_manager = RequestManager(redis_host=self._redis_host,
redis_port=self._redis_port,
redis_db=self._redis_db)
state_manager = ConnectionStateManager(
seconds_before_contact_check=SECONDS_BEFORE_CONTACT_CHECK,
seconds_before_disconnect=SECONDS_BEFORE_UNREGISTER)
poller_timeout = max([
int(
min([SECONDS_BEFORE_CONTACT_CHECK, SECONDS_BEFORE_UNREGISTER])
/ 10.0 * 1000), 500
])
loop = True
while loop:
sockets = dict(poller.poll(poller_timeout))
# get a new client request
if self._client_socket in sockets and sockets[
self._client_socket] == zmq.POLLIN:
message = self._client_socket.recv_multipart()
client_request = ClientRequest(
source=extract_source_from_message(message),
content=extract_content_from_message(message))
request_manager.append(client_request)
# register endpoints of a worker or mark the worker as waiting
if self._worker_control_socket in sockets and sockets[
self._worker_control_socket] == zmq.POLLIN:
message = self._worker_control_socket.recv_multipart()
source = extract_source_from_message(message)
content = extract_content_from_message(message)
assert len(source) == 1
worker_id = source[0]
content = decode(msgpack.unpackb(content))
state_manager.contact_from(worker_id)
if 'endpoints' in content:
# first connection to register endpoints
request_manager.register(worker_id, content['endpoints'])
if 'ping' in content and content['ping']:
# ping to check if broker is still available
logger.debug('Received ping from {}'.format(worker_id))
if worker_id in request_manager.registered_workers:
# only respond to registered workers, otherwise disconnect
self._worker_control_socket.send_multipart([
worker_id, b'',
msgpack.packb(
ControlRequest(ControlRequest.PONG).content)
])
else:
# we don't know this worker, so remove it
self.send_kick(worker_id)
state_manager.disconnect(worker_id)
if 'pong' in content and content['pong']:
logger.debug('Received pong from {}'.format(worker_id))
if 'disconnect' in content and content['disconnect']:
# disconnect worker
request_manager.unregister(worker_id)
state_manager.disconnect(worker_id)
if 'r' in content:
# acknowledge reception of task
pass
if 'w' in content and content['w']:
# signal that task is done
request_manager.worker_available(worker_id)
# receive responses and send them back to the client
if self._worker_response_socket in sockets and sockets[
self._worker_response_socket] == zmq.POLLIN:
message = self._worker_response_socket.recv_multipart()
source = extract_source_from_message(message)
response = extract_content_from_message(message)
response = msgpack.unpackb(response)
worker_response = source + [b''] + [b'OK']
self._worker_response_socket.send_multipart(worker_response)
request = request_manager[response[b'id'].decode()]
del request_manager[response[b'id'].decode()]
del response[b'id']
client_response = request.source + [b''] + [
msgpack.packb(response)
]
self._client_socket.send_multipart(client_response)
# send requests to workers
for worker_id, request in request_manager():
logger.info('Sending {} to {}'.format(request.id, worker_id))
self.send_request(worker_id, request)
# send ping messages to workers
for worker_id in state_manager.get_connections_to_ping():
logger.debug('Pinging {}'.format(worker_id))
self.send_ping(worker_id)
# send kick messages to non-responsive workers
for worker_id in state_manager.get_connections_to_disconnect():
logger.info('Kicking {}'.format(worker_id))
self.send_kick(worker_id)
request_manager.unregister(worker_id)