X-Git-Url: https://git.piment-noir.org/?a=blobdiff_plain;ds=sidebyside;f=src%2Fworker%2Fabstract-worker.ts;h=56a7696f0638fabaaccd7eb4100b7d19e358f0cd;hb=1e3214b63e262557aadebf3c57e8388de6a4bbe4;hp=64627efb8823198c6cb09c81fbdded3bef046f01;hpb=15a8bf08edd98c1f9ef4a4d3f89d2cffbe30852e;p=poolifier.git

diff --git a/src/worker/abstract-worker.ts b/src/worker/abstract-worker.ts
index 64627efb..56a7696f 100644
--- a/src/worker/abstract-worker.ts
+++ b/src/worker/abstract-worker.ts
@@ -84,7 +84,11 @@ export abstract class AbstractWorker<
        * The maximum time to keep this worker active while idle.
        * The pool automatically checks and terminates this worker when the time expires.
        */
-      maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME
+      maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME,
+      /**
+       * The function to call when the worker is killed.
+       */
+      killHandler: EMPTY_FUNCTION
     }
   ) {
     super(type)
@@ -100,6 +104,7 @@ export abstract class AbstractWorker<
     this.opts.maxInactiveTime =
       opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME
     delete this.opts.async
+    this.opts.killHandler = opts.killHandler ?? EMPTY_FUNCTION
   }
 
   /**
@@ -119,7 +124,7 @@ export abstract class AbstractWorker<
       this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
       this.taskFunctions.set(
         typeof taskFunctions.name === 'string' &&
-          taskFunctions.name.trim().length > 0
+        taskFunctions.name.trim().length > 0
           ? taskFunctions.name
           : 'fn1',
         boundFn
@@ -293,17 +298,17 @@ export abstract class AbstractWorker<
    * @param message - The received message.
    */
   protected messageListener (message: MessageValue<Data>): void {
-    if (message.workerId != null && message.workerId !== this.id) {
-      throw new Error('Message worker id does not match worker id')
+    if (this.isMain) {
+      throw new Error('Cannot handle message to worker in main worker')
+    } else if (message.workerId != null && message.workerId !== this.id) {
+      throw new Error('Message worker id does not match the worker id')
     } else if (message.workerId === this.id) {
       if (message.statistics != null) {
         // Statistics message received
         this.statistics = message.statistics
       } else if (message.checkActive != null) {
         // Check active message received
-        !this.isMain && message.checkActive
-          ? this.startCheckActive()
-          : this.stopCheckActive()
+        message.checkActive ? this.startCheckActive() : this.stopCheckActive()
       } else if (message.taskId != null && message.data != null) {
         // Task message received
         this.run(message)
@@ -320,8 +325,31 @@ export abstract class AbstractWorker<
    * @param message - The kill message.
    */
   protected handleKillMessage (message: MessageValue<Data>): void {
-    !this.isMain && this.stopCheckActive()
-    this.emitDestroy()
+    this.stopCheckActive()
+    if (isAsyncFunction(this.opts.killHandler)) {
+      (this.opts.killHandler?.() as Promise<void>)
+        .then(() => {
+          this.sendToMainWorker({ kill: 'success', workerId: this.id })
+          return null
+        })
+        .catch(() => {
+          this.sendToMainWorker({ kill: 'failure', workerId: this.id })
+        })
+        .finally(() => {
+          this.emitDestroy()
+        })
+        .catch(EMPTY_FUNCTION)
+    } else {
+      try {
+        // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+        this.opts.killHandler?.() as void
+        this.sendToMainWorker({ kill: 'success', workerId: this.id })
+      } catch (error) {
+        this.sendToMainWorker({ kill: 'failure', workerId: this.id })
+      } finally {
+        this.emitDestroy()
+      }
+    }
   }
 
   /**
@@ -395,9 +423,6 @@ export abstract class AbstractWorker<
    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
    */
   protected run (task: Task<Data>): void {
-    if (this.isMain) {
-      throw new Error('Cannot run a task in the main worker')
-    }
     const fn = this.getTaskFunction(task.name)
     if (isAsyncFunction(fn)) {
       this.runInAsyncScope(this.runAsync.bind(this), this, fn, task)
@@ -454,7 +479,7 @@ export abstract class AbstractWorker<
   ): void {
     let taskPerformance = this.beginTaskPerformance(task.name)
     fn(task.data)
-      .then(res => {
+      .then((res) => {
         taskPerformance = this.endTaskPerformance(taskPerformance)
         this.sendToMainWorker({
           data: res,
@@ -464,7 +489,7 @@ export abstract class AbstractWorker<
         })
         return null
       })
-      .catch(e => {
+      .catch((e) => {
         const errorMessage = this.handleError(e as Error | string)
         this.sendToMainWorker({
           taskError: {
@@ -529,7 +554,7 @@ export abstract class AbstractWorker<
   }
 
   private updateLastTaskTimestamp (): void {
-    if (!this.isMain && this.activeInterval != null) {
+    if (this.activeInterval != null) {
       this.lastTaskTimestamp = performance.now()
     }
   }