From 4672f93a4c7ad4dcea482e7e05f12fde4f2879c9 Mon Sep 17 00:00:00 2001
From: Adam Sage <30732752+asagedev@users.noreply.github.com>
Date: Sun, 15 Sep 2024 15:13:25 -0400
Subject: [PATCH] Add websocket keepalive

---
 quartz/cli/handlers.js | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/quartz/cli/handlers.js b/quartz/cli/handlers.js
index 12e7e8ec9..95666d470 100644
--- a/quartz/cli/handlers.js
+++ b/quartz/cli/handlers.js
@@ -317,14 +317,10 @@ export async function handleBuild(argv) {
   }
 
   if (argv.serve) {
-    const connections = []
-    const clientRefresh = () => connections.forEach((conn) => conn.send("rebuild"))
-
     if (argv.baseDir !== "" && !argv.baseDir.startsWith("/")) {
       argv.baseDir = "/" + argv.baseDir
     }
 
-    await build(clientRefresh)
     const server = http.createServer(async (req, res) => {
       if (argv.baseDir && !req.url?.startsWith(argv.baseDir)) {
         console.log(
@@ -410,8 +406,28 @@ export async function handleBuild(argv) {
       return serve()
     })
     server.listen(argv.port)
+    function heartbeat() {
+        this.isAlive = true;
+    }
     const wss = new WebSocketServer({ port: argv.wsPort })
-    wss.on("connection", (ws) => connections.push(ws))
+    wss.on("connection", function connection(ws) {
+        ws.isAlive = true;
+        ws.on('error', console.error);
+        ws.on('pong', heartbeat);
+    })
+    const interval = setInterval(function ping() {
+        wss.clients.forEach(function each(ws) {
+            if (ws.isAlive === false) return ws.terminate();
+
+            ws.isAlive = false;
+            ws.ping();
+        });
+    }, 30000);
+    wss.on('close', function close() {
+        clearInterval(interval);
+    });
+    const clientRefresh = () => wss.clients.forEach((conn) => conn.send("rebuild"))
+    await build(clientRefresh)
     console.log(
       chalk.cyan(
         `Started a Quartz server listening at http://localhost:${argv.port}${argv.baseDir}`,