@outputai/core 0.12.0 → 0.12.1-next.1243f78.0

package/src/worker/interruption.js

@@ -1,33 +1,72 @@
+import { serializeError } from '#helpers/error_serializer';
 import { createChildLogger } from '#logger';
 
 const log = createChildLogger( 'Interruption' );
 
-const FORCE_QUIT_GRACE_MS = 1000;
+const DOUBLE_SIGNAL_IGNORE_TIME = 1000;
+const KILL_AFTER_UNCAUGHT_TIME = 60_000;
 const INTERRUPTION_SIGNALS = [ 'SIGTERM', 'SIGINT', 'SIGUSR2' ];
+const UNCAUGHT_ERROR_TYPES = [ 'uncaughtException', 'unhandledRejection' ];
 
-export const setupInterruptionHandler = cb => {
-  const state = { interruptionReceivedAt: null };
+export class KillSignError extends Error {
+  name = 'KillSignError';
+};
+
+export class UncaughtError extends Error {
+  name = 'UncaughtError';
+};
+
+const state = {
+  lastSignalDate: null,
+  setupCompleted: false,
+  uncaughtHandled: false,
+  abortController: null
+};
 
-  const handle = signal => {
-    log.info( 'Signal Received', { signal } );
+/** Creates a countdown to kill the code */
+const createTerminationWatchdog = ( message, timeoutMs ) =>
+  setTimeout( () => {
+    log.warn( message );
+    process.exit( 1 );
+  }, timeoutMs ).unref();
 
-    if ( state.interruptionReceivedAt ) {
-      const elapsed = Date.now() - state.interruptionReceivedAt;
+/** Handles a signal interruption */
+const handleSignal = signal => {
+  log.info( 'Signal Received', { signal } );
 
-      // If running with npx, 2 kill signals are received in rapid succession,
-      // this ignores the second interruption when it is right after the first.
-      if ( elapsed < FORCE_QUIT_GRACE_MS ) {
-        return;
-      }
+  if ( state.lastSignalDate ) {
+    const elapsed = Date.now() - state.lastSignalDate;
+
+    // If running with npx, 2 kill signals are received in rapid succession,
+    // this ignores the second interruption when it is right after the first.
+    if ( elapsed > DOUBLE_SIGNAL_IGNORE_TIME ) {
       log.warn( 'Force quitting...' );
       process.exit( 1 );
-      return;
     }
+    return;
+  }
+
+  state.lastSignalDate = Date.now();
+  state.abortController.abort( new KillSignError( signal ) );
+};
 
-    state.interruptionReceivedAt = Date.now();
-    log.warn( 'Initiating shutdown...' );
-    cb();
-  };
+/** handles uncaught exceptions, unhandled promises */
+const handleUncaught = ( error, type ) => {
+  log.warn( 'Uncaught exception', { type, error: serializeError( error ) } );
+
+  if ( !state.uncaughtHandled ) {
+    state.abortController.abort( new UncaughtError( type, { cause: error } ) );
+    // after detecting an uncaught exception, starts a countdown, if the code doesn't finish in time, force quit
+    createTerminationWatchdog( 'Uncaught exception handling timed out, force quitting...', KILL_AFTER_UNCAUGHT_TIME );
+    state.uncaughtHandled = true;
+  }
+};
 
-  INTERRUPTION_SIGNALS.forEach( signal => process.on( signal, () => handle( signal ) ) );
+export const setupInterruptionHandler = abortController => {
+  if ( !state.setupCompleted ) {
+    state.abortController = abortController;
+    INTERRUPTION_SIGNALS.forEach( signal => process.on( signal, () => handleSignal( signal ) ) );
+    UNCAUGHT_ERROR_TYPES.forEach( type => process.on( type, error => handleUncaught( error, type ) ) );
+    state.setupCompleted = true;
+  }
 };

package/src/worker/interruption.spec.js

@@ -1,98 +1,177 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-import { setupInterruptionHandler } from './interruption.js';
 
-const { mockLog } = vi.hoisted( () => ( {
-  mockLog: { info: vi.fn(), warn: vi.fn() }
+const DOUBLE_SIGNAL_IGNORE_TIME = 1000;
+const KILL_AFTER_UNCAUGHT_TIME = 60_000;
+
+const { mockLog, serializeErrorMock } = vi.hoisted( () => ( {
+  mockLog: { info: vi.fn(), warn: vi.fn(), error: vi.fn() },
+  serializeErrorMock: vi.fn( () => 'serialized' )
 } ) );
 
 vi.mock( '#logger', () => ( { createChildLogger: () => mockLog } ) );
+vi.mock( '#helpers/error_serializer', () => ( { serializeError: serializeErrorMock } ) );
+
+// The module keeps its attachment state at module scope, so every test needs a fresh copy.
+const loadModule = async () => {
+  vi.resetModules();
+  return import( './interruption.js' );
+};
 
 describe( 'setupInterruptionHandler', () => {
-  const onHandlers = {};
-  const callback = vi.fn();
-  const exitMock = vi.fn();
-  const originalOn = process.on;
-  const originalExit = process.exit;
+  const handlers = {};
+  const spies = {};
 
   beforeEach( () => {
     vi.clearAllMocks();
-    Object.keys( onHandlers ).forEach( key => delete onHandlers[key] );
-    process.on = vi.fn( ( event, handler ) => {
-      onHandlers[event] = handler;
+    Object.keys( handlers ).forEach( key => delete handlers[key] );
+
+    spies.on = vi.spyOn( process, 'on' ).mockImplementation( ( event, handler ) => {
+      handlers[event] = handler;
+      return process;
     } );
-    process.exit = exitMock;
+    spies.exit = vi.spyOn( process, 'exit' ).mockImplementation( () => undefined );
   } );
 
   afterEach( () => {
     vi.useRealTimers();
-    process.on = originalOn;
-    process.exit = originalExit;
+    vi.restoreAllMocks();
   } );
 
-  it( 'registers interruption signal handlers', () => {
-    setupInterruptionHandler( callback );
+  it( 'attaches a handler for every signal and uncaught error type', async () => {
+    const { setupInterruptionHandler } = await loadModule();
+
+    setupInterruptionHandler( new AbortController() );
 
-    expect( process.on ).toHaveBeenCalledWith( 'SIGTERM', expect.any( Function ) );
-    expect( process.on ).toHaveBeenCalledWith( 'SIGINT', expect.any( Function ) );
-    expect( process.on ).toHaveBeenCalledWith( 'SIGUSR2', expect.any( Function ) );
+    expect( Object.keys( handlers ) ).toEqual( [
+      'SIGTERM',
+      'SIGINT',
+      'SIGUSR2',
+      'uncaughtException',
+      'unhandledRejection'
+    ] );
   } );
 
-  it( 'logs and invokes callback on first SIGTERM', () => {
-    setupInterruptionHandler( callback );
+  it( 'attaches once and keeps the first abort controller', async () => {
+    const { setupInterruptionHandler } = await loadModule();
+    const first = new AbortController();
+    const second = new AbortController();
 
-    onHandlers.SIGTERM();
+    setupInterruptionHandler( first );
+    const attachedCalls = spies.on.mock.calls.length;
+    setupInterruptionHandler( second );
 
-    expect( mockLog.info ).toHaveBeenCalledWith( 'Signal Received', { signal: 'SIGTERM' } );
-    expect( mockLog.warn ).toHaveBeenCalledWith( 'Initiating shutdown...' );
-    expect( callback ).toHaveBeenCalledOnce();
-    expect( exitMock ).not.toHaveBeenCalled();
+    expect( spies.on ).toHaveBeenCalledTimes( attachedCalls );
+
+    handlers.SIGTERM();
+
+    expect( first.signal.aborted ).toBe( true );
+    expect( second.signal.aborted ).toBe( false );
   } );
 
-  it( 'logs and invokes callback on first SIGINT', () => {
-    setupInterruptionHandler( callback );
+  describe( 'signals', () => {
+    it( 'aborts with a kill signal error naming the signal', async () => {
+      const { setupInterruptionHandler, KillSignError } = await loadModule();
+      const controller = new AbortController();
 
-    onHandlers.SIGINT();
+      setupInterruptionHandler( controller );
+      handlers.SIGTERM();
 
-    expect( mockLog.info ).toHaveBeenCalledWith( 'Signal Received', { signal: 'SIGINT' } );
-    expect( mockLog.warn ).toHaveBeenCalledWith( 'Initiating shutdown...' );
-    expect( callback ).toHaveBeenCalledOnce();
-    expect( exitMock ).not.toHaveBeenCalled();
-  } );
+      expect( mockLog.info ).toHaveBeenCalledWith( 'Signal Received', { signal: 'SIGTERM' } );
+      expect( controller.signal.reason ).toBeInstanceOf( KillSignError );
+      expect( controller.signal.reason.message ).toBe( 'SIGTERM' );
+      expect( spies.exit ).not.toHaveBeenCalled();
+    } );
 
-  it( 'logs and invokes callback on first SIGUSR2', () => {
-    setupInterruptionHandler( callback );
+    it( 'ignores a second signal received within the grace window', async () => {
+      vi.useFakeTimers();
+      const { setupInterruptionHandler } = await loadModule();
+      const controller = new AbortController();
 
-    onHandlers.SIGUSR2();
+      setupInterruptionHandler( controller );
+      handlers.SIGINT();
+      vi.advanceTimersByTime( DOUBLE_SIGNAL_IGNORE_TIME - 1 );
+      handlers.SIGINT();
 
-    expect( mockLog.info ).toHaveBeenCalledWith( 'Signal Received', { signal: 'SIGUSR2' } );
-    expect( mockLog.warn ).toHaveBeenCalledWith( 'Initiating shutdown...' );
-    expect( callback ).toHaveBeenCalledOnce();
-    expect( exitMock ).not.toHaveBeenCalled();
-  } );
+      expect( mockLog.warn ).not.toHaveBeenCalledWith( 'Force quitting...' );
+      expect( spies.exit ).not.toHaveBeenCalled();
+    } );
 
-  it( 'ignores a second signal received within the grace period', () => {
-    vi.useFakeTimers();
-    setupInterruptionHandler( callback );
+    it( 'force quits on a second signal received after the grace window', async () => {
+      vi.useFakeTimers();
+      const { setupInterruptionHandler } = await loadModule();
+      const controller = new AbortController();
 
-    onHandlers.SIGTERM();
-    onHandlers.SIGINT();
+      setupInterruptionHandler( controller );
+      handlers.SIGINT();
+      vi.advanceTimersByTime( DOUBLE_SIGNAL_IGNORE_TIME + 1 );
+      handlers.SIGINT();
 
-    expect( callback ).toHaveBeenCalledOnce();
-    expect( mockLog.warn ).toHaveBeenCalledTimes( 1 );
-    expect( mockLog.warn ).toHaveBeenCalledWith( 'Initiating shutdown...' );
-    expect( exitMock ).not.toHaveBeenCalled();
+      expect( mockLog.warn ).toHaveBeenCalledWith( 'Force quitting...' );
+      expect( spies.exit ).toHaveBeenCalledWith( 1 );
+    } );
   } );
 
-  it( 'force quits on a second signal after the grace period', () => {
-    vi.useFakeTimers();
-    setupInterruptionHandler( callback );
+  describe( 'uncaught errors', () => {
+    it( 'aborts with an uncaught error wrapping the original', async () => {
+      const { setupInterruptionHandler, UncaughtError } = await loadModule();
+      const controller = new AbortController();
+      const error = new TypeError( 'boom' );
+
+      setupInterruptionHandler( controller );
+      handlers.uncaughtException( error );
+
+      // Recorded here because a prior abort would leave the reason handler with nothing to report.
+      expect( mockLog.warn ).toHaveBeenCalledWith( 'Uncaught exception', { type: 'uncaughtException', error: 'serialized' } );
+      expect( serializeErrorMock ).toHaveBeenCalledWith( error );
+      expect( mockLog.error ).not.toHaveBeenCalled();
+      expect( controller.signal.reason ).toBeInstanceOf( UncaughtError );
+      expect( controller.signal.reason.message ).toBe( 'uncaughtException' );
+      expect( controller.signal.reason.cause ).toBe( error );
+    } );
+
+    it( 'wraps non error rejection values for unhandled rejections', async () => {
+      const { setupInterruptionHandler, UncaughtError } = await loadModule();
+      const controller = new AbortController();
 
-    onHandlers.SIGTERM();
-    vi.advanceTimersByTime( 1001 );
-    onHandlers.SIGINT();
+      setupInterruptionHandler( controller );
+      handlers.unhandledRejection( 'not an error' );
 
-    expect( callback ).toHaveBeenCalledOnce();
-    expect( mockLog.warn ).toHaveBeenCalledWith( 'Force quitting...' );
-    expect( exitMock ).toHaveBeenCalledWith( 1 );
+      expect( mockLog.warn ).toHaveBeenCalledWith( 'Uncaught exception', { type: 'unhandledRejection', error: 'serialized' } );
+      expect( serializeErrorMock ).toHaveBeenCalledWith( 'not an error' );
+      expect( mockLog.error ).not.toHaveBeenCalled();
+      expect( controller.signal.reason ).toBeInstanceOf( UncaughtError );
+      expect( controller.signal.reason.message ).toBe( 'unhandledRejection' );
+      expect( controller.signal.reason.cause ).toBe( 'not an error' );
+    } );
+
+    it( 'forces an exit when the shutdown outlives the watchdog', async () => {
+      vi.useFakeTimers();
+      const { setupInterruptionHandler } = await loadModule();
+
+      setupInterruptionHandler( new AbortController() );
+      handlers.uncaughtException( new Error( 'boom' ) );
+
+      vi.advanceTimersByTime( KILL_AFTER_UNCAUGHT_TIME - 1 );
+
+      expect( spies.exit ).not.toHaveBeenCalled();
+
+      vi.advanceTimersByTime( 1 );
+
+      expect( mockLog.warn ).toHaveBeenCalledWith( 'Uncaught exception handling timed out, force quitting...' );
+      expect( spies.exit ).toHaveBeenCalledWith( 1 );
+    } );
+
+    it( 'keeps the watchdog from holding the process open', async () => {
+      const { setupInterruptionHandler } = await loadModule();
+      setupInterruptionHandler( new AbortController() );
+
+      const timer = { unref: vi.fn() };
+      const setTimeoutSpy = vi.spyOn( globalThis, 'setTimeout' ).mockReturnValue( timer );
+
+      handlers.uncaughtException( new Error( 'boom' ) );
+
+      expect( setTimeoutSpy ).toHaveBeenCalledWith( expect.any( Function ), KILL_AFTER_UNCAUGHT_TIME );
+      expect( timer.unref ).toHaveBeenCalledOnce();
+    } );
   } );
 } );

package/src/worker/shutdown.js

@@ -0,0 +1,53 @@
+import { serializeError } from '#helpers/error_serializer';
+import { createChildLogger } from '#logger';
+
+const log = createChildLogger( 'Worker Shutdown' );
+
+/**
+ * Gracefully shuts every running service down, in order.
+ * A service that fails to stop never interrupts the sequence, it is reported back instead.
+ * @returns {Promise<Array<{ service: string, error: Error }>>} one entry per service that failed to stop
+ */
+export const shutdownServices = async ( { workerRunner, connectionMonitor, catalogPublisher, connection } ) => {
+  const shutdownJobs = [
+    {
+      service: 'worker',
+      label: 'Stopping Worker',
+      isActive: () => workerRunner?.running,
+      stop: () => workerRunner.stop()
+    },
+    {
+      service: 'connection_monitor',
+      label: 'Stopping Connection Monitor',
+      isActive: () => connectionMonitor?.running,
+      stop: () => connectionMonitor.stop()
+    },
+    {
+      service: 'catalog_publisher',
+      label: 'Interrupting Catalog Publisher',
+      isActive: () => catalogPublisher?.running,
+      stop: () => catalogPublisher.interrupt()
+    },
+    {
+      service: 'connection',
+      label: 'Closing Connection',
+      isActive: () => connection,
+      stop: () => connection.close()
+    }
+  ];
+
+  const failures = [];
+  for ( const { service, label, isActive, stop } of shutdownJobs ) {
+    if ( isActive() ) {
+      log.info( `${label}...` );
+      // a try covers a service that throws instead of rejecting, which a catch on the promise would miss
+      try {
+        await stop();
+      } catch ( error ) {
+        failures.push( { service, error } );
+        log.warn( `${label} error`, { error: serializeError( error ) } );
+      }
+    }
+  }
+  return failures;
+};

package/src/worker/shutdown.spec.js

@@ -0,0 +1,197 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+import { shutdownServices } from './shutdown.js';
+
+const { mockLog, serializeErrorMock } = vi.hoisted( () => ( {
+  mockLog: { info: vi.fn(), warn: vi.fn() },
+  serializeErrorMock: vi.fn( error => ( { name: error.name, message: error.message } ) )
+} ) );
+
+vi.mock( '#logger', () => ( { createChildLogger: () => mockLog } ) );
+vi.mock( '#helpers/error_serializer', () => ( { serializeError: serializeErrorMock } ) );
+
+/** Stands in for the services registry built during startup, with every service active. */
+const createServices = () => ( {
+  workerRunner: { running: true, stop: vi.fn().mockResolvedValue( undefined ) },
+  connectionMonitor: { running: true, stop: vi.fn().mockResolvedValue( undefined ) },
+  catalogPublisher: { running: true, interrupt: vi.fn().mockResolvedValue( undefined ) },
+  connection: { close: vi.fn().mockResolvedValue( undefined ) }
+} );
+
+/** The registry as it looks before startup assigns anything. */
+const createEmptyServices = () => ( {
+  workerRunner: null,
+  connectionMonitor: null,
+  catalogPublisher: null,
+  connection: null
+} );
+
+describe( 'shutdownServices', () => {
+  beforeEach( () => {
+    vi.clearAllMocks();
+  } );
+
+  it( 'stops every active service, worker first and connection last', async () => {
+    const services = createServices();
+
+    await shutdownServices( services );
+
+    expect( services.workerRunner.stop ).toHaveBeenCalledOnce();
+    expect( services.connectionMonitor.stop ).toHaveBeenCalledOnce();
+    expect( services.catalogPublisher.interrupt ).toHaveBeenCalledOnce();
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+
+    expect( services.workerRunner.stop.mock.invocationCallOrder[0] )
+      .toBeLessThan( services.connectionMonitor.stop.mock.invocationCallOrder[0] );
+    expect( services.connectionMonitor.stop.mock.invocationCallOrder[0] )
+      .toBeLessThan( services.catalogPublisher.interrupt.mock.invocationCallOrder[0] );
+    expect( services.catalogPublisher.interrupt.mock.invocationCallOrder[0] )
+      .toBeLessThan( services.connection.close.mock.invocationCallOrder[0] );
+  } );
+
+  it( 'announces each step it takes', async () => {
+    await shutdownServices( createServices() );
+
+    expect( mockLog.info ).toHaveBeenCalledWith( 'Stopping Worker...' );
+    expect( mockLog.info ).toHaveBeenCalledWith( 'Stopping Connection Monitor...' );
+    expect( mockLog.info ).toHaveBeenCalledWith( 'Interrupting Catalog Publisher...' );
+    expect( mockLog.info ).toHaveBeenCalledWith( 'Closing Connection...' );
+    expect( mockLog.warn ).not.toHaveBeenCalled();
+  } );
+
+  it( 'waits for a service to settle before stopping the next one', async () => {
+    const deferred = {};
+    const services = createServices();
+    services.workerRunner.stop.mockReturnValue( new Promise( resolve => {
+      deferred.resolve = resolve;
+    } ) );
+
+    const stopped = shutdownServices( services );
+    await Promise.resolve();
+
+    expect( services.connectionMonitor.stop ).not.toHaveBeenCalled();
+
+    deferred.resolve();
+    await stopped;
+
+    expect( services.connectionMonitor.stop ).toHaveBeenCalledOnce();
+  } );
+
+  it( 'skips a worker runner that is no longer running', async () => {
+    const services = createServices();
+    services.workerRunner.running = false;
+
+    await shutdownServices( services );
+
+    expect( services.workerRunner.stop ).not.toHaveBeenCalled();
+    expect( mockLog.info ).not.toHaveBeenCalledWith( 'Stopping Worker...' );
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+  } );
+
+  it( 'skips a connection monitor that already exited', async () => {
+    const services = createServices();
+    services.connectionMonitor.running = false;
+
+    await shutdownServices( services );
+
+    expect( services.connectionMonitor.stop ).not.toHaveBeenCalled();
+    expect( mockLog.info ).not.toHaveBeenCalledWith( 'Stopping Connection Monitor...' );
+  } );
+
+  it( 'skips a catalog publisher that already finished', async () => {
+    const services = createServices();
+    services.catalogPublisher.running = false;
+
+    await shutdownServices( services );
+
+    expect( services.catalogPublisher.interrupt ).not.toHaveBeenCalled();
+    expect( mockLog.info ).not.toHaveBeenCalledWith( 'Interrupting Catalog Publisher...' );
+  } );
+
+  it( 'reads the state at each turn, so a service stopped by an earlier step is skipped', async () => {
+    const services = createServices();
+    services.workerRunner.stop.mockImplementation( async () => {
+      services.connectionMonitor.running = false;
+    } );
+
+    await shutdownServices( services );
+
+    expect( services.connectionMonitor.stop ).not.toHaveBeenCalled();
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+  } );
+
+  it( 'does nothing when startup never built a service', async () => {
+    await expect( shutdownServices( createEmptyServices() ) ).resolves.toEqual( [] );
+
+    expect( mockLog.info ).not.toHaveBeenCalled();
+    expect( mockLog.warn ).not.toHaveBeenCalled();
+  } );
+
+  it( 'closes the connection when startup failed before anything else existed', async () => {
+    const services = { ...createEmptyServices(), connection: { close: vi.fn().mockResolvedValue( undefined ) } };
+
+    await shutdownServices( services );
+
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+    expect( mockLog.info ).toHaveBeenCalledWith( 'Closing Connection...' );
+  } );
+
+  it( 'reports nothing when every service stops cleanly', async () => {
+    await expect( shutdownServices( createServices() ) ).resolves.toEqual( [] );
+  } );
+
+  it( 'warns, reports the failure and keeps going when a service fails to stop', async () => {
+    const error = new Error( 'drain failed' );
+    const services = createServices();
+    services.workerRunner.stop.mockRejectedValue( error );
+
+    await expect( shutdownServices( services ) ).resolves.toEqual( [ { service: 'worker', error } ] );
+
+    expect( serializeErrorMock ).toHaveBeenCalledWith( error );
+    expect( mockLog.warn ).toHaveBeenCalledWith( 'Stopping Worker error', {
+      error: { name: 'Error', message: 'drain failed' }
+    } );
+    expect( services.connectionMonitor.stop ).toHaveBeenCalledOnce();
+    expect( services.catalogPublisher.interrupt ).toHaveBeenCalledOnce();
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+  } );
+
+  it( 'warns and reports when closing the connection fails', async () => {
+    const error = new Error( 'close failed' );
+    const services = createServices();
+    services.connection.close.mockRejectedValue( error );
+
+    await expect( shutdownServices( services ) ).resolves.toEqual( [ { service: 'connection', error } ] );
+
+    expect( mockLog.warn ).toHaveBeenCalledWith( 'Closing Connection error', {
+      error: { name: 'Error', message: 'close failed' }
+    } );
+  } );
+
+  it( 'warns and reports a service that throws instead of rejecting', async () => {
+    const error = new Error( 'shutdown blew up' );
+    const services = createServices();
+    services.workerRunner.stop.mockImplementation( () => {
+      throw error;
+    } );
+
+    await expect( shutdownServices( services ) ).resolves.toEqual( [ { service: 'worker', error } ] );
+
+    expect( mockLog.warn ).toHaveBeenCalledWith( 'Stopping Worker error', {
+      error: { name: 'Error', message: 'shutdown blew up' }
+    } );
+    expect( services.connection.close ).toHaveBeenCalledOnce();
+  } );
+
+  it( 'reports every service that failed, in the order they were stopped', async () => {
+    const drainError = new Error( 'drain failed' );
+    const closeError = new Error( 'close failed' );
+    const services = createServices();
+    services.workerRunner.stop.mockRejectedValue( drainError );
+    services.connection.close.mockRejectedValue( closeError );
+
+    await expect( shutdownServices( services ) ).resolves.toEqual( [
+      { service: 'worker', error: drainError },
+      { service: 'connection', error: closeError }
+    ] );
+  } );
+} );

package/src/worker/worker_runner.js

@@ -0,0 +1,66 @@
+/**
+ * Runs a Temporal worker and owns its graceful drain.
+ *
+ * NOTE:
+ * Temporal worker shutdown is a bit odd.
+ * worker.run() is an async job that only resolves when calling worker.shutdown().
+ * But worker.shutdown() is not async and returns nothing, so there is no way to await it.
+ * All code that needs to run after shutdown needs to be after `await worker.run()`.
+ */
+export class WorkerRunner {
+  #worker = null;
+  #signal = null;
+  #running = false;
+  #execution = null;
+
+  /**
+   * @param {object} options
+   * @param {import('@temporalio/worker').Worker} options.worker - already created Temporal worker
+   * @param {AbortSignal} options.signal - drains the worker when aborted
+   */
+  constructor( { worker, signal } ) {
+    this.#worker = worker;
+    this.#signal = signal;
+  }
+
+  /** Returns whether the worker is polling */
+  get running() {
+    return this.#running;
+  }
+
+  /**
+   * Starts polling the task queue. Aborting the signal drains the worker.
+   * @returns {Promise<void>} resolves when the worker has fully stopped, rejects when it fails
+   */
+  start() {
+    if ( this.#execution ) {
+      return this.#execution;
+    }
+    if ( this.#signal.aborted ) {
+      return Promise.resolve();
+    }
+
+    // the listener must not return the drain, Node rethrows a rejection returned by an event listener
+    this.#signal.addEventListener( 'abort', () => {
+      this.stop();
+    }, { once: true } );
+
+    this.#running = true;
+    this.#execution = this.#worker.run().finally( () => {
+      this.#running = false;
+    } );
+
+    return this.#execution;
+  }
+
+  /**
+   * Triggers the graceful drain.
+   * @returns {Promise<void>} resolves when the worker has fully stopped, rejects when the drain fails
+   */
+  stop() {
+    if ( this.#worker.getStatus().runState === 'RUNNING' ) {
+      this.#worker.shutdown();
+    }
+    return this.#execution ?? Promise.resolve();
+  }
+};