class Controller(object):
def __init__(self):
self._uid = utils.uid()
self._addr = utils.get_addr()
# pipes of this controller
self._imanager = ControllerPipeStorage()
self._omanager = ControllerPipeStorage()
# context and poller
self._context = zmq.Context()
self._context.sndhwm = _configs.CTL_CTL_HWM
self._context.rcvhwm = _configs.CTL_CTL_HWM
# socket pools
self._ns_socket = None
# control router and dispatcher
self._control_router = None
self._control_router_port = 0
self._control_dispatcher = CallbackRegistry()
# queue of ControlMessage
self._control_mqueue = queue.Queue()
self._control_poller = zmq.Poller()
# peers respect to the controller
# map uid => ControllerPeer
self._controller_peers = dict()
# the peers of input pipes (i.e. the output pipes)
self._input_from = dict()
self._input_cache = dict()
self._output_to = dict()
self._output_to_pipe = collections.defaultdict(
dict) # the peers of output pipes
self._output_to_id = dict()
self._output_cache = dict() # map pipe_name => cache
self._data_poller = zmq.Poller()
# threads and stop-event
self._all_socks = set()
self._all_threads = []
self._stop_event = threading.Event()
def socket(self, socket_type):
sock = self._context.socket(socket_type)
self._all_socks.add(sock)
return sock
def close_socket(self, sock):
utils.graceful_close(sock)
self._all_socks.remove(sock)
return self
def initialize(self, pipes=None):
pipes = pipes or []
for pipe in pipes:
if pipe.direction == 'IN':
self._imanager.put(pipe)
else:
assert pipe.direction == 'OUT'
self._omanager.put(pipe)
pipe.set_controller(self)
# setup ns socket
self._ns_socket = self.socket(zmq.REQ)
self._ns_socket.connect(
os.getenv(
'JAC_NAME_SERVER',
'{}://localhost:{}'.format(_configs.NS_CTL_PROTOCAL,
_configs.NS_CTL_PORT)))
self._control_poller.register(self._ns_socket, zmq.POLLIN)
# setup router socket
self._control_router = self.socket(zmq.ROUTER)
self._control_router_port = self._control_router.bind_to_random_port(
'tcp://*')
self._control_poller.register(self._control_router, zmq.POLLIN)
# register on the name-server
response = utils.req_send_and_recv(
self._ns_socket, {
'action': _configs.Actions.NS_REGISTER_CTL_REQ,
'uid': self._uid,
'ctl_protocal': 'tcp',
'ctl_addr': self._addr,
'ctl_port': self._control_router_port,
'meta': {}
})
assert response['action'] == _configs.Actions.NS_REGISTER_CTL_REP
# register pipes on name-server
response = utils.req_send_and_recv(
self._ns_socket, {
'action': _configs.Actions.NS_REGISTER_OUTPUTS_REQ,
'uid': self._uid,
'outputs': list(self._omanager.keys())
})
assert response['action'] == _configs.Actions.NS_REGISTER_OUTPUTS_REP
# query name-server for ipipes
response = utils.req_send_and_recv(
self._ns_socket, {
'action': _configs.Actions.NS_REGISTER_INPUTS_REQ,
'uid': self._uid,
'inputs': list(self._imanager.keys())
})
assert response['action'] == _configs.Actions.NS_REGISTER_INPUTS_REP
logger.info('IPipes query {}.'.format(response['results']))
self._initialize_recv_peers(response['results'])
# setup dispatcher
self._control_dispatcher.register(_configs.Actions.NS_HEARTBEAT_REP,
lambda msg: None)
self._control_dispatcher.register(_configs.Actions.CTL_CONNECT_REQ,
self._on_ctl_connect_req)
self._control_dispatcher.register(_configs.Actions.CTL_CONNECT_REP,
self._on_ctl_connect_rep)
self._control_dispatcher.register(_configs.Actions.CTL_CONNECTED_REQ,
self._on_ctl_connected_req)
self._control_dispatcher.register(_configs.Actions.CTL_CONNECTED_REP,
lambda msg: None)
self._control_dispatcher.register(_configs.Actions.NS_NOTIFY_OPEN_REQ,
self._on_ctl_notify_open_req)
self._control_dispatcher.register(_configs.Actions.NS_NOTIFY_OPEN_REP,
lambda msg: None)
self._control_dispatcher.register(_configs.Actions.NS_NOTIFY_CLOSE_REQ,
self._on_ctl_notify_close_req)
self._control_dispatcher.register(_configs.Actions.NS_NOTIFY_CLOSE_REP,
lambda msg: None)
# run threads
self._all_threads.append(
threading.Thread(target=self._main, name='ctl-main'))
self._all_threads.append(
threading.Thread(target=self._main_heartbeat,
name='ctl-main-ns-heartbeat'))
for i in self._all_threads:
i.start()
def finalize(self):
self._stop_event.set()
for i in self._all_threads:
i.join()
for sock in self._all_socks:
utils.graceful_close(sock)
def _main(self):
wait = 0
while True:
if self._stop_event.wait(wait / 1000):
break
nr_done = 0
socks = dict(self._control_poller.poll(0))
nr_done += self._main_do_control_recv(socks)
nr_done += self._main_do_control_send()
socks = dict(self._data_poller.poll(0))
nr_done += self._main_do_data_recv(socks)
nr_done += self._main_do_data_send()
if nr_done > 0:
wait = wait / 2 if wait > 1 else 0
else:
wait = wait + 1 if wait < 50 else 50
def _main_heartbeat(self):
while True:
self._control_mqueue.put(
ControlMessage(self._ns_socket,
None, {
'action': _configs.Actions.NS_HEARTBEAT_REQ,
'uid': self._uid
},
countdown=0))
if self._stop_event.wait(_configs.NS_HEARTBEAT_INTERVAL):
break
def _main_do_control_recv(self, socks):
nr_done = 0
# ns
if self._ns_socket in socks:
for msg in utils.iter_recv(utils.req_recv_json, self._ns_socket):
self._control_dispatcher.dispatch(msg['action'], msg)
nr_done += 1
# router
if self._control_router in socks:
for identifier, msg in utils.iter_recv(utils.router_recv_json,
self._control_router):
self._control_dispatcher.dispatch(msg['action'], identifier,
msg)
nr_done += 1
for info, csock in self._controller_peers.values():
if csock in socks:
for msg in utils.iter_recv(utils.req_recv_json, csock):
self._control_dispatcher.dispatch(msg['action'], msg)
nr_done += 1
return nr_done
def _main_do_control_send(self):
nr_scheduled = self._control_mqueue.qsize()
nr_done = 0
for i in range(nr_scheduled):
job = self._control_mqueue.get()
if job.identifier is not None:
rc = utils.router_send_json(job.sock,
job.identifier,
job.payload,
flag=zmq.NOBLOCK)
else:
rc = utils.req_send_json(job.sock,
job.payload,
flag=zmq.NOBLOCK)
if not rc:
if job.countdown > 0:
self._control_mqueue.put(
ControlMessage(job[0], job[1], job[2],
job.countdown - 1))
else:
nr_done += 1
return nr_done
def _main_do_data_recv(self, in_socks):
nr_done = 0
for name in self._input_from:
cache = self._input_cache.pop(name, None)
if cache is None:
peer = self._input_from[name]
if peer.dsock in in_socks:
msg = utils.pull_pyobj(peer.dsock)
cache = (msg['name'], msg['from_identifier'],
msg.get('to_identifier', None), msg['data'])
if cache is None:
continue
nr_done_this = 0
if cache[2] is None: # is broadcast
for p in self._imanager.filter_notfull(cache[0]):
p.raw_queue.put_nowait(
BroadcastMessage(cache[1], cache[-1]))
nr_done_this += 1
else:
for p in self._imanager.filter_notfull(cache[0]):
if p.identifier == cache[2]:
p.raw_queue.put_nowait(
UnicastMessage(cache[1], cache[2], cache[-1]))
nr_done_this += 1
if nr_done_this > 0:
nr_done += nr_done_this
else:
self._input_cache[name] = cache
return nr_done
def _main_do_data_send(self):
nr_done = 0
for name in self._omanager.keys():
cache = self._output_cache.get(name, None)
if cache is None:
pipes = self._omanager.filter_notempty(name)
if len(pipes) != 0:
pipe = random.choice(pipes)
cache = pipe.raw_queue.get_nowait()
self._output_cache[name] = cache
if cache is None:
continue
nr_done_this = 0
if isinstance(cache, BroadcastMessage):
for peer in self._output_to_pipe[name].values():
nr_done_this += utils.push_pyobj(peer.dsock, {
'uid': self._uid,
'name': name,
'from_identifier': cache.from_identifier,
'data': cache.payload
},
flag=zmq.NOBLOCK)
elif isinstance(cache, UnicastMessage):
if (name, cache.to_identifier) in self._output_to_id:
peer = self._output_to_id[(name, cache.to_identifier)]
nr_done_this += utils.push_pyobj(peer.dsock, {
'uid': self._uid,
'name': name,
'from_identifier': cache.from_identifier,
'to_identifier': cache.to_identifier,
'data': cache.payload
},
flag=zmq.NOBLOCK)
else:
raise TypeError('Unknown message type: {}.'.format(
type(cache)))
if nr_done_this > 0:
self._output_cache[name] = None
nr_done += nr_done_this
return nr_done
# BEGIN:: Connection
def _initialize_recv_peers(self, results):
for peers in results.values():
for info in peers:
uid = info['uid']
if uid not in self._controller_peers:
self._controller_peers[uid] = ControllerPeer(info, None)
self._do_setup_ctl_peer(uid)
def _do_setup_ctl_peer(self, uid):
info, sock = self._controller_peers[uid]
if sock is not None:
return
sock = self.socket(zmq.REQ)
sock.connect('{}://{}:{}'.format(info['ctl_protocal'],
info['ctl_addr'], info['ctl_port']))
self._control_poller.register(sock, zmq.POLLIN)
self._control_mqueue.put(
ControlMessage(sock,
None, {
'action': _configs.Actions.CTL_CONNECT_REQ,
'uid': self._uid,
'inputs': self._imanager.pipe_info()
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
self._controller_peers[uid] = ControllerPeer(info, sock)
logger.info('Connecting to "{}".'.format(uid))
def _on_ctl_connect_req(self, identifier, msg):
uid, pipes = msg['uid'], msg['inputs']
flag = False
for name, _ in pipes:
if name in self._omanager:
flag = True
break
if flag:
response = {}
if uid in self._output_to:
port = self._output_to[uid].port
else:
sock = self.socket(zmq.PUSH)
port = sock.bind_to_random_port('{}://{}'.format(
_configs.CTL_DAT_PROTOCAL, _configs.CTL_DAT_HOST))
pipes_rec = {p[0] for p in pipes}
ids_rec = {tuple(p) for p in pipes}
peer = PipePeer(self._addr, port, sock, pipes_rec, ids_rec)
self._output_to[uid] = peer
for p in pipes_rec:
self._output_to_pipe[p][uid] = peer
for i in ids_rec:
self._output_to_id[i] = peer
logger.info('Connection opened for "{}": port={}.'.format(
uid, port))
if port > 0:
response = {
'dat_protocal': _configs.CTL_DAT_PROTOCAL,
'dat_addr': self._addr,
'dat_port': port
}
self._control_mqueue.put(
ControlMessage(self._control_router,
identifier, {
'action': _configs.Actions.CTL_CONNECT_REP,
'uid': self._uid,
'conn': response
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
def _on_ctl_connect_rep(self, msg):
uid, conn = msg['uid'], msg['conn']
if len(conn) and uid not in self._input_from:
sock = self.socket(zmq.PULL)
sock.connect('{}://{}:{}'.format(conn['dat_protocal'],
conn['dat_addr'],
conn['dat_port']))
self._data_poller.register(sock, zmq.POLLIN)
self._input_from[uid] = PipePeer(conn['dat_addr'],
conn['dat_port'], sock, None,
None)
logger.info(
'Connection established to "{}": remote_port={}.'.format(
uid, conn['dat_port']))
self._control_mqueue.put(
ControlMessage(self._controller_peers[uid].csock,
None, {
'action': _configs.Actions.CTL_CONNECTED_REQ,
'uid': self._uid
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
def _on_ctl_connected_req(self, identifier, msg):
self._control_mqueue.put(
ControlMessage(self._control_router,
identifier, {
'action': _configs.Actions.CTL_CONNECTED_REP,
'uid': self._uid
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
logger.info('Connection established for "{}".'.format(msg['uid']))
# END:: Connection
def _on_ctl_notify_open_req(self, identifier, msg):
uid = msg['uid']
if uid not in self._controller_peers:
self._controller_peers[uid] = ControllerPeer(msg['info'], None)
self._do_setup_ctl_peer(uid)
self._control_mqueue.put(
ControlMessage(self._control_router,
identifier, {
'action': _configs.Actions.NS_NOTIFY_OPEN_REP,
'uid': self._uid
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
logger.info('Found new controller: "{}".'.format(uid))
def _on_ctl_notify_close_req(self, identifier, msg):
uid = msg['uid']
if uid in self._controller_peers:
peer = self._controller_peers.pop(uid)
self._control_poller.unregister(peer.csock)
self.close_socket(peer.csock)
if uid in self._input_from:
peer = self._input_from.pop(uid)
self._data_poller.unregister(peer.dsock)
self.close_socket(peer.dsock)
if uid in self._input_cache:
del self._input_cache[uid]
if uid in self._output_to:
peer = self._output_to.pop(uid)
self.close_socket(peer.dsock)
for k in peer.pipes:
self._output_to_pipe[k].pop(uid)
for k in peer.ids:
self._output_to_id.pop(k)
if uid in self._output_cache:
del self._output_cache[uid]
self._control_mqueue.put(
ControlMessage(self._control_router,
identifier, {
'action': _configs.Actions.NS_NOTIFY_CLOSE_REP,
'uid': self._uid
},
countdown=_configs.CTL_CTL_SND_COUNTDOWN))
logger.info('Close timeout controller: "{}".'.format(uid))
class NameServer(object):
def __init__(self, host, port, protocal):
self.storage = NameServerControllerStorage()
self._addr = '{}://{}:{}'.format(protocal, host, port)
self._context_lock = threading.Lock()
self._context = zmq.Context()
self._router = self._context.socket(zmq.ROUTER)
self._poller = zmq.Poller()
self._dispatcher = CallbackRegistry()
self._req_socks = set()
self._all_threads = list()
self._control_send_queue = queue.Queue()
def mainloop(self):
self.initialize()
try:
self._all_threads.append(
threading.Thread(target=self.main, name='name-server-main'))
self._all_threads.append(
threading.Thread(target=self.main_cleanup,
name='name-server-cleanup'))
for i in self._all_threads:
i.start()
finally:
self.finalize()
def initialize(self):
self._router.bind(self._addr)
self._poller.register(self._router, zmq.POLLIN)
self._dispatcher.register(_configs.Actions.NS_REGISTER_CTL_REQ,
self._on_ns_register_controller_req)
self._dispatcher.register(_configs.Actions.NS_REGISTER_OUTPUTS_REQ,
self._on_ns_register_outputs_req)
self._dispatcher.register(_configs.Actions.NS_REGISTER_INPUTS_REQ,
self._on_ns_register_inputs_req)
self._dispatcher.register(_configs.Actions.NS_HEARTBEAT_REQ,
self._on_ns_heartbeat_req)
self._dispatcher.register(_configs.Actions.NS_NOTIFY_OPEN_REP,
lambda msg: None)
self._dispatcher.register(_configs.Actions.NS_NOTIFY_CLOSE_REP,
lambda msg: None)
def finalize(self):
for i in self._all_threads:
i.join()
for sock in self._req_socks:
utils.graceful_close(sock)
utils.graceful_close(self._router)
if not self._context.closed:
self._context.destroy(0)
def main_cleanup(self):
while True:
with self._context_lock:
now = time.time()
for k, v in list(self.storage.items()):
if (now - v['last_heartbeat']) > _configs.NS_CLEANUP_WAIT:
info, req_sock = self.storage.unregister(k)
self._poller.unregister(req_sock)
utils.graceful_close(req_sock)
self._req_socks.remove(req_sock)
# TODO:: use controller's heartbeat
all_peers_to_inform = set()
for i in info['inputs']:
all_peers_to_inform = all_peers_to_inform.union(
self.storage.get_outputs(i))
for i in info['outputs']:
all_peers_to_inform = all_peers_to_inform.union(
self.storage.get_inputs(i))
logger.debug('Inform died: {}.'.format(
str(all_peers_to_inform)))
for peer in all_peers_to_inform:
self._control_send_queue.put({
'sock':
self.storage.get_req_sock(peer),
'countdown':
_configs.CTL_CTL_SND_COUNTDOWN,
'payload': {
'action':
_configs.Actions.NS_NOTIFY_CLOSE_REQ,
'uid': k
},
})
logger.info(
'Unregister timeout controller {}.'.format(k))
time.sleep(_configs.NS_CLEANUP_WAIT)
def main(self):
while True:
with self._context_lock:
socks = dict(self._poller.poll(50))
self._main_do_send()
self._main_do_recv(socks)
def _main_do_send(self):
nr_send = self._control_send_queue.qsize()
for i in range(nr_send):
job = self._control_send_queue.get()
rc = utils.req_send_json(job['sock'],
job['payload'],
flag=zmq.NOBLOCK)
if not rc:
job['countdown'] -= 1
if job['countdown'] >= 0:
self._control_send_queue.put(job)
else:
logger.warning('Drop job: {}.'.format(str(job)))
def _main_do_recv(self, socks):
if self._router in socks and socks[self._router] == zmq.POLLIN:
for identifier, msg in utils.iter_recv(utils.router_recv_json,
self._router):
self._dispatcher.dispatch(msg['action'], identifier, msg)
for k in socks:
if k in self._req_socks and socks[k] == zmq.POLLIN:
for msg in utils.iter_recv(utils.req_recv_json, k):
self._dispatcher.dispatch(msg['action'], msg)
def _on_ns_register_controller_req(self, identifier, msg):
req_sock = self._context.socket(zmq.REQ)
req_sock.connect('{}://{}:{}'.format(msg['ctl_protocal'],
msg['ctl_addr'], msg['ctl_port']))
self.storage.register(msg, req_sock)
self._req_socks.add(req_sock)
self._poller.register(req_sock, zmq.POLLIN)
utils.router_send_json(
self._router, identifier,
{'action': _configs.Actions.NS_REGISTER_CTL_REP})
logger.info('Controller registered: {}.'.format(msg['uid']))
def _on_ns_register_outputs_req(self, identifier, msg):
self.storage.register_outputs(msg)
all_peers_to_inform = set()
for i in msg['outputs']:
all_peers_to_inform = all_peers_to_inform.union(
self.storage.get_inputs(i))
for peer in all_peers_to_inform:
self._control_send_queue.put({
'sock':
self.storage.get_req_sock(peer),
'countdown':
_configs.CTL_CTL_SND_COUNTDOWN,
'payload': {
'action': _configs.Actions.NS_NOTIFY_OPEN_REQ,
'uid': msg['uid'],
'info': self.storage.get(msg['uid'])
},
})
utils.router_send_json(
self._router, identifier,
{'action': _configs.Actions.NS_REGISTER_OUTPUTS_REP})
logger.info('Controller pipes registered: out={} (uid="{}").'.format(
msg['outputs'], msg['uid']))
def _on_ns_register_inputs_req(self, identifier, msg):
self.storage.register_inputs(msg)
res = {}
for name in msg['inputs']:
all_pipes = self.storage.get_outputs(name)
all_pipes = list(map(self.storage.get, all_pipes))
res[name] = all_pipes
utils.router_send_json(self._router, identifier, {
'action': _configs.Actions.NS_REGISTER_INPUTS_REP,
'results': res
})
def _on_ns_heartbeat_req(self, identifier, msg):
if self.storage.contains(msg['uid']):
self.storage.get(msg['uid'])['last_heartbeat'] = time.time()
logger.debug('Heartbeat {}: time={}.'.format(
msg['uid'], time.time()))
utils.router_send_json(
self._router, identifier,
{'action': _configs.Actions.NS_HEARTBEAT_REP})