Пример #1
0
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()
Пример #2
0
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()
Пример #3
0
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
Пример #4
0
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')
Пример #5
0
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())