From 6f2f2c7c092970f20357a1846de6ea15b9e91bca Mon Sep 17 00:00:00 2001
From: crosstyan <crosstyan@outlook.com>
Date: Wed, 24 Jun 2026 15:44:15 +0800
Subject: [PATCH] Fix shutdown with active WebSocket clients

---
 server.js | 48 ++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 44 insertions(+), 4 deletions(-)

diff --git a/server.js b/server.js
index 555d077..ef0be3e 100644
--- a/server.js
+++ b/server.js
@@ -11,6 +11,7 @@ const PORT = Number.parseInt(process.env.PORT || '3000', 10);
 const HOST = process.env.HOST || '127.0.0.1';
 const TAIL_LINES = Number.parseInt(process.env.TAIL_LINES || '200', 10);
 const WATCH_INTERVAL_MS = Number.parseInt(process.env.WATCH_INTERVAL_MS || '500', 10);
+const SHUTDOWN_TIMEOUT_MS = Number.parseInt(process.env.SHUTDOWN_TIMEOUT_MS || '3000', 10);
 const PUBLIC_DIR = path.join(__dirname, 'public');
 const LOG_FILE = path.resolve(process.env.LOG_FILE || process.argv[2] || 'sample.log');
 
@@ -275,14 +276,53 @@ async function main() {
     console.log(`Streaming ${LOG_FILE}`);
   });
 
-  function shutdown() {
+  let shuttingDown = false;
+
+  function shutdown(signal) {
+    if (shuttingDown) {
+      console.error(`${signal} received again; forcing shutdown.`);
+      process.exit(1);
+    }
+
+    shuttingDown = true;
+    console.log(`${signal} received; shutting down.`);
     fs.unwatchFile(LOG_FILE);
+
+    const shutdownTimer = setTimeout(() => {
+      console.error(`Shutdown exceeded ${SHUTDOWN_TIMEOUT_MS}ms; forcing exit.`);
+      process.exit(1);
+    }, SHUTDOWN_TIMEOUT_MS);
+    shutdownTimer.unref();
+
+    const terminateTimer = setTimeout(() => {
+      for (const client of wss.clients) {
+        if (client.readyState !== WebSocket.CLOSED) {
+          client.terminate();
+        }
+      }
+    }, Math.min(500, SHUTDOWN_TIMEOUT_MS));
+    terminateTimer.unref();
+
+    for (const client of wss.clients) {
+      client.close(1001, 'Server shutting down');
+    }
+
     wss.close();
-    server.close(() => process.exit(0));
+
+    server.close((error) => {
+      clearTimeout(shutdownTimer);
+
+      if (error) {
+        console.error(error);
+        process.exit(1);
+      }
+
+      process.exit(0);
+    });
   }
 
-  process.on('SIGINT', shutdown);
-  process.on('SIGTERM', shutdown);
+  process.on('SIGINT', () => shutdown('SIGINT'));
+  process.on('SIGTERM', () => shutdown('SIGTERM'));
 }
 
 main().catch((error) => {