feat: add task function worker node affinity (#126)

* feat: add task function worker node affinity

Port task function worker node affinity from poolifier/poolifier PR #2269.
Allows restricting task execution to specific worker nodes by their indices
via a workerNodeKeys property on task function objects.

* [autofix.ci] apply automated fixes

* fix: add workerNodeKeys validation and remove misleading JSDoc

* [autofix.ci] apply automated fixes

* fix: add missing affinity checks in least-elu strategy

* ci: migrate npm publish to OIDC trusted publisher

* fix: add missing nextWorkerNodeKey assignment in weighted round robin and align workerNodeKeys JSDoc with upstream

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>

Author: jerome-benoit

.github/workflows/release-please.yml

@@ -83,6 +83,10 @@ jobs:
     needs: build-release
     runs-on: ubuntu-latest
 
+    permissions:
+      contents: read
+      id-token: write
+
     steps:
       - name: Checkout
         uses: actions/checkout@v6
@@ -108,23 +112,15 @@ jobs:
       - name: Publish Release
         if: ${{ contains(github.ref_name, '-') == false }}
         run: cd dist && npm publish
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
 
       - name: Publish Release Candidate
         if: ${{ contains(github.ref_name, '-rc') == true }}
         run: cd dist && npm publish --tag next
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
 
       - name: Publish Beta Release
         if: ${{ contains(github.ref_name, '-beta') == true }}
         run: cd dist && npm publish --tag beta
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}
 
       - name: Publish Alpha Release
         if: ${{ contains(github.ref_name, '-alpha') == true }}
         run: cd dist && npm publish --tag alpha
-        env:
-          NODE_AUTH_TOKEN: ${{ secrets.NPM_TOKEN }}

README.md

@@ -58,6 +58,7 @@ Please consult our [general guidelines](#general-guidelines).
 - Support for abortable task function ✔
 - Support for multiple task functions with per task function queuing priority
   and tasks distribution strategy ✔
+- Support for worker node affinity per task function ✔
 - Support for task functions
   [CRUD](https://en.wikipedia.org/wiki/Create,_read,_update_and_delete)
   operations at runtime ✔

docs/api.md

@@ -91,8 +91,22 @@ This method is available on both pool implementations and returns a boolean.
 `name` (mandatory) The task function name.\
 `fn` (mandatory) The task function
 `(data?: Data) => Response | Promise<Response>` or task function object
-`{ taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy }`.
-Priority range is the same as Unix nice levels.
+`{ taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy, workerNodeKeys?: number[] }`.
+Priority range is the same as Unix nice levels. `workerNodeKeys` is an array of
+worker node keys to restrict task execution to specific workers (worker node
+affinity).
+
+#### Worker Node Affinity Notes
+
+- Worker node keys are validated at registration time against the pool's maximum
+  size (`maximumNumberOfWorkers ?? minimumNumberOfWorkers`).
+- The number of worker node keys cannot exceed the pool's maximum size
+  (`maximumNumberOfWorkers ?? minimumNumberOfWorkers`).
+- In dynamic pools, you can reference worker node keys up to the maximum pool
+  size. Workers that don't exist yet are automatically created when a task
+  targeting them is executed.
+- At execution time, if no specified worker is ready, selection retries until
+  one becomes available or retries are exhausted.
 
 This method is available on both pool implementations and returns a boolean
 promise.
@@ -225,9 +239,12 @@ An object with these properties:
 ### `class YourWorker extends ThreadWorker`
 
 `taskFunctions` (mandatory) The task function or task functions object
-`Record<string, (data?: Data) => Response | Promise<Response> | { taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy }>`
+`Record<string, (data?: Data) => Response | Promise<Response> | { taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy, workerNodeKeys?: number[] }>`
 that you want to execute on the worker. Priority range is the same as Unix nice
-levels.\
+levels. `workerNodeKeys` is an array of worker node keys to restrict task
+execution to specific workers (worker node affinity). See
+[Worker Node Affinity Notes](#worker-node-affinity-notes) above for validation
+behavior.\
 `opts` (optional) An object with these properties:
 
 - `killBehavior` (optional) - Dictates if your worker will be deleted in case a
@@ -270,8 +287,11 @@ This method is available on both worker implementations and returns
 `name` (mandatory) The task function name.\
 `fn` (mandatory) The task function
 `(data?: Data) => Response | Promise<Response>` or task function object
-`{ taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy }`.
-Priority range is the same as Unix nice levels.
+`{ taskFunction: (data?: Data) => Response | Promise<Response>, priority?: number, strategy?: WorkerChoiceStrategy, workerNodeKeys?: number[] }`.
+Priority range is the same as Unix nice levels. `workerNodeKeys` is an array of
+worker node keys to restrict task execution to specific workers (worker node
+affinity). See [Worker Node Affinity Notes](#worker-node-affinity-notes) above
+for validation behavior.
 
 This method is available on both worker implementations and returns
 `{ status: boolean, error?: Error }`.

src/pools/abstract-pool.ts

@@ -45,6 +45,7 @@ import {
   checkValidPriority,
   checkValidTasksQueueOptions,
   checkValidWorkerChoiceStrategy,
+  checkValidWorkerNodeKeys,
   getDefaultTasksQueueOptions,
   updateEluWorkerUsage,
   updateRunTimeWorkerUsage,
@@ -842,8 +843,8 @@ export abstract class AbstractPool<
   private async sendTaskFunctionOperationToWorkers(
     message: MessageValue<Data>,
   ): Promise<boolean> {
-    const targetWorkerNodeKeys = [...this.workerNodes.keys()]
-    if (targetWorkerNodeKeys.length === 0) {
+    const targetWorkerNodeCount = this.workerNodes.length
+    if (targetWorkerNodeCount === 0) {
       return true
     }
     const responsesReceived: MessageValue<Response>[] = []
@@ -853,14 +854,14 @@ export abstract class AbstractPool<
       reject: (reason?: unknown) => void,
     ): void => {
       this.checkMessageWorkerId(message)
+      const workerNodeKey = this.getWorkerNodeKeyByWorkerId(message.workerId)
       if (
         message.taskFunctionOperationStatus != null &&
-        targetWorkerNodeKeys.includes(
-          this.getWorkerNodeKeyByWorkerId(message.workerId),
-        )
+        workerNodeKey >= 0 &&
+        workerNodeKey < targetWorkerNodeCount
       ) {
         responsesReceived.push(message)
-        if (responsesReceived.length >= targetWorkerNodeKeys.length) {
+        if (responsesReceived.length >= targetWorkerNodeCount) {
           if (
             responsesReceived.every(
               (msg) => msg.taskFunctionOperationStatus === true,
@@ -883,19 +884,20 @@ export abstract class AbstractPool<
       }
     }
     let listener: ((message: MessageValue<Response>) => void) | undefined
+    const workerNodeKeys = [...this.workerNodes.keys()]
     try {
       return await new Promise<boolean>((resolve, reject) => {
         listener = (message: MessageValue<Response>) => {
           taskFunctionOperationsListener(message, resolve, reject)
         }
-        for (const workerNodeKey of targetWorkerNodeKeys) {
+        for (const workerNodeKey of workerNodeKeys) {
           this.registerWorkerMessageListener(workerNodeKey, listener)
           this.sendToWorker(workerNodeKey, message)
         }
       })
     } finally {
       if (listener != null) {
-        for (const workerNodeKey of targetWorkerNodeKeys) {
+        for (const workerNodeKey of workerNodeKeys) {
           this.deregisterWorkerMessageListener(workerNodeKey, listener)
         }
       }
@@ -928,6 +930,10 @@ export abstract class AbstractPool<
     }
     checkValidPriority(fn.priority)
     checkValidWorkerChoiceStrategy(fn.strategy)
+    checkValidWorkerNodeKeys(
+      fn.workerNodeKeys,
+      this.maximumNumberOfWorkers ?? this.minimumNumberOfWorkers,
+    )
     const opResult = await this.sendTaskFunctionOperationToWorkers({
       taskFunctionOperation: 'add',
       taskFunctionProperties: buildTaskFunctionProperties(name, fn),
@@ -1057,6 +1063,26 @@ export abstract class AbstractPool<
     )?.strategy
   }
 
+  /**
+   * Gets task function worker node keys affinity set, if any.
+   * @param name - The task function name.
+   * @returns The task function worker node keys affinity set, or `undefined` if not defined.
+   */
+  private readonly getTaskFunctionWorkerNodeKeysSet = (
+    name?: string,
+  ): ReadonlySet<number> | undefined => {
+    name = name ?? DEFAULT_TASK_NAME
+    const taskFunctionsProperties = this.listTaskFunctionsProperties()
+    if (name === DEFAULT_TASK_NAME) {
+      name = taskFunctionsProperties[1]?.name
+    }
+    const workerNodeKeys = taskFunctionsProperties.find(
+      (taskFunctionProperties: TaskFunctionProperties) =>
+        taskFunctionProperties.name === name,
+    )?.workerNodeKeys
+    return workerNodeKeys != null ? new Set(workerNodeKeys) : undefined
+  }
+
   /**
    * Gets worker node task function priority, if any.
    *
@@ -1549,7 +1575,20 @@ export abstract class AbstractPool<
    * @returns The chosen worker node key.
    */
   private chooseWorkerNode(name?: string): number {
-    if (this.shallCreateDynamicWorker()) {
+    const workerNodeKeysSet = this.getTaskFunctionWorkerNodeKeysSet(name)
+    if (workerNodeKeysSet != null) {
+      const maxPoolSize = this.maximumNumberOfWorkers ??
+        this.minimumNumberOfWorkers
+      const targetSize = max(...workerNodeKeysSet) + 1
+      while (
+        this.started &&
+        !this.destroying &&
+        this.workerNodes.length < targetSize &&
+        this.workerNodes.length < maxPoolSize
+      ) {
+        this.createAndSetupDynamicWorkerNode()
+      }
+    } else if (this.shallCreateDynamicWorker()) {
       const workerNodeKey = this.createAndSetupDynamicWorkerNode()
       if (
         this.workerChoiceStrategiesContext?.getPolicy().dynamicWorkerUsage ===
@@ -1560,6 +1599,7 @@ export abstract class AbstractPool<
     }
     return this.workerChoiceStrategiesContext!.execute(
       this.getTaskFunctionWorkerChoiceStrategy(name),
+      workerNodeKeysSet,
     )
   }
 
@@ -2172,6 +2212,14 @@ export abstract class AbstractPool<
     if (ready == null || !ready) {
       throw new Error(`Worker ${workerId?.toString()} failed to initialize`)
     }
+    const maxPoolSize = this.maximumNumberOfWorkers ??
+      this.minimumNumberOfWorkers
+    for (const taskFunctionProperties of taskFunctionsProperties ?? []) {
+      checkValidWorkerNodeKeys(
+        taskFunctionProperties.workerNodeKeys,
+        maxPoolSize,
+      )
+    }
     const workerNodeKey = this.getWorkerNodeKeyByWorkerId(workerId)
     const workerNode = this.workerNodes[workerNodeKey]
     workerNode.info.ready = ready

src/pools/selection-strategies/abstract-worker-choice-strategy.ts

@@ -101,7 +101,9 @@ export abstract class AbstractWorkerChoiceStrategy<
   public abstract update(workerNodeKey: number): boolean
 
   /** @inheritDoc */
-  public abstract choose(): number | undefined
+  public abstract choose(
+    workerNodeKeysSet?: ReadonlySet<number>,
+  ): number | undefined
 
   /** @inheritDoc */
   public abstract remove(workerNodeKey: number): boolean
@@ -144,6 +146,47 @@ export abstract class AbstractWorkerChoiceStrategy<
     return workerNodeKey
   }
 
+  /**
+   * Gets the next worker node key in a round-robin fashion.
+   *
+   * @returns The next worker node key.
+   */
+  protected getRoundRobinNextWorkerNodeKey(): number {
+    return this.nextWorkerNodeKey === this.pool.workerNodes.length - 1
+      ? 0
+      : (this.nextWorkerNodeKey ?? this.previousWorkerNodeKey) + 1
+  }
+
+  /**
+   * Gets the worker node key from a single-element affinity set.
+   *
+   * @param workerNodeKeysSet - The worker node keys affinity set.
+   * @returns The worker node key if ready, `undefined` otherwise.
+   */
+  protected getSingleWorkerNodeKey(
+    workerNodeKeysSet: ReadonlySet<number>,
+  ): number | undefined {
+    const [workerNodeKey] = workerNodeKeysSet
+    return this.isWorkerNodeReady(workerNodeKey) ? workerNodeKey : undefined
+  }
+
+  /**
+   * Whether the worker node is eligible for selection (ready and in affinity set).
+   *
+   * @param workerNodeKey - The worker node key.
+   * @param workerNodeKeysSet - The worker node keys affinity set. If undefined, all workers are eligible.
+   * @returns Whether the worker node is eligible.
+   */
+  protected isWorkerNodeEligible(
+    workerNodeKey: number,
+    workerNodeKeysSet?: ReadonlySet<number>,
+  ): boolean {
+    return (
+      this.isWorkerNodeReady(workerNodeKey) &&
+      (workerNodeKeysSet == null || workerNodeKeysSet.has(workerNodeKey))
+    )
+  }
+
   /**
    * Gets the worker node task runtime.
    * If the task statistics require the average runtime, the average runtime is returned.

src/pools/selection-strategies/fair-share-worker-choice-strategy.ts

@@ -77,9 +77,9 @@ export class FairShareWorkerChoiceStrategy<
   }
 
   /** @inheritDoc */
-  public choose(): number | undefined {
+  public choose(workerNodeKeysSet?: ReadonlySet<number>): number | undefined {
     this.setPreviousWorkerNodeKey(this.nextWorkerNodeKey)
-    this.nextWorkerNodeKey = this.fairShareNextWorkerNodeKey()
+    this.nextWorkerNodeKey = this.fairShareNextWorkerNodeKey(workerNodeKeysSet)
     return this.nextWorkerNodeKey
   }
 
@@ -96,27 +96,30 @@ export class FairShareWorkerChoiceStrategy<
     return true
   }
 
-  private fairShareNextWorkerNodeKey(): number | undefined {
+  private fairShareNextWorkerNodeKey(
+    workerNodeKeysSet?: ReadonlySet<number>,
+  ): number | undefined {
+    if (workerNodeKeysSet?.size === 0) {
+      return undefined
+    }
+    if (workerNodeKeysSet?.size === 1) {
+      return this.getSingleWorkerNodeKey(workerNodeKeysSet)
+    }
     const chosenWorkerNodeKey = this.pool.workerNodes.reduce(
       (minWorkerNodeKey: number, workerNode, workerNodeKey, workerNodes) => {
-        if (!this.isWorkerNodeReady(workerNodeKey)) {
+        if (!this.isWorkerNodeEligible(workerNodeKey, workerNodeKeysSet)) {
           return minWorkerNodeKey
         }
-        if (minWorkerNodeKey === -1) {
-          workerNode.strategyData = {
-            ...workerNode.strategyData,
-            virtualTaskEndTimestamp: this
-              .computeWorkerNodeVirtualTaskEndTimestamp(workerNodeKey),
-          }
-          return workerNodeKey
-        }
         if (workerNode.strategyData?.virtualTaskEndTimestamp == null) {
           workerNode.strategyData = {
             ...workerNode.strategyData,
             virtualTaskEndTimestamp: this
               .computeWorkerNodeVirtualTaskEndTimestamp(workerNodeKey),
           }
         }
+        if (minWorkerNodeKey === -1) {
+          return workerNodeKey
+        }
         return workerNode.strategyData.virtualTaskEndTimestamp! <
             workerNodes[minWorkerNodeKey].strategyData!.virtualTaskEndTimestamp!
           ? workerNodeKey

src/pools/selection-strategies/interleaved-weighted-round-robin-worker-choice-strategy.ts

@@ -85,7 +85,13 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
   }
 
   /** @inheritDoc */
-  public choose(): number | undefined {
+  public choose(workerNodeKeysSet?: ReadonlySet<number>): number | undefined {
+    if (workerNodeKeysSet?.size === 0) {
+      return undefined
+    }
+    if (workerNodeKeysSet?.size === 1) {
+      return this.getSingleWorkerNodeKey(workerNodeKeysSet)
+    }
     for (
       let roundIndex = this.roundId;
       roundIndex < this.roundWeights.length;
@@ -106,7 +112,7 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
         }
         const workerWeight = this.opts!.weights![workerNodeKey]
         if (
-          this.isWorkerNodeReady(workerNodeKey) &&
+          this.isWorkerNodeEligible(workerNodeKey, workerNodeKeysSet) &&
           workerWeight >= this.roundWeights[roundIndex] &&
           this.workerNodeVirtualTaskExecutionTime < workerWeight
         ) {
@@ -118,6 +124,7 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
           return this.nextWorkerNodeKey
         }
       }
+      this.workerNodeId = 0
     }
     this.interleavedWeightedRoundRobinNextWorkerNodeId()
     return undefined