fix(worker): avoid a race in shutdown hooks (#918)

Author: mjameswh

packages/test/src/test-worker-lifecycle.ts

@@ -5,6 +5,7 @@
  * @module
  */
 import { Worker } from '@temporalio/worker';
+import { Runtime } from '@temporalio/worker';
 import test from 'ava';
 import { RUN_INTEGRATION_TESTS } from './helpers';
 import { defaultOptions, isolateFreeWorker } from './mock-native-worker';
@@ -27,23 +28,59 @@ if (RUN_INTEGRATION_TESTS) {
     t.is(worker.getState(), 'STOPPED');
     await t.throwsAsync(worker.run(), { message: 'Poller was already started' });
   });
+
+  test.serial('Worker shuts down gracefully if interrupted before running', async (t) => {
+    const worker = await Worker.create({
+      ...defaultOptions,
+      shutdownGraceTime: '500ms',
+      taskQueue: 'shutdown-test',
+    });
+    t.is(worker.getState(), 'INITIALIZED');
+    process.emit('SIGINT', 'SIGINT');
+    const p = worker.run();
+    t.is(worker.getState(), 'RUNNING');
+    await p;
+    t.is(worker.getState(), 'STOPPED');
+  });
 }
 
 test.serial('Mocked run shuts down gracefully', async (t) => {
-  const worker = isolateFreeWorker({
-    shutdownGraceTime: '500ms',
-    taskQueue: 'shutdown-test',
-  });
-  t.is(worker.getState(), 'INITIALIZED');
-  const p = worker.run();
-  t.is(worker.getState(), 'RUNNING');
-  process.emit('SIGINT', 'SIGINT');
-  await p;
-  t.is(worker.getState(), 'STOPPED');
-  await t.throwsAsync(worker.run(), { message: 'Poller was already started' });
+  try {
+    const worker = isolateFreeWorker({
+      shutdownGraceTime: '500ms',
+      taskQueue: 'shutdown-test',
+    });
+    t.is(worker.getState(), 'INITIALIZED');
+    const p = worker.run();
+    t.is(worker.getState(), 'RUNNING');
+    process.emit('SIGINT', 'SIGINT');
+    await p;
+    t.is(worker.getState(), 'STOPPED');
+    await t.throwsAsync(worker.run(), { message: 'Poller was already started' });
+  } finally {
+    if (Runtime._instance) await Runtime._instance.shutdown();
+  }
+});
+
+test.serial('Mocked run shuts down gracefully if interrupted before running', async (t) => {
+  try {
+    const worker = isolateFreeWorker({
+      shutdownGraceTime: '500ms',
+      taskQueue: 'shutdown-test',
+    });
+    // worker.native.initiateShutdown = () => new Promise(() => undefined);
+    t.is(worker.getState(), 'INITIALIZED');
+    process.emit('SIGINT', 'SIGINT');
+    const p = worker.run();
+    t.is(worker.getState(), 'RUNNING');
+    await p;
+    t.is(worker.getState(), 'STOPPED');
+  } finally {
+    if (Runtime._instance) await Runtime._instance.shutdown();
+  }
 });
 
-test('Mocked run throws if not shut down gracefully', async (t) => {
+test.serial('Mocked run throws if not shut down gracefully', async (t) => {
   const worker = isolateFreeWorker({
     shutdownGraceTime: '5ms',
     taskQueue: 'shutdown-test',

packages/worker/src/runtime.ts

@@ -101,6 +101,7 @@ export class Runtime {
   protected readonly logPollPromise: Promise<void>;
   public readonly logger: Logger;
   protected readonly shutdownSignalCallbacks = new Set<() => void>();
+  protected state: 'RUNNING' | 'SHUTTING_DOWN' = 'RUNNING';
 
   static _instance?: Runtime;
   static instantiator?: 'install' | 'instance';
@@ -381,7 +382,11 @@ export class Runtime {
    * @hidden
    */
   public registerShutdownSignalCallback(callback: () => void): void {
-    this.shutdownSignalCallbacks.add(callback);
+    if (this.state === 'RUNNING') {
+      this.shutdownSignalCallbacks.add(callback);
+    } else {
+      queueMicrotask(callback);
+    }
   }
 
   /**
@@ -416,6 +421,7 @@ export class Runtime {
    * Bound to `this` for use with `process.on` and `process.off`
    */
   protected startShutdownSequence = (): void => {
+    this.state = 'SHUTTING_DOWN';
     this.teardownShutdownHook();
     for (const callback of this.shutdownSignalCallbacks) {
       queueMicrotask(callback); // Run later