implement worker pool for high-concurrency NSFW detection

Author: pratham-developer

nsfwDetection/nsfwProcess.js

@@ -1,41 +1,90 @@
-// nsfwDetection/nsfwProcess.js
 import { fork } from 'child_process';
 import path from 'path';
+import os from 'os';
 
 const nsfwServicePath = path.resolve('nsfwDetection/nsfwService.cjs');
-const nsfwProcess = fork(nsfwServicePath, [], { stdio: ['pipe', 'pipe', 'pipe', 'ipc'] });
+const WORKER_COUNT = os.cpus().length; // Use number of CPU cores
+const WORKER_TIMEOUT = 5000; // 5 seconds per request
 
-let ready = false;
-const pending = new Map();
+let workers = [];
+let readyWorkers = new Set();
+let pending = new Map(); // Map: id -> { resolve, reject, timeout, workerIdx }
+let rrIndex = 0; // Round-robin index
 
-nsfwProcess.on('message', (msg) => {
-  if (msg.type === 'ready') ready = true;
-  if (msg.type === 'result' && pending.has(msg.id)) {
-    // Log predictions from the child process
-    if (msg.predictions) {
-      console.log('NSFW Predictions:', msg.predictions);
+function spawnWorker(idx) {
+  const worker = fork(nsfwServicePath, [], { stdio: ['pipe', 'pipe', 'pipe', 'ipc'] });
+  worker.on('message', (msg) => {
+    if (msg.type === 'ready') {
+      readyWorkers.add(idx);
     }
-    pending.get(msg.id)(msg.safe);
-    pending.delete(msg.id);
-  }
-  if (msg.type === 'error' && pending.has(msg.id)) {
-    pending.get(msg.id)(false, msg.error);
-    pending.delete(msg.id);
-  }
-});
+    if (msg.type === 'result' && pending.has(msg.id)) {
+      const { resolve, timeout } = pending.get(msg.id);
+      clearTimeout(timeout);
+      resolve(msg.safe);
+      pending.delete(msg.id);
+    }
+    if (msg.type === 'error' && pending.has(msg.id)) {
+      const { reject, timeout } = pending.get(msg.id);
+      clearTimeout(timeout);
+      reject(new Error(msg.error));
+      pending.delete(msg.id);
+    }
+  });
+  worker.on('exit', () => {
+    readyWorkers.delete(idx);
+    // Clean up any pending tasks for this worker
+    for (const [id, task] of pending.entries()) {
+      if (task.workerIdx === idx) {
+        clearTimeout(task.timeout);
+        task.reject(new Error('NSFW worker crashed'));
+        pending.delete(id);
+      }
+    }
+    // Restart worker
+    setTimeout(() => {
+      workers[idx] = spawnWorker(idx);
+    }, 2000);
+  });
+  worker.on('error', () => {
+    readyWorkers.delete(idx);
+  });
+  return worker;
+}
+
+// Initialize worker pool
+for (let i = 0; i < WORKER_COUNT; i++) {
+  workers.push(spawnWorker(i));
+}
 
 export function isNsfwReady() {
-  return ready;
+  return readyWorkers.size > 0;
 }
 
 export function checkNsfw(imagePath) {
   return new Promise((resolve, reject) => {
-    if (!ready) return reject(new Error('NSFW service not ready'));
+    if (readyWorkers.size === 0) return reject(new Error('NSFW service not ready'));
+    // Pick next ready worker in round-robin
+    let tries = 0;
+    let workerIdx = rrIndex;
+    while (!readyWorkers.has(workerIdx) && tries < WORKER_COUNT) {
+      workerIdx = (workerIdx + 1) % WORKER_COUNT;
+      tries++;
+    }
+    if (!readyWorkers.has(workerIdx)) {
+      return reject(new Error('No NSFW workers ready'));
+    }
+    rrIndex = (workerIdx + 1) % WORKER_COUNT;
     const id = Date.now() + Math.random();
-    nsfwProcess.send({ type: 'check', imagePath, id });
-    pending.set(id, (safe, error) => {
-      if (error) return reject(new Error(error));
-      resolve(safe);
-    });
+    try {
+      workers[workerIdx].send({ type: 'check', imagePath, id });
+      const timeout = setTimeout(() => {
+        pending.delete(id);
+        reject(new Error('NSFW check timed out'));
+      }, WORKER_TIMEOUT);
+      pending.set(id, { resolve, reject, timeout, workerIdx });
+    } catch (err) {
+      pending.delete(id);
+      reject(new Error('Failed to send image to NSFW service'));
+    }
   });
-}
+}

nsfwDetection/nsfwService.cjs

@@ -1,49 +1,39 @@
-// nsfwDetection/nsfwService.cjs
 const nsfw = require('nsfwjs');
 const tf = require('@tensorflow/tfjs-node');
 const fs = require('fs');
 
 let model;
 
 function isUnsafe(predictions) {
-  // Initialize probabilities with fallbacks
   const prob = { Porn: 0, Hentai: 0, Sexy: 0, Neutral: 0, Drawing: 0 };
   predictions.forEach(p => prob[p.className] = p.probability);
 
-  // Pre-calculate frequently used values
   const maxNsfw = Math.max(prob.Porn, prob.Hentai, prob.Sexy);
   const maxSafe = Math.max(prob.Neutral, prob.Drawing);
   const topPrediction = predictions.reduce((a, b) =>
     a.probability > b.probability ? a : b
   );
 
-  // 1. Explicit content threshold
   if (prob.Porn > 0.35 || prob.Hentai > 0.35) return true;
-
-  // 2. Weighted NSFW index
   if ((0.5 * prob.Porn + 0.3 * prob.Sexy + 0.2 * prob.Hentai) > 0.4) return true;
-
-  // 3. Relative risk ratio (with uncertainty buffer)
   if (maxNsfw >= 0.2 && maxSafe > 0 && maxNsfw / maxSafe > 2.5) return true;
-
-  // 4. Top class enforcement
-  if (['Porn', 'Hentai'].includes(topPrediction.className) && topPrediction.probability > 0.25) {
-    return true;
-  }
-
-  // 5. Sexy content dominance
+  if (['Porn', 'Hentai'].includes(topPrediction.className) && topPrediction.probability > 0.25) return true;
   if (prob.Sexy > 0.6 && prob.Sexy > (prob.Neutral + prob.Drawing)) return true;
 
-  // 6. [NEW] Uncertainty check - reject ambiguous predictions
   const entropy = -Object.values(prob).reduce((sum, p) => sum + (p * Math.log2(p || 1e-10)), 0);
-  if (entropy > 0.8 && maxSafe < 0.4) return true; // High uncertainty + low safe confidence
+  if (entropy > 0.8 && maxSafe < 0.4) return true;
 
   return false;
 }
 
 async function loadModel() {
-  model = await nsfw.load();
-  process.send && process.send({ type: 'ready' });
+  try {
+    model = await nsfw.load();
+    process.send && process.send({ type: 'ready' });
+  } catch (err) {
+    console.error('Model loading failed:', err);
+    process.exit(1);
+  }
 }
 
 loadModel();
@@ -62,4 +52,15 @@ process.on('message', async (msg) => {
       process.send({ type: 'error', error: err.message, id: msg.id });
     }
   }
-});
+});
+
+// Handle unexpected crashes
+process.on('uncaughtException', (err) => {
+  console.error('Uncaught Exception:', err);
+  process.exit(1);
+});
+
+process.on('unhandledRejection', (reason) => {
+  console.error('Unhandled Rejection:', reason);
+  process.exit(1);
+});