Test that activity is cancelled on worker shutdown (#439)

Author: bergundy

packages/test/src/activities/index.ts

@@ -1,7 +1,7 @@
 /* eslint-disable @typescript-eslint/no-empty-function */
 /* eslint-disable @typescript-eslint/explicit-module-boundary-types */
 import { Context } from '@temporalio/activity';
-import { WorkflowClient, WorkflowHandle } from '@temporalio/client';
+import { Client, WorkflowHandle } from '@temporalio/client';
 import { ApplicationFailure, QueryDefinition } from '@temporalio/common';
 import { ProtoActivityInput, ProtoActivityResult } from '../../protos/root';
 import { cancellableFetch as cancellableFetchInner } from './cancellable-fetch';
@@ -66,11 +66,11 @@ export async function waitForCancellation(): Promise<void> {
 function getSchedulingWorkflowHandle(): WorkflowHandle {
   const { info, connection, dataConverter } = getContext();
   const { workflowExecution } = info;
-  const client = new WorkflowClient({ connection, namespace: info.workflowNamespace, dataConverter });
-  return client.getHandle(workflowExecution.workflowId, workflowExecution.runId);
+  const client = new Client({ connection, namespace: info.workflowNamespace, dataConverter });
+  return client.workflow.getHandle(workflowExecution.workflowId, workflowExecution.runId);
 }
 
-async function signalSchedulingWorkflow(signalName: string) {
+export async function signalSchedulingWorkflow(signalName: string) {
   const handle = getSchedulingWorkflowHandle();
   await handle.signal(signalName);
 }

packages/test/src/activities/interceptors.ts

@@ -1,10 +1,10 @@
 import * as activity from '@temporalio/activity';
-import { Connection } from '@temporalio/client';
+import { ConnectionLike } from '@temporalio/client';
 import { defaultDataConverter, LoadedDataConverter } from '@temporalio/common';
 import { ActivityExecuteInput, ActivityInboundCallsInterceptor, Next } from '@temporalio/worker';
 
 export class ConnectionInjectorInterceptor implements ActivityInboundCallsInterceptor {
-  constructor(public readonly connection: Connection, public readonly dataConverter = defaultDataConverter) {}
+  constructor(public readonly connection: ConnectionLike, public readonly dataConverter = defaultDataConverter) {}
   async execute(input: ActivityExecuteInput, next: Next<ActivityInboundCallsInterceptor, 'execute'>): Promise<unknown> {
     Object.assign(activity.Context.current(), {
       connection: this.connection,
@@ -18,7 +18,7 @@ export class ConnectionInjectorInterceptor implements ActivityInboundCallsInterc
  * Extend the basic activity Context
  */
 export interface Context extends activity.Context {
-  connection: Connection;
+  connection: ConnectionLike;
   dataConverter: LoadedDataConverter;
 }
 

packages/test/src/helpers.ts

@@ -73,6 +73,7 @@ export const bundlerOptions = {
   // workflow module, add it to this list:
   ignoreModules: [
     '@temporalio/common/lib/internal-non-workflow',
+    '@temporalio/activity',
     '@temporalio/client',
     '@temporalio/testing',
     '@temporalio/worker',

packages/test/src/test-activities-are-cancelled-on-worker-shutdown.ts

@@ -0,0 +1,51 @@
+import crypto from 'crypto';
+import * as activity from '@temporalio/activity';
+import * as testing from '@temporalio/testing';
+import * as workflow from '@temporalio/workflow';
+import { bundlerOptions, test, Worker } from './helpers';
+
+export async function testWorkflow(): Promise<void> {
+  await workflow.proxyActivities({ startToCloseTimeout: '1m' }).testActivity();
+}
+
+test('Worker cancels heartbeating activities after shutdown has been requested', async (t) => {
+  const env = await testing.TestWorkflowEnvironment.createLocal();
+  try {
+    let properlyCancelled = false;
+    const taskQueue = 'activity-heartbeat-after-shutdown';
+    const worker = await Worker.create({
+      connection: env.nativeConnection,
+      taskQueue,
+      workflowsPath: __filename,
+      bundlerOptions,
+      activities: {
+        async testActivity() {
+          const ctx = activity.Context.current();
+          worker.shutdown();
+          ctx.heartbeat();
+          try {
+            await ctx.cancelled;
+          } catch (err) {
+            if (err instanceof activity.CancelledFailure && err.message === 'WORKER_SHUTDOWN') {
+              properlyCancelled = true;
+            }
+            throw err;
+          }
+        },
+      },
+    });
+    const handle = await env.client.workflow.start(testWorkflow, {
+      workflowId: crypto.randomUUID(),
+      taskQueue,
+    });
+    try {
+      // If worker completes within graceful shutdown period, the activity has successfully been cancelled
+      await worker.run();
+    } finally {
+      await handle.terminate();
+    }
+    t.true(properlyCancelled);
+  } finally {
+    await env.teardown();
+  }
+});

packages/testing/src/index.ts

@@ -24,7 +24,7 @@ import {
 } from '@temporalio/client';
 import { ActivityFunction, CancelledFailure } from '@temporalio/common';
 import { msToTs, tsToMs } from '@temporalio/common/lib/time';
-import { NativeConnection, Runtime, appendDefaultInterceptors } from '@temporalio/worker';
+import { NativeConnection, Runtime } from '@temporalio/worker';
 import {
   EphemeralServer,
   EphemeralServerConfig,

packages/worker/src/workflow/vm.ts

@@ -2,12 +2,10 @@ import assert from 'assert';
 import { AsyncLocalStorage } from 'async_hooks';
 import vm from 'vm';
 import { gte } from 'semver';
-import { WorkflowInfo } from '@temporalio/workflow';
 import { IllegalStateError } from '@temporalio/common';
 import { Workflow, WorkflowCreateOptions, WorkflowCreator } from './interface';
 import { WorkflowBundleWithSourceMapAndFilename } from './workflow-worker-thread/input';
 import {
-  ActivationContext,
   BaseVMWorkflow,
   globalHandlers,
   injectConsole,