fix(worker): queue job completions

Author: RomainLanz

src/job_pool.ts

@@ -3,7 +3,7 @@ import type { AcquiredJob } from './contracts/adapter.js'
 /**
  * Entry representing an active job in the pool.
  */
-interface PoolEntry {
+export interface PoolEntry {
   /** Promise that resolves when the job completes */
   promise: Promise<void>
   /** The acquired job data */
@@ -12,6 +12,11 @@ interface PoolEntry {
   queue: string
 }
 
+interface CompletedEntry {
+  id: string
+  entry: PoolEntry
+}
+
 /**
  * Manages concurrent job execution with a fixed pool size.
  *
@@ -31,6 +36,9 @@ interface PoolEntry {
  */
 export class JobPool {
   #activeJobs = new Map<string, PoolEntry>()
+  #completedEntries: CompletedEntry[] = []
+  #completedHead = 0
+  #completionAvailable?: PromiseWithResolvers<void>
 
   /** Number of currently running jobs */
   get size() {
@@ -69,8 +77,12 @@ export class JobPool {
    * @param promise - Promise that resolves when the job completes
    */
   add(job: AcquiredJob, queue: string, promise: Promise<void>) {
-    void promise.catch(() => {})
-    this.#activeJobs.set(job.id, { promise, job, queue })
+    const entry = { promise, job, queue }
+    this.#activeJobs.set(job.id, entry)
+    void promise.then(
+      () => this.#enqueueCompletion(job.id, entry),
+      () => this.#enqueueCompletion(job.id, entry)
+    )
   }
 
   /**
@@ -100,27 +112,26 @@ export class JobPool {
   /**
    * Wait for the next job to complete and return it.
    *
-   * Uses `Promise.race()` internally, so the fastest job wins.
-   * The completed job is removed from the pool.
+   * Completions are queued in settlement order and remain available until a
+   * consumer asks for them. The completed job is removed from the pool.
    *
    * @returns The first job to complete (success or failure)
    */
   async waitForNextCompletion(): Promise<PoolEntry> {
-    const completedJobId = await Promise.race(
-      [...this.#activeJobs.entries()].map(async ([id, { promise }]) => {
-        try {
-          await promise
-        } catch {
-          // Errors are handled in Worker#execute
-        }
-        return id
-      })
-    )
+    while (true) {
+      while (this.#completedHead >= this.#completedEntries.length) {
+        this.#completionAvailable ??= Promise.withResolvers<void>()
+        await this.#completionAvailable.promise
+      }
+
+      const completed = this.#completedEntries[this.#completedHead++]!
+      this.#compactCompletedJobs()
 
-    const completed = this.#activeJobs.get(completedJobId)!
-    this.#activeJobs.delete(completedJobId)
+      if (this.#activeJobs.get(completed.id) !== completed.entry) continue
 
-    return completed
+      this.#activeJobs.delete(completed.id)
+      return completed.entry
+    }
   }
 
   /**
@@ -140,5 +151,24 @@ export class JobPool {
 
     await Promise.all(promises)
     this.#activeJobs.clear()
+    this.#completedEntries = []
+    this.#completedHead = 0
+  }
+
+  #enqueueCompletion(jobId: string, entry: PoolEntry): void {
+    if (this.#activeJobs.get(jobId) !== entry) return
+
+    this.#completedEntries.push({ id: jobId, entry })
+    this.#completionAvailable?.resolve()
+    this.#completionAvailable = undefined
+  }
+
+  #compactCompletedJobs(): void {
+    if (this.#completedHead < 1_024 || this.#completedHead * 2 < this.#completedEntries.length) {
+      return
+    }
+
+    this.#completedEntries = this.#completedEntries.slice(this.#completedHead)
+    this.#completedHead = 0
   }
 }

src/worker_session.ts

@@ -2,7 +2,7 @@ import { setTimeout } from 'node:timers/promises'
 import debug from './debug.js'
 import * as errors from './exceptions.js'
 import { parse } from './utils.js'
-import { JobPool } from './job_pool.js'
+import { JobPool, type PoolEntry } from './job_pool.js'
 import { WorkerHeartbeat, type HeartbeatOperationWrapper } from './worker_heartbeat.js'
 import { Locator } from './locator.js'
 import type { SingleJobDispatchRequest } from './job_dispatch_runtime.js'
@@ -71,6 +71,9 @@ export class WorkerSession {
   #stopOperation?: Promise<void>
   #pool = new JobPool()
   #fillOperation?: Promise<PoolFillResult>
+  #completionOperation?: Promise<PoolEntry>
+  #completedEntry?: PoolEntry
+  #completionDelayController?: AbortController
   #lastStalledCheck = 0
   #delayController?: AbortController
 
@@ -308,19 +311,12 @@ export class WorkerSession {
           }
 
           const hasCapacity = this.#pool.hasCapacity(this.#settings.concurrency)
-          const result = await Promise.race([
-            this.#pool
-              .waitForNextCompletion()
-              .then((completed) => ({ kind: 'completed' as const, completed })),
-            ...(hasCapacity
-              ? [setTimeout(this.#settings.idleDelay).then(() => ({ kind: 'tick' as const }))]
-              : []),
-          ])
+          const completed = await this.#waitForCompletion(hasCapacity)
 
           if (!this.#running) break
-          if (result.kind === 'tick') continue
+          if (!completed) continue
 
-          yield { type: 'completed', queue: result.completed.queue, job: result.completed.job }
+          yield { type: 'completed', queue: completed.queue, job: completed.job }
         } catch (error) {
           if (!this.#running) break
 
@@ -338,6 +334,56 @@ export class WorkerSession {
     }
   }
 
+  async #waitForCompletion(hasCapacity: boolean): Promise<PoolEntry | null> {
+    const completionOperation = this.#getCompletionOperation()
+
+    if (!hasCapacity || this.#completedEntry) {
+      return this.#consumeCompletion(completionOperation)
+    }
+
+    const controller = new AbortController()
+    this.#completionDelayController = controller
+
+    try {
+      await setTimeout(this.#settings.idleDelay, undefined, { signal: controller.signal })
+    } catch (error) {
+      if (!controller.signal.aborted) {
+        throw error
+      }
+    } finally {
+      if (this.#completionDelayController === controller) {
+        this.#completionDelayController = undefined
+      }
+    }
+
+    return this.#completedEntry ? this.#consumeCompletion(completionOperation) : null
+  }
+
+  #getCompletionOperation(): Promise<PoolEntry> {
+    if (this.#completionOperation) return this.#completionOperation
+
+    const completionOperation = this.#pool.waitForNextCompletion().then((completed) => {
+      if (this.#completionOperation === completionOperation) {
+        this.#completedEntry = completed
+        this.#completionDelayController?.abort()
+      }
+      return completed
+    })
+    this.#completionOperation = completionOperation
+    return completionOperation
+  }
+
+  async #consumeCompletion(completionOperation: Promise<PoolEntry>): Promise<PoolEntry> {
+    const completed = this.#completedEntry ?? (await completionOperation)
+
+    if (this.#completionOperation === completionOperation) {
+      this.#completionOperation = undefined
+      this.#completedEntry = undefined
+    }
+
+    return completed
+  }
+
   async #fillPool(): Promise<PoolFillResult> {
     const slotsAvailable = this.#settings.concurrency - this.#pool.size
 

tests/job_pool.spec.ts

@@ -38,13 +38,13 @@ test.group('JobPool', () => {
     cleanup,
   }) => {
     const execution = Promise.withResolvers<void>()
-    const originalCatch = execution.promise.catch.bind(execution.promise)
+    const originalThen = execution.promise.then.bind(execution.promise)
     let observed = false
 
-    execution.promise.catch = ((...args: Parameters<Promise<void>['catch']>) => {
-      observed = true
-      return originalCatch(...args)
-    }) as Promise<void>['catch']
+    execution.promise.then = ((onFulfilled, onRejected) => {
+      observed = onRejected !== undefined
+      return originalThen(onFulfilled, onRejected)
+    }) as Promise<void>['then']
 
     const pool = new JobPool()
     cleanup(async () => {
@@ -57,6 +57,41 @@ test.group('JobPool', () => {
     assert.isTrue(observed)
   })
 
+  test('attaches one settlement handler while a completion wait remains pending', async ({
+    assert,
+    cleanup,
+  }) => {
+    const execution = Promise.withResolvers<void>()
+    let attachments = 0
+    const observablePromise = {
+      then(onFulfilled, onRejected) {
+        attachments++
+        return execution.promise.then(onFulfilled, onRejected)
+      },
+      catch(onRejected) {
+        return this.then(undefined, onRejected)
+      },
+    } as Promise<void>
+
+    const pool = new JobPool()
+    cleanup(async () => {
+      execution.resolve()
+      await pool.drain()
+    })
+
+    pool.add(createJob('observed-once'), 'default', observablePromise)
+    const completion = pool.waitForNextCompletion()
+
+    for (let tick = 0; tick < 10; tick++) {
+      await setTimeout(0)
+    }
+
+    assert.equal(attachments, 1)
+
+    execution.resolve()
+    await completion
+  })
+
   test('should check capacity correctly', ({ assert }) => {
     const pool = new JobPool()
 
@@ -85,6 +120,81 @@ test.group('JobPool', () => {
     assert.equal(pool.size, 1)
   })
 
+  test('returns a completion queued before waiting begins', async ({ assert }) => {
+    const pool = new JobPool()
+    pool.add(createJob('already-completed'), 'default', Promise.resolve())
+
+    await setTimeout(0)
+
+    const completed = await pool.waitForNextCompletion()
+    assert.equal(completed.job.id, 'already-completed')
+    assert.isTrue(pool.isEmpty())
+  })
+
+  test('returns a newer job that settles after waiting begins', async ({ assert }) => {
+    const pool = new JobPool()
+    const longExecution = Promise.withResolvers<void>()
+    pool.add(createJob('long-running'), 'default', longExecution.promise)
+
+    const completion = pool.waitForNextCompletion()
+    pool.add(createJob('newer-fast-job'), 'default', Promise.resolve())
+
+    assert.equal((await completion).job.id, 'newer-fast-job')
+    longExecution.resolve()
+    await pool.drain()
+  })
+
+  test('ignores a stale completion after an active job id is replaced', async ({ assert }) => {
+    const pool = new JobPool()
+    const staleExecution = Promise.withResolvers<void>()
+    const currentExecution = Promise.withResolvers<void>()
+
+    pool.add(createJob('reused-id'), 'stale', staleExecution.promise)
+    staleExecution.resolve()
+    await setTimeout(0)
+    pool.add(createJob('reused-id'), 'current', currentExecution.promise)
+
+    let delivered = false
+    const completion = pool.waitForNextCompletion().then((entry) => {
+      delivered = true
+      return entry
+    })
+
+    await setTimeout(0)
+    assert.isFalse(delivered)
+
+    currentExecution.resolve()
+    assert.equal((await completion).queue, 'current')
+    assert.isTrue(pool.isEmpty())
+  })
+
+  test('returns every job once in settlement order', async ({ assert }) => {
+    const pool = new JobPool()
+    const first = Promise.withResolvers<void>()
+    const second = Promise.withResolvers<void>()
+    const third = Promise.withResolvers<void>()
+
+    pool.add(createJob('first'), 'default', first.promise)
+    pool.add(createJob('second'), 'default', second.promise)
+    pool.add(createJob('third'), 'default', third.promise)
+
+    second.resolve()
+    await setTimeout(0)
+    first.resolve()
+    await setTimeout(0)
+    third.resolve()
+
+    assert.deepEqual(
+      [
+        (await pool.waitForNextCompletion()).job.id,
+        (await pool.waitForNextCompletion()).job.id,
+        (await pool.waitForNextCompletion()).job.id,
+      ],
+      ['second', 'first', 'third']
+    )
+    assert.isTrue(pool.isEmpty())
+  })
+
   test('should remove job from pool after completion', async ({ assert }) => {
     const pool = new JobPool()
 
@@ -173,4 +283,22 @@ test.group('JobPool', () => {
 
     assert.isTrue(pool.isEmpty())
   })
+
+  test('drain clears completions that settled before they were consumed', async ({ assert }) => {
+    const pool = new JobPool()
+    const running = Promise.withResolvers<void>()
+
+    pool.add(createJob('settled'), 'default', Promise.resolve())
+    pool.add(createJob('running'), 'default', running.promise)
+    await setTimeout(0)
+
+    const drain = pool.drain()
+    running.resolve()
+    await drain
+
+    assert.isTrue(pool.isEmpty())
+
+    pool.add(createJob('after-drain'), 'default', Promise.resolve())
+    assert.equal((await pool.waitForNextCompletion()).job.id, 'after-drain')
+  })
 })