feat: support events and std streams from an abort handler

josdejong · Dec 12, 2024 · ea38666 · ea38666
1 parent 6bf949e
commit ea38666
Show file tree

Hide file tree

Showing 3 changed files with 78 additions and 11 deletions.
diff --git a/src/WorkerHandler.js b/src/WorkerHandler.js
@@ -220,13 +220,11 @@ function objectToError (obj) {
 
 function handleEmittedStdPayload(handler, payload) {
   // TODO: refactor if parallel task execution gets added
-  if (Object.keys(handler.processing).length !== 1) {
-    return;
-  }
-  var task = Object.values(handler.processing)[0]
-  if (task.options && typeof task.options.on === 'function') {
-    task.options.on(payload);
-  }
+  Object.values(handler.processing)
+    .forEach(task => task?.options?.on(payload));
+
+  Object.values(handler.tracking)
+    .forEach(task => task?.options?.on(payload)); 
 }
 
 /**
@@ -299,6 +297,16 @@ function WorkerHandler(script, _options) {
             task.resolver.resolve(response.result);
           }
         }
+      } else {
+        // if the task is not the current, it might be tracked for cleanup
+        var task = me.tracking[id];
+        if (task !== undefined) {
+          if (response.isEvent) {
+            if (task.options && typeof task.options.on === 'function') {
+              task.options.on(response.payload);
+            }
+          }
+        } 
       }
 
       if (response.method === CLEANUP_METHOD_ID) {
@@ -422,7 +430,8 @@ WorkerHandler.prototype.exec = function(method, params, resolver, options) {
     if (error instanceof Promise.CancellationError || error instanceof Promise.TimeoutError) {
       me.tracking[id] = {
         id,
-        resolver: Promise.defer()
+        resolver: Promise.defer(),
+        options: options,
       };
 
       // remove this task from the queue. It is already rejected (hence this

diff --git a/test/Pool.test.js b/test/Pool.test.js
@@ -1356,7 +1356,7 @@ describe('Pool', function () {
         maxWorkers: 1,
         onCreateWorker: () => {
           workerCount += 1;
-        }
+        },
       });
 
       let task = pool.exec('asyncTimeout', [],  {});
@@ -1561,6 +1561,40 @@ describe('Pool', function () {
         });
       });
     });
+
+    it('should trigger event stdout in abort handler', function (done) {
+      var pool = createPool(__dirname + '/workers/cleanup-abort.js', {
+        maxWorkers: 1,
+        workerType: 'process',
+        emitStdStreams: true, 
+        workerTerminateTimeout: 1000,
+      });
+
+      pool.exec('stdoutStreamOnAbort', [], {
+        on: function (payload) {
+          assert.strictEqual(payload.stdout.trim(), "Hello, world!");
+          pool.terminate();
+          done();
+        }
+      }).timeout(50);
+    });
+
+    it('should trigger event in abort handler', function (done) {
+      var pool = createPool(__dirname + '/workers/cleanup-abort.js', {
+        maxWorkers: 1,
+        workerType: 'process',
+        emitStdStreams: true, 
+        workerTerminateTimeout: 1000,
+      });
+
+      pool.exec('eventEmitOnAbort', [], {
+        on: function (payload) {
+          assert.strictEqual(payload.status, 'cleanup_success');
+          pool.terminate();
+          done();
+        }
+      }).timeout(50);
+    });
   });
 
   describe('validate', () => {

diff --git a/test/workers/cleanup-abort.js b/test/workers/cleanup-abort.js
@@ -5,8 +5,7 @@ function asyncTimeout() {
   return new Promise(function (resolve) {
     let timeout = setTimeout(() => {
         resolve();
-    }, 5000); 
-
+    }, 5000);
     me.worker.addAbortListener(async function () {
         clearTimeout(timeout);
         resolve();
@@ -34,11 +33,36 @@ function asyncAbortHandlerNeverResolves() {
   });
 }
 
+function stdoutStreamOnAbort() {
+  var me = this;
+  return new Promise(function (resolve) {
+    me.worker.addAbortListener(async function () {
+        console.log("Hello, world!");
+        resolve();
+    });
+  });
+}
+
+function eventEmitOnAbort() {
+  var me = this;
+  return new Promise(function (resolve) {
+    me.worker.addAbortListener(async function () {
+        workerpool.workerEmit({
+          status: 'cleanup_success',
+        });
+        resolve();
+    });
+  });
+}
+
+
 // create a worker and register public functions
 workerpool.worker(
   {
     asyncTimeout: asyncTimeout,
     asyncAbortHandlerNeverResolves: asyncAbortHandlerNeverResolves,
+    stdoutStreamOnAbort: stdoutStreamOnAbort,
+    eventEmitOnAbort: eventEmitOnAbort,
   },
   {
     abortListenerTimeout: 1000