fix(worker): Wait for worker shutdown if runUntil promise throws (#943)

Author: bergundy

packages/test/src/test-replay.ts

@@ -24,22 +24,12 @@ async function getHistories(fname: string): Promise<History> {
   }
 }
 
-function historator(histories: Array<History>, goSlow?: boolean) {
-  return {
-    timesCalled: 0,
-    async *[Symbol.asyncIterator]() {
-      for (const history of histories) {
-        this.timesCalled++;
-        yield { workflowId: 'fake', history };
-        if (goSlow) {
-          // This matters because the exception propagation from the worker takes a long time
-          // compared to this generator. This sleep makes it more realistic for a
-          // stream-from-net-or-disk situation
-          await new Promise((resolve) => setTimeout(resolve, 1000));
-        }
-      }
-    },
-  };
+function historator(histories: Array<History>) {
+  return (async function* () {
+    for (const history of histories) {
+      yield { workflowId: 'fake', history };
+    }
+  })();
 }
 
 const test = anyTest as TestInterface<Context>;
@@ -125,9 +115,7 @@ test('multiple-histories-replay', async (t) => {
     },
     { histories }
   );
-  t.deepEqual(histories.timesCalled, 2);
-  t.deepEqual(res.errors.length, 0);
-  t.pass();
+  t.is(res.errors.length, 0);
 });
 
 test('multiple-histories-replay-fails-fast', async (t) => {
@@ -136,7 +124,7 @@ test('multiple-histories-replay-fails-fast', async (t) => {
   // change workflow type to break determinism
   hist1.events[0].workflowExecutionStartedEventAttributes!.workflowType!.name = 'http';
   hist2.events[0].workflowExecutionStartedEventAttributes!.workflowType!.name = 'http';
-  const histories = historator([hist1, hist2], true);
+  const histories = historator([hist1, hist2]);
 
   await t.throwsAsync(
     Worker.runReplayHistories(
@@ -147,7 +135,6 @@ test('multiple-histories-replay-fails-fast', async (t) => {
       { histories }
     )
   );
-  t.deepEqual(histories.timesCalled, 1);
 });
 
 test('multiple-histories-replay-fails-slow', async (t) => {
@@ -165,6 +152,5 @@ test('multiple-histories-replay-fails-slow', async (t) => {
     },
     { histories }
   );
-  t.deepEqual(histories.timesCalled, 2);
-  t.deepEqual(res.errors.length, 1);
+  t.is(res.errors.length, 1);
 });

packages/test/src/test-testenvironment.ts

@@ -1,7 +1,7 @@
 import { WorkflowFailedError } from '@temporalio/client';
 import { TestWorkflowEnvironment, workflowInterceptorModules } from '@temporalio/testing';
 import { Connection } from '@temporalio/testing/lib/connection';
-import { Runtime, Worker } from '@temporalio/worker';
+import { Worker } from '@temporalio/worker';
 import anyTest, { TestInterface } from 'ava';
 import { v4 as uuid4 } from 'uuid';
 import {
@@ -19,9 +19,6 @@ interface Context {
 const test = anyTest as TestInterface<Context>;
 
 test.before(async (t) => {
-  Runtime.install({
-    telemetryOptions: { tracingFilter: 'DEBUG' },
-  });
   t.context = {
     testEnv: await TestWorkflowEnvironment.createTimeSkipping(),
   };
@@ -153,17 +150,15 @@ test.serial('Workflow code can run assertions', async (t) => {
     },
   });
 
-  await worker.runUntil(async () => {
-    const err: WorkflowFailedError = await t.throwsAsync(
+  const err: WorkflowFailedError = await t.throwsAsync(
+    worker.runUntil(
       client.workflow.execute(assertFromWorkflow, {
         workflowId: uuid4(),
         taskQueue: 'test',
         args: [6],
-      }),
-      {
-        instanceOf: WorkflowFailedError,
-      }
-    );
-    t.is(err.cause?.message, 'Expected values to be strictly equal:\n\n6 !== 7\n');
-  });
+      })
+    ),
+    { instanceOf: WorkflowFailedError }
+  );
+  t.is(err.cause?.message, 'Expected values to be strictly equal:\n\n6 !== 7\n');
 });

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

@@ -99,3 +99,21 @@ test.serial('Mocked run throws if not shut down gracefully', async (t) => {
   t.is(worker.getState(), 'FAILED');
   await t.throwsAsync(worker.run(), { message: 'Poller was already started' });
 });
+
+test.serial('Mocked throws combined error in runUntil', async (t) => {
+  const worker = isolateFreeWorker({
+    shutdownGraceTime: '5ms',
+    taskQueue: 'shutdown-test',
+  });
+  worker.native.initiateShutdown = () => new Promise(() => undefined);
+  const err = await t.throwsAsync(
+    worker.runUntil(async () => {
+      throw new Error('inner');
+    })
+  );
+  t.is(worker.getState(), 'FAILED');
+  t.is(err.message, 'Worker terminated with fatal error in `runUntil`');
+  const { workerError, innerError } = (err as any).cause;
+  t.is(workerError.message, 'Timed out while waiting for worker to shutdown gracefully');
+  t.is(innerError.message, 'inner');
+});

packages/worker/src/index.ts

@@ -28,7 +28,15 @@ export * from './interceptors';
 export * from './logger';
 export { History, Runtime, RuntimeOptions } from './runtime';
 export * from './sinks';
-export { DataConverter, defaultPayloadConverter, errors, State, Worker } from './worker';
+export {
+  CombinedWorkerRunError,
+  CombinedWorkerRunErrorCause,
+  DataConverter,
+  defaultPayloadConverter,
+  errors,
+  State,
+  Worker,
+} from './worker';
 export {
   appendDefaultInterceptors,
   CompiledWorkerOptions,

packages/worker/src/worker.ts

@@ -140,6 +140,33 @@ export type ActivityTaskWithContext = ContextAware<{
 
 type CompiledWorkerOptionsWithBuildId = CompiledWorkerOptions & { buildId: string };
 
+/**
+ * Combined error information for {@link Worker.runUntil}
+ */
+export interface CombinedWorkerRunErrorCause {
+  /**
+   * Error thrown by a Worker
+   */
+  workerError: unknown;
+  /**
+   * Error thrown by the wrapped promise or function
+   */
+  innerError: unknown;
+}
+
+/**
+ * Error thrown by {@link Worker.runUntil}
+ */
+export class CombinedWorkerRunError extends Error {
+  public readonly name = 'CombinedWorkerRunError';
+  public readonly cause: CombinedWorkerRunErrorCause;
+
+  constructor(message: string, { cause }: { cause: CombinedWorkerRunErrorCause }) {
+    super(message);
+    this.cause = cause;
+  }
+}
+
 export interface NativeWorkerLike {
   type: 'Worker';
   initiateShutdown: Promisify<OmitFirstParam<typeof native.workerInitiateShutdown>>;
@@ -1644,21 +1671,45 @@ export class Worker {
    * @returns the result of `fnOrPromise`
    *
    * Throws on fatal Worker errors.
+   *
+   * **SDK versions `< 1.5.0`**:
+   * This method would not wait for worker to complete shutdown if the inner `fnOrPromise` threw an error.
+   *
+   * **SDK versions `>=1.5.0`**:
+   * This method always waits for both worker shutdown and inner `fnOrPromise` to resolve.
+   * If one of worker run -or- the inner promise throw an error, that error is rethrown.
+   * If both throw an error, a {@link CombinedWorkerRunError} with a `cause` attribute containing both errors.
    */
   async runUntil<R>(fnOrPromise: Promise<R> | (() => Promise<R>)): Promise<R> {
-    const runAndShutdown = async () => {
+    const workerRunPromise = this.run();
+    const innerPromise = (async () => {
       try {
-        if (typeof fnOrPromise === 'function') {
-          return await fnOrPromise();
-        } else {
-          return await fnOrPromise;
-        }
+        const p = typeof fnOrPromise === 'function' ? fnOrPromise() : fnOrPromise;
+        return await p;
       } finally {
-        this.shutdown();
+        if (this.state === 'RUNNING') {
+          this.shutdown();
+        }
       }
-    };
-    const [_, ret] = await Promise.all([this.run(), runAndShutdown()]);
-    return ret;
+    })();
+    const [innerResult, workerRunResult] = await Promise.allSettled([innerPromise, workerRunPromise]);
+
+    if (workerRunResult.status === 'rejected') {
+      if (innerResult.status === 'rejected') {
+        throw new CombinedWorkerRunError('Worker terminated with fatal error in `runUntil`', {
+          cause: {
+            workerError: workerRunResult.reason,
+            innerError: innerResult.reason,
+          },
+        });
+      } else {
+        throw workerRunResult.reason;
+      }
+    } else if (innerResult.status === 'rejected') {
+      throw innerResult.reason;
+    } else {
+      return innerResult.value;
+    }
   }
 
   /**