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

server/index.js

@@ -33,9 +33,17 @@ const serviceNames = uniq(
 
 // The number of processes to run, defaults to the number of CPU's available
 // for the web service, and 1 for collaboration during the beta period.
-const processCount = serviceNames.includes("collaboration")
-  ? 1
-  : env.WEB_CONCURRENCY || undefined;
+let processCount = env.WEB_CONCURRENCY || undefined;
+
+if (serviceNames.includes("collaboration")) {
+  if (env.WEB_CONCURRENCY !== 1) {
+    Logger.info(
+      "lifecycle",
+      "Note: Restricting process count to 1 due to use of collaborative service"
+    );
+  }
+  processCount = 1;
+}
 
 // This function will only be called once in the original process
 function master() {
@@ -72,15 +80,6 @@ async function start(id: string, disconnect: () => void) {
   router.get("/_health", (ctx) => (ctx.body = "OK"));
   app.use(router.routes());
 
-  if (
-    serviceNames.includes("websockets") &&
-    serviceNames.includes("collaboration")
-  ) {
-    throw new Error(
-      "Cannot run websockets and collaboration services in the same process"
-    );
-  }
-
   // loop through requested services at startup
   for (const name of serviceNames) {
     if (!Object.keys(services).includes(name)) {

server/presenters/env.js

@@ -6,7 +6,7 @@ export default function present(env: Object): Object {
   return {
     URL: env.URL.replace(/\/$/, ""),
     CDN_URL: (env.CDN_URL || "").replace(/\/$/, ""),
-    COLLABORATION_URL: (env.COLLABORATION_URL || "")
+    COLLABORATION_URL: (env.COLLABORATION_URL || env.URL)
       .replace(/\/$/, "")
       .replace(/^http/, "ws"),
     DEPLOYMENT: env.DEPLOYMENT,

server/presenters/team.js

@@ -1,5 +1,4 @@
 // @flow
-import env from "../env";
 import { Team } from "../models";
 
 export default function present(team: Team) {
@@ -8,9 +7,7 @@ export default function present(team: Team) {
     name: team.name,
     avatarUrl: team.logoUrl,
     sharing: team.sharing,
-    collaborativeEditing: !!(
-      team.collaborativeEditing && env.COLLABORATION_URL
-    ),
+    collaborativeEditing: team.collaborativeEditing,
     documentEmbeds: team.documentEmbeds,
     guestSignin: team.guestSignin,
     subdomain: team.subdomain,

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);
   });