From b13db3124bfd53c3b44a90e1af4175e77e64aa10 Mon Sep 17 00:00:00 2001
From: Asher <ash@coder.com>
Date: Thu, 28 Jan 2021 12:47:50 -0600
Subject: [PATCH] Add health websocket

This is used by some of our services.
---
 src/node/routes/health.ts | 17 +++++++++++++++++
 src/node/routes/index.ts  |  1 +
 src/node/wsRouter.ts      |  3 +++
 3 files changed, 21 insertions(+)

diff --git a/src/node/routes/health.ts b/src/node/routes/health.ts
index 20dab71a..f38bb0ab 100644
--- a/src/node/routes/health.ts
+++ b/src/node/routes/health.ts
@@ -1,4 +1,5 @@
 import { Router } from "express"
+import { wss, Router as WsRouter } from "../wsRouter"
 
 export const router = Router()
 
@@ -8,3 +9,19 @@ router.get("/", (req, res) => {
     lastHeartbeat: req.heart.lastHeartbeat,
   })
 })
+
+export const wsRouter = WsRouter()
+
+wsRouter.ws("/", async (req) => {
+  wss.handleUpgrade(req, req.socket, req.head, (ws) => {
+    ws.on("message", () => {
+      ws.send(
+        JSON.stringify({
+          event: "health",
+          status: req.heart.alive() ? "alive" : "expired",
+          lastHeartbeat: req.heart.lastHeartbeat,
+        }),
+      )
+    })
+  })
+})
diff --git a/src/node/routes/index.ts b/src/node/routes/index.ts
index 1c4483ff..5a84ff6e 100644
--- a/src/node/routes/index.ts
+++ b/src/node/routes/index.ts
@@ -133,6 +133,7 @@ export const register = async (
   wsApp.use("/vscode", vscode.wsRouter.router)
 
   app.use("/healthz", health.router)
+  wsApp.use("/healthz", health.wsRouter.router)
 
   if (args.auth === AuthType.Password) {
     app.use("/login", login.router)
diff --git a/src/node/wsRouter.ts b/src/node/wsRouter.ts
index 1ff13aaf..d829d082 100644
--- a/src/node/wsRouter.ts
+++ b/src/node/wsRouter.ts
@@ -1,6 +1,7 @@
 import * as express from "express"
 import * as expressCore from "express-serve-static-core"
 import * as http from "http"
+import Websocket from "ws"
 import * as pluginapi from "../../typings/pluginapi"
 
 export const handleUpgrade = (app: express.Express, server: http.Server): void => {
@@ -48,3 +49,5 @@ export class WebsocketRouter {
 export function Router(): WebsocketRouter {
   return new WebsocketRouter()
 }
+
+export const wss = new Websocket.Server({ noServer: true })