refactor(websocket): streamline WebSocket server setup and client instantiation

- Removed the request validation logic from the WebSocket connection handler.
- Added a cleanup function to close the WebSocket server.
- Introduced a singleton pattern for the WebSocket client to manage connections more efficiently.
This commit is contained in:
Mauricio Siu
2025-04-04 01:55:29 -06:00
parent d43b098a7a
commit 36172491a4
2 changed files with 26 additions and 17 deletions

View File

@@ -1,5 +1,4 @@
import type http from "node:http"; import type http from "node:http";
import { validateRequest } from "@dokploy/server/index";
import { applyWSSHandler } from "@trpc/server/adapters/ws"; import { applyWSSHandler } from "@trpc/server/adapters/ws";
import { WebSocketServer } from "ws"; import { WebSocketServer } from "ws";
import { appRouter } from "../api/root"; import { appRouter } from "../api/root";
@@ -13,11 +12,13 @@ export const setupDrawerLogsWebSocketServer = (
path: "/drawer-logs", path: "/drawer-logs",
}); });
// Set up tRPC WebSocket handler
applyWSSHandler({ applyWSSHandler({
wss: wssTerm, wss: wssTerm,
router: appRouter, router: appRouter,
createContext: createTRPCContext as any, createContext: createTRPCContext as any,
}); });
server.on("upgrade", (req, socket, head) => { server.on("upgrade", (req, socket, head) => {
const { pathname } = new URL(req.url || "", `http://${req.headers.host}`); const { pathname } = new URL(req.url || "", `http://${req.headers.host}`);
@@ -31,13 +32,8 @@ export const setupDrawerLogsWebSocketServer = (
} }
}); });
wssTerm.on("connection", async (ws, req) => { // Return cleanup function
const _url = new URL(req.url || "", `http://${req.headers.host}`); return () => {
const { user, session } = await validateRequest(req); wssTerm.close();
};
if (!user || !session) {
ws.close();
return;
}
});
}; };

View File

@@ -27,15 +27,28 @@ const getWsUrl = () => {
const protocol = window.location.protocol === "https:" ? "wss:" : "ws:"; const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
const host = window.location.host; const host = window.location.host;
// Use the base URL for all tRPC WebSocket connections
return `${protocol}${host}/drawer-logs`; return `${protocol}${host}/drawer-logs`;
}; };
const wsClient = // Singleton WebSocket client instance
typeof window !== "undefined" let wsClientInstance: ReturnType<typeof createWSClient> | null = null;
? createWSClient({
url: getWsUrl() || "", const getWsClient = () => {
}) if (typeof window === "undefined") return null;
: null;
if (!wsClientInstance) {
wsClientInstance = createWSClient({
url: getWsUrl() || "",
onClose: () => {
// Reset the instance when connection closes so it can be recreated
wsClientInstance = null;
},
});
}
return wsClientInstance;
};
/** A set of type-safe react-query hooks for your tRPC API. */ /** A set of type-safe react-query hooks for your tRPC API. */
export const api = createTRPCNext<AppRouter>({ export const api = createTRPCNext<AppRouter>({
@@ -57,7 +70,7 @@ export const api = createTRPCNext<AppRouter>({
splitLink({ splitLink({
condition: (op) => op.type === "subscription", condition: (op) => op.type === "subscription",
true: wsLink({ true: wsLink({
client: wsClient!, client: getWsClient()!,
}), }),
false: splitLink({ false: splitLink({
condition: (op) => op.input instanceof FormData, condition: (op) => op.input instanceof FormData,