From b9ec7dc5a4712a466a2ee0839a1e6e6de267fd13 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E3=81=BE=E3=81=A3=E3=81=A1=E3=82=83=E3=81=A8=E3=83=BC?=
 =?UTF-8?q?=E3=81=AB=E3=82=85?=
 <17376330+u1-liquid@users.noreply.github.com>
Date: Fri, 25 Aug 2023 15:55:25 +0900
Subject: [PATCH] =?UTF-8?q?enhance:=20=E3=82=B7=E3=83=A3=E3=83=83=E3=83=88?=
 =?UTF-8?q?=E3=83=80=E3=82=A6=E3=83=B3=E6=99=82=E3=80=81DB=E6=8E=A5?=
 =?UTF-8?q?=E7=B6=9A=E3=81=8C=E5=88=87=E3=82=8C=E3=81=9F=E3=82=89=E7=A2=BA?=
 =?UTF-8?q?=E5=AE=9F=E3=81=AB=E7=B5=82=E4=BA=86=E3=81=95=E3=81=9B=E3=82=8B?=
 =?UTF-8?q?=20(MisskeyIO#159)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 packages/backend/src/GlobalModule.ts |  5 ++++
 packages/backend/src/boot/index.ts   | 43 ++++++++++++++++------------
 2 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/packages/backend/src/GlobalModule.ts b/packages/backend/src/GlobalModule.ts
index 9f1ee9fcaa..666461c1bb 100644
--- a/packages/backend/src/GlobalModule.ts
+++ b/packages/backend/src/GlobalModule.ts
@@ -4,6 +4,7 @@
  */
 
 import { setTimeout } from 'node:timers/promises';
+import process from 'node:process';
 import { Global, Inject, Module } from '@nestjs/common';
 import * as Redis from 'ioredis';
 import { DataSource } from 'typeorm';
@@ -103,5 +104,9 @@ export class GlobalModule implements OnApplicationShutdown {
 
 	async onApplicationShutdown(signal: string): Promise<void> {
 		await this.dispose();
+		process.emitWarning('Misskey is shutting down', {
+			code: 'MISSKEY_SHUTDOWN',
+			detail: `Application received ${signal} signal`,
+		});
 	}
 }
diff --git a/packages/backend/src/boot/index.ts b/packages/backend/src/boot/index.ts
index 2b1f609e9b..28bcdb5ff4 100644
--- a/packages/backend/src/boot/index.ts
+++ b/packages/backend/src/boot/index.ts
@@ -9,6 +9,7 @@
 
 import cluster from 'node:cluster';
 import { EventEmitter } from 'node:events';
+import process from 'node:process';
 import chalk from 'chalk';
 import Xev from 'xev';
 import Logger from '@/logger.js';
@@ -29,6 +30,8 @@ const ev = new Xev();
 
 //#region Events
 
+let isShuttingDown = false;
+
 if (cluster.isPrimary && !envOption.disableClustering) {
 	// Listen new workers
 	cluster.on('fork', worker => {
@@ -41,25 +44,22 @@ if (cluster.isPrimary && !envOption.disableClustering) {
 	});
 
 	// Listen for dying workers
-	cluster.on('exit', (worker, code, signal?) => {
+	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`));
-		}
+		clusterLogger.error(chalk.red(`[${worker.id}] died (${signal || code})`));
+		if (!isShuttingDown) cluster.fork();
+		else clusterLogger.info(chalk.yellow('Worker respawn disabled because of shutdown'));
+	});
+
+	process.on('SIGINT', () => {
+		logger.warn(chalk.yellow('Process received SIGINT'));
+		isShuttingDown = true;
+	});
+
+	process.on('SIGTERM', () => {
+		logger.warn(chalk.yellow('Process received SIGTERM'));
+		isShuttingDown = true;
 	});
 }
 
@@ -78,7 +78,14 @@ process.on('uncaughtException', err => {
 
 // Dying away...
 process.on('exit', code => {
-	logger.info(`The process is going to exit with code ${code}`);
+	logger.warn(chalk.yellow(`The process is going to exit with code ${code}`));
+});
+
+process.on('warning', warning => {
+	if ((warning as never)['code'] !== 'MISSKEY_SHUTDOWN') return;
+	logger.warn(chalk.yellow(`${warning.message}: ${(warning as never)['detail']}`));
+	for (const id in cluster.workers) cluster.workers[id]?.process.kill('SIGTERM');
+	process.exit();
 });
 
 //#endregion