fix(cli): hotswap should wait for lambda's updateFunctionCode to complete (#18536)

There are [upcoming changes](https://aws.amazon.com/blogs/compute/coming-soon-expansion-of-aws-lambda-states-to-all-functions/)
that will rollout Lambda states to all Lambda Functions. Prior to
this update (current functionality) when you made an
`updateFunctionCode` request the function was immediately available for
both invocation and future updates. Once this change is rolled out this
will no longer be the case. With Lambda states, when you make an update
to a Lambda Function, it will not be available for future updates until
the `LastUpdateStatus` returns `Successful`.

This PR introduces a custom waiter that will wait for the update to
complete before proceeding. The waiter will wait until the
`State=Active` and the `LastUpdateStatus=Successful`.

The `State` controls whether or not the function can be invoked, and the
`LastUpdateStatus` controls whether the function can be updated. Based
on this, I am not considering a deployment complete until both are
successful. To see a more in depth analysis of the different values see #18386.

In my testing I found that the time it took for a function to go from
`LastUpdateStatus=InProgress` to `LastUpdateStatus=Successful` was:

- ~1 second for a zip Function not in a VPC
- ~25 seconds for a container Function or a Function in a VPC
- ~2 minutes to deploy a VPC function (best proxy for StateReasonCode=Restoring)

There are a couple of built in waiters that could have been used for
this, namely
[functionUpdated](https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/AWS/Lambda.html#functionUpdated-waiter).
This waiter uses `getFunctionConfiguration` which has a quota of 15
requests/second. In addition the waiter polls every 5 seconds and this
cannot be configured. Because hotswapping is sensitive to any latency
that is introduced, I created a custom waiter that uses `getFunction`.
`getFunction` has a quota of 100 requests/second and the custom waiter
can be configured to poll every 1 second or every 5 seconds depending on
what type of function is being updated.

fixes #18386


----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: corymhall

packages/aws-cdk/lib/api/hotswap/lambda-functions.ts

@@ -1,5 +1,6 @@
 import { Writable } from 'stream';
 import * as archiver from 'archiver';
+import * as AWS from 'aws-sdk';
 import { flatMap } from '../../util';
 import { ISDK } from '../aws-auth';
 import { CfnEvaluationException, EvaluateCloudFormationTemplate } from '../evaluate-cloudformation-template';
@@ -232,25 +233,18 @@ class LambdaFunctionHotswapOperation implements HotswapOperation {
     const operations: Promise<any>[] = [];
 
     if (resource.code !== undefined) {
-      const updateFunctionCodePromise = lambda.updateFunctionCode({
+      const updateFunctionCodeResponse = await lambda.updateFunctionCode({
         FunctionName: this.lambdaFunctionResource.physicalName,
         S3Bucket: resource.code.s3Bucket,
         S3Key: resource.code.s3Key,
         ImageUri: resource.code.imageUri,
         ZipFile: resource.code.functionCodeZip,
       }).promise();
 
+      await this.waitForLambdasCodeUpdateToFinish(updateFunctionCodeResponse, lambda);
+
       // only if the code changed is there any point in publishing a new Version
       if (this.lambdaFunctionResource.publishVersion) {
-        // we need to wait for the code update to be done before publishing a new Version
-        await updateFunctionCodePromise;
-        // if we don't wait for the Function to finish updating,
-        // we can get a "The operation cannot be performed at this time. An update is in progress for resource:"
-        // error when publishing a new Version
-        await lambda.waitFor('functionUpdated', {
-          FunctionName: this.lambdaFunctionResource.physicalName,
-        }).promise();
-
         const publishVersionPromise = lambda.publishVersion({
           FunctionName: this.lambdaFunctionResource.physicalName,
         }).promise();
@@ -269,8 +263,6 @@ class LambdaFunctionHotswapOperation implements HotswapOperation {
         } else {
           operations.push(publishVersionPromise);
         }
-      } else {
-        operations.push(updateFunctionCodePromise);
       }
     }
 
@@ -304,6 +296,53 @@ class LambdaFunctionHotswapOperation implements HotswapOperation {
     // run all of our updates in parallel
     return Promise.all(operations);
   }
+
+  /**
+   * After a Lambda Function is updated, it cannot be updated again until the
+   * `State=Active` and the `LastUpdateStatus=Successful`.
+   *
+   * Depending on the configuration of the Lambda Function this could happen relatively quickly
+   * or very slowly. For example, Zip based functions _not_ in a VPC can take ~1 second whereas VPC
+   * or Container functions can take ~25 seconds (and 'idle' VPC functions can take minutes).
+   */
+  private async waitForLambdasCodeUpdateToFinish(currentFunctionConfiguration: AWS.Lambda.FunctionConfiguration, lambda: AWS.Lambda): Promise<void> {
+    const functionIsInVpcOrUsesDockerForCode = currentFunctionConfiguration.VpcConfig?.VpcId ||
+        currentFunctionConfiguration.PackageType === 'Image';
+
+    // if the function is deployed in a VPC or if it is a container image function
+    // then the update will take much longer and we can wait longer between checks
+    // otherwise, the update will be quick, so a 1-second delay is fine
+    const delaySeconds = functionIsInVpcOrUsesDockerForCode ? 5 : 1;
+
+    // configure a custom waiter to wait for the function update to complete
+    (lambda as any).api.waiters.updateFunctionCodeToFinish = {
+      name: 'UpdateFunctionCodeToFinish',
+      operation: 'getFunction',
+      // equates to 1 minute for zip function not in a VPC and
+      // 5 minutes for container functions or function in a VPC
+      maxAttempts: 60,
+      delay: delaySeconds,
+      acceptors: [
+        {
+          matcher: 'path',
+          argument: "Configuration.LastUpdateStatus == 'Successful' && Configuration.State == 'Active'",
+          expected: true,
+          state: 'success',
+        },
+        {
+          matcher: 'path',
+          argument: 'Configuration.LastUpdateStatus',
+          expected: 'Failed',
+          state: 'failure',
+        },
+      ],
+    };
+
+    const updateFunctionCodeWaiter = new (AWS as any).ResourceWaiter(lambda, 'updateFunctionCodeToFinish');
+    await updateFunctionCodeWaiter.wait({
+      FunctionName: this.lambdaFunctionResource.physicalName,
+    }).promise();
+  }
 }
 
 /**

packages/aws-cdk/test/api/hotswap/hotswap-deployments.test.ts

@@ -8,7 +8,7 @@ let mockGetEndpointSuffix: () => string;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({});
   mockUpdateMachineDefinition = jest.fn();
   mockGetEndpointSuffix = jest.fn(() => 'amazonaws.com');
   hotswapMockSdkProvider.stubLambda({

packages/aws-cdk/test/api/hotswap/hotswap-test-setup.ts

@@ -83,8 +83,29 @@ export class HotswapMockSdkProvider {
     });
   }
 
-  public stubLambda(stubs: SyncHandlerSubsetOf<AWS.Lambda>) {
-    this.mockSdkProvider.stubLambda(stubs);
+  public stubLambda(
+    stubs: SyncHandlerSubsetOf<AWS.Lambda>,
+    serviceStubs?: SyncHandlerSubsetOf<AWS.Service>,
+    additionalProperties: { [key: string]: any } = {},
+  ): void {
+    this.mockSdkProvider.stubLambda(stubs, {
+      api: {
+        waiters: {},
+      },
+      makeRequest() {
+        return {
+          promise: () => Promise.resolve({}),
+          response: {},
+          addListeners: () => {},
+        };
+      },
+      ...serviceStubs,
+      ...additionalProperties,
+    });
+  }
+
+  public getLambdaApiWaiters(): { [key: string]: any } {
+    return (this.mockSdkProvider.sdk.lambda() as any).api.waiters;
   }
 
   public setUpdateProjectMock(mockUpdateProject: (input: codebuild.UpdateProjectInput) => codebuild.UpdateProjectOutput) {

packages/aws-cdk/test/api/hotswap/lambda-functions-docker-hotswap-deployments.test.ts

@@ -5,16 +5,26 @@ let mockUpdateLambdaCode: (params: Lambda.Types.UpdateFunctionCodeRequest) => La
 let mockTagResource: (params: Lambda.Types.TagResourceRequest) => {};
 let mockUntagResource: (params: Lambda.Types.UntagResourceRequest) => {};
 let hotswapMockSdkProvider: setup.HotswapMockSdkProvider;
+let mockMakeRequest: (operation: string, params: any) => AWS.Request<any, AWS.AWSError>;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({
+    PackageType: 'Image',
+  });
   mockTagResource = jest.fn();
   mockUntagResource = jest.fn();
+  mockMakeRequest = jest.fn().mockReturnValue({
+    promise: () => Promise.resolve({}),
+    response: {},
+    addListeners: () => {},
+  });
   hotswapMockSdkProvider.stubLambda({
     updateFunctionCode: mockUpdateLambdaCode,
     tagResource: mockTagResource,
     untagResource: mockUntagResource,
+  }, {
+    makeRequest: mockMakeRequest,
   });
 });
 
@@ -65,3 +75,53 @@ test('calls the updateLambdaCode() API when it receives only a code difference i
     ImageUri: 'new-image',
   });
 });
+
+test('calls the getFunction() API with a delay of 5', async () => {
+  // GIVEN
+  setup.setCurrentCfnStackTemplate({
+    Resources: {
+      Func: {
+        Type: 'AWS::Lambda::Function',
+        Properties: {
+          Code: {
+            ImageUri: 'current-image',
+          },
+          FunctionName: 'my-function',
+        },
+        Metadata: {
+          'aws:asset:path': 'old-path',
+        },
+      },
+    },
+  });
+  const cdkStackArtifact = setup.cdkStackArtifactOf({
+    template: {
+      Resources: {
+        Func: {
+          Type: 'AWS::Lambda::Function',
+          Properties: {
+            Code: {
+              ImageUri: 'new-image',
+            },
+            FunctionName: 'my-function',
+          },
+          Metadata: {
+            'aws:asset:path': 'new-path',
+          },
+        },
+      },
+    },
+  });
+
+  // WHEN
+  await hotswapMockSdkProvider.tryHotswapDeployment(cdkStackArtifact);
+
+  // THEN
+  expect(mockMakeRequest).toHaveBeenCalledWith('getFunction', { FunctionName: 'my-function' });
+  expect(hotswapMockSdkProvider.getLambdaApiWaiters()).toEqual(expect.objectContaining({
+    updateFunctionCodeToFinish: expect.objectContaining({
+      name: 'UpdateFunctionCodeToFinish',
+      delay: 5,
+    }),
+  }));
+});