Add WebSocket JSON RPC feature

Signed-off-by: YAMADA Hideki <yamada.hideki@po.ntts.co.jp>
Signed-off-by: FUJITA Tomonori <fujita.tomonori@lab.ntt.co.jp>

1 files changed, 44 insertions, 0 deletions

diff --git a/ryu/app/wsgi.py b/ryu/app/wsgi.py
index 406fbab9..68707319 100644
--- a/ryu/app/wsgi.py
+++ b/ryu/app/wsgi.py
@@ -23,6 +23,11 @@ from ryu import cfg
 from ryu.lib import hub
 from routes import Mapper
 from routes.util import URLGenerator
+from tinyrpc.server import RPCServer
+from tinyrpc.dispatch import RPCDispatcher
+from tinyrpc.dispatch import public as rpc_public
+from tinyrpc.protocols.jsonrpc import JSONRPCProtocol
+from tinyrpc.transports import ServerTransport
 
 CONF = cfg.CONF
 CONF.register_cli_opts([
@@ -115,6 +120,45 @@ class ControllerBase(object):
         return getattr(self, action)(req, **kwargs)
 
 
+class WebSocketDisconnectedError(Exception):
+    pass
+
+
+class WebSocketServerTransport(ServerTransport):
+    def __init__(self, ws):
+        self.ws = ws
+
+    def receive_message(self):
+        message = self.ws.wait()
+        if message is None:
+            raise WebSocketDisconnectedError()
+        context = None
+        return (context, message)
+
+    def send_reply(self, context, reply):
+        self.ws.send(unicode(reply))
+
+
+class WebSocketRPCServer(RPCServer):
+    def __init__(self, ws, rpc_callback):
+        dispatcher = RPCDispatcher()
+        dispatcher.register_instance(rpc_callback)
+        super(WebSocketRPCServer, self).__init__(
+            WebSocketServerTransport(ws),
+            JSONRPCProtocol(),
+            dispatcher,
+        )
+
+    def serve_forever(self):
+        try:
+            super(WebSocketRPCServer, self).serve_forever()
+        except WebSocketDisconnectedError:
+            return
+
+    def _spawn(self, func, *args, **kwargs):
+        hub.spawn(func, *args, **kwargs)
+
+
 class wsgify_hack(webob.dec.wsgify):
     def __call__(self, environ, start_response):
         self.kwargs['start_response'] = start_response