From 0a2d0818e797e18cb2265ba83da35b81743c7de0 Mon Sep 17 00:00:00 2001
From: "Navy.gif" <navydotgif@gmail.com>
Date: Fri, 10 Feb 2023 23:56:26 +0200
Subject: [PATCH] handle sigterm for graceful shutdown

---
 src/controller/Controller.js | 12 ++++++++++++
 src/controller/Shard.js      | 20 ++++++++++++++++++--
 src/server/Server.js         | 19 +++++++++++--------
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/src/controller/Controller.js b/src/controller/Controller.js
index a60a811..4be2b56 100644
--- a/src/controller/Controller.js
+++ b/src/controller/Controller.js
@@ -38,6 +38,8 @@ class Controller extends EventEmitter {
             this.logger.error(ex.stack);
         });
 
+        process.on('SIGINT', this.#shutdown.bind(this));
+
         this.on('built', () => {
             this.logger.info(`API Controller built`);
             this._built = true;
@@ -90,6 +92,16 @@ class Controller extends EventEmitter {
         shard.on('message', (msg) => this._handleMessage(shard, msg));
     }
 
+    async #shutdown () {
+        this.logger.info('Received SIGINT, shutting down');
+        const promises = this.shards.map(shard => shard.awaitShutdown());
+        await Promise.all(promises);
+        this.logger.info(`Shutdown completed, goodbye`);
+
+        // eslint-disable-next-line no-process-exit
+        process.exit();
+    }
+
 }
 
 module.exports = Controller;
\ No newline at end of file
diff --git a/src/controller/Shard.js b/src/controller/Shard.js
index 02261e7..7c6699f 100644
--- a/src/controller/Shard.js
+++ b/src/controller/Shard.js
@@ -2,7 +2,8 @@ const ChildProcess = require('node:child_process');
 const { EventEmitter } = require('node:events');
 
 const { Util } = require('../util');
-
+// Give subprocess 90s to shut down before being forcibly killed
+const KillTO = 90 * 1000;
 class Shard extends EventEmitter {
 
     constructor (controller, id, options = {}) {
@@ -31,6 +32,8 @@ class Shard extends EventEmitter {
         // Set in the spawn method
         this.spawnedAt = null;
 
+        this._awaitingShutdown = null;
+
     }
 
     async spawn (waitForReady = false) {
@@ -85,7 +88,7 @@ class Shard extends EventEmitter {
                 const to = setTimeout(() => {
                     this.process.kill();
                     resolve();
-                }, 5000);
+                }, KillTO);
                 // Gracefully handle exit
                 this.process.once('exit', (code, signal) => {
                     clearTimeout(to);
@@ -118,6 +121,13 @@ class Shard extends EventEmitter {
         });
     }
 
+    awaitShutdown () {
+        this._respawn = false;
+        return new Promise((resolve) => {
+            this._awaitingShutdown = resolve;
+        });
+    }
+
     _handleMessage (message) {
         if (message) {
             if (message._ready) {
@@ -125,6 +135,10 @@ class Shard extends EventEmitter {
                 this.emit('ready');
                 return;
             } else if (message._shutdown) {
+                setTimeout(() => {
+                    if (this.process) 
+                        this.process.kill('SIGKILL');
+                }, KillTO);
                 this.ready = false;
                 this.emit('shutdown');
                 return;
@@ -150,6 +164,8 @@ class Shard extends EventEmitter {
             this.process.removeAllListeners();
         
         this.emit('death');
+        if (this._awaitingShutdown) 
+            this._awaitingShutdown();
 
         if (code !== 0) 
             this.crashes.push(Date.now() - this.spawnedAt);
diff --git a/src/server/Server.js b/src/server/Server.js
index 6b4236a..db1d4fa 100644
--- a/src/server/Server.js
+++ b/src/server/Server.js
@@ -155,6 +155,7 @@ class Server extends EventEmitter {
         this.app.use(this.#ready.bind(this)); // denies requests before the server is ready
 
         process.on('message', this._handleMessage.bind(this));
+        process.on('SIGINT', this.shutdown.bind(this));
 
     }
 
@@ -209,17 +210,19 @@ class Server extends EventEmitter {
 
     async shutdown () {
         this.logger.info(`Received shutdown command, initiating graceful shutdown`);
+        // Tells the manager the shard is shutting down, will set a timeout to kill the shard if it hangs
         process.send({ _shutdown: true });
 
-        this._ready = false; // stops any new connections
+        this._ready = false;
+        this.server.close(async () => {
+            await this.mongodb.close();
+            await this.mariadb.close();
+            this.logger.status('DB shutdowns complete.');
 
-        await this.mongodb.close();
-        await this.mariadb.close();
-        this.logger.status('DB shutdowns complete.');
-
-        this.logger.status('Shutdown complete. Goodbye');
-        // eslint-disable-next-line no-process-exit
-        process.exit();
+            this.logger.status('Shutdown complete. Goodbye');
+            // eslint-disable-next-line no-process-exit
+            process.exit();
+        });
 
     }