diff --git a/src/core/dev-server/worker.ts b/src/core/dev-server/worker.ts index b763f7e0c6..ce6a82ca70 100644 --- a/src/core/dev-server/worker.ts +++ b/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) => { diff --git a/src/presets/_nitro/runtime/nitro-dev.ts b/src/presets/_nitro/runtime/nitro-dev.ts index 587c28e75a..2c948a7059 100644 --- a/src/presets/_nitro/runtime/nitro-dev.ts +++ b/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(); 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), diff --git a/test/dev-shutdown/.gitignore b/test/dev-shutdown/.gitignore new file mode 100644 index 0000000000..0b7102254a --- /dev/null +++ b/test/dev-shutdown/.gitignore @@ -0,0 +1,2 @@ +.nitro +.output diff --git a/test/dev-shutdown/dev-shutdown.test.ts b/test/dev-shutdown/dev-shutdown.test.ts new file mode 100644 index 0000000000..f330979342 --- /dev/null +++ b/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()), + isTest: false, + isCI: false, +})); + +describe("nitro:dev:shutdown", () => { + const rootDir = fileURLToPath(new URL("fixture", import.meta.url)); + let closeNitro: (() => Promise) | 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((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((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); +}); diff --git a/test/dev-shutdown/fixture/nitro.config.ts b/test/dev-shutdown/fixture/nitro.config.ts new file mode 100644 index 0000000000..959d2859cc --- /dev/null +++ b/test/dev-shutdown/fixture/nitro.config.ts @@ -0,0 +1,7 @@ +import { defineNitroConfig } from "nitropack/config"; + +export default defineNitroConfig({ + experimental: { + websocket: true, + }, +}); diff --git a/test/dev-shutdown/fixture/routes/_ws.ts b/test/dev-shutdown/fixture/routes/_ws.ts new file mode 100644 index 0000000000..1e8fd8f121 --- /dev/null +++ b/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()); + }, +});