From b2be2fc8411e22fa5e811a8098ab6ea20c0ce777 Mon Sep 17 00:00:00 2001
From: zkldi <20380519+zkldi@users.noreply.github.com>
Date: Fri, 12 Nov 2021 03:39:17 +0000
Subject: [PATCH] Add graceful sigterm handling to workers

---
 server/src/lib/handlers/sigterm.ts           | 35 ++++++++++++++++++++
 server/src/lib/score-import/worker/worker.ts |  3 ++
 server/src/main.ts                           | 20 ++---------
 3 files changed, 40 insertions(+), 18 deletions(-)
 create mode 100644 server/src/lib/handlers/sigterm.ts

diff --git a/server/src/lib/handlers/sigterm.ts b/server/src/lib/handlers/sigterm.ts
new file mode 100644
index 000000000..84b3f235a
--- /dev/null
+++ b/server/src/lib/handlers/sigterm.ts
@@ -0,0 +1,35 @@
+import { monkDB } from "external/mongo/db";
+import { CloseRedisConnection } from "external/redis/redis";
+import http from "http";
+import https from "https";
+import { rootLogger } from "lib/logger/logger";
+import { CloseScoreImportQueue } from "lib/score-import/worker/queue";
+
+const logger = rootLogger;
+
+export function HandleSIGTERMGracefully(instance?: http.Server | https.Server) {
+	logger.info("SIGTERM Received, closing program.", { shutdownInfo: true });
+
+	if (instance) {
+		instance.close(() => {
+			CloseEverythingElse();
+		});
+	} else {
+		CloseEverythingElse();
+	}
+}
+
+async function CloseEverythingElse() {
+	logger.info("Closing Mongo Database.", { shutdownInfo: true });
+	await monkDB.close();
+
+	logger.info("Closing Redis Connection.", { shutdownInfo: true });
+	CloseRedisConnection();
+
+	logger.info("Closing Score Import Queue.", { shutdownInfo: true });
+	CloseScoreImportQueue();
+
+	logger.info("Everything closed. Waiting for process to exit naturally.", {
+		shutdownInfo: true,
+	});
+}
diff --git a/server/src/lib/score-import/worker/worker.ts b/server/src/lib/score-import/worker/worker.ts
index acb730512..74c152e1a 100644
--- a/server/src/lib/score-import/worker/worker.ts
+++ b/server/src/lib/score-import/worker/worker.ts
@@ -12,6 +12,7 @@
 // Explicitly set this before importing anything!
 process.env.IS_SCORE_WORKER_SERVER = "true";
 
+import { HandleSIGTERMGracefully } from "lib/handlers/sigterm";
 import CreateLogCtx from "lib/logger/logger";
 import { ImportTypes } from "tachi-common";
 import { FormatUserDoc, GetUserWithID } from "utils/user";
@@ -67,3 +68,5 @@ ScoreImportQueue.process(async <I extends ImportTypes>(job: ScoreImportJob<I>) =
 
 	return importDocument;
 });
+
+process.on("SIGTERM", HandleSIGTERMGracefully);
diff --git a/server/src/main.ts b/server/src/main.ts
index 6bf04750b..bd84c2fbd 100644
--- a/server/src/main.ts
+++ b/server/src/main.ts
@@ -19,6 +19,7 @@ import fetch from "utils/fetch";
 import http from "http";
 import { CloseRedisConnection } from "external/redis/redis";
 import { CloseScoreImportQueue } from "lib/score-import/worker/queue";
+import { HandleSIGTERMGracefully } from "lib/handlers/sigterm";
 
 const logger = CreateLogCtx(__filename);
 
@@ -74,24 +75,7 @@ if (ServerConfig.ENABLE_SERVER_HTTPS) {
 	logger.info(`HTTP Listening on port ${Environment.port}`, { bootInfo: true });
 }
 
-process.on("SIGTERM", () => {
-	logger.info("SIGTERM Received, closing program.", { shutdownInfo: true });
-
-	instance.close(async () => {
-		logger.info("Closing Mongo Database.", { shutdownInfo: true });
-		await monkDB.close();
-
-		logger.info("Closing Redis Connection.", { shutdownInfo: true });
-		CloseRedisConnection();
-
-		logger.info("Closing Score Import Queue.", { shutdownInfo: true });
-		CloseScoreImportQueue();
-
-		logger.info("Everything closed. Waiting for process to exit naturally.", {
-			shutdownInfo: true,
-		});
-	});
-});
+process.on("SIGTERM", HandleSIGTERMGracefully);
 
 if (process.env.INVOKE_JOB_RUNNER) {
 	logger.info(`Spawning a tachi-server job runner inline.`, { bootInfo: true });