From daf429eb849d688031e07ea13422fa7e67a7b7fe Mon Sep 17 00:00:00 2001
From: riku6460 <17585784+riku6460@users.noreply.github.com>
Date: Tue, 22 Aug 2023 21:42:21 +0900
Subject: [PATCH] Graceful Shutdown (MisskeyIO#156)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---------

Co-authored-by: まっちゃとーにゅ <17376330+u1-liquid@users.noreply.github.com>
---
 Dockerfile                                    |  2 +-
 package.json                                  |  2 +
 packages/backend/src/boot/index.ts            | 46 +++++++++++++------
 packages/backend/src/server/ServerService.ts  |  2 +-
 .../server/api/StreamingApiServerService.ts   |  7 ++-
 5 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 45ca57bc21..787e248cb1 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -86,4 +86,4 @@ ENV MALLOC_CONF=background_thread:true,metadata_thp:auto,dirty_decay_ms:30000,mu
 ENV NODE_ENV=production
 HEALTHCHECK --interval=5s --retries=20 CMD ["/bin/bash", "/misskey/healthcheck.sh"]
 ENTRYPOINT ["/usr/bin/tini", "--"]
-CMD ["pnpm", "run", "migrateandstart"]
+CMD ["pnpm", "run", "migrateandstart:docker"]
diff --git a/package.json b/package.json
index b385fcaeaa..e531e1a76c 100644
--- a/package.json
+++ b/package.json
@@ -18,11 +18,13 @@
 		"build": "pnpm build-pre && pnpm -r build && pnpm gulp",
 		"build-storybook": "pnpm --filter frontend build-storybook",
 		"start": "pnpm check:connect && cd packages/backend && node ./built/boot/index.js",
+		"start:docker": "pnpm check:connect && cd packages/backend && exec node ./built/boot/index.js",
 		"start:test": "cd packages/backend && cross-env NODE_ENV=test node ./built/boot/index.js",
 		"init": "pnpm migrate",
 		"migrate": "cd packages/backend && pnpm migrate",
 		"check:connect": "cd packages/backend && pnpm check:connect",
 		"migrateandstart": "pnpm migrate && pnpm start",
+		"migrateandstart:docker": "pnpm migrate && exec pnpm start:docker",
 		"gulp": "pnpm exec gulp build",
 		"watch": "pnpm dev",
 		"dev": "node ./scripts/dev.mjs",
diff --git a/packages/backend/src/boot/index.ts b/packages/backend/src/boot/index.ts
index fc8fc2ffb4..2b1f609e9b 100644
--- a/packages/backend/src/boot/index.ts
+++ b/packages/backend/src/boot/index.ts
@@ -29,23 +29,39 @@ const ev = new Xev();
 
 //#region Events
 
-// Listen new workers
-cluster.on('fork', worker => {
-	clusterLogger.debug(`Process forked: [${worker.id}]`);
-});
+if (cluster.isPrimary && !envOption.disableClustering) {
+	// Listen new workers
+	cluster.on('fork', worker => {
+		clusterLogger.debug(`Process forked: [${worker.id}]`);
+	});
 
-// Listen online workers
-cluster.on('online', worker => {
-	clusterLogger.debug(`Process is now online: [${worker.id}]`);
-});
+	// Listen online workers
+	cluster.on('online', worker => {
+		clusterLogger.debug(`Process is now online: [${worker.id}]`);
+	});
 
-// Listen for dying workers
-cluster.on('exit', worker => {
-	// Replace the dead worker,
-	// we're not sentimental
-	clusterLogger.error(chalk.red(`[${worker.id}] died :(`));
-	cluster.fork();
-});
+	// Listen for dying workers
+	cluster.on('exit', (worker, code, signal?) => {
+		// Replace the dead worker,
+		// we're not sentimental
+		if (signal) {
+			switch (signal) {
+				case 'SIGINT':
+				case 'SIGTERM':
+					console.log(chalk.green(`[${worker.id}] exited by signal: ${signal}`));
+					break;
+				default:
+					console.error(chalk.red(`[${worker.id}] killed by signal: ${signal}`));
+					cluster.fork();
+					break;
+			}
+		} else if (code !== 0) {
+			console.error(chalk.red(`[${worker.id}] exited with error code: ${code}`));
+		} else {
+			console.log(chalk.green(`[${worker.id}] exited normally`));
+		}
+	});
+}
 
 // Display detail of unhandled promise rejection
 if (!envOption.quiet) {
diff --git a/packages/backend/src/server/ServerService.ts b/packages/backend/src/server/ServerService.ts
index 4a61ed5212..5719189c6b 100644
--- a/packages/backend/src/server/ServerService.ts
+++ b/packages/backend/src/server/ServerService.ts
@@ -254,7 +254,7 @@ export class ServerService implements OnApplicationShutdown {
 
 	@bindThis
 	public async dispose(): Promise<void> {
-		await this.streamingApiServerService.detach();
+		this.streamingApiServerService.detach();
 		await this.#fastify.close();
 	}
 
diff --git a/packages/backend/src/server/api/StreamingApiServerService.ts b/packages/backend/src/server/api/StreamingApiServerService.ts
index afa0bc5826..b8022dd765 100644
--- a/packages/backend/src/server/api/StreamingApiServerService.ts
+++ b/packages/backend/src/server/api/StreamingApiServerService.ts
@@ -173,13 +173,12 @@ export class StreamingApiServerService {
 	}
 
 	@bindThis
-	public detach(): Promise<void> {
+	public detach(): void {
 		if (this.#cleanConnectionsIntervalId) {
 			clearInterval(this.#cleanConnectionsIntervalId);
 			this.#cleanConnectionsIntervalId = null;
 		}
-		return new Promise((resolve) => {
-			this.#wss.close(() => resolve());
-		});
+		this.#wss.close();
+		this.#wss.clients.forEach(client => client.terminate());
 	}
 }