chore: Allow websockets and collaboration service to run in the same process (#2674)

server/services/collaboration.js

@@ -1,16 +1,17 @@
 // @flow
 import http from "http";
+import url from "url";
 import { Server } from "@hocuspocus/server";
 import Koa from "koa";
-import websocket from "koa-easy-ws";
-import Router from "koa-router";
+import WebSocket from "ws";
 import AuthenticationExtension from "../collaboration/authentication";
 import LoggerExtension from "../collaboration/logger";
 import PersistenceExtension from "../collaboration/persistence";
 import TracingExtension from "../collaboration/tracing";
 
 export default function init(app: Koa, server: http.Server) {
-  const router = new Router();
+  const path = "/collaboration";
+  const wss = new WebSocket.Server({ noServer: true });
 
   const hocuspocus = Server.configure({
     extensions: [
@@ -21,22 +22,16 @@ export default function init(app: Koa, server: http.Server) {
     ],
   });
 
-  // Websockets for collaborative editing
-  router.get("/collaboration/:documentName", async (ctx) => {
-    let { documentName } = ctx.params;
+  server.on("upgrade", function (req, socket, head) {
+    if (req.url.indexOf(path) > -1) {
+      const documentName = url.parse(req.url).pathname?.split("/").pop();
 
-    if (ctx.ws) {
-      const ws = await ctx.ws();
-      hocuspocus.handleConnection(ws, ctx.request, documentName);
+      wss.handleUpgrade(req, socket, Buffer.alloc(0), (client) => {
+        hocuspocus.handleConnection(client, req, documentName);
+      });
     }
-
-    ctx.response.status = 101;
   });
 
-  app.use(websocket());
-  app.use(router.routes());
-  app.use(router.allowedMethods());
-
   server.on("shutdown", () => {
     hocuspocus.destroy();
   });

server/services/websockets.js

@@ -16,13 +16,28 @@ import { getUserForJWT } from "../utils/jwt";
 const { can } = policy;
 
 export default function init(app: Koa, server: http.Server) {
+  const path = "/realtime";
+
   // Websockets for events and non-collaborative documents
   const io = IO(server, {
-    path: "/realtime",
+    path,
     serveClient: false,
     cookie: false,
   });
 
+  // Remove the upgrade handler that we just added when registering the IO engine
+  // And re-add it with a check to only handle the realtime path, this allows
+  // collaboration websockets to exist in the same process as engine.io.
+  const listeners = server.listeners("upgrade");
+  const ioHandleUpgrade = listeners.pop();
+  server.removeListener("upgrade", ioHandleUpgrade);
+
+  server.on("upgrade", function (req, socket, head) {
+    if (req.url.indexOf(path) > -1) {
+      ioHandleUpgrade(req, socket, head);
+    }
+  });
+
   server.on("shutdown", () => {
     Metrics.gaugePerInstance("websockets.count", 0);
   });