From 8daca59ca61dac8fd7f68f3135d69092c70e599f Mon Sep 17 00:00:00 2001
From: popkirby <popkirby@gmail.com>
Date: Sat, 8 Jul 2023 07:27:26 +0900
Subject: [PATCH] perf(backend): use mutex for nsfw model loading (#11109)

Co-authored-by: tamaina <tamaina@hotmail.co.jp>
---
 CHANGELOG.md                           |  1 +
 packages/backend/package.json          |  1 +
 packages/backend/src/core/AiService.ts | 10 +++++++++-
 pnpm-lock.yaml                         |  9 +++++++++
 4 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 288ab07963..434420522d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -33,6 +33,7 @@
 
 ### Server
 - JSON.parse の回数を削減することで、ストリーミングのパフォーマンスを向上しました
+- nsfwjs のモデルロードを排他することで、重複ロードによってメモリ使用量が増加しないように
 - 連合の配送ジョブのパフォーマンスを向上（ロック機構の見直し、Redisキャッシュの活用）
 
 ## 13.13.2
diff --git a/packages/backend/package.json b/packages/backend/package.json
index 288a8ea7bc..87f7f9ddf7 100644
--- a/packages/backend/package.json
+++ b/packages/backend/package.json
@@ -75,6 +75,7 @@
 		"accepts": "1.3.8",
 		"ajv": "8.12.0",
 		"archiver": "5.3.1",
+		"async-mutex": "^0.4.0",
 		"autwh": "0.1.0",
 		"bcryptjs": "2.4.3",
 		"blurhash": "2.0.5",
diff --git a/packages/backend/src/core/AiService.ts b/packages/backend/src/core/AiService.ts
index 02501b832b..c0596446dd 100644
--- a/packages/backend/src/core/AiService.ts
+++ b/packages/backend/src/core/AiService.ts
@@ -4,6 +4,7 @@ import { dirname } from 'node:path';
 import { Inject, Injectable } from '@nestjs/common';
 import * as nsfw from 'nsfwjs';
 import si from 'systeminformation';
+import { Mutex } from 'async-mutex';
 import type { Config } from '@/config.js';
 import { DI } from '@/di-symbols.js';
 import { bindThis } from '@/decorators.js';
@@ -17,6 +18,7 @@ let isSupportedCpu: undefined | boolean = undefined;
 @Injectable()
 export class AiService {
 	private model: nsfw.NSFWJS;
+	private modelLoadMutex: Mutex = new Mutex();
 
 	constructor(
 		@Inject(DI.config)
@@ -39,7 +41,13 @@ export class AiService {
 
 			const tf = await import('@tensorflow/tfjs-node');
 
-			if (this.model == null) this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 });
+			if (this.model == null) {
+				await this.modelLoadMutex.runExclusive(async () => {
+					if (this.model == null) {
+						this.model = await nsfw.load(`file://${_dirname}/../../nsfw-model/`, { size: 299 });
+					}
+				});
+			}
 
 			const buffer = await fs.promises.readFile(path);
 			const image = await tf.node.decodeImage(buffer, 3) as any;
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
index 69d25e807d..0d99cede2e 100644
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@@ -140,6 +140,9 @@ importers:
       archiver:
         specifier: 5.3.1
         version: 5.3.1
+      async-mutex:
+        specifier: ^0.4.0
+        version: 0.4.0
       autwh:
         specifier: 0.1.0
         version: 0.1.0
@@ -9468,6 +9471,12 @@ packages:
     resolution: {integrity: sha512-csOlWGAcRFJaI6m+F2WKdnMKr4HhdhFVBk0H/QbJFMCr+uO2kwohwXQPxw/9OCxp05r5ghVBFSyioixx3gfkNQ==}
     dev: true
 
+  /async-mutex@0.4.0:
+    resolution: {integrity: sha512-eJFZ1YhRR8UN8eBLoNzcDPcy/jqjsg6I1AP+KvWQX80BqOSW1oJPJXDylPUEeMr2ZQvHgnQ//Lp6f3RQ1zI7HA==}
+    dependencies:
+      tslib: 2.6.0
+    dev: false
+
   /async-settle@1.0.0:
     resolution: {integrity: sha512-VPXfB4Vk49z1LHHodrEQ6Xf7W4gg1w0dAPROHngx7qgDjqmIQ+fXmwgGXTW/ITLai0YLSvWepJOP9EVpMnEAcw==}
     engines: {node: '>= 0.10'}