Merge pull request #73 from nibble-4bits/feature/task-state-timeout

Feature/task state timeout

__tests__/stateActions/TaskStateAction.test.ts

@@ -1,6 +1,8 @@
 import type { TaskState } from '../../src/typings/TaskState';
 import { LambdaClient } from '../../src/aws/LambdaClient';
 import { TaskStateAction } from '../../src/stateMachine/stateActions/TaskStateAction';
+import { StatesTimeoutError } from '../../src/error/predefined/StatesTimeoutError';
+import { sleep } from '../../src/util';
 
 afterEach(() => {
   jest.clearAllMocks();
@@ -70,4 +72,48 @@ describe('Task State', () => {
     expect(mockInvokeFunction).not.toHaveBeenCalled();
     expect(stateResult).toBe(8);
   });
+
+  test('should throw `States.Timeout` error if action runs longer than the value specified in `TimeoutSeconds` field', async () => {
+    const definition: TaskState = {
+      Type: 'Task',
+      Resource: 'mock-arn',
+      TimeoutSeconds: 1,
+      End: true,
+    };
+    const stateName = 'TaskState';
+    const input = { num1: 5, num2: 3 };
+    const context = {};
+
+    const localHandlerFn = async () => {
+      await sleep(1100);
+      return 1;
+    };
+    const options = { overrideFn: localHandlerFn, awsConfig: undefined };
+
+    const taskStateAction = new TaskStateAction(definition, stateName);
+
+    await expect(() => taskStateAction.execute(input, context, options)).rejects.toThrow(StatesTimeoutError);
+  });
+
+  test('should throw `States.Timeout` error if action runs longer than the value specified in `TimeoutSecondsPath` field', async () => {
+    const definition: TaskState = {
+      Type: 'Task',
+      Resource: 'mock-arn',
+      TimeoutSecondsPath: '$.taskTimeout',
+      End: true,
+    };
+    const stateName = 'TaskState';
+    const input = { num1: 5, num2: 3, taskTimeout: 1 };
+    const context = {};
+
+    const localHandlerFn = async () => {
+      await sleep(1100);
+      return 1;
+    };
+    const options = { overrideFn: localHandlerFn, awsConfig: undefined };
+
+    const taskStateAction = new TaskStateAction(definition, stateName);
+
+    await expect(() => taskStateAction.execute(input, context, options)).rejects.toThrow(StatesTimeoutError);
+  });
 });

docs/feature-support.md

@@ -38,6 +38,8 @@ The following features all come from the [Amazon States Language specification](
     - [x] `TimestampPath`
   - [x] Task
     - [x] `Resource` (only Lambda functions supported)
+    - [x] `TimeoutSeconds`
+    - [x] `TimeoutSecondsPath`
   - [x] Parallel
     - [x] `Branches`
   - [x] Map
@@ -106,9 +108,7 @@ The following features all come from the [Amazon States Language specification](
 
 - States
   - Task
-    - [ ] `TimeoutSeconds`
     - [ ] `HeartbeatSeconds`
-    - [ ] `TimeoutSecondsPath`
     - [ ] `HeartbeatSecondsPath`
     - [ ] `Credentials`
   - Map

src/stateMachine/stateActions/TaskStateAction.ts

@@ -2,25 +2,69 @@ import type { TaskState } from '../../typings/TaskState';
 import type { JSONValue } from '../../typings/JSONValue';
 import type { ActionResult, TaskStateActionOptions } from '../../typings/StateActions';
 import type { Context } from '../../typings/Context';
+import { StatesTimeoutError } from '../../error/predefined/StatesTimeoutError';
 import { BaseStateAction } from './BaseStateAction';
 import { LambdaClient } from '../../aws/LambdaClient';
+import { jsonPathQuery } from '../jsonPath/JsonPath';
+import { IntegerConstraint } from '../jsonPath/constraints/IntegerConstraint';
 
 class TaskStateAction extends BaseStateAction<TaskState> {
+  private timeoutAbortController: AbortController;
+
   constructor(stateDefinition: TaskState, stateName: string) {
     super(stateDefinition, stateName);
+    this.timeoutAbortController = new AbortController();
+  }
+
+  private createTimeoutPromise(input: JSONValue, context: Context): Promise<never> | undefined {
+    const state = this.stateDefinition;
+
+    if (!state.TimeoutSeconds && !state.TimeoutSecondsPath) return;
+
+    let timeout: number;
+    if (state.TimeoutSeconds) {
+      timeout = state.TimeoutSeconds;
+    } else if (state.TimeoutSecondsPath) {
+      timeout = jsonPathQuery<number>(state.TimeoutSecondsPath, input, context, {
+        constraints: [IntegerConstraint.greaterThanOrEqual(1)],
+      });
+    }
+
+    return new Promise<never>((_, reject) => {
+      const handleTimeoutAbort = () => clearTimeout(timeoutId);
+
+      const timeoutId = setTimeout(() => {
+        this.timeoutAbortController.signal.removeEventListener('abort', handleTimeoutAbort);
+        reject(new StatesTimeoutError());
+      }, timeout * 1000);
+
+      this.timeoutAbortController.signal.addEventListener('abort', handleTimeoutAbort, { once: true });
+    });
   }
 
   override async execute(input: JSONValue, context: Context, options: TaskStateActionOptions): Promise<ActionResult> {
     const state = this.stateDefinition;
+    const racingPromises = [];
+    const timeoutPromise = this.createTimeoutPromise(input, context);
 
-    // If local override for task resource is defined, use that
     if (options.overrideFn) {
-      const result = await options.overrideFn(input);
-      return this.buildExecutionResult(result);
+      // If local override for task resource is defined, use that
+      const resultPromise = options.overrideFn(input);
+      racingPromises.push(resultPromise);
+    } else {
+      // Else, call Lambda in AWS using SDK
+      const lambdaClient = new LambdaClient(options.awsConfig);
+      const resultPromise = lambdaClient.invokeFunction(state.Resource, input);
+      racingPromises.push(resultPromise);
     }
 
-    const lambdaClient = new LambdaClient(options.awsConfig);
-    const result = await lambdaClient.invokeFunction(state.Resource, input);
+    if (timeoutPromise) {
+      racingPromises.push(timeoutPromise);
+    }
+
+    const result = await Promise.race(racingPromises);
+    this.timeoutAbortController.abort();
+
     return this.buildExecutionResult(result);
   }
 }

src/typings/TaskState.ts

@@ -21,6 +21,8 @@ interface BaseTaskState
     CatchableState {
   Type: 'Task';
   Resource: string;
+  TimeoutSeconds?: number;
+  TimeoutSecondsPath?: string;
 }
 
 export type TaskState = (IntermediateState | TerminalState) & BaseTaskState;