1# -*- coding: utf-8 -*- 2__doc__ = """ 3WSGI entities to support WebSocket from within gevent. 4 5Its usage is rather simple: 6 7.. code-block: python 8 9 from gevent import monkey; monkey.patch_all() 10 from ws4py.websocket import EchoWebSocket 11 from ws4py.server.geventserver import WSGIServer 12 from ws4py.server.wsgiutils import WebSocketWSGIApplication 13 14 server = WSGIServer(('localhost', 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket)) 15 server.serve_forever() 16 17""" 18import logging 19import sys 20 21import gevent 22from gevent.pywsgi import WSGIHandler, WSGIServer as _WSGIServer 23from gevent.pool import Pool 24 25from ws4py import format_addresses 26from ws4py.server.wsgiutils import WebSocketWSGIApplication 27 28logger = logging.getLogger('ws4py') 29 30__all__ = ['WebSocketWSGIHandler', 'WSGIServer', 31 'GEventWebSocketPool'] 32 33class WebSocketWSGIHandler(WSGIHandler): 34 """ 35 A WSGI handler that will perform the :rfc:`6455` 36 upgrade and handshake before calling the WSGI application. 37 38 If the incoming request doesn't have a `'Upgrade'` header, 39 the handler will simply fallback to the gevent builtin's handler 40 and process it as per usual. 41 """ 42 43 def run_application(self): 44 upgrade_header = self.environ.get('HTTP_UPGRADE', '').lower() 45 if upgrade_header: 46 try: 47 # Build and start the HTTP response 48 self.environ['ws4py.socket'] = self.socket or self.environ['wsgi.input'].rfile._sock 49 self.result = self.application(self.environ, self.start_response) or [] 50 self.process_result() 51 except: 52 raise 53 else: 54 del self.environ['ws4py.socket'] 55 self.socket = None 56 self.rfile.close() 57 58 ws = self.environ.pop('ws4py.websocket') 59 if ws: 60 self.server.pool.track(ws) 61 else: 62 gevent.pywsgi.WSGIHandler.run_application(self) 63 64class GEventWebSocketPool(Pool): 65 """ 66 Simple pool of bound websockets. 67 Internally it uses a gevent group to track 68 the websockets. The server should call the ``clear`` 69 method to initiate the closing handshake when the 70 server is shutdown. 71 """ 72 73 def track(self, websocket): 74 logger.info("Managing websocket %s" % format_addresses(websocket)) 75 return self.spawn(websocket.run) 76 77 def clear(self): 78 logger.info("Terminating server and all connected websockets") 79 for greenlet in self: 80 try: 81 websocket = greenlet._run.im_self 82 if websocket: 83 websocket.close(1001, 'Server is shutting down') 84 except: 85 pass 86 finally: 87 self.discard(greenlet) 88 89class WSGIServer(_WSGIServer): 90 handler_class = WebSocketWSGIHandler 91 92 def __init__(self, *args, **kwargs): 93 """ 94 WSGI server that simply tracks websockets 95 and send them a proper closing handshake 96 when the server terminates. 97 98 Other than that, the server is the same 99 as its :class:`gevent.pywsgi.WSGIServer` 100 base. 101 """ 102 _WSGIServer.__init__(self, *args, **kwargs) 103 self.pool = GEventWebSocketPool() 104 105 def stop(self, *args, **kwargs): 106 self.pool.clear() 107 _WSGIServer.stop(self, *args, **kwargs) 108 109if __name__ == '__main__': 110 import os 111 112 from ws4py import configure_logger 113 configure_logger() 114 115 from ws4py.websocket import EchoWebSocket 116 server = WSGIServer(('127.0.0.1', 9000), 117 WebSocketWSGIApplication(handler_cls=EchoWebSocket)) 118 server.serve_forever() 119