Working prototype with Multi-process support for collab endpoint

Depends on y-crdt/pycrdt-websocket#128

Signed-off-by: Alexis Jeandet <alexis.jeandet@member.fsf.org>

Author: jeandet

docker/Dockerfile

@@ -12,6 +12,7 @@ ENV SPEASY_CACHE_PATH=/data \
     SPEASY_PROXY_SERVER_URL="" \
     SPEASY_PROXY_LOG_LEVEL="WARN"\
     SPEASY_CORE_HTTP_REWRITE_RULES=""\
+    WITH_COLLAB="true" \
     PORT=$PORT\
     SPEASY_PROXY_LOG_CONFIG_FILE="/home/speasy/speasy_proxy/logging.yaml"
 

docker/entry_point.sh

@@ -1,4 +1,19 @@
 #!/usr/bin/env bash
 
 #~/.local/bin/gunicorn  speasy_proxy:app -w 32 -k uvicorn.workers.UvicornWorker --ws-max-queue=128  --backlog=32 --limit-max-requests=1024 --root-path=$SPEASY_PROXY_PREFIX --host="0.0.0.0" --port=$PORT --proxy-headers --limit-concurrency=10000 --log-config=/home/speasy/speasy_proxy/logging.yaml
-~/.local/bin/gunicorn  speasy_proxy:app --timeout 600 --max-requests 10000 --backlog 2048  --worker-connections 1000 -w 32 --threads=8 -k speasy_proxy.UvicornWorker.SpeasyUvicornWorker
+if [ "$SPEASY_PROXY_COLLAB_ENDPOINT_ENABLE" = "true" ]; then
+  echo "SPEASY_PROXY_COLLAB_ENDPOINT_ENABLE= is set, starting collab server"
+  python3 -m speasy_proxy.collab_server &
+  COLLAB_PID=$!
+fi
+
+PYTHON_PATH=$(dirname "$(which python3)")
+
+"$PYTHON_PATH"/gunicorn  speasy_proxy:app --timeout 600 --max-requests 10000 --backlog 2048  --worker-connections 1000 -w 32 --threads=8 -k speasy_proxy.UvicornWorker.SpeasyUvicornWorker
+
+if [ "$SPEASY_PROXY_COLLAB_ENDPOINT_ENABLE" = "true" ]; then
+  echo "killing collab server with PID $COLLAB_PID"
+  kill $COLLAB_PID
+  wait $COLLAB_PID
+  echo "collab server with PID $COLLAB_PID has been killed"
+fi

pyproject.toml

@@ -31,7 +31,10 @@ classifiers = [
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
 ]
-dependencies = ['apscheduler < 4.0.0', 'humanize', 'speasy>=1.5.0', 'pycdfpp>=0.7.3', 'diskcache', 'jinja2', 'pyzstd', 'bokeh>=3.0.0', 'matplotlib', 'fastapi[all]', 'uvicorn', 'pycrdt-websocket >=0.15.4,<0.16.0']
+dependencies = ['apscheduler < 4.0.0', 'humanize', 'speasy>=1.5.0', 'pycdfpp>=0.7.3', 'diskcache', 'jinja2', 'pyzstd', 'bokeh>=3.0.0', 'matplotlib', 'fastapi[all]', 'uvicorn']
+[project.optional-dependencies]
+collab = ['pycrdt-websocket >=0.15.4,<0.16.0', 'hypercorn', 'httpx_ws']
+
 [project.urls]
 homepage = "https://github.com/SciQLop/speasy_proxy"
 

speasy_proxy/UvicornWorker.py

@@ -1,4 +1,4 @@
-from uvicorn.workers import UvicornWorker
+from uvicorn_worker import UvicornWorker
 import os
 
 

speasy_proxy/api/v1/ws_collaboration.py

@@ -5,12 +5,63 @@
 
 if collab_endpoint.enable():
     from asyncio import create_task
+    from anyio import get_cancelled_exc_class, Lock
+    from contextlib import asynccontextmanager
     from fastapi import WebSocket
-    from pycrdt_websocket import WebsocketServer
-    from pycrdt_websocket.websocket import HttpxWebsocket
+    from httpx_ws import aconnect_ws
+    from pycrdt import Array, Doc, Provider
+    from pycrdt.websocket import WebsocketServer
+    from pycrdt.websocket.websocket import HttpxWebsocket
     from .routes import router
 
 
+    class Websocket:
+        def __init__(self, websocket, path: str):
+            self._websocket = websocket
+            self._path = path
+            self._send_lock = Lock()
+
+        @property
+        def path(self) -> str:
+            return self._path
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self) -> bytes:
+            try:
+                message = await self.recv()
+            except Exception:
+                raise StopAsyncIteration()
+            return message
+
+        async def send(self, message: bytes):
+            async with self._send_lock:
+                await self._websocket.send_bytes(message)
+
+        async def recv(self) -> bytes:
+            b = await self._websocket.receive_bytes()
+            return bytes(b)
+
+
+    @asynccontextmanager
+    async def aprovider_factory(port, room_name, ydoc=None, log=None):
+        ydoc = Doc() if ydoc is None else ydoc
+        server_websocket = None
+        connect = aconnect_ws(f"http://localhost:{port}/{room_name}")
+        try:
+            async with connect as websocket:
+                websocket_provider = Provider(ydoc, Websocket(websocket, room_name), log)
+                async with websocket_provider as websocket_provider:
+                    yield ydoc, server_websocket
+        except get_cancelled_exc_class():
+            pass
+
+
+    def provider_factory(path, doc, log):
+        return aprovider_factory(collab_endpoint.port(), path, ydoc=doc, log=log)
+
+
     @router.websocket("/collaboration/{path:path}")
     async def websocket_endpoint(path: str, websocket: WebSocket):
         await websocket.accept()
@@ -21,7 +72,7 @@ async def websocket_endpoint(path: str, websocket: WebSocket):
     async def get_websocket_server():
         global WEBSOCKET_SERVER
         if WEBSOCKET_SERVER is None:
-            WEBSOCKET_SERVER = WebsocketServer()
+            WEBSOCKET_SERVER = WebsocketServer(provider_factory=provider_factory)
             create_task(WEBSOCKET_SERVER.start())
             await WEBSOCKET_SERVER.started.wait()
         return WEBSOCKET_SERVER

speasy_proxy/config/__init__.py

@@ -4,4 +4,6 @@
 collab_endpoint = ConfigSection("PROXY_COLLAB_ENDPOINT",
                                 enable={"default": False,
                                         "type_ctor": lambda x: {'true': True,
-                                                                'false': False}.get(x.lower(), False)})
+                                                                'false': False}.get(x.lower(), False)},
+                                port ={"default": 9999, "type_ctor": int},
+                                )

speasy_proxy/sciqlop_collab_server.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+import asyncio
+from hypercorn import Config
+from hypercorn.asyncio import serve
+from pycrdt.websocket import ASGIServer, WebsocketServer
+from speasy_proxy.config import collab_endpoint
+import os
+
+
+async def main():
+    websocket_server = WebsocketServer()
+    app = ASGIServer(websocket_server)
+    config = Config()
+    config.bind = [f"localhost:{collab_endpoint.port()}"]
+    async with websocket_server:
+        await serve(app, config, mode="asgi")
+
+asyncio.run(main())