Merge pull request #57 from ServerlessLife/do-not-kill-node-worker-if-working

fix: Do not kill node worker if working

Author: ServerlessLife

src/nodeWorker.ts

@@ -5,7 +5,14 @@ import { Configuration } from './configuration.js';
 import { getModuleDirname, getProjectDirname } from './getDirname.js';
 import { Logger } from './logger.js';
 
-const workers = new Map<string, Worker>();
+interface MyWorker extends Worker {
+  used?: boolean;
+  toKill?: boolean;
+  onMessage?: (msg: any) => void;
+  onError?: (err: any) => void;
+}
+
+const workers = new Map<string, MyWorker>();
 
 /**
  * Run the function in a Node.js Worker Thread
@@ -22,7 +29,9 @@ async function runInWorker(input: {
   const func = await Configuration.getLambda(input.fuctionRequest.functionId);
 
   return new Promise<void>((resolve, reject) => {
-    let worker = workers.get(input.fuctionRequest.workerId);
+    let worker: MyWorker | undefined = workers.get(
+      input.fuctionRequest.workerId,
+    );
 
     if (!worker) {
       worker = startWorker({
@@ -33,31 +42,41 @@ async function runInWorker(input: {
         environment: input.environment,
         verbose: Configuration.config.verbose,
       });
+      worker.used = false;
+      worker.toKill = false;
     } else {
       Logger.verbose(
         `[Function ${input.fuctionRequest.functionId}] [Worker ${input.fuctionRequest.workerId}] Reusing worker`,
       );
     }
 
-    worker.on('message', (msg) => {
+    worker.onMessage = (msg) => {
       Logger.verbose(
         `[Function ${input.fuctionRequest.functionId}] [Worker ${input.fuctionRequest.workerId}] Worker message`,
         JSON.stringify(msg),
       );
+
+      worker.used = false;
       if (msg?.errorType) {
         reject(msg);
       } else {
         resolve(msg);
       }
-    });
-    worker.on('error', (err) => {
+
+      if (worker.toKill) {
+        worker.toKill = false;
+        void worker.terminate();
+      }
+    };
+    worker.onError = (err) => {
       Logger.error(
         `[Function ${input.fuctionRequest.functionId}] [Worker ${input.fuctionRequest.workerId}] Error`,
         err,
       );
       reject(err);
-    });
+    };
 
+    worker.used = true;
     worker.postMessage({
       env: input.fuctionRequest.env,
       event: input.fuctionRequest.event,
@@ -89,7 +108,7 @@ function startWorker(input: WorkerRequest) {
 
   const localProjectDir = getProjectDirname();
 
-  const worker = new Worker(
+  const worker: MyWorker = new Worker(
     path.resolve(path.join(getModuleDirname(), `./nodeWorkerRunner.mjs`)),
     {
       env: {
@@ -118,8 +137,16 @@ function startWorker(input: WorkerRequest) {
     );
     workers.delete(input.workerId);
   });
+
   workers.set(input.workerId, worker);
 
+  worker.on('message', (msg) => {
+    worker?.onMessage?.(msg);
+  });
+  worker.on('error', (err) => {
+    worker?.onError?.(err);
+  });
+
   return worker;
 }
 
@@ -130,7 +157,18 @@ async function stopAllWorkers() {
   Logger.verbose('Stopping all workers');
   const promises: Promise<any>[] = [];
   for (const worker of workers.values()) {
-    promises.push(worker.terminate());
+    if (worker.used) {
+      worker.toKill = true;
+      // set timout for 5 minutes and kill the worker if it is still running
+      setTimeout(() => {
+        if (worker.toKill) {
+          worker.toKill = false;
+          void worker.terminate();
+        }
+      }, 300000);
+    } else {
+      promises.push(worker.terminate());
+    }
   }
   workers.clear();
   await Promise.all(promises);