Remove environment checks from persistent variables (#115)

Author: apteryxxyz

.changeset/pretty-candles-knock.md

@@ -0,0 +1,5 @@
+---
+"next-ws": minor
+---
+
+Remove environment checks for getting persistent variables

examples/custom-server/package.json

@@ -3,7 +3,7 @@
   "private": true,
   "scripts": {
     "build": "next build",
-    "start": "NODE_ENV=production tsx server.ts",
+    "start": "NODE_ENV=production tsx --require=\"./global.js\" server.ts",
     "dev": "tsx --require=\"./global.js\" server.ts",
     "prepare": "next-ws patch"
   },

src/server/persistent.ts

@@ -1,82 +1,30 @@
-import * as logger from 'next/dist/build/output/log.js';
-
-function getEnvironmentMeta() {
-  const isCustomServer = !process.title.startsWith('next-');
-  const isMainProcess = process.env.NEXT_WS_MAIN_PROCESS === '1';
-  const isDevelopment = process.env.NODE_ENV === 'development';
-  return { isCustomServer, isMainProcess, isDevelopment };
-}
-
-function mainProcessOnly(callerName: string) {
-  if (process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK === '1') return;
-
-  const meta = getEnvironmentMeta();
-  if (!meta.isCustomServer && !meta.isMainProcess) {
-    throw new Error(
-      `[next-ws] Attempt to call '${callerName}' outside the main process.
-You may be attempting to interact with the WebSocket server outside of an UPGRADE handler. This will fail in production, as Next.js employs a worker process for routing, which do not have access to the WebSocket server on the main process.
-You can resolve this by using a custom server.`,
-    );
-  } else if (!meta.isCustomServer) {
-    logger.warnOnce(
-      `[next-ws] The function '${callerName}' was called without a custom server.
-This could lead to unintended behaviour, especially if you're attempting to interact with the WebSocket server outside of an UPGRADE handler.
-Please note, while such configurations might function during development, they will fail in production. This is because Next.js employs a worker process for routing in production, which do not have access to the WebSocket server on the main process.
-You can resolve this by using a custom server.`,
-    );
-  }
-}
-
-//
-
-import type { Server as HttpServer } from 'node:http';
-export const kHttpServer = Symbol.for('kHttpServer');
-
-export function getHttpServer() {
-  mainProcessOnly('getHttpServer');
-  return Reflect.get(globalThis, kHttpServer) as HttpServer | undefined;
-}
-
-export function setHttpServer<T extends HttpServer | undefined>(server: T) {
-  mainProcessOnly('setHttpServer');
-  Reflect.set(globalThis, kHttpServer, server);
-  return getHttpServer() as T;
-}
-
-export function useHttpServer<T extends HttpServer | undefined>(
-  server: () => T,
-): T {
-  mainProcessOnly('useHttpServer');
-  const existing = getHttpServer();
-  if (existing) return existing as T;
-  return setHttpServer(server());
-}
-
-//
-
-import type { WebSocketServer } from 'ws';
-export const kWebSocketServer = Symbol.for('kWebSocketServer');
-
-export function getWebSocketServer() {
-  mainProcessOnly('getWebSocketServer');
-  return Reflect.get(globalThis, kWebSocketServer) as
-    | WebSocketServer
-    | undefined;
-}
-
-export function setWebSocketServer<T extends WebSocketServer | undefined>(
-  server: T,
-) {
-  mainProcessOnly('setWebSocketServer');
-  Reflect.set(globalThis, kWebSocketServer, server);
-  return getWebSocketServer() as T;
-}
-
-export function useWebSocketServer<T extends WebSocketServer | undefined>(
-  server: () => T,
-): T {
-  mainProcessOnly('useWebSocketServer');
-  const existing = getWebSocketServer();
-  if (existing) return existing as T;
-  return setWebSocketServer(server()) as T;
-}
+function useGlobal<T>(key: PropertyKey) {
+  return [
+    function get() {
+      return Reflect.get(globalThis, key) as T | undefined;
+    },
+    function set(value: T) {
+      return Reflect.set(globalThis, key, value);
+    },
+    function use(getter: () => T) {
+      const existing = Reflect.get(globalThis, key);
+      if (existing) return existing as T;
+      Reflect.set(globalThis, key, getter());
+      return Reflect.get(globalThis, key) as T;
+    },
+  ] as const;
+}
+
+// ===== HTTP Server ===== //
+
+export const [getHttpServer, setHttpServer, useHttpServer] = //
+  useGlobal<import('node:http').Server>(
+    Symbol.for('next-ws.http-server'), //
+  );
+
+// ===== WebSocket Server ===== //
+
+export const [getWebSocketServer, setWebSocketServer, useWebSocketServer] =
+  useGlobal<import('ws').WebSocketServer>(
+    Symbol.for('next-ws.websocket-server'),
+  );

src/server/setup.ts

@@ -7,24 +7,18 @@ import { toNextRequest } from './helpers/request.js';
 import { useHttpServer, useWebSocketServer } from './persistent.js';
 
 export function setupWebSocketServer(nextServer: NextNodeServer) {
-  process.env.NEXT_WS_MAIN_PROCESS = String(1);
-
-  process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK = String(1);
   const httpServer = //
     // @ts-expect-error - serverOptions is protected
     useHttpServer(() => nextServer.serverOptions?.httpServer);
   if (!httpServer)
     return logger.error('[next-ws] was not able to find the HTTP server');
   const wsServer = //
     useWebSocketServer(() => new WebSocketServer({ noServer: true }));
-  if (!wsServer)
-    return logger.error('[next-ws] was not able to find the WebSocket server');
-  process.env.NEXT_WS_SKIP_ENVIRONMENT_CHECK = String(0);
 
   logger.ready('[next-ws] has started the WebSocket server');
 
   // Prevent double-attaching
-  const kInstalled = Symbol.for('kInstalled');
+  const kInstalled = Symbol.for('next-ws.http-server.attached');
   if (Reflect.has(httpServer, kInstalled)) return;
   Reflect.set(httpServer, kInstalled, true);