trying to improve our gevent api to support chaussette's usage

Lawouach · Lawouach · commit eddfa5d3f00c · 2013-08-07T21:15:50.000+02:00
diff --git a/example/echo_chaussette_server.py b/example/echo_chaussette_server.py
@@ -0,0 +1,41 @@
+# -*- coding: utf-8 -*-
+import logging
+
+from gevent import monkey; monkey.patch_all()
+import gevent
+from ws4py.server.geventserver import WebSocketWSGIApplication, \
+     WebSocketWSGIHandler, GEventWebSocketPool
+from ws4py.websocket import EchoWebSocket
+
+from chaussette.backend._gevent import Server as GeventServer
+
+class ws4pyServer(GeventServer):
+    handler_class = WebSocketWSGIHandler
+
+    def __init__(self, *args, **kwargs):
+        GeventServer.__init__(self, *args, **kwargs)
+        self.pool = GEventWebSocketPool()
+
+    def stop(self, *args, **kwargs):
+        self.pool.clear()
+        self.pool = None
+        GeventServer.stop(self, *args, **kwargs)
+
+if __name__ == '__main__':
+    import os, socket, sys
+    from ws4py import configure_logger
+    logger = configure_logger()
+
+    from chaussette.backend import register
+    register('ws4py', ws4pyServer)
+
+    from chaussette.server import make_server
+    server = make_server(app=WebSocketWSGIApplication(handler_cls=EchoWebSocket),
+                         host='unix:///%s/ws.sock' % os.getcwd(),
+                         address_family=socket.AF_UNIX,
+                         backend='ws4py',
+                         logger=logger)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        sys.exit(0)
diff --git a/ws4py/__init__.py b/ws4py/__init__.py
@@ -57,7 +57,11 @@ def configure_logger(stdout=True, filepath=None, level=logging.INFO):
     return logger
 
 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)
-    
+    me = ws.local_address
+    peer = ws.peer_address
+    if isinstance(me, tuple) and isinstance(peer, tuple):
+        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)
+
+    return "[Bound to '%s']" % me
diff --git a/ws4py/server/geventserver.py b/ws4py/server/geventserver.py
@@ -27,7 +27,8 @@
 
 logger = logging.getLogger('ws4py')
 
-__all__ = ['WebSocketWSGIHandler', 'WSGIServer']
+__all__ = ['WebSocketWSGIHandler', 'WSGIServer',
+           'GEventWebSocketPool']
 
 class WebSocketWSGIHandler(WSGIHandler):
     """
@@ -56,10 +57,35 @@ def run_application(self):
 
                 ws = self.environ.pop('ws4py.websocket')
                 if ws:
-                    self.server.link_websocket_to_server(ws)
+                    self.server.pool.track(ws)
         else:
             gevent.pywsgi.WSGIHandler.run_application(self)
 
+class GEventWebSocketPool(Group):
+    """
+    Simple pool of bound websockets.
+    Internally it uses a gevent group to track
+    the websockets. The server should call the ``clear``
+    method to initiate the closing handshake when the
+    server is shutdown.
+    """
+
+    def track(self, websocket):
+        logger.info("Managing websocket %s" % format_addresses(websocket))
+        return self.spawn(websocket.run)
+
+    def clear(self):
+        logger.info("Terminating server and all connected websockets")
+        for greenlet in self:
+            try:
+                websocket = greenlet._run.im_self
+                if websocket:
+                    websocket.close(1001, 'Server is shutting down')
+            except:
+                pass
+            finally:
+                self.discard(greenlet)
+                
 class WSGIServer(_WSGIServer):
     handler_class = WebSocketWSGIHandler
 
@@ -74,24 +100,15 @@ def __init__(self, *args, **kwargs):
         base.
         """
         _WSGIServer.__init__(self, *args, **kwargs)
-        self._websockets = Group()
-
-    def link_websocket_to_server(self, websocket):
-        logger.info("Managing websocket %s" % format_addresses(websocket))
-        self._websockets.spawn(websocket.run)
+        self.pool = GEventWebSocketPool()
 
     def stop(self, *args, **kwargs):
-        logger.info("Terminating server and all connected websockets")
-        for greenlet in self._websockets:
-            try:
-                websocket = greenlet._run.im_self
-                if websocket:
-                    websocket.close(1001, 'Server is shutting down')
-            except:
-                pass
+        self.pool.clear()
         _WSGIServer.stop(self, *args, **kwargs)
 
 if __name__ == '__main__':
+    import os
+
     from ws4py import configure_logger
     configure_logger()
 
