class EventLoop:
def __init__(self, server_id, handler: Handler):
self.select = DefaultSelector()
self.server_id = server_id
self.handler = handler
def register(self, fd: Socket, events: SelectorKey, data=None):
fd.setblocking(False)
self.select.register(fd, events, data)
def modify(self, fd: Socket, events: SelectorKey, data=None):
self.select.modify(fd, events, data)
def unregister(self, fd):
self.select.unregister(fd)
def run(self):
self.register(self.server_id, selectors.EVENT_READ)
while True:
events = self.select.select()
for key, mask in events:
if mask == EVENT_READ:
if key.fileobj == self.server_id:
self.accept(key.fileobj)
else:
self.read(key.fileobj)
elif mask == EVENT_WRITE:
print(f"EventLoop, key={key}, data={key.data}")
self.write(key.fileobj, key.data)
def accept(self, sock):
"""sock 要处理的 fd, mask 为事件类型 READ or WRITE"""
conn = self.handler.accept(sock) # Should be ready
self.register(conn, EVENT_READ, self.read)
def read(self, conn):
for data in self.handler.read(conn):
if data and data not in ('q', 'Q'):
resp = self.handler.handle(conn, data)
self.modify(conn, EVENT_WRITE, resp)
else:
self.unregister(conn)
self.handler.close(conn)
def write(self, conn, data):
self.handler.send(conn, data)
self.modify(conn, EVENT_READ)
def close(self):
self.select.close()
class Selector:
def __init__(self):
self._sel = DefaultSelector()
self._fds = {}
def register_read(self, task, fileno):
return self._register(task, fileno, EVENT_READ)
def register_write(self, task, fileno):
return self._register(task, fileno, EVENT_WRITE)
def _unregister(self, fd):
fileno = fd.fileno()
if fileno in self._fds:
LOG.debug('unregister fd %r', fd)
del self._fds[fileno]
self._sel.unregister(fileno)
def _register(self, task, fileno, events):
fd = self._fds.get(fileno, None)
if fd is None:
fd = KernelFd(self, fileno, task)
self._fds[fileno] = fd
if fd.task is not task:
raise RuntimeError(
'file descriptor already registered by other task')
if fd.selector_key is None:
LOG.debug('register fd %r, events=%r', fd, events)
fd.selector_key = self._sel.register(fd, events)
elif fd.events != events:
LOG.debug('modify fd %r, events=%r', fd, events)
fd.selector_key = self._sel.modify(fd, events)
return fd
def poll(self, timeout):
io_events = self._sel.select(timeout=timeout)
for key, mask in io_events:
fd = key.fileobj
if fd.task.is_alive and fd.events & mask:
LOG.debug('task %r wakeup by fd %r', fd.task, fd)
yield fd.task
def close(self):
self._sel.close()
class SimpleServer(object):
""" This is a simple server receiving and sending utf-8 plain text lines
over TCP/IP. It should use epoll or other efficient mechanisms if
available to support a large number of clients without cpu hogging.
All data sent to connections will be utf-8 encoded. All received lines
from clients are expected to be utf-8 encoded.
Lines longer than 10KB will be quietly dropped.
"""
def __init__(self, port, callback):
logging.basicConfig(level=logging.INFO)
self.clients_kept_for_sending = []
self.selector = DefaultSelector()
self.callback = callback
self.server_socket = socket.socket(family=socket.AF_INET6,
type=socket.SOCK_STREAM)
self.server_socket.bind(("::0", port))
self.server_socket.listen(5)
self.selector.register(self.server_socket, EVENT_READ, data=None)
self.server_socket.setblocking(False)
def close(self):
self.server_socket.close()
def _remove_socket(self, socket):
try:
self.selector.unregister(socket)
except KeyError:
pass
def _handle_new_client(self, socket, addr):
client = SimpleClient(socket, self, addr)
self.selector.register(socket, EVENT_READ, data=client)
self.callback(client)
def _mark_client_for_sending(self, client, marked):
if marked:
self.selector.modify(client.socket, EVENT_READ | EVENT_WRITE,
data=client)
else:
self.selector.modify(client.socket, EVENT_READ, data=client)
def tick(self, timeout=0):
""" This will check for any server events and process them.
If timeout is greater than zero, it will wait at most that
amount of time in seconds until something happens.
(This is useful if you want to save some cpu time and you have
nothing else to do)
For debugging purposes, it returns the amount of events processed.
"""
# get rid of old connections we kept around for sending:
entry_removed = True
while entry_removed:
entry_removed = False
for entry in self.clients_kept_for_sending:
if entry[1] <= time.time() or len(entry[0].send_buffer) == 0:
try:
entry[0].close()
except:
pass
self.clients_kept_for_sending.remove(entry)
entry_removed = True
# wait for socket events:
events = self.selector.select(timeout=timeout)
events_processed = 0
for key, mask in events:
if mask & EVENT_READ != 0:
events_processed = events_processed + 1
# handle accept for server socket:
if key.data == None:
conn, addr = self.server_socket.accept()
self._handle_new_client(conn, addr)
continue
# handle client socket receive:
key.data._do_recv()
continue
if mask & EVENT_WRITE != 0:
events_processed = events_processed + 1
# handle client socket send:
key.data._do_outstanding_send()
continue
logging.debugging("unknown socket event happening, bug?")
return events_processed
class EventLoop:
"""event loop, handles IO"""
def __init__(self):
self.co_result = {}
self.selector = DefaultSelector()
def run_until_complete(self, coroutine):
self._continue(coroutine, None)
while coroutine not in self.co_result:
self._tick()
result = self.co_result[coroutine]
del self.co_result[coroutine]
return result
def run_forever(self):
while True:
self._tick()
def start_coroutine(self, co):
self._continue(co, None)
def _continue(self, co, awaited):
try:
cmd = co.send(awaited)
except StopIteration as e:
self.co_result[co] = e.value
else:
if isinstance(cmd, EventLoopCommand.StreamAccept):
self._register(EVENT_READ, cmd.sock, (cmd, co))
elif isinstance(cmd, EventLoopCommand.StreamConnect):
try:
cmd.sock.connect(cmd.addr)
except BlockingIOError:
self._register(EVENT_WRITE, cmd.sock, (cmd, co))
except:
raise
else:
# connect finished immediately
self._continue(co, None)
elif isinstance(cmd, EventLoopCommand.StreamRecv):
self._register(EVENT_READ, cmd.sock, (cmd, co))
elif isinstance(cmd, EventLoopCommand.StreamSend):
self._register(EVENT_WRITE, cmd.sock, (cmd, co))
else:
raise RuntimeError('unknown cmd')
def _register(self, event, sock, data):
try:
key = self.selector.get_key(sock)
except KeyError:
self.selector.register(sock, event, {event: data})
else:
assert key.events & event == 0
key.data[event] = data
self.selector.modify(sock, key.events | event, key.data)
def _unregister(self, event, sock):
key = self.selector.get_key(sock)
assert key.events & event != 0
new_mask = key.events & ~event
del key.data[event]
if new_mask == 0:
self.selector.unregister(sock)
else:
self.selector.modify(sock, new_mask, key.data)
def _tick(self):
for key, mask in self.selector.select(1.0):
if mask & EVENT_WRITE != 0:
cmd, co = key.data[EVENT_WRITE]
self._unregister(EVENT_WRITE, cmd.sock)
if isinstance(cmd, EventLoopCommand.StreamConnect):
self._continue(co, cmd.sock.connect(cmd.addr))
elif isinstance(cmd, EventLoopCommand.StreamSend):
self._continue(co, cmd.sock.send(cmd.buf))
else:
raise RuntimeError('unknown cmd')
elif mask & EVENT_READ != 0:
cmd, co = key.data[EVENT_READ]
self._unregister(EVENT_READ, cmd.sock)
if isinstance(cmd, EventLoopCommand.StreamAccept):
self._continue(co, cmd.sock.accept())
elif isinstance(cmd, EventLoopCommand.StreamRecv):
self._continue(co, cmd.sock.recv(cmd.buf_len))
else:
raise RuntimeError('unknown cmd')
class MQueue:
def __init__(self, id, host, port):
self.local_node = Node(self, host, port, id=id)
self._sock = ObjectSocket()
self._lock = RLock()
self._msg_cond = Condition(self._lock)
self._nodes_cond = Condition(self._lock)
self._event_thread = Thread(target=self.event_loop, daemon=True)
self._sel = Selector()
self._msgbuf = []
self._message_handlers = {}
self.nodes = {id: self.local_node}
self.size = None
self.timestamp = 0
self.running = True
self.controller = None
def log(self, msg, *args):
# with open(os.path.join(os.path.dirname(sys.argv[0]), 'output-{}.log'.format(self.local_node.id)), 'a') as f:
with open(os.path.join(os.path.dirname(sys.argv[0]), 'debug.log'),
'a') as f:
t = time.time()
date = time.strftime(
'%Y-%m-%d %H:%M:%S',
time.localtime(t)) + '{:.3f}'.format(t - int(t))[1:]
if isinstance(msg, str):
msg = msg.format(*args)
else:
msg = str(msg)
print('{} [{}]: {}'.format(date, self.local_node.id, msg),
file=f,
flush=True)
def event_loop(self):
while True:
# self.log('Waiting for event...')
for key, events in self._sel.select():
try:
# self.log('Event in event_log: {}', key)
if 'handler' in key.data:
key.data['handler'](key.fileobj, key.data)
except EOFError:
self._sel.unregister(key.fileobj)
except SystemExit:
self.log('EXITING')
with self._lock:
self.running = False
self._msg_cond.notify_all()
except:
self.log('Exception in event_loop (key: {}): {}', key,
format_exc())
self._sel.unregister(key.fileobj)
def on_connection(self, sock, data):
c, a = self._sock.accept()
self.log('New connection from {}', a)
self._sel.register(c, EVENT_READ, {'handler': self.on_raw_data})
def _extract_msg(self, sock, data):
if not 'msg' in data:
msg = Message()
data['msg'] = msg
else:
msg = data['msg']
if not msg.iscomplete():
msg.feed(sock)
return msg
def on_raw_data(self, sock, data):
msg = self._extract_msg(sock, data)
if msg.iscomplete():
data = msg.getdata()
self.log('Message received: {} from {}', data, sock.getpeername())
if 'nodes' in data and 'total-node-count' in data:
nodes = data['nodes']
count = data['total-node-count']
self.log('Controller: Nodes: {}, total count: {}', nodes,
count)
with self._lock:
self.size = count
self.log('Setting self.size = {}', self.size)
self._nodes_cond.notify_all()
for id, host, port in nodes:
node = Node(self, host, port, id=id)
node.connect()
self._sel.modify(sock, EVENT_READ,
{'handler': self.on_control})
self.controller = sock
elif 'node-info' in data:
id, host, port = map(lambda k: data['node-info'][k],
['id', 'host', 'port'])
with self._lock:
self.nodes[id] = Node(self, host, port, sock=sock, id=id)
self.log(
'New node (connection from {}): {} (total nodes: {} {})',
sock.getpeername(), id, len(self.nodes), self.nodes)
self._sel.modify(sock, EVENT_READ, {
'handler': self.on_message,
'node': self.nodes[id]
})
self._nodes_cond.notify_all()
def on_message(self, sock, data):
try:
msg = self._extract_msg(sock, data)
except EOFError:
self.log('Node {} disconnected', data['node'])
with self._lock:
del self.nodes[data['node'].id]
raise
except OSError as ose:
self.log('Node {} disconnected due to an error', data['node'])
raise
if msg.iscomplete():
msg = msg.getdata()
del data['msg']
if isinstance(msg, dict) and 'timestamp' in msg and 'data' in msg:
with self._lock:
self.timestamp = max(self.timestamp, msg['timestamp']) + 1
self.log('Message from {}: {} (ts {}). New timestamp: {}',
data['node'].id, msg, msg['timestamp'],
self.timestamp)
msg['node'] = data['node']
if self.controller:
try:
self.controller.sendobj({
'timestamp':
msg['timestamp'],
'from':
msg['node'].id,
'to':
self.local_node.id,
'data':
msg['data'],
'in-channel':
True
})
except:
self.log(
'Exception while sending message {} to controller: {}',
m, format_exc())
for filter, handler in self._message_handlers.items():
if filter(msg):
handler(msg)
return
self._msgbuf.append(msg)
self._msg_cond.notify_all()
def on_control(self, sock, data):
msg = self._extract_msg(sock, data)
if msg.iscomplete():
del data['msg']
msg = msg.getdata()
self.log('Control message: {}', msg)
try:
exec(msg)
except SystemExit:
raise
except:
pass
def register(self, node):
with self._lock:
self.nodes[node.id] = node
self._sel.register(node.sock, EVENT_READ, {
'handler': self.on_message,
'node': node
})
self.log(
'New node (connection to {} using {}): {} (total nodes: {} {})',
node.sock.getpeername(), node.sock.getsockname(), node.id,
len(self.nodes), self.nodes)
self._nodes_cond.notify_all()
def send(self, id, msg):
assert isinstance(msg, dict), 'Message should be dict'
if id != self.local_node.id:
with self._lock:
self.timestamp += 1
self.nodes[id].sock.sendobj({
'timestamp': self.timestamp,
'data': msg
})
def broadcast(self, msg):
with self._lock:
self.timestamp += 1
for n in self.nodes.values():
if n.id != self.local_node.id:
n.sock.sendobj({'timestamp': self.timestamp, 'data': msg})
def recv(self, filter, timeout=None):
with self._lock:
while True:
if not self.running:
self.log('MAIN THREAD TERMINATING')
exit()
for m in self._msgbuf:
if filter(m):
self._msgbuf.remove(m)
if self.controller:
try:
self.controller.sendobj({
'timestamp':
m['timestamp'],
'from':
m['node'].id,
'to':
self.local_node.id,
'data':
m['data'],
'in-channel':
False
})
except:
self.log(
'Exception while sending message {} to controller: {}',
m, format_exc())
return m
if timeout is not None:
t = _time()
if not self._msg_cond.wait(timeout):
raise TimeoutError
if timeout is not None:
timeout -= _time() - t
if timeout <= 0:
raise TimeoutError
def become_coord(self):
with self._lock:
self.timestamp += 1
self._msgbuf.insert(
0, {
'node': self.local_node,
'data': {
'type': MSG_REQUEST,
'change-state': 'coord'
}
})
self._msg_cond.notify_all()
def set_message_handler(self, filter, handler):
with self._lock:
self._message_handlers[filter] = handler
def clear_message_handler(self, filter):
with self._lock:
if filter in self._message_handlers:
del self._message_handlers[filter]
def wait_for_init(self):
with self._lock:
while self.size is None or len(self.nodes) < self.size:
self.log('In wait: self.size: {}, len(self.nodes): {}',
self.size, len(self.nodes))
self._nodes_cond.wait()
def start(self):
# self.log('Starting MQueue...')
self._sock.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
self._sock.bind((self.local_node.host, self.local_node.port))
self._sock.listen(5)
self._sel.register(self._sock, EVENT_READ,
{'handler': self.on_connection})
self._event_thread.start()
def loop(self):
try:
self.local_node.loop()
except:
self.mqueue.log('Exception in the main loop: {}', format_exc())