fix(dev): close upgraded sockets and resolve on worker shutdown timeout

src/core/dev-server/worker.ts

@@ -151,6 +151,7 @@ export class NodeDevWorker implements DevWorker {
           if (process.env.DEBUG) {
             consola.warn(`force closing dev worker...`);
           }
+          resolve();
         }, gracefulShutdownTimeoutSec * 1000);
 
         this.#worker?.on("message", (message) => {

src/presets/_nitro/runtime/nitro-dev.ts

@@ -6,6 +6,7 @@ import { trapUnhandledNodeErrors } from "nitropack/runtime/internal";
 import { startScheduleRunner } from "nitropack/runtime/internal";
 import { scheduledTasks, tasks } from "#nitro-internal-virtual/tasks";
 import { Server } from "node:http";
+import type { Duplex } from "node:stream";
 import { join } from "node:path";
 import nodeCrypto from "node:crypto";
 import { parentPort, threadId } from "node:worker_threads";
@@ -39,6 +40,7 @@ const nitroApp = useNitroApp();
 
 const server = new Server(toNodeListener(nitroApp.h3App));
 let listener: Server | undefined;
+const upgradedSockets = new Set<Duplex>();
 
 listen()
   .catch(() => listen(true /* use random port */))
@@ -51,7 +53,11 @@ listen()
 // https://crossws.unjs.io/adapters/node
 if (import.meta._websocket) {
   const { handleUpgrade } = wsAdapter(nitroApp.h3App.websocket);
-  server.on("upgrade", handleUpgrade);
+  server.on("upgrade", (req, socket, head) => {
+    upgradedSockets.add(socket);
+    socket.once("close", () => upgradedSockets.delete(socket));
+    return handleUpgrade(req, socket, head);
+  });
 }
 
 // Register tasks handlers
@@ -136,6 +142,9 @@ function getSocketAddress() {
 
 async function shutdown() {
   server.closeAllConnections?.();
+  for (const socket of upgradedSockets) {
+    socket.destroy();
+  }
   await Promise.all([
     new Promise((resolve) => listener?.close(resolve)),
     nitroApp.hooks.callHook("close").catch(console.error),

test/dev-shutdown/.gitignore

@@ -0,0 +1,2 @@
+.nitro
+.output

test/dev-shutdown/dev-shutdown.test.ts

@@ -0,0 +1,54 @@
+import { fileURLToPath } from "node:url";
+import { build, createDevServer, createNitro, prepare } from "../../src/core";
+import { afterAll, describe, expect, it, vi } from "vitest";
+
+vi.mock("std-env", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("std-env")>()),
+  isTest: false,
+  isCI: false,
+}));
+
+describe("nitro:dev:shutdown", () => {
+  const rootDir = fileURLToPath(new URL("fixture", import.meta.url));
+  let closeNitro: (() => Promise<void>) | undefined;
+
+  afterAll(() => closeNitro?.());
+
+  it("closes promptly with an open websocket connection", async () => {
+    const nitro = await createNitro({
+      rootDir,
+      dev: true,
+      preset: "nitro-dev",
+      buildDir: fileURLToPath(new URL(".nitro", import.meta.url)),
+      output: { dir: fileURLToPath(new URL(".output", import.meta.url)) },
+    });
+    closeNitro = () => nitro.close();
+    const devServer = createDevServer(nitro);
+    const server = await devServer.listen(0, { hostname: "127.0.0.1" });
+    await prepare(nitro);
+    const ready = new Promise<void>((resolve) => {
+      nitro.hooks.hook("dev:reload", () => resolve());
+    });
+    await build(nitro);
+    await ready;
+
+    const ws = new WebSocket(server.url.replace(/^http/, "ws") + "_ws");
+    const echo = await new Promise<string>((resolve, reject) => {
+      ws.addEventListener("open", () => ws.send("ping"));
+      ws.addEventListener("message", (e) => resolve(String(e.data)));
+      ws.addEventListener("error", reject);
+    });
+    expect(echo).toBe("ping");
+
+    const start = performance.now();
+    await Promise.race([
+      devServer.close(),
+      new Promise((_, reject) =>
+        setTimeout(() => reject(new Error("dev server close timed out")), 8000)
+      ),
+    ]);
+    expect(performance.now() - start).toBeLessThan(2000);
+    ws.close();
+    await server.close();
+  }, 30_000);
+});

test/dev-shutdown/fixture/nitro.config.ts

@@ -0,0 +1,7 @@
+import { defineNitroConfig } from "nitropack/config";
+
+export default defineNitroConfig({
+  experimental: {
+    websocket: true,
+  },
+});

test/dev-shutdown/fixture/routes/_ws.ts

@@ -0,0 +1,7 @@
+import { defineWebSocketHandler } from "h3";
+
+export default defineWebSocketHandler({
+  message(peer, message) {
+    peer.send(message.text());
+  },
+});