def unregister_worker(self, wid):
"""Unregister the worker with the given id and stop all timers for the worker.
If the worker id is not registered, nothing happens.
:param wid: the worker id.
:type wid: str
:rtype: None
"""
try:
wtracker = self._workers[wid]
except KeyError:
# not registered, ignore
return
# remove this workers' data from the map
wtracker.shutdown()
service = wtracker.service
if service in self._services:
wq, wr = self._services[service]
wq.remove(wid)
del self._workers[wid]
self.reset_node_info(wid)
_LOG.info("Worker with id: '%s' was removed from the pool." %
bytes_to_hexstring(wid))
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
"""
_LOG.debug("Received a new reply from worker: %s." % rp)
cmd = msg.pop(0)
if cmd in self._worker_cmds:
fnc = self._worker_cmds[cmd]
fnc(rp, msg)
else:
# ignore unknown command
# DISCONNECT worker
_LOG.info(
"Worker with id: '%s' is trying to use an unknown command." %
bytes_to_hexstring(rp[0]))
self.disconnect(rp[0])
return
def register_worker_info(self,
wid,
service=SERVICE_BROKER,
status=WORKER_ONLINE_STATUS,
worker_type=WK_TYPE_BROKER,
address=get_current_ip(DEFAULT_INTEFACE),
protocols=None):
"""Update the worker info list.
:param wid: the worker id.
:type wid: str
:param service: the service name.
:type service: str
:param status: the current network status of the worker.
:type status: str
:param worker_type: the specific worker type.
:type worker_type: str
:param address: the ipv4 or ipv6 of this worker.
:type address: str
:param protocols: the routing protocols reported by the worker.
:type protocols: str
:rtype: None
"""
service_list = set([])
service_list.add(service)
existing_worker = next((x for x in self._workers_info
if self._workers_info[x]['ip'] == address),
None)
if existing_worker:
service_list.update(self._workers_info[existing_worker]['service'])
del self._workers_info[existing_worker]
if isinstance(protocols, str):
protocols = protocols.split(' ')
worker_info = {
'id': bytes_to_hexstring(wid),
'status': status,
'ip': address,
'type': worker_type,
'service': list(service_list),
'protocols': protocols
}
self._workers_info[wid] = worker_info
# broker will store a list of protocols supported by all clients
if wid != self.main_stream.getsockopt(zmq.IDENTITY):
if self._workers_info[self.main_stream.getsockopt(
zmq.IDENTITY)]['protocols'] is None:
self._workers_info[self.main_stream.getsockopt(
zmq.IDENTITY)]['protocols'] = protocols
else:
self._workers_info[self.main_stream.getsockopt(zmq.IDENTITY)]['protocols'] = \
list(set(self._workers_info[self.main_stream.getsockopt(zmq.IDENTITY)]['protocols']).intersection(
set(protocols)))
return
def reset_node_info(self, wid, status=WORKER_OFFLINE_STATUS):
worker_info = {
'id': bytes_to_hexstring(wid),
'status': status,
'ip': self._workers_info[wid]['ip'],
'type': self._workers_info[wid]['type'],
'service': self._workers_info[wid]['service'],
'protocols': self._workers_info[wid]['protocols']
}
try:
self._workers_info[wid] = worker_info
except KeyError:
pass
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]
wtracker = self._workers.get(ret_id)
if not wtracker:
# worker not found, ignore message
return
service = wtracker.service
# make worker available again
try:
wq, wr = self._services[service]
cp, msg = split_address(msg)
if cp[0] == 'BROKER':
self.update_worker_info(ret_id, msg)
return
cmd = msg.pop(0)
self.client_response(cp, service, cmd,
bytes_to_hexstring(wtracker.id), msg)
wq.put(wtracker.id)
self.change_worker_status(rp[0], WORKER_ONLINE_STATUS)
if wr:
proto, rp, msg = wr.pop(0)
self.on_client(proto, rp, msg)
except KeyError:
# unknown service
_LOG.info("Worker with id: '%s' reports an unknown service." %
bytes_to_hexstring(ret_id))
self.disconnect(ret_id)
return
def on_timer(self):
"""Method called on timer expiry.
Checks which workers are dead and unregisters them.
:rtype: None
"""
# use list to avoid size change during iteration error
for wtracker in list(self._workers.values()):
if not wtracker.is_alive():
_LOG.debug("Worker with id: '%s' timed out." %
bytes_to_hexstring(wtracker.id))
self.unregister_worker(wtracker.id)
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]
_LOG.debug("Worker with id: '%s' wants to disconnect." %
bytes_to_hexstring(wid))
self.change_worker_status(wid, WORKER_INACTIVE_STATUS)
self.unregister_worker(wid)
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
"""
try:
wtracker = self._workers[wid]
except KeyError:
# not registered, ignore
return
_LOG.info("Requesting from worker with id: '%s' to disconnect." %
bytes_to_hexstring(wid))
to_send = [wid, b'', WORKER_PROTO, MSG_DISCONNECT]
if self.main_stream.closed():
self.shutdown()
self.main_stream.send_multipart(to_send)
self.unregister_worker(wid)
return
def register_worker(self, wid, service, worker_type, address, protocols):
"""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
:param worker_type: the type of the worker.
:type worker_type: str
:param address: the ipv4 or upv6 address of the worker.
:type address: str
:param protocols: the routing protocols reported by the worker.
:type protocols: str
:rtype: None
"""
if wid in self._workers:
return
self._workers[wid] = WorkerTracker(WORKER_PROTO, wid, service,
self.main_stream)
# If service exists then add this worker to its workers queue, if not create it.
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, [])
self.register_worker_info(wid, service, WORKER_ONLINE_STATUS,
worker_type, address, protocols)
_LOG.info(
"New worker of type: '%s' registered with id: '%s', for service: '%s' and can be found at '%s'."
% (worker_type, bytes_to_hexstring(wid), service, address))
return