feat: add retriers

Author: sam-n-johnston

src/stateTasks/executors/TaskExecutor.ts

@@ -8,6 +8,7 @@ import type { TaskStateDefinition } from '../../types/State';
 import { StateInfoHandler } from '../../StateInfoHandler';
 import { StateProcessor } from '../../StateProcessor';
 import { Context } from '../../Context/Context';
+import { Retriers } from '../../types/Retriers';
 
 export class TaskExecutor extends StateTypeExecutor {
   public async execute(
@@ -27,7 +28,16 @@ export class TaskExecutor extends StateTypeExecutor {
     const functionLambda = await import(`${lambdaPath}`);
 
     this.injectEnvVarsLambdaSpecific(stateInfo.environment);
-    const output = await functionLambda[stateInfo.handlerName](input, context);
+
+    // eslint-disable-next-line @typescript-eslint/ban-types
+    let output: any;
+    if (stateDefinition.Retry) {
+      const retrier = Retriers.create(stateDefinition.Retry);
+      output = await retrier.retry(functionLambda[stateInfo.handlerName].bind(this, input, context));
+    } else {
+      output = await functionLambda[stateInfo.handlerName](input, context);
+    }
+
     this.removeEnvVarsLambdaSpecific(stateInfo.environment);
 
     const outputJson = this.processOutput(input, output, stateDefinition);

src/types/Retrier.ts

@@ -0,0 +1,54 @@
+import { TaskRetryRule } from './State';
+
+export enum StatesErrors {
+  DataLimitExceeded = 'States.DataLimitExceeded',
+  Runtime = 'States.Runtime',
+  Timeout = 'States.Timeout',
+  TaskFailed = 'States.TaskFailed',
+  Permissions = 'States.Permissions',
+}
+
+export class Retrier {
+  private currentNumberOfRetries = 0;
+  private currentIntervalSeconds: number;
+
+  private constructor(
+    private readonly _ErrorEquals: StatesErrors[],
+    private readonly _IntervalSeconds: number = 1,
+    private readonly _MaxAttempts: number = 3,
+    private readonly _BackoffRate: number = 2.0,
+  ) {
+    this.currentIntervalSeconds = _IntervalSeconds;
+  }
+
+  public static create(taskRetryRule: TaskRetryRule): Retrier {
+    return new Retrier(
+      taskRetryRule.ErrorEquals,
+      taskRetryRule.IntervalSeconds,
+      taskRetryRule.MaxAttempts,
+      taskRetryRule.BackoffRate,
+    );
+  }
+
+  async retry(fn: () => any): Promise<any> {
+    try {
+      return await fn();
+    } catch (error) {
+      if (this.currentNumberOfRetries < this._MaxAttempts) {
+        this.currentNumberOfRetries++;
+        return await new Promise((resolve) => {
+          setTimeout(async () => {
+            this.currentIntervalSeconds = this.currentIntervalSeconds * this._BackoffRate;
+            return resolve(await this.retry(fn));
+          }, this.currentIntervalSeconds * 1000);
+        });
+      } else {
+        throw error;
+      }
+    }
+  }
+
+  get ErrorEquals(): StatesErrors[] {
+    return this._ErrorEquals;
+  }
+}

src/types/Retriers.ts

@@ -0,0 +1,35 @@
+import { Retrier, StatesErrors } from './Retrier';
+import { TaskRetryRule } from './State';
+
+export class Retriers {
+  private constructor(private readonly _retriers: Retrier[]) {}
+
+  public static create(taskRetryRules: TaskRetryRule[]): Retriers {
+    const retriers = taskRetryRules.map((taskRetryRule) => {
+      return Retrier.create(taskRetryRule);
+    });
+    return new Retriers(retriers);
+  }
+
+  private getRetrierBasedOn(statesError: StatesErrors): Retrier | undefined {
+    const retrier = this._retriers.find((retrier) => {
+      return retrier.ErrorEquals.includes(statesError);
+    });
+    if (!retrier) {
+      return;
+    }
+    return retrier;
+  }
+
+  // TODO: Add the number of retries to the context object
+  async retry(fn: () => any): Promise<any> {
+    // TODO: Add timeout error catching
+    const retrier = this.getRetrierBasedOn(StatesErrors.TaskFailed);
+
+    if (retrier) {
+      return await retrier.retry(fn);
+    } else {
+      return fn();
+    }
+  }
+}

src/types/State.ts

@@ -1,5 +1,6 @@
 import type { StateMachineDefinition } from './StateMachine';
 import { StateType } from '../stateTasks/StateType';
+import { StatesErrors } from './Retrier';
 
 export type StateInfo = {
   handlerPath: string;
@@ -84,23 +85,15 @@ export type PassStateDefinition = CommonStateDefinition & {
   Parameters?: string; // TODO: TBD
 };
 
-export type TaskErrorName =
-  | 'State.ALL'
-  | 'States.DataLimitExceeded'
-  | 'States.Runtime'
-  | 'States.Timeout'
-  | 'States.TaskFailed'
-  | 'States.Permissions';
-
 export type TaskRetryRule = {
-  ErrorEquals: TaskErrorName[];
+  ErrorEquals: StatesErrors[];
   IntervalSeconds?: number;
   MaxAttempts?: number;
   BackoffRate?: number;
 };
 
 export type TaskCatchRule = {
-  ErrorEquals: TaskErrorName[];
+  ErrorEquals: StatesErrors[];
   Next: string;
   ResultPath?: string;
 };

tests/src/Retrier.test.ts

@@ -0,0 +1,51 @@
+import { Retrier, StatesErrors } from '../../src/types/Retrier';
+
+describe('Retrier', () => {
+  describe('when the function succeeds', () => {
+    it('should call it once', () => {
+      const retrier = Retrier.create({ ErrorEquals: [StatesErrors.TaskFailed] });
+      const retriedFunction = jest.fn();
+      retrier.retry(retriedFunction);
+      expect(retriedFunction).toHaveBeenCalledTimes(1);
+    });
+
+    describe('when the function is async', () => {
+      it('should call it once', () => {
+        const retrier = Retrier.create({ ErrorEquals: [StatesErrors.TaskFailed] });
+        const retriedFunction = jest.fn().mockImplementation(() => Promise.resolve);
+        retrier.retry(retriedFunction);
+        expect(retriedFunction).toHaveBeenCalledTimes(1);
+      });
+    });
+  });
+
+  describe('when the function fails', () => {
+    it('should call it 3 times', () => {
+      const retrier = Retrier.create({ ErrorEquals: [StatesErrors.TaskFailed] });
+      const retriedFunction = jest.fn().mockImplementation(() => {
+        throw new Error('MyError');
+      });
+
+      try {
+        retrier.retry(retriedFunction);
+      } catch (error) {
+        expect(error.message).toEqual('MyError');
+        expect(retriedFunction).toHaveBeenCalledTimes(3);
+      }
+    });
+
+    it('should call it 5 times', () => {
+      const retrier = Retrier.create({ ErrorEquals: [StatesErrors.TaskFailed], MaxAttempts: 5 });
+      const retriedFunction = jest.fn().mockImplementation(() => {
+        throw new Error('MyError');
+      });
+
+      try {
+        retrier.retry(retriedFunction);
+      } catch (error) {
+        expect(error.message).toEqual('MyError');
+        expect(retriedFunction).toHaveBeenCalledTimes(5);
+      }
+    });
+  });
+});