diff --git a/src/node/routes/health.ts b/src/node/routes/health.ts index f38bb0ab..faf1f9c7 100644 --- a/src/node/routes/health.ts +++ b/src/node/routes/health.ts @@ -13,8 +13,8 @@ router.get("/", (req, res) => { export const wsRouter = WsRouter() wsRouter.ws("/", async (req) => { - wss.handleUpgrade(req, req.socket, req.head, (ws) => { - ws.on("message", () => { + wss.handleUpgrade(req, req.ws, req.head, (ws) => { + ws.addEventListener("message", () => { ws.send( JSON.stringify({ event: "health", @@ -23,5 +23,6 @@ wsRouter.ws("/", async (req) => { }), ) }) + req.ws.resume() }) }) diff --git a/test/health.test.ts b/test/health.test.ts new file mode 100644 index 00000000..4eae9c60 --- /dev/null +++ b/test/health.test.ts @@ -0,0 +1,40 @@ +import * as httpserver from "./httpserver" +import * as integration from "./integration" + +describe("health", () => { + let codeServer: httpserver.HttpServer | undefined + + afterEach(async () => { + if (codeServer) { + await codeServer.close() + codeServer = undefined + } + }) + + it("/healthz", async () => { + ;[, , codeServer] = await integration.setup(["--auth=none"], "") + const resp = await codeServer.fetch("/healthz") + expect(resp.status).toBe(200) + const json = await resp.json() + expect(json).toStrictEqual({ lastHeartbeat: 0, status: "expired" }) + }) + + it("/healthz (websocket)", async () => { + ;[, , codeServer] = await integration.setup(["--auth=none"], "") + const ws = codeServer.ws("/healthz") + const message = await new Promise((resolve, reject) => { + ws.on("error", console.error) + ws.on("message", (message) => { + try { + const j = JSON.parse(message.toString()) + resolve(j) + } catch (error) { + reject(error) + } + }) + ws.on("open", () => ws.send(JSON.stringify({ event: "health" }))) + }) + ws.terminate() + expect(message).toStrictEqual({ event: "health", status: "expired", lastHeartbeat: 0 }) + }) +}) diff --git a/test/httpserver.ts b/test/httpserver.ts index 4fe54f88..399ccda1 100644 --- a/test/httpserver.ts +++ b/test/httpserver.ts @@ -1,5 +1,6 @@ import * as express from "express" import * as http from "http" +import * as net from "net" import * as nodeFetch from "node-fetch" import Websocket from "ws" import * as util from "../src/common/util" @@ -8,13 +9,21 @@ import { handleUpgrade } from "../src/node/wsRouter" // Perhaps an abstraction similar to this should be used in app.ts as well. export class HttpServer { - private hs = http.createServer() + private readonly sockets = new Set() + private cleanupTimeout?: NodeJS.Timeout - public constructor(hs?: http.Server) { - // See usage in test/integration.ts - if (hs) { - this.hs = hs - } + // See usage in test/integration.ts + public constructor(private readonly hs = http.createServer()) { + this.hs.on("connection", (socket) => { + this.sockets.add(socket) + socket.on("close", () => { + this.sockets.delete(socket) + if (this.cleanupTimeout && this.sockets.size === 0) { + clearTimeout(this.cleanupTimeout) + this.cleanupTimeout = undefined + } + }) + }) } /** @@ -54,6 +63,8 @@ export class HttpServer { */ public close(): Promise { return new Promise((res, rej) => { + // Close will not actually close anything; it just waits until everything + // is closed. this.hs.close((err) => { if (err) { rej(err) @@ -61,6 +72,19 @@ export class HttpServer { } res() }) + + // If there are sockets remaining we might need to force close them or + // this promise might never resolve. + if (this.sockets.size > 0) { + // Give sockets a chance to close up shop. + this.cleanupTimeout = setTimeout(() => { + this.cleanupTimeout = undefined + for (const socket of this.sockets.values()) { + console.warn("a socket was left hanging") + socket.destroy() + } + }, 1000) + } }) } diff --git a/test/test-plugin/src/index.ts b/test/test-plugin/src/index.ts index 592ad372..772b59d8 100644 --- a/test/test-plugin/src/index.ts +++ b/test/test-plugin/src/index.ts @@ -28,7 +28,8 @@ export const plugin: cs.Plugin = { wsRouter() { const wr = cs.WsRouter() wr.ws("/test-app", (req) => { - cs.wss.handleUpgrade(req, req.socket, req.head, (ws) => { + cs.wss.handleUpgrade(req, req.ws, req.head, (ws) => { + req.ws.resume() ws.send("hello") }) })