Cluster pool is now working

[poolifier.git] / src / pools / abstract-pool.ts

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 6e4465edd7a262af497cf8e9b84942ca9b80d4f6..879e3e4f78f2b3ab60a5918a865fa7c711e3d882 100644 (file)
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -24,6 +24,7 @@ export interface IWorker {
   on(event: 'error', handler: ErrorHandler<this>): void
   on(event: 'online', handler: OnlineHandler<this>): void
   on(event: 'exit', handler: ExitHandler<this>): void
+  once(event: 'exit', handler: ExitHandler<this>): void
 }
 
 /**
@@ -82,7 +83,7 @@ export abstract class AbstractPool<
 
   /**
    * - `key`: The `Worker`
-   * - `value`: Number of tasks that has been assigned to that worker since it started
+   * - `value`: Number of tasks currently in progress on the worker.
    */
   public readonly tasks: Map<Worker, number> = new Map<Worker, number>()
 
@@ -119,7 +120,6 @@ export abstract class AbstractPool<
     if (!this.filePath) {
       throw new Error('Please specify a file with a worker implementation')
     }
-
     this.setupHook()
 
     for (let i = 1; i <= this.numberOfWorkers; i++) {
@@ -161,9 +161,7 @@ export abstract class AbstractPool<
   }
 
   public async destroy (): Promise<void> {
-    for (const worker of this.workers) {
-      await this.destroyWorker(worker)
-    }
+    await Promise.all(this.workers.map(worker => this.destroyWorker(worker)))
   }
 
   /**
@@ -200,6 +198,20 @@ export abstract class AbstractPool<
     }
   }
 
+  /**
+   * Increase the number of tasks that the given workers has done.
+   *
+   * @param worker Workers whose tasks are increased.
+   */
+  protected decreaseWorkersTasks (worker: Worker): void {
+    const numberOfTasksTheWorkerHas = this.tasks.get(worker)
+    if (numberOfTasksTheWorkerHas !== undefined) {
+      this.tasks.set(worker, numberOfTasksTheWorkerHas - 1)
+    } else {
+      throw Error('Worker could not be found in tasks map')
+    }
+  }
+
   /**
    * Removes the given worker from the pool.
    *
@@ -221,8 +233,10 @@ export abstract class AbstractPool<
    */
   protected chooseWorker (): Worker {
     const chosenWorker = this.workers[this.nextWorkerIndex]
-    this.nextWorkerIndex++
-    this.nextWorkerIndex %= this.workers.length
+    this.nextWorkerIndex =
+      this.workers.length - 1 === this.nextWorkerIndex
+        ? 0
+        : this.nextWorkerIndex + 1
     return chosenWorker
   }
 
@@ -237,15 +251,13 @@ export abstract class AbstractPool<
     message: MessageValue<Data>
   ): void
 
-  protected abstract registerWorkerMessageListener (
-    port: Worker,
-    listener: (message: MessageValue<Response>) => void
-  ): void
+  protected abstract registerWorkerMessageListener<
+    Message extends Data | Response
+  > (worker: Worker, listener: (message: MessageValue<Message>) => void): void
 
-  protected abstract unregisterWorkerMessageListener (
-    port: Worker,
-    listener: (message: MessageValue<Response>) => void
-  ): void
+  protected abstract unregisterWorkerMessageListener<
+    Message extends Data | Response
+  > (worker: Worker, listener: (message: MessageValue<Message>) => void): void
 
   protected internalExecute (
     worker: Worker,
@@ -255,7 +267,7 @@ export abstract class AbstractPool<
       const listener: (message: MessageValue<Response>) => void = message => {
         if (message.id === messageId) {
           this.unregisterWorkerMessageListener(worker, listener)
-          this.increaseWorkersTask(worker)
+          this.decreaseWorkersTasks(worker)
           if (message.error) reject(message.error)
           else resolve(message.data as Response)
         }
@@ -288,8 +300,8 @@ export abstract class AbstractPool<
 
     worker.on('error', this.opts.errorHandler ?? (() => {}))
     worker.on('online', this.opts.onlineHandler ?? (() => {}))
-    // TODO handle properly when a worker exit
     worker.on('exit', this.opts.exitHandler ?? (() => {}))
+    worker.once('exit', () => this.removeWorker(worker))
 
     this.workers.push(worker)