Source code for ws4py.server.geventserver

# -*- coding: utf-8 -*-
__doc__ = """
WSGI entities to support WebSocket from within gevent.

Its usage is rather simple:

.. code-block: python

    from gevent import monkey; monkey.patch_all()
    from ws4py.websocket import EchoWebSocket
    from ws4py.server.geventserver import WSGIServer
    from ws4py.server.wsgiutils import WebSocketWSGIApplication

    server = WSGIServer(('localhost', 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket))
    server.serve_forever()

"""
import logging
import sys

import gevent
from gevent.pywsgi import WSGIHandler, WSGIServer as _WSGIServer
from gevent.pool import Pool

from ws4py import format_addresses
from ws4py.server.wsgiutils import WebSocketWSGIApplication

logger = logging.getLogger('ws4py')

__all__ = ['WebSocketWSGIHandler', 'WSGIServer',
           'GEventWebSocketPool']

[docs]class WebSocketWSGIHandler(WSGIHandler): """ A WSGI handler that will perform the :rfc:`6455` upgrade and handshake before calling the WSGI application. If the incoming request doesn't have a `'Upgrade'` header, the handler will simply fallback to the gevent builtin's handler and process it as per usual. """
[docs] def run_application(self): upgrade_header = self.environ.get('HTTP_UPGRADE', '').lower() if upgrade_header: try: # Build and start the HTTP response self.environ['ws4py.socket'] = self.socket or self.environ['wsgi.input'].rfile._sock self.result = self.application(self.environ, self.start_response) or [] self.process_result() except: raise else: del self.environ['ws4py.socket'] self.socket = None self.rfile.close() ws = self.environ.pop('ws4py.websocket') if ws: self.server.pool.track(ws) else: gevent.pywsgi.WSGIHandler.run_application(self)
[docs]class GEventWebSocketPool(Pool): """ 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. """
[docs] def track(self, websocket): logger.info("Managing websocket %s" % format_addresses(websocket)) return self.spawn(websocket.run)
[docs] 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)
[docs]class WSGIServer(_WSGIServer): handler_class = WebSocketWSGIHandler def __init__(self, *args, **kwargs): """ WSGI server that simply tracks websockets and send them a proper closing handshake when the server terminates. Other than that, the server is the same as its :class:`gevent.pywsgi.WSGIServer` base. """ _WSGIServer.__init__(self, *args, **kwargs) self.pool = GEventWebSocketPool()
[docs] def stop(self, *args, **kwargs): self.pool.clear() _WSGIServer.stop(self, *args, **kwargs)
if __name__ == '__main__': import os from ws4py import configure_logger configure_logger() from ws4py.websocket import EchoWebSocket server = WSGIServer(('127.0.0.1', 9000), WebSocketWSGIApplication(handler_cls=EchoWebSocket)) server.serve_forever()