Cloned DMOJ

event_socket_server/engines/__init__.py

@@ -0,0 +1,17 @@
+import select
+
+__author__ = 'Quantum'
+engines = {}
+
+from .select_server import SelectServer  # noqa: E402, import not at top for consistency
+engines['select'] = SelectServer
+
+if hasattr(select, 'poll'):
+    from .poll_server import PollServer
+    engines['poll'] = PollServer
+
+if hasattr(select, 'epoll'):
+    from .epoll_server import EpollServer
+    engines['epoll'] = EpollServer
+
+del select

event_socket_server/engines/epoll_server.py

@@ -0,0 +1,17 @@
+import select
+__author__ = 'Quantum'
+
+if not hasattr(select, 'epoll'):
+    raise ImportError('System does not support epoll')
+
+from .poll_server import PollServer  # noqa: E402, must be imported here
+
+
+class EpollServer(PollServer):
+    poll = select.epoll
+    WRITE = select.EPOLLIN | select.EPOLLOUT | select.EPOLLERR | select.EPOLLHUP
+    READ = select.EPOLLIN | select.EPOLLERR | select.EPOLLHUP
+    POLLIN = select.EPOLLIN
+    POLLOUT = select.EPOLLOUT
+    POLL_CLOSE = select.EPOLLHUP | select.EPOLLERR
+    NEED_CLOSE = True

event_socket_server/engines/poll_server.py

@@ -0,0 +1,97 @@
+import errno
+import logging
+import select
+import threading
+
+from ..base_server import BaseServer
+
+logger = logging.getLogger('event_socket_server')
+
+if not hasattr(select, 'poll'):
+    raise ImportError('System does not support poll')
+
+
+class PollServer(BaseServer):
+    poll = select.poll
+    WRITE = select.POLLIN | select.POLLOUT | select.POLLERR | select.POLLHUP
+    READ = select.POLLIN | select.POLLERR | select.POLLHUP
+    POLLIN = select.POLLIN
+    POLLOUT = select.POLLOUT
+    POLL_CLOSE = select.POLLERR | select.POLLHUP
+    NEED_CLOSE = False
+
+    def __init__(self, *args, **kwargs):
+        super(PollServer, self).__init__(*args, **kwargs)
+        self._poll = self.poll()
+        self._fdmap = {}
+        self._server_fds = {sock.fileno(): sock for sock in self._servers}
+        self._close_lock = threading.RLock()
+
+    def _register_write(self, client):
+        logger.debug('On write mode: %s', client.client_address)
+        self._poll.modify(client.fileno(), self.WRITE)
+
+    def _register_read(self, client):
+        logger.debug('On read mode: %s', client.client_address)
+        self._poll.modify(client.fileno(), self.READ)
+
+    def _clean_up_client(self, client, finalize=False):
+        logger.debug('Taking close lock: cleanup')
+        with self._close_lock:
+            logger.debug('Cleaning up client: %s, finalize: %d', client.client_address, finalize)
+            fd = client.fileno()
+            try:
+                self._poll.unregister(fd)
+            except IOError as e:
+                if e.errno != errno.ENOENT:
+                    raise
+            except KeyError:
+                pass
+            del self._fdmap[fd]
+            super(PollServer, self)._clean_up_client(client, finalize)
+
+    def _serve(self):
+        for fd, sock in self._server_fds.items():
+            self._poll.register(fd, self.POLLIN)
+            sock.listen(16)
+        try:
+            while not self._stop.is_set():
+                for fd, event in self._poll.poll(self._dispatch_event()):
+                    if fd in self._server_fds:
+                        client = self._accept(self._server_fds[fd])
+                        logger.debug('Accepting: %s', client.client_address)
+                        fd = client.fileno()
+                        self._poll.register(fd, self.READ)
+                        self._fdmap[fd] = client
+                    elif event & self.POLL_CLOSE:
+                        logger.debug('Client closed: %s', self._fdmap[fd].client_address)
+                        self._clean_up_client(self._fdmap[fd])
+                    else:
+                        logger.debug('Taking close lock: event loop')
+                        with self._close_lock:
+                            try:
+                                client = self._fdmap[fd]
+                            except KeyError:
+                                pass
+                            else:
+                                logger.debug('Client active: %s, read: %d, write: %d',
+                                             client.client_address,
+                                             event & self.POLLIN,
+                                             event & self.POLLOUT)
+                                if event & self.POLLIN:
+                                    logger.debug('Non-blocking read on client: %s', client.client_address)
+                                    self._nonblock_read(client)
+                                # Might be closed in the read handler.
+                                if event & self.POLLOUT and fd in self._fdmap:
+                                    logger.debug('Non-blocking write on client: %s', client.client_address)
+                                    self._nonblock_write(client)
+        finally:
+            logger.info('Shutting down server')
+            self.on_shutdown()
+            for client in self._clients:
+                self._clean_up_client(client, True)
+            for fd, sock in self._server_fds.items():
+                self._poll.unregister(fd)
+                sock.close()
+            if self.NEED_CLOSE:
+                self._poll.close()

event_socket_server/engines/select_server.py

@@ -0,0 +1,49 @@
+import select
+
+from ..base_server import BaseServer
+
+
+class SelectServer(BaseServer):
+    def __init__(self, *args, **kwargs):
+        super(SelectServer, self).__init__(*args, **kwargs)
+        self._reads = set(self._servers)
+        self._writes = set()
+
+    def _register_write(self, client):
+        self._writes.add(client)
+
+    def _register_read(self, client):
+        self._writes.remove(client)
+
+    def _clean_up_client(self, client, finalize=False):
+        self._writes.discard(client)
+        self._reads.remove(client)
+        super(SelectServer, self)._clean_up_client(client, finalize)
+
+    def _serve(self, select=select.select):
+        for server in self._servers:
+            server.listen(16)
+        try:
+            while not self._stop.is_set():
+                r, w, x = select(self._reads, self._writes, self._reads, self._dispatch_event())
+                for s in r:
+                    if s in self._servers:
+                        self._reads.add(self._accept(s))
+                    else:
+                        self._nonblock_read(s)
+
+                for client in w:
+                    self._nonblock_write(client)
+
+                for s in x:
+                    s.close()
+                    if s in self._servers:
+                        raise RuntimeError('Server is in exceptional condition')
+                    else:
+                        self._clean_up_client(s)
+        finally:
+            self.on_shutdown()
+            for client in self._clients:
+                self._clean_up_client(client, True)
+            for server in self._servers:
+                server.close()