perf(backend): use mutex for nsfw model loading

MisskeyIO · Jul 5, 2023 · d4a9eba · d4a9eba
1 parent 3e93450
commit d4a9eba
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 78 deletions.
diff --git 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
@@ -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)
@@ -31,16 +33,22 @@ export class AiService {
  const cpuFlags = await this.getCpuFlags();
  isSupportedCpu = REQUIRED_CPU_FLAGS.every(required => cpuFlags.includes(required));
  }
- 
+
  if (!isSupportedCpu) {
  console.error('This CPU cannot use TensorFlow.');
  return null;
  }
- 
+
  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;
  try {