trying out an event based solution to handle connections. This ought to allow for a lot more websocket connections at once

Lawouach · Lawouach · commit c3aba8b51fb2 · 2013-02-02T12:21:17.000+01:00
diff --git a/ws4py/manager.py b/ws4py/manager.py
@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+import logging
+import select
+import threading
+import time
+
+logger = logging.getLogger('ws4py')
+
+def format_addresses(ws):
+    me_ip, me_port = ws.local_address
+    peer_ip, peer_port = ws.peer_address
+    return "[Local => %s:%d | Remote => %s:%d]" % (me_ip, me_port, peer_ip, peer_port)
+    
+class SelectPoller(object):
+    def __init__(self, timeout=0.1):
+        self._fds = []
+        self.timeout = timeout
+
+    def release(self):
+        self._fds = []
+
+    def register(self, fd):
+        if fd not in self._fds:
+            self._fds.append(fd)
+
+    def unregister(self, fd):
+        if fd in self._fds:
+            self._fds.remove(fd)
+
+    def poll(self):
+        r, w, x = select.select(self._fds, [], [], self.timeout)
+        return r
+
+class EPollPoller(object):
+    def __init__(self, timeout=0.1):
+        self.poller = select.epoll()
+        self.timeout = timeout
+
+    def release(self):
+        self.poller.close()
+
+    def release(self):
+        self.poller.close()
+
+    def register(self, fd):
+        try:
+            self.poller.register(fd, select.EPOLLIN | select.EPOLLPRI)
+        except IOError:
+            pass
+
+    def unregister(self, fd):
+        self.poller.unregister(fd)
+
+    def poll(self):
+        events = self.poller.poll(timeout=self.timeout)
+        for fd, event in events:
+            if event | select.EPOLLIN | select.EPOLLPRI:
+                yield fd
+
+class WebSocketManager(threading.Thread):
+    def __init__(self, poller=None):
+        threading.Thread.__init__(self)
+        self.lock = threading.Lock()
+        self.websockets = {}
+
+        if hasattr(select, "epoll"):
+            self.poller = EPollPoller()
+            logger.info("Using epoll")
+        else:
+            self.poller = SelectPoller()
+            logger.info("Using select as epoll is not available")
+
+    def __len__(self):
+        return len(self.websockets)
+
+    def add(self, websocket):
+        logger.info("Managing websocket %s" % format_addresses(websocket))
+        websocket.opened()
+        with self.lock:
+            fd = websocket.sock.fileno()
+            self.websockets[fd] = websocket
+            self.poller.register(fd)
+
+    def remove(self, websocket):
+        logger.info("Removing websocket %s" % format_addresses(websocket))
+        with self.lock:
+            fd = websocket.sock.fileno()
+            self.websockets.pop(fd)
+            self.poller.unregister(fd)
+
+    def stop(self):
+        self.running = False
+        with self.lock:
+            self.websockets.clear()
+            self.poller.release()
+
+    def run(self):
+        self.running = True
+        while self.running:
+            with self.lock:
+                polled = self.poller.poll()
+
+            if not self.running:
+                break
+
+            for fd in polled:
+                if not self.running:
+                    break
+
+                ws = self.websockets.get(fd)
+
+                if ws:
+                    if not ws.terminated and not ws.once():
+                        with self.lock:
+                            fd = ws.sock.fileno()
+                            self.websockets.pop(fd)
+                            self.poller.unregister(fd)
+
+                        if not ws.terminated:
+                            logger.info("Terminating websocket %s" % format_addresses(ws))
+                            ws.terminate()
+
+    def close_all(self, code=1001, message='Server is shutting down'):
+        with self.lock:
+            logger.info("Closing all websockets with [%d] '%s'" % (code, message))
+            for ws in self.websockets.itervalues():
+                ws.close(code=1001, reason=message)
+
+    def broadcast(self, message, binary=False):
+        with self.lock:
+            websockets = self.websockets.copy()
+
+        for ws in websockets.itervalues():
+            if not ws.terminated:
+                try:
+                    ws.send(message, binary)
+                except Exception, ex:
+                    pass
diff --git a/ws4py/server/cherrypyserver.py b/ws4py/server/cherrypyserver.py
@@ -77,6 +77,7 @@ def ws(self):
 from ws4py.exc import HandshakeError
 from ws4py.websocket import WebSocket
 from ws4py.compat import py3k, enc, dec, get_connection, detach_connection
+from ws4py.manager import WebSocketManager
 
 __all__ = ['WebSocketTool', 'WebSocketPlugin']
 
@@ -280,18 +281,17 @@ def _set_internal_flags(self):
 class WebSocketPlugin(plugins.SimplePlugin):
     def __init__(self, bus):
         plugins.SimplePlugin.__init__(self, bus)
-        self.pool = {}
+        self.manager = WebSocketManager()
 
     def start(self):
         cherrypy.log("Starting WebSocket processing")
-        self.bus.subscribe('main', self.monitor)
         self.bus.subscribe('stop', self.cleanup)
         self.bus.subscribe('handle-websocket', self.handle)
         self.bus.subscribe('websocket-broadcast', self.broadcast)
+        self.manager.start()
 
     def stop(self):
         cherrypy.log("Terminating WebSocket processing")
-        self.bus.unsubscribe('main', self.monitor)
         self.bus.unsubscribe('stop', self.cleanup)
         self.bus.unsubscribe('handle-websocket', self.handle)
         self.bus.unsubscribe('websocket-broadcast', self.broadcast)
@@ -303,39 +303,15 @@ def handle(self, ws_handler, peer_addr):
         :param ws_handler: websocket handler instance
         :param peer_addr: remote peer address for tracing purpose
         """
-        cherrypy.log("Managing WebSocket connection from %s:%d" % (peer_addr[0], peer_addr[1]))
-        th = threading.Thread(target=ws_handler.run, name="WebSocket client at %s:%d" % (peer_addr[0], peer_addr[1]))
-        th.daemon = True
-        self.pool[ws_handler] = (th, peer_addr)
-        th.start()
-
-    def monitor(self):
-        """
-        Called within the engine's mainloop to drop connections
-        that have terminated since last iteration.
-        """
-        if py3k:
-            handlers = list(self.pool.keys())
-        else:
-            handlers = self.pool.keys()[:]
-
-        for handler in handlers:
-            if handler.terminated:
-                th, addr = self.pool[handler]
-                cherrypy.log("Removing WebSocket connection %s:%d" % (addr[0], addr[1]))
-                th.join()
-                del self.pool[handler]
+        self.manager.add(ws_handler)
 
     def cleanup(self):
         """
         Terminate all connections and clear the pool. Executed when the engine stops.
         """
-        cherrypy.log("Closing %d WebSocket connections" % len(self.pool))
-        for handler in self.pool:
-            handler.close(code=1001, reason='Server is shutting down')
-            th, addr = self.pool[handler]
-            th.join()
-        self.pool.clear()
+        self.manager.close_all()
+        self.manager.stop()
+        self.manager.join()
 
     def broadcast(self, message, binary=False):
         """
@@ -346,12 +322,7 @@ def broadcast(self, message, binary=False):
           of the connected handler.
         :param binary: whether or not the message is a binary one
         """
-        for ws_handler in self.pool:
-            try:
-                if not ws_handler.terminated:
-                    ws_handler.send(message, binary)
-            except:
-                cherrypy.log(traceback=True)
+        self.manager.broadcast(message, binary)
 
 if __name__ == '__main__':
     import random
diff --git a/ws4py/websocket.py b/ws4py/websocket.py
