fix(custom-resources): inactive lambda functions fail on invoke

packages/@aws-cdk/custom-resources/lib/provider-framework/runtime/outbound.ts

@@ -43,7 +43,31 @@ async function defaultInvokeFunction(req: AWS.Lambda.InvocationRequest): Promise
     lambda = new AWS.Lambda(awsSdkConfig);
   }
 
-  return lambda.invoke(req).promise();
+  try {
+    /**
+     * Try an initial invoke.
+     *
+     * When you try to invoke a function that is inactive, the invocation fails and Lambda sets
+     * the function to pending state until the function resources are recreated.
+     * If Lambda fails to recreate the resources, the function is set to the inactive state.
+     *
+     * We're using invoke first because `waitFor` doesn't trigger an inactive function to do anything,
+     * it just runs `getFunction` and checks the state.
+     */
+    return await lambda.invoke(req).promise();
+  } catch (error) {
+
+    /**
+     * The status of the Lambda function is checked every second for up to 300 seconds.
+     * Exits the loop on 'Active' state and throws an error on 'Inactive' or 'Failed'.
+     *
+     * And now we wait.
+     */
+    await lambda.waitFor('functionActiveV2', {
+      FunctionName: req.FunctionName,
+    }).promise();
+    return await lambda.invoke(req).promise();
+  }
 }
 
 export let startExecution = defaultStartExecution;

packages/@aws-cdk/custom-resources/test/provider-framework/outbound.test.ts

@@ -0,0 +1,174 @@
+import * as aws from 'aws-sdk';
+import { invokeFunction } from '../../lib/provider-framework/runtime/outbound';
+
+jest.mock('aws-sdk', () => {
+  return {
+    Lambda: class {
+      public invoke() {
+        return { promise: () => mockInvoke() };
+      }
+
+      public waitFor() {
+        return { promise: () => mockWaitFor() };
+      }
+    },
+  };
+});
+
+let mockInvoke: () => Promise<aws.Lambda.InvocationResponse>;
+
+const req: aws.Lambda.InvocationRequest = {
+  FunctionName: 'Whatever',
+  Payload: {
+    IsThisATest: 'Yes, this is a test',
+    AreYouSure: 'Yes, I am sure',
+  },
+};
+
+let invokeCount: number = 0;
+let expectedFunctionStates: string[] = [];
+let receivedFunctionStates: string[] = [];
+
+const mockWaitFor = async (): Promise<aws.Lambda.GetFunctionResponse> => {
+  let state = expectedFunctionStates.pop();
+  while (state !== 'Active') {
+    receivedFunctionStates.push(state!);
+    // If it goes back to inactive it's failed
+    if (state === 'Inactive') throw new Error('Not today');
+    // If failed... it's failed
+    if (state === 'Failed') throw new Error('Broken');
+    // If pending, continue the loop, no other valid options
+    if (state !== 'Pending') throw new Error('State is confused');
+    state = expectedFunctionStates.pop();
+  }
+  receivedFunctionStates.push(state);
+  return {
+    Configuration: {
+      State: 'Active',
+    },
+  };
+};
+
+describe('invokeFunction tests', () => {
+  afterEach(() => {
+    invokeCount = 0;
+    expectedFunctionStates = [];
+    receivedFunctionStates = [];
+  });
+
+  // Success cases
+  test('Inactive function that reactivates does not throw error', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 0) {
+        invokeCount++;
+        throw new Error('Better luck next time');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('Active');
+    expectedFunctionStates.push('Pending');
+
+    expect(await invokeFunction(req)).toEqual({ Payload: req.Payload });
+    expect(invokeCount).toEqual(2);
+    expect(receivedFunctionStates).toEqual(['Pending', 'Active']);
+  });
+
+  test('Active function does not run waitFor or retry invoke', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 1) {
+        invokeCount++;
+        throw new Error('This should not happen in this test');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('Active');
+
+    expect(await invokeFunction(req)).toEqual({ Payload: req.Payload });
+    expect(invokeCount).toEqual(1);
+    expect(receivedFunctionStates).toEqual([]);
+  });
+
+  // Failure cases
+  test('Inactive function that goes back to inactive throws error', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 0) {
+        invokeCount++;
+        throw new Error('Better luck next time');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('Inactive');
+    expectedFunctionStates.push('Pending');
+    expectedFunctionStates.push('Pending');
+
+    await expect(invokeFunction(req)).rejects.toThrowError(new Error('Not today'));
+    expect(invokeCount).toEqual(1);
+    expect(receivedFunctionStates).toEqual(['Pending', 'Pending', 'Inactive']);
+  });
+
+  test('Inactive function that goes to failed throws error', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 0) {
+        invokeCount++;
+        throw new Error('Better luck next time');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('Failed');
+    expectedFunctionStates.push('Pending');
+    expectedFunctionStates.push('Pending');
+
+    await expect(invokeFunction(req)).rejects.toThrowError(new Error('Broken'));
+    expect(invokeCount).toEqual(1);
+    expect(receivedFunctionStates).toEqual(['Pending', 'Pending', 'Failed']);
+  });
+
+  test('Inactive function that returns other value throws error', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 0) {
+        invokeCount++;
+        throw new Error('Better luck next time');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('NewFunctionWhoDis');
+    expectedFunctionStates.push('Pending');
+    expectedFunctionStates.push('Pending');
+
+    await expect(invokeFunction(req)).rejects.toThrowError(new Error('State is confused'));
+    expect(invokeCount).toEqual(1);
+    expect(receivedFunctionStates).toEqual(['Pending', 'Pending', 'NewFunctionWhoDis']);
+  });
+
+  test('Wait for stops on terminal responses', async () => {
+    mockInvoke = async () => {
+      if (invokeCount == 0) {
+        invokeCount++;
+        throw new Error('Better luck next time');
+      }
+      invokeCount++;
+      return { Payload: req.Payload };
+    };
+
+    expectedFunctionStates.push('SomethingElse');
+    expectedFunctionStates.push('Pending');
+    expectedFunctionStates.push('Inactive');
+    expectedFunctionStates.push('Pending');
+    expectedFunctionStates.push('Pending');
+
+    await expect(invokeFunction(req)).rejects.toThrowError(new Error('Not today'));
+    expect(invokeCount).toEqual(1);
+    expect(receivedFunctionStates).toEqual(['Pending', 'Pending', 'Inactive']);
+  });
+});
+