refactor(worker): extract heartbeat lifecycle

Author: RomainLanz

src/worker_heartbeat.ts

@@ -0,0 +1,88 @@
+import debug from './debug.js'
+import type { Adapter } from './contracts/adapter.js'
+import type { JobPool } from './job_pool.js'
+
+export type HeartbeatOperationWrapper = <T>(operation: () => Promise<T>) => Promise<T>
+
+export interface WorkerHeartbeatOptions {
+  workerId: string
+  queues: readonly string[]
+  interval: number
+  pool: JobPool
+  adapter: Adapter
+  wrapInternal: HeartbeatOperationWrapper
+}
+
+/**
+ * Renews the leases of Jobs owned by one WorkerSession.
+ */
+export class WorkerHeartbeat {
+  readonly #workerId: string
+  readonly #queues: readonly string[]
+  readonly #interval: number
+  readonly #pool: JobPool
+  readonly #adapter: Adapter
+  readonly #wrapInternal: HeartbeatOperationWrapper
+
+  #timer?: NodeJS.Timeout
+  #renewal?: Promise<void>
+
+  constructor(options: WorkerHeartbeatOptions) {
+    this.#workerId = options.workerId
+    this.#queues = Object.freeze([...options.queues])
+    this.#interval = options.interval
+    this.#pool = options.pool
+    this.#adapter = options.adapter
+    this.#wrapInternal = options.wrapInternal
+  }
+
+  start(): void {
+    this.#stopTimer()
+
+    this.#timer = setInterval(() => {
+      if (this.#renewal) return
+
+      const renewal = this.#renewActiveJobs()
+      this.#renewal = renewal
+
+      const clearRenewal = () => {
+        if (this.#renewal === renewal) {
+          this.#renewal = undefined
+        }
+      }
+
+      void renewal.then(clearRenewal, clearRenewal)
+    }, this.#interval)
+
+    this.#timer.unref?.()
+  }
+
+  async stop(): Promise<void> {
+    this.#stopTimer()
+    await this.#renewal
+  }
+
+  #stopTimer(): void {
+    if (!this.#timer) return
+
+    clearInterval(this.#timer)
+    this.#timer = undefined
+  }
+
+  async #renewActiveJobs(): Promise<void> {
+    if (this.#pool.isEmpty()) return
+
+    const jobIdsByQueue = this.#pool.activeJobIdsByQueue()
+
+    for (const queue of this.#queues) {
+      const jobIds = jobIdsByQueue.get(queue)
+      if (!jobIds || jobIds.length === 0) continue
+
+      try {
+        await this.#wrapInternal(() => this.#adapter.renewJobs(queue, jobIds))
+      } catch (error) {
+        debug('worker %s: failed to renew jobs on queue %s: %O', this.#workerId, queue, error)
+      }
+    }
+  }
+}

src/worker_session.ts

@@ -3,6 +3,7 @@ import debug from './debug.js'
 import * as errors from './exceptions.js'
 import { parse } from './utils.js'
 import { JobPool } 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'
 import type { Adapter, AcquiredJob } from './contracts/adapter.js'
@@ -18,7 +19,7 @@ interface PoolFillResult {
   errors: unknown[]
 }
 
-export type InternalOperationWrapper = <T>(operation: () => Promise<T>) => Promise<T>
+export type InternalOperationWrapper = HeartbeatOperationWrapper
 export type JobExecutor = Pick<JobExecutionRuntime, 'execute'>
 export interface ScheduleDispatcher {
   dispatch(request: SingleJobDispatchRequest): Promise<unknown>
@@ -58,6 +59,7 @@ export class WorkerSession {
   readonly #scheduleDispatcher: ScheduleDispatcher
   readonly #wrapInternal: InternalOperationWrapper
   readonly #settings: WorkerSessionSettings
+  readonly #heartbeat: WorkerHeartbeat
 
   #mode?: SessionMode
   #running = false
@@ -70,8 +72,6 @@ export class WorkerSession {
   #pool = new JobPool()
   #fillOperation?: Promise<PoolFillResult>
   #lastStalledCheck = 0
-  #heartbeatTimer?: NodeJS.Timeout
-  #heartbeatRenewal?: Promise<void>
   #delayController?: AbortController
 
   constructor(options: WorkerSessionOptions) {
@@ -82,6 +82,14 @@ export class WorkerSession {
     this.#scheduleDispatcher = options.scheduleDispatcher
     this.#wrapInternal = options.wrapInternal
     this.#settings = options.settings
+    this.#heartbeat = new WorkerHeartbeat({
+      workerId: options.workerId,
+      queues: this.#queues,
+      interval: Math.max(Math.floor(options.settings.stalledThreshold / 2), 1),
+      pool: this.#pool,
+      adapter: options.adapter,
+      wrapInternal: options.wrapInternal,
+    })
   }
 
   assertQueues(queues: readonly string[]): void {
@@ -192,8 +200,7 @@ export class WorkerSession {
       await this.#pool.drain()
     }
 
-    this.#stopHeartbeat()
-    await this.#heartbeatRenewal
+    await this.#heartbeat.stop()
     await this.#closeCycleGenerator()
 
     if (this.#startCompletion) {
@@ -260,7 +267,7 @@ export class WorkerSession {
   }
 
   async *#cycles(): AsyncGenerator<WorkerCycle, void, unknown> {
-    this.#startHeartbeat()
+    this.#heartbeat.start()
 
     try {
       while (this.#running) {
@@ -326,7 +333,7 @@ export class WorkerSession {
       }
     } finally {
       if (!this.#stopping) {
-        this.#stopHeartbeat()
+        await this.#heartbeat.stop()
       }
     }
   }
@@ -468,53 +475,6 @@ export class WorkerSession {
     }
   }
 
-  #startHeartbeat(): void {
-    this.#stopHeartbeat()
-
-    const interval = Math.max(Math.floor(this.#settings.stalledThreshold / 2), 1)
-
-    this.#heartbeatTimer = setInterval(() => {
-      if (this.#heartbeatRenewal) return
-
-      const renewal = this.#renewActiveJobs()
-      this.#heartbeatRenewal = renewal
-
-      const clearRenewal = () => {
-        if (this.#heartbeatRenewal === renewal) {
-          this.#heartbeatRenewal = undefined
-        }
-      }
-
-      void renewal.then(clearRenewal, clearRenewal)
-    }, interval)
-
-    this.#heartbeatTimer.unref?.()
-  }
-
-  #stopHeartbeat(): void {
-    if (this.#heartbeatTimer) {
-      clearInterval(this.#heartbeatTimer)
-      this.#heartbeatTimer = undefined
-    }
-  }
-
-  async #renewActiveJobs(): Promise<void> {
-    if (this.#pool.isEmpty()) return
-
-    const jobIdsByQueue = this.#pool.activeJobIdsByQueue()
-
-    for (const queue of this.#queues) {
-      const jobIds = jobIdsByQueue.get(queue)
-      if (!jobIds || jobIds.length === 0) continue
-
-      try {
-        await this.#wrapInternal(() => this.#adapter.renewJobs(queue, jobIds))
-      } catch (error) {
-        debug('worker %s: failed to renew jobs on queue %s: %O', this.#workerId, queue, error)
-      }
-    }
-  }
-
   async #dispatchDueSchedules(): Promise<void> {
     while (this.#running) {
       const schedule = await this.#wrapInternal(() => this.#adapter.claimDueSchedule())

tests/worker_heartbeat.spec.ts

@@ -0,0 +1,153 @@
+import { setTimeout } from 'node:timers/promises'
+import { test } from '@japa/runner'
+import { JobPool } from '../src/job_pool.js'
+import { WorkerHeartbeat } from '../src/worker_heartbeat.js'
+import { ControllableAdapter } from './_mocks/controllable_adapter.js'
+import { trackPromise } from './_utils/track_promise.js'
+
+function createHeartbeat(adapter: ControllableAdapter, pool: JobPool, interval = 10) {
+  return new WorkerHeartbeat({
+    workerId: 'test-worker',
+    queues: ['default'],
+    interval,
+    pool,
+    adapter,
+    wrapInternal: (operation) => operation(),
+  })
+}
+
+async function acquireIntoPool(
+  adapter: ControllableAdapter,
+  pool: JobPool,
+  id: string,
+  execution: Promise<void>
+): Promise<void> {
+  adapter.setWorkerId('test-worker')
+  await adapter.pushOn('default', {
+    id,
+    name: 'TestJob',
+    payload: {},
+    attempts: 0,
+    priority: 0,
+  })
+
+  const job = await adapter.popFrom('default')
+  if (!job) throw new Error('Expected an acquired Job')
+  pool.add(job, 'default', execution)
+}
+
+test.group('WorkerHeartbeat', () => {
+  test('waits for an in-flight renewal before stopping', async ({ assert, cleanup }) => {
+    const adapter = new ControllableAdapter()
+    const pool = new JobPool()
+    const execution = Promise.withResolvers<void>()
+    await acquireIntoPool(adapter, pool, 'renewed-job', execution.promise)
+    adapter.renewals.block(1)
+    const heartbeat = createHeartbeat(adapter, pool, 1)
+
+    cleanup(async () => {
+      adapter.releaseAll()
+      execution.resolve()
+      await heartbeat.stop()
+      await pool.drain()
+    })
+
+    heartbeat.start()
+    await adapter.renewals.waitForStarted()
+
+    const stop = trackPromise(heartbeat.stop())
+    assert.isFalse(stop.settled)
+
+    adapter.renewals.release(1)
+    await stop.promise
+  })
+
+  test('does not overlap renewals', async ({ assert, cleanup }) => {
+    const adapter = new ControllableAdapter()
+    const pool = new JobPool()
+    const execution = Promise.withResolvers<void>()
+    await acquireIntoPool(adapter, pool, 'renewed-job', execution.promise)
+    adapter.renewals.block(1)
+    const heartbeat = createHeartbeat(adapter, pool, 1)
+
+    cleanup(async () => {
+      adapter.releaseAll()
+      execution.resolve()
+      await heartbeat.stop()
+      await pool.drain()
+    })
+
+    heartbeat.start()
+    await adapter.renewals.waitForStarted()
+    await setTimeout(20)
+
+    assert.equal(adapter.renewals.calls, 1)
+
+    adapter.renewals.release(1)
+    await heartbeat.stop()
+  })
+
+  test('absorbs renewal errors', async ({ assert, cleanup }) => {
+    const adapter = new ControllableAdapter()
+    const pool = new JobPool()
+    const execution = Promise.withResolvers<void>()
+    await acquireIntoPool(adapter, pool, 'renewed-job', execution.promise)
+    adapter.renewals.fail(1, new Error('renewal failed'))
+    const heartbeat = createHeartbeat(adapter, pool, 1)
+
+    cleanup(async () => {
+      execution.resolve()
+      await heartbeat.stop()
+      await pool.drain()
+    })
+
+    heartbeat.start()
+    await adapter.renewals.waitForSettled(1)
+    await heartbeat.stop()
+
+    assert.isAtLeast(adapter.renewals.calls, 1)
+  })
+
+  test('prevents an active Job from being recovered as stalled', async ({ assert, cleanup }) => {
+    const adapter = new ControllableAdapter()
+    const pool = new JobPool()
+    const execution = Promise.withResolvers<void>()
+    await acquireIntoPool(adapter, pool, 'long-running-job', execution.promise)
+    const heartbeat = createHeartbeat(adapter, pool, 20)
+
+    cleanup(async () => {
+      execution.resolve()
+      await heartbeat.stop()
+      await pool.drain()
+    })
+
+    heartbeat.start()
+    await adapter.renewals.waitForStarted()
+    await setTimeout(60)
+
+    const recovered = await adapter.recoverStalledJobs('default', 40, 1)
+    assert.equal(recovered, 0)
+  })
+
+  test('stops renewing after stop resolves', async ({ assert, cleanup }) => {
+    const adapter = new ControllableAdapter()
+    const pool = new JobPool()
+    const execution = Promise.withResolvers<void>()
+    await acquireIntoPool(adapter, pool, 'heartbeat-job', execution.promise)
+    const heartbeat = createHeartbeat(adapter, pool)
+
+    cleanup(async () => {
+      execution.resolve()
+      await heartbeat.stop()
+      await pool.drain()
+    })
+
+    heartbeat.start()
+    await adapter.renewals.waitForStarted()
+    await heartbeat.stop()
+
+    const renewalsAtStop = adapter.renewals.calls
+    await setTimeout(40)
+    assert.equal(adapter.renewals.calls, renewalsAtStop)
+  })
+})