feat(webapp): add support for running web services (api, engine, webapp) in cluster mode for better perf (#2472)

* feat(webapp): add support for running web services (api, engine, webapp) in cluster mode for better perf

* cleaned up signal handling and resolved some valid 🐇 issues

apps/webapp/app/eventLoopMonitor.server.ts

@@ -5,6 +5,7 @@ import { env } from "./env.server";
 import { context, Context } from "@opentelemetry/api";
 import { performance } from "node:perf_hooks";
 import { logger } from "./services/logger.server";
+import { signalsEmitter } from "./services/signals.server";
 
 const THRESHOLD_NS = env.EVENT_LOOP_MONITOR_THRESHOLD_MS * 1e6;
 
@@ -110,6 +111,13 @@ function startEventLoopUtilizationMonitoring() {
     lastEventLoopUtilization = currentEventLoopUtilization;
   }, env.EVENT_LOOP_MONITOR_UTILIZATION_INTERVAL_MS);
 
+  signalsEmitter.on("SIGTERM", () => {
+    clearInterval(interval);
+  });
+  signalsEmitter.on("SIGINT", () => {
+    clearInterval(interval);
+  });
+
   return () => {
     clearInterval(interval);
   };

apps/webapp/app/services/realtime/relayRealtimeStreams.server.ts

@@ -1,5 +1,6 @@
 import { AuthenticatedEnvironment } from "../apiAuth.server";
 import { logger } from "../logger.server";
+import { signalsEmitter } from "../signals.server";
 import { StreamIngestor, StreamResponder } from "./types";
 import { LineTransformStream } from "./utils.server";
 import { v1RealtimeStreams } from "./v1StreamsGlobal.server";
@@ -243,12 +244,17 @@ export class RelayRealtimeStreams implements StreamIngestor, StreamResponder {
 }
 
 function initializeRelayRealtimeStreams() {
-  return new RelayRealtimeStreams({
+  const service = new RelayRealtimeStreams({
     ttl: 1000 * 60 * 5, // 5 minutes
     cleanupInterval: 1000 * 60, // 1 minute
     fallbackIngestor: v1RealtimeStreams,
     fallbackResponder: v1RealtimeStreams,
   });
+
+  signalsEmitter.on("SIGTERM", service.close.bind(service));
+  signalsEmitter.on("SIGINT", service.close.bind(service));
+
+  return service;
 }
 
 export const relayRealtimeStreams = singleton(

apps/webapp/app/services/runsReplicationInstance.server.ts

@@ -3,8 +3,8 @@ import invariant from "tiny-invariant";
 import { env } from "~/env.server";
 import { singleton } from "~/utils/singleton";
 import { provider } from "~/v3/tracer.server";
-import { logger } from "./logger.server";
 import { RunsReplicationService } from "./runsReplicationService.server";
+import { signalsEmitter } from "./signals.server";
 
 export const runsReplicationInstance = singleton(
   "runsReplicationInstance",
@@ -80,8 +80,8 @@ function initializeRunsReplicationInstance() {
         });
       });
 
-    process.on("SIGTERM", service.shutdown.bind(service));
-    process.on("SIGINT", service.shutdown.bind(service));
+    signalsEmitter.on("SIGTERM", service.shutdown.bind(service));
+    signalsEmitter.on("SIGINT", service.shutdown.bind(service));
   }
 
   return service;

apps/webapp/app/services/runsReplicationService.server.ts

@@ -204,6 +204,8 @@ export class RunsReplicationService {
   }
 
   public async shutdown() {
+    if (this._isShuttingDown) return;
+
     this._isShuttingDown = true;
 
     this.logger.info("Initiating shutdown of runs replication service");

apps/webapp/app/services/signals.server.ts

@@ -0,0 +1,32 @@
+import { EventEmitter } from "events";
+import { singleton } from "~/utils/singleton";
+
+export type SignalsEvents = {
+  SIGTERM: [
+    {
+      time: Date;
+      signal: NodeJS.Signals;
+    }
+  ];
+  SIGINT: [
+    {
+      time: Date;
+      signal: NodeJS.Signals;
+    }
+  ];
+};
+
+export type SignalsEventArgs<T extends keyof SignalsEvents> = SignalsEvents[T];
+
+export type SignalsEmitter = EventEmitter<SignalsEvents>;
+
+function initializeSignalsEmitter() {
+  const emitter = new EventEmitter<SignalsEvents>();
+
+  process.on("SIGTERM", () => emitter.emit("SIGTERM", { time: new Date(), signal: "SIGTERM" }));
+  process.on("SIGINT", () => emitter.emit("SIGINT", { time: new Date(), signal: "SIGINT" }));
+
+  return emitter;
+}
+
+export const signalsEmitter = singleton("signalsEmitter", initializeSignalsEmitter);

apps/webapp/app/v3/dynamicFlushScheduler.server.ts

@@ -1,6 +1,7 @@
 import { Logger } from "@trigger.dev/core/logger";
 import { nanoid } from "nanoid";
 import pLimit from "p-limit";
+import { signalsEmitter } from "~/services/signals.server";
 
 export type DynamicFlushSchedulerConfig<T> = {
   batchSize: number;
@@ -22,6 +23,7 @@ export class DynamicFlushScheduler<T> {
   private readonly BATCH_SIZE: number;
   private readonly FLUSH_INTERVAL: number;
   private flushTimer: NodeJS.Timeout | null;
+  private metricsReporterTimer: NodeJS.Timeout | undefined;
   private readonly callback: (flushId: string, batch: T[]) => Promise<void>;
 
   // New properties for dynamic scaling
@@ -41,6 +43,7 @@ export class DynamicFlushScheduler<T> {
     droppedEvents: 0,
     droppedEventsByKind: new Map<string, number>(),
   };
+  private isShuttingDown: boolean = false;
 
   // New properties for load shedding
   private readonly loadSheddingThreshold: number;
@@ -75,6 +78,7 @@ export class DynamicFlushScheduler<T> {
 
     this.startFlushTimer();
     this.startMetricsReporter();
+    this.setupShutdownHandlers();
   }
 
   addToBatch(items: T[]): void {
@@ -119,8 +123,8 @@ export class DynamicFlushScheduler<T> {
     this.currentBatch.push(...itemsToAdd);
     this.totalQueuedItems += itemsToAdd.length;
 
-    // Check if we need to create a batch
-    if (this.currentBatch.length >= this.currentBatchSize) {
+    // Check if we need to create a batch (if we are shutting down, create a batch immediately because the flush timer is stopped)
+    if (this.currentBatch.length >= this.currentBatchSize || this.isShuttingDown) {
       this.createBatch();
     }
 
@@ -137,6 +141,23 @@ export class DynamicFlushScheduler<T> {
     this.resetFlushTimer();
   }
 
+  private setupShutdownHandlers(): void {
+    signalsEmitter.on("SIGTERM", () =>
+      this.shutdown().catch((error) => {
+        this.logger.error("Error shutting down dynamic flush scheduler", {
+          error,
+        });
+      })
+    );
+    signalsEmitter.on("SIGINT", () =>
+      this.shutdown().catch((error) => {
+        this.logger.error("Error shutting down dynamic flush scheduler", {
+          error,
+        });
+      })
+    );
+  }
+
   private startFlushTimer(): void {
     this.flushTimer = setInterval(() => this.checkAndFlush(), this.FLUSH_INTERVAL);
   }
@@ -145,6 +166,9 @@ export class DynamicFlushScheduler<T> {
     if (this.flushTimer) {
       clearInterval(this.flushTimer);
     }
+
+    if (this.isShuttingDown) return;
+
     this.startFlushTimer();
   }
 
@@ -226,7 +250,7 @@ export class DynamicFlushScheduler<T> {
   }
 
   private lastConcurrencyAdjustment: number = Date.now();
-  
+
   private adjustConcurrency(backOff: boolean = false): void {
     const currentConcurrency = this.limiter.concurrency;
     let newConcurrency = currentConcurrency;
@@ -281,7 +305,7 @@ export class DynamicFlushScheduler<T> {
 
   private startMetricsReporter(): void {
     // Report metrics every 30 seconds
-    setInterval(() => {
+    this.metricsReporterTimer = setInterval(() => {
       const droppedByKind: Record<string, number> = {};
       this.metrics.droppedEventsByKind.forEach((count, kind) => {
         droppedByKind[kind] = count;
@@ -356,10 +380,18 @@ export class DynamicFlushScheduler<T> {
 
   // Graceful shutdown
   async shutdown(): Promise<void> {
+    if (this.isShuttingDown) return;
+
+    this.isShuttingDown = true;
+
     if (this.flushTimer) {
       clearInterval(this.flushTimer);
     }
 
+    if (this.metricsReporterTimer) {
+      clearInterval(this.metricsReporterTimer);
+    }
+
     // Flush any remaining items
     if (this.currentBatch.length > 0) {
       this.createBatch();

apps/webapp/app/v3/marqs/index.server.ts

@@ -24,6 +24,7 @@ import z from "zod";
 import { env } from "~/env.server";
 import { AuthenticatedEnvironment } from "~/services/apiAuth.server";
 import { logger } from "~/services/logger.server";
+import { signalsEmitter } from "~/services/signals.server";
 import { singleton } from "~/utils/singleton";
 import { legacyRunEngineWorker } from "../legacyRunEngineWorker.server";
 import { concurrencyTracker } from "../services/taskRunConcurrencyTracker.server";
@@ -112,6 +113,7 @@ export class MarQS {
   private queueDequeueCooloffPeriod: Map<string, number> = new Map();
   private queueDequeueCooloffCounts: Map<string, number> = new Map();
   private clearCooloffPeriodInterval: NodeJS.Timeout;
+  isShuttingDown: boolean = false;
 
   constructor(private readonly options: MarQSOptions) {
     this.redis = options.redis;
@@ -151,11 +153,14 @@ export class MarQS {
   }
 
   #setupShutdownHandlers() {
-    process.on("SIGTERM", () => this.shutdown("SIGTERM"));
-    process.on("SIGINT", () => this.shutdown("SIGINT"));
+    signalsEmitter.on("SIGTERM", () => this.shutdown("SIGTERM"));
+    signalsEmitter.on("SIGINT", () => this.shutdown("SIGINT"));
   }
 
   async shutdown(signal: NodeJS.Signals) {
+    if (this.isShuttingDown) return;
+    this.isShuttingDown = true;
+
     console.log("👇 Shutting down marqs", this.name, signal);
     clearInterval(this.clearCooloffPeriodInterval);
     this.#rebalanceWorkers.forEach((worker) => worker.stop());

apps/webapp/app/v3/tracing.server.ts

@@ -41,33 +41,14 @@ export async function startSpanWithEnv<T>(
   fn: (span: Span) => Promise<T>,
   options?: SpanOptions
 ): Promise<T> {
-  return startSpan(
-    tracer,
-    name,
-    async (span) => {
-      try {
-        return await fn(span);
-      } catch (e) {
-        if (e instanceof Error) {
-          span.recordException(e);
-        } else {
-          span.recordException(new Error(String(e)));
-        }
-
-        throw e;
-      } finally {
-        span.end();
-      }
+  return startSpan(tracer, name, fn, {
+    ...options,
+    attributes: {
+      ...attributesFromAuthenticatedEnv(env),
+      ...options?.attributes,
     },
-    {
-      attributes: {
-        ...attributesFromAuthenticatedEnv(env),
-        ...options?.attributes,
-      },
-      kind: SpanKind.SERVER,
-      ...options,
-    }
-  );
+    kind: SpanKind.SERVER,
+  });
 }
 
 export async function emitDebugLog(

apps/webapp/server.ts

@@ -12,168 +12,257 @@ import type { Server as IoServer } from "socket.io";
 import { WebSocketServer } from "ws";
 import { RateLimitMiddleware } from "~/services/apiRateLimit.server";
 import { type RunWithHttpContextFunction } from "~/services/httpAsyncStorage.server";
+import cluster from "node:cluster";
+import os from "node:os";
 
-const app = express();
+const ENABLE_CLUSTER = process.env.ENABLE_CLUSTER === "1";
+const cpuCount = os.availableParallelism();
+const WORKERS =
+  Number.parseInt(process.env.WEB_CONCURRENCY || process.env.CLUSTER_WORKERS || "", 10) || cpuCount;
 
-if (process.env.DISABLE_COMPRESSION !== "1") {
-  app.use(compression());
+function forkWorkers() {
+  for (let i = 0; i < WORKERS; i++) {
+    cluster.fork();
+  }
 }
 
-// http://expressjs.com/en/advanced/best-practice-security.html#at-a-minimum-disable-x-powered-by-header
-app.disable("x-powered-by");
+function installPrimarySignalHandlers() {
+  let didHandleSigterm = false;
+  let didHandleSigint = false;
+  let didGracefulExit = false;
 
-// Remix fingerprints its assets so we can cache forever.
-app.use("/build", express.static("public/build", { immutable: true, maxAge: "1y" }));
-
-// Everything else (like favicon.ico) is cached for an hour. You may want to be
-// more aggressive with this caching.
-app.use(express.static("public", { maxAge: "1h" }));
-
-app.use(morgan("tiny"));
-
-process.title = "node webapp-server";
-
-const MODE = process.env.NODE_ENV;
-const BUILD_DIR = path.join(process.cwd(), "build");
-const build = require(BUILD_DIR);
-
-const port = process.env.REMIX_APP_PORT || process.env.PORT || 3000;
-
-if (process.env.HTTP_SERVER_DISABLED !== "true") {
-  const socketIo: { io: IoServer } | undefined = build.entry.module.socketIo;
-  const wss: WebSocketServer | undefined = build.entry.module.wss;
-  const apiRateLimiter: RateLimitMiddleware = build.entry.module.apiRateLimiter;
-  const engineRateLimiter: RateLimitMiddleware = build.entry.module.engineRateLimiter;
-  const runWithHttpContext: RunWithHttpContextFunction = build.entry.module.runWithHttpContext;
-
-  app.use((req, res, next) => {
-    // helpful headers:
-    res.set("Strict-Transport-Security", `max-age=${60 * 60 * 24 * 365 * 100}`);
-
-    // Add X-Robots-Tag header for test-cloud.trigger.dev
-    if (req.hostname !== "cloud.trigger.dev") {
-      res.set("X-Robots-Tag", "noindex, nofollow");
+  const forward = (signal: NodeJS.Signals) => {
+    for (const id in cluster.workers) {
+      const w = cluster.workers[id];
+      if (w?.process?.pid) {
+        try {
+          process.kill(w.process.pid, signal);
+        } catch {}
+      }
     }
+  };
 
-    // /clean-urls/ -> /clean-urls
-    if (req.path.endsWith("/") && req.path.length > 1) {
-      const query = req.url.slice(req.path.length);
-      const safepath = req.path.slice(0, -1).replace(/\/+/g, "/");
-      res.redirect(301, safepath + query);
-      return;
-    }
-    next();
+  const gracefulExit = () => {
+    if (didGracefulExit) return;
+    didGracefulExit = true;
+
+    const timeoutMs = Number(process.env.GRACEFUL_SHUTDOWN_TIMEOUT || 30_000);
+    // wait for workers to exit, then exit the primary too
+    const maybeExit = () => {
+      const alive = Object.values(cluster.workers || {}).some((w) => w && !w.isDead());
+      if (!alive) process.exit(0);
+    };
+    setInterval(maybeExit, 1000);
+    setTimeout(() => process.exit(0), timeoutMs);
+  };
+
+  process.on("SIGTERM", () => {
+    if (didHandleSigterm) return;
+    didHandleSigterm = true;
+    forward("SIGTERM");
+    gracefulExit();
+  });
+  process.on("SIGINT", () => {
+    if (didHandleSigint) return;
+    didHandleSigint = true;
+    forward("SIGINT");
+    gracefulExit();
+  });
+}
+
+if (ENABLE_CLUSTER && cluster.isPrimary) {
+  process.title = `node webapp-server primary`;
+  console.log(`[cluster] Primary ${process.pid} is starting with ${WORKERS} workers`);
+  forkWorkers();
+
+  cluster.on("exit", (worker, code, signal) => {
+    const intentional =
+      // If we sent "shutdown", the worker will exit with code 0 after closing.
+      code === 0 || worker.exitedAfterDisconnect;
+    console.log(
+      `[cluster] worker ${worker.process.pid} exited (code=${code}, signal=${signal}, intentional=${intentional})`
+    );
+    // If it wasn't during a shutdown, replace the worker.
+    if (!intentional) cluster.fork();
   });
 
-  app.use((req, res, next) => {
-    // Generate a unique request ID for each request
-    const requestId = nanoid();
+  installPrimarySignalHandlers();
+} else {
+  const app = express();
 
-    runWithHttpContext({ requestId, path: req.url, host: req.hostname, method: req.method }, next);
-  });
+  if (process.env.DISABLE_COMPRESSION !== "1") {
+    app.use(compression());
+  }
 
-  if (process.env.DASHBOARD_AND_API_DISABLED !== "true") {
-    if (process.env.ALLOW_ONLY_REALTIME_API === "true") {
-      // Block all requests that do not start with /realtime
-      app.use((req, res, next) => {
-        // Make sure /healthcheck is still accessible
-        if (!req.url.startsWith("/realtime") && req.url !== "/healthcheck") {
-          res.status(404).send("Not Found");
-          return;
-        }
+  // http://expressjs.com/en/advanced/best-practice-security.html#at-a-minimum-disable-x-powered-by-header
+  app.disable("x-powered-by");
 
-        next();
+  // Remix fingerprints its assets so we can cache forever.
+  app.use("/build", express.static("public/build", { immutable: true, maxAge: "1y" }));
+
+  // Everything else (like favicon.ico) is cached for an hour. You may want to be
+  // more aggressive with this caching.
+  app.use(express.static("public", { maxAge: "1h" }));
+
+  app.use(morgan("tiny"));
+
+  process.title = ENABLE_CLUSTER
+    ? `node webapp-worker-${cluster.isWorker ? cluster.worker?.id : "solo"}`
+    : "node webapp-server";
+
+  const MODE = process.env.NODE_ENV;
+  const BUILD_DIR = path.join(process.cwd(), "build");
+  const build = require(BUILD_DIR);
+
+  const port = process.env.REMIX_APP_PORT || process.env.PORT || 3000;
+
+  if (process.env.HTTP_SERVER_DISABLED !== "true") {
+    const socketIo: { io: IoServer } | undefined = build.entry.module.socketIo;
+    const wss: WebSocketServer | undefined = build.entry.module.wss;
+    const apiRateLimiter: RateLimitMiddleware = build.entry.module.apiRateLimiter;
+    const engineRateLimiter: RateLimitMiddleware = build.entry.module.engineRateLimiter;
+    const runWithHttpContext: RunWithHttpContextFunction = build.entry.module.runWithHttpContext;
+
+    app.use((req, res, next) => {
+      // helpful headers:
+      res.set("Strict-Transport-Security", `max-age=${60 * 60 * 24 * 365 * 100}`);
+
+      // Add X-Robots-Tag header for test-cloud.trigger.dev
+      if (req.hostname !== "cloud.trigger.dev") {
+        res.set("X-Robots-Tag", "noindex, nofollow");
+      }
+
+      // /clean-urls/ -> /clean-urls
+      if (req.path.endsWith("/") && req.path.length > 1) {
+        const query = req.url.slice(req.path.length);
+        const safepath = req.path.slice(0, -1).replace(/\/+/g, "/");
+        res.redirect(301, safepath + query);
+        return;
+      }
+      next();
+    });
+
+    app.use((req, res, next) => {
+      // Generate a unique request ID for each request
+      const requestId = nanoid();
+
+      runWithHttpContext(
+        { requestId, path: req.url, host: req.hostname, method: req.method },
+        next
+      );
+    });
+
+    if (process.env.DASHBOARD_AND_API_DISABLED !== "true") {
+      if (process.env.ALLOW_ONLY_REALTIME_API === "true") {
+        // Block all requests that do not start with /realtime
+        app.use((req, res, next) => {
+          // Make sure /healthcheck is still accessible
+          if (!req.url.startsWith("/realtime") && req.url !== "/healthcheck") {
+            res.status(404).send("Not Found");
+            return;
+          }
+
+          next();
+        });
+      }
+
+      app.use(apiRateLimiter);
+      app.use(engineRateLimiter);
+
+      app.all(
+        "*",
+        // @ts-ignore
+        createRequestHandler({
+          build,
+          mode: MODE,
+        })
+      );
+    } else {
+      // we need to do the health check here at /healthcheck
+      app.get("/healthcheck", (req, res) => {
+        res.status(200).send("OK");
       });
     }
 
-    app.use(apiRateLimiter);
-    app.use(engineRateLimiter);
+    const server = app.listen(port, () => {
+      console.log(
+        `✅ server ready: http://localhost:${port} [NODE_ENV: ${MODE}]${
+          ENABLE_CLUSTER && cluster.isWorker ? ` [worker ${cluster.worker?.id}/${process.pid}]` : ""
+        }`
+      );
 
-    app.all(
-      "*",
-      // @ts-ignore
-      createRequestHandler({
-        build,
-        mode: MODE,
-      })
-    );
-  } else {
-    // we need to do the health check here at /healthcheck
-    app.get("/healthcheck", (req, res) => {
-      res.status(200).send("OK");
-    });
-  }
-
-  const server = app.listen(port, () => {
-    console.log(`✅ server ready: http://localhost:${port} [NODE_ENV: ${MODE}]`);
-
-    if (MODE === "development") {
-      broadcastDevReady(build)
-        .then(() => logDevReady(build))
-        .catch(console.error);
-    }
-  });
-
-  server.keepAliveTimeout = 65 * 1000;
-  // Mitigate against https://github.com/triggerdotdev/trigger.dev/security/dependabot/128
-  // by not allowing 2000+ headers to be sent and causing a DoS
-  // headers will instead be limited by the maxHeaderSize
-  server.maxHeadersCount = 0;
-
-  process.on("SIGTERM", () => {
-    server.close((err) => {
-      if (err) {
-        console.error("Error closing express server:", err);
-      } else {
-        console.log("Express server closed gracefully.");
+      if (MODE === "development") {
+        broadcastDevReady(build)
+          .then(() => logDevReady(build))
+          .catch(console.error);
       }
     });
-  });
 
-  socketIo?.io.attach(server);
-  server.removeAllListeners("upgrade"); // prevent duplicate upgrades from listeners created by io.attach()
+    server.keepAliveTimeout = 65 * 1000;
+    // Mitigate against https://github.com/triggerdotdev/trigger.dev/security/dependabot/128
+    // by not allowing 2000+ headers to be sent and causing a DoS
+    // headers will instead be limited by the maxHeaderSize
+    server.maxHeadersCount = 0;
 
-  server.on("upgrade", async (req, socket, head) => {
-    console.log(
-      `Attemping to upgrade connection at url ${req.url} with headers: ${JSON.stringify(
-        req.headers
-      )}`
-    );
+    let didCloseServer = false;
 
-    socket.on("error", (err) => {
-      console.error("Connection upgrade error:", err);
-    });
+    function closeServer(signal: NodeJS.Signals) {
+      if (didCloseServer) return;
+      didCloseServer = true;
 
-    const url = new URL(req.url ?? "", "http://localhost");
-
-    // Upgrade socket.io connection
-    if (url.pathname.startsWith("/socket.io/")) {
-      console.log(`Socket.io client connected, upgrading their connection...`);
-
-      // https://github.com/socketio/socket.io/issues/4693
-      (socketIo?.io.engine as EngineServer).handleUpgrade(req, socket, head);
-      return;
+      server.close((err) => {
+        if (err) {
+          console.error("Error closing express server:", err);
+        } else {
+          console.log("Express server closed gracefully.");
+        }
+      });
     }
 
-    // Only upgrade the connecting if the path is `/ws`
-    if (url.pathname !== "/ws") {
-      // Setting the socket.destroy() error param causes an error event to be emitted which needs to be handled with socket.on("error") to prevent uncaught exceptions.
-      socket.destroy(
-        new Error(
-          "Cannot connect because of invalid path: Please include `/ws` in the path of your upgrade request."
-        )
-      );
-      return;
-    }
+    process.on("SIGTERM", closeServer);
+    process.on("SIGINT", closeServer);
 
-    console.log(`Client connected, upgrading their connection...`);
+    socketIo?.io.attach(server);
+    server.removeAllListeners("upgrade"); // prevent duplicate upgrades from listeners created by io.attach()
 
-    // Handle the WebSocket connection
-    wss?.handleUpgrade(req, socket, head, (ws) => {
-      wss?.emit("connection", ws, req);
+    server.on("upgrade", async (req, socket, head) => {
+      console.log(`Attemping to upgrade connection at url ${req.url}`);
+
+      socket.on("error", (err) => {
+        console.error("Connection upgrade error:", err);
+      });
+
+      const url = new URL(req.url ?? "", "http://localhost");
+
+      // Upgrade socket.io connection
+      if (url.pathname.startsWith("/socket.io/")) {
+        console.log(`Socket.io client connected, upgrading their connection...`);
+
+        // https://github.com/socketio/socket.io/issues/4693
+        (socketIo?.io.engine as EngineServer).handleUpgrade(req, socket, head);
+        return;
+      }
+
+      // Only upgrade the connecting if the path is `/ws`
+      if (url.pathname !== "/ws") {
+        // Setting the socket.destroy() error param causes an error event to be emitted which needs to be handled with socket.on("error") to prevent uncaught exceptions.
+        socket.destroy(
+          new Error(
+            "Cannot connect because of invalid path: Please include `/ws` in the path of your upgrade request."
+          )
+        );
+        return;
+      }
+
+      console.log(`Client connected, upgrading their connection...`);
+
+      // Handle the WebSocket connection
+      wss?.handleUpgrade(req, socket, head, (ws) => {
+        wss?.emit("connection", ws, req);
+      });
     });
-  });
-} else {
-  require(BUILD_DIR);
-  console.log(`✅ app ready (skipping http server)`);
+  } else {
+    require(BUILD_DIR);
+    console.log(`✅ app ready (skipping http server)`);
+  }
 }