feat: activity override and cancellation

apps/tests/aws-runtime/test/test-service.ts

@@ -11,6 +11,8 @@ import {
   workflow,
   heartbeat,
   HeartbeatTimeout,
+  completeActivity,
+  failActivity,
 } from "@eventual/core";
 import { SendMessageCommand, SQSClient } from "@aws-sdk/client-sqs";
 import { AsyncWriterTestEvent } from "./async-writer-handler.js";
@@ -148,11 +150,13 @@ export const childWorkflow = workflow(
   }
 );
 
-const slowActivity = activity(
-  "slowAct",
-  { timeoutSeconds: 5 },
-  () => new Promise((resolve) => setTimeout(resolve, 10 * 1000))
-);
+const delay = (seconds: number) =>
+  new Promise((resolve) => setTimeout(resolve, seconds * 1000));
+
+const slowActivity = activity("slowAct", { timeoutSeconds: 5 }, async () => {
+  await delay(10);
+  return "done finally";
+});
 
 const slowWf = workflow("slowWorkflow", { timeoutSeconds: 5 }, () =>
   sleepFor(10)
@@ -332,3 +336,144 @@ const sendFinishEvent = activity("sendFinish", async (executionId: string) => {
     proxy: true,
   });
 });
+
+/**
+ * An event which is raised by the activityOverrideActivity.
+ * Can provide a value to the activity via activity token or defer the responsibility to the workflow via signal.
+ */
+const activityOverrideEvent = event<{
+  executionId: string;
+  token: string;
+  location: "handler" | "signal";
+  type: "complete" | "fail";
+}>("activityOverrideEvent");
+
+/**
+ * An async activity which will be cancelled using it's activity token.
+ */
+const activityOverrideActivity = activity(
+  "eventPublish",
+  async ({
+    type,
+    location,
+    executionId,
+  }: {
+    executionId: string;
+    location: "handler" | "signal";
+    type: "complete" | "fail";
+  }) => {
+    return asyncResult(
+      async (token) =>
+        await activityOverrideEvent.publish({
+          token,
+          location,
+          type,
+          executionId,
+        })
+    );
+  }
+);
+
+/**
+ * A signal called by the activityOverrideEvent handler to pass the activity token to the workflow.
+ * Used to test activity completion from the workflow.
+ */
+const activityOverrideSignal = new Signal<{
+  token: string;
+  type: "complete" | "fail";
+}>("activityOverrideSignal");
+
+activityOverrideEvent.on(async ({ token, location, type, executionId }) => {
+  if (location === "handler") {
+    if (type === "complete") {
+      await completeActivity(token, "from the event handler!");
+    } else {
+      await failActivity(token, new Error("WHY!!!"));
+    }
+  } else {
+    await activityOverrideSignal.send(executionId, { token, type });
+  }
+});
+
+/**
+ * Activity which waits to be closed/cancelled.
+ * If it is closed, it signals the workflow with "complete", if not it signals with "fail".
+ */
+const activityOverrideAwareActivity = activity(
+  "overrideAware",
+  async ({ executionId }: { executionId: string }) => {
+    let n = 0;
+    while (n < 10) {
+      await delay(1);
+      const { closed } = await heartbeat();
+      if (closed) {
+        await activityOverrideSignal.send(executionId, {
+          token: "",
+          type: "complete",
+        });
+        return;
+      }
+    }
+
+    await activityOverrideSignal.send(executionId, { token: "", type: "fail" });
+  }
+);
+
+/**
+ * Activity Override Tests.
+ *
+ * 1. cancel an activity from the workflow and record the result
+ * 2. Use the activity token to cancel, fail, or complete an async token from within an event handler or workflow
+ * 3. Cancel an activity and use a signal to signify that it knows it was cancelled and record the result.
+ */
+export const overrideWorkflow = workflow(
+  "override",
+  async (_, { execution: { id: executionId } }) => {
+    const act = slowActivity();
+    act.cancel("because");
+
+    const results1 = await Promise.allSettled([act]);
+
+    const signalHandler = activityOverrideSignal.on(async ({ token, type }) => {
+      if (type === "complete") {
+        await completeActivity(token, "from the signal handler!");
+      } else {
+        await failActivity(token, new Error("BECAUSE!!!"));
+      }
+    });
+
+    const results2 = await Promise.allSettled([
+      activityOverrideActivity({
+        location: "handler",
+        type: "complete",
+        executionId,
+      }),
+      activityOverrideActivity({
+        location: "handler",
+        type: "fail",
+        executionId,
+      }),
+      activityOverrideActivity({
+        location: "signal",
+        type: "complete",
+        executionId,
+      }),
+      activityOverrideActivity({
+        location: "signal",
+        type: "fail",
+        executionId,
+      }),
+    ]);
+
+    const aware = activityOverrideAwareActivity({ executionId });
+    await sleepFor(1);
+    aware.cancel("because");
+
+    signalHandler.dispose();
+
+    // the override activity SHOULD send a signal when it realizes it is cancelled.
+    const result = await expectSignal(activityOverrideSignal);
+
+    return [results1, results2, result];
+  }
+);

apps/tests/aws-runtime/test/tester.test.ts

@@ -1,4 +1,8 @@
-import { HeartbeatTimeout } from "@eventual/core";
+import {
+  ActivityCancelled,
+  EventualError,
+  HeartbeatTimeout,
+} from "@eventual/core";
 import { eventualRuntimeTestHarness } from "./runtime-test-harness.js";
 import {
   eventDrivenWorkflow,
@@ -10,6 +14,7 @@ import {
   workflow2,
   workflow3,
   workflow4,
+  overrideWorkflow,
 } from "./test-service.js";
 
 jest.setTimeout(100 * 1000);
@@ -31,7 +36,7 @@ eventualRuntimeTestHarness(({ testCompletion }) => {
     { status: "fulfilled", value: ["HELLO SAM", "HELLO CHRIS", "HELLO SAM"] },
     { status: "fulfilled", value: ["hello sam", "hello chris", "hello sam"] },
     { status: "fulfilled", value: "hello sam" },
-    { status: "rejected", reason: "Error" },
+    { status: "rejected", reason: { name: "Error", message: "failed" } },
   ]);
 
   testCompletion("parent-child", parentWorkflow, "done");
@@ -45,7 +50,10 @@ eventualRuntimeTestHarness(({ testCompletion }) => {
 
   testCompletion("asyncActivities", asyncWorkflow, [
     "hello from the async writer!",
-    "AsyncWriterError",
+    {
+      name: "AsyncWriterError",
+      message: "I was told to fail this activity, sorry.",
+    },
   ]);
 
   testCompletion("heartbeat", heartbeatWorkflow, 10, [
@@ -62,4 +70,21 @@ eventualRuntimeTestHarness(({ testCompletion }) => {
   ]);
 
   testCompletion("event-driven", eventDrivenWorkflow, "done!");
+
+  testCompletion("overrideActivities", overrideWorkflow, [
+    [{ status: "rejected", reason: new ActivityCancelled("because").toJSON() }],
+    [
+      { status: "fulfilled", value: "from the event handler!" },
+      {
+        status: "rejected",
+        reason: new EventualError("Error", "WHY!!!").toJSON(),
+      },
+      { status: "fulfilled", value: "from the signal handler!" },
+      {
+        status: "rejected",
+        reason: new EventualError("Error", "BECAUSE!!!").toJSON(),
+      },
+    ],
+    { token: "", type: "complete" },
+  ]);
 });

packages/@eventual/aws-cdk/src/activities.ts

@@ -87,10 +87,12 @@ export class Activities extends Construct implements IActivities, IGrantable {
 
   public configureCompleteActivity(func: Function) {
     this.props.workflows.configureSendWorkflowEvent(func);
+    this.configureUpdateActivity(func);
   }
 
   public grantCompleteActivity(grantable: IGrantable) {
     this.props.workflows.grantSendWorkflowEvent(grantable);
+    this.grantUpdateActivity(grantable);
   }
 
   public configureUpdateActivity(func: Function) {

packages/@eventual/aws-runtime/src/clients/activity-runtime-client.ts

@@ -67,7 +67,7 @@ export class AWSActivityRuntimeClient implements ActivityRuntimeClient {
     executionId: string,
     seq: number,
     heartbeatTime: string
-  ): Promise<{ cancelled: boolean }> {
+  ): Promise<{ closed: boolean }> {
     const item = await this.props.dynamo.send(
       new UpdateItemCommand({
         Key: {
@@ -86,27 +86,36 @@ export class AWSActivityRuntimeClient implements ActivityRuntimeClient {
     );
 
     return {
-      cancelled:
-        (item.Attributes as ActivityExecutionRecord).cancelled?.BOOL ?? false,
+      closed:
+        (item.Attributes as ActivityExecutionRecord).closed?.BOOL ?? false,
     };
   }
 
-  async cancelActivity(executionId: string, seq: number) {
-    await this.props.dynamo.send(
+  /**
+   * An activity execution is closed when it already has a result (completed or failed).
+   *
+   * Close the activity to prevent others from emitting events for it and report to the activity worker
+   * that the activity is no longer able to report a result.
+   */
+  async closeActivity(executionId: string, seq: number) {
+    const item = await this.props.dynamo.send(
       new UpdateItemCommand({
         Key: {
           pk: { S: ActivityExecutionRecord.key(executionId, seq) },
         },
         UpdateExpression:
-          "SET cancelled=:cancelled, executionId = :executionId, seq = :seq",
+          "SET closed=:closed, executionId = :executionId, seq = :seq",
         ExpressionAttributeValues: {
-          ":cancelled": { BOOL: true },
+          ":closed": { BOOL: true },
           ":executionId": { S: executionId },
           ":seq": { N: `${seq}` },
         },
         TableName: this.props.activityTableName,
+        ReturnValues: ReturnValue.UPDATED_OLD,
       })
     );
+
+    return { alreadyClosed: item.Attributes?.["closed"]?.BOOL ?? false };
   }
 
   async getActivity(
@@ -132,7 +141,7 @@ export interface ActivityExecutionRecord {
   executionId: AttributeValue.SMember;
   seq: AttributeValue.NMember;
   heartbeatTime?: AttributeValue.SMember;
-  cancelled?: AttributeValue.BOOLMember;
+  closed?: AttributeValue.BOOLMember;
 }
 
 export namespace ActivityExecutionRecord {
@@ -148,7 +157,7 @@ function createActivityFromRecord(
   return {
     executionId: activityRecord.executionId.S,
     seq: Number(activityRecord.seq.N),
-    cancelled: Boolean(activityRecord.cancelled?.BOOL ?? false),
+    closed: Boolean(activityRecord.closed?.BOOL ?? false),
     heartbeatTime: activityRecord?.heartbeatTime?.S,
   };
 }

packages/@eventual/aws-runtime/src/handlers/orchestrator.ts

@@ -6,6 +6,7 @@ import type { SQSEvent, SQSRecord } from "aws-lambda";
 import { logger, loggerMiddlewares } from "../logger.js";
 import { AWSMetricsClient } from "../clients/metrics-client.js";
 import {
+  createActivityRuntimeClient,
   createEventClient,
   createExecutionHistoryClient,
   createTimerClient,
@@ -24,6 +25,7 @@ const orchestrate = createOrchestrator({
   workflowRuntimeClient: createWorkflowRuntimeClient(),
   workflowClient: createWorkflowClient(),
   eventClient: createEventClient(),
+  activityRuntimeClient: createActivityRuntimeClient(),
   metricsClient: AWSMetricsClient,
   logger,
 });