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

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,17 @@ 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();
 
-      // 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();
+      await this.waitForLastUpdateStatusComplete(updateFunctionCodeResponse, lambda);
 
+      if (this.lambdaFunctionResource.publishVersion) {
         const publishVersionPromise = lambda.publishVersion({
           FunctionName: this.lambdaFunctionResource.physicalName,
         }).promise();
@@ -269,8 +262,6 @@ class LambdaFunctionHotswapOperation implements HotswapOperation {
         } else {
           operations.push(publishVersionPromise);
         }
-      } else {
-        operations.push(updateFunctionCodePromise);
       }
     }
 
@@ -304,6 +295,58 @@ 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 waitForLastUpdateStatusComplete(currentFunctionConfiguration: AWS.Lambda.FunctionConfiguration, lambda: AWS.Lambda): Promise<void> {
+    let delay = 1; // in seconds
+    // 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
+    if (currentFunctionConfiguration.VpcConfig || currentFunctionConfiguration.PackageType === 'Image') {
+      delay = 5;
+    }
+
+    // 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,
+      acceptors: [
+        {
+          state: 'success',
+          matcher: 'path',
+          argument: "(Configuration.LastUpdateStatus == 'Successful' && Configuration.State == 'Active')",
+          expected: true,
+        },
+        {
+          state: 'failure',
+          matcher: 'path',
+          argument: 'Configuration.LastUpdateStatus',
+          expected: 'Failed',
+        },
+        {
+          state: 'retry',
+          matcher: 'path',
+          argument: 'Configuration.LastUpdateStatus',
+          expected: 'InProgress',
+        },
+      ],
+    };
+
+    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,11 +8,23 @@ let mockGetEndpointSuffix: () => string;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({});
   mockUpdateMachineDefinition = jest.fn();
   mockGetEndpointSuffix = jest.fn(() => 'amazonaws.com');
   hotswapMockSdkProvider.stubLambda({
     updateFunctionCode: mockUpdateLambdaCode,
+  }, {
+    // these are needed for the waiter API that the Lambda service hotswap uses
+    api: {
+      waiters: {},
+    },
+    makeRequest() {
+      return {
+        promise: () => Promise.resolve({}),
+        response: {},
+        addListeners: () => {},
+      };
+    },
   });
   hotswapMockSdkProvider.setUpdateStateMachineMock(mockUpdateMachineDefinition);
   hotswapMockSdkProvider.stubGetEndpointSuffix(mockGetEndpointSuffix);

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

@@ -83,8 +83,8 @@ export class HotswapMockSdkProvider {
     });
   }
 
-  public stubLambda(stubs: SyncHandlerSubsetOf<AWS.Lambda>) {
-    this.mockSdkProvider.stubLambda(stubs);
+  public stubLambda(stubs: SyncHandlerSubsetOf<AWS.Lambda>, additionalProperties: { [key: string]: any } = {}) {
+    this.mockSdkProvider.stubLambda(stubs, additionalProperties);
   }
 
   public setUpdateProjectMock(mockUpdateProject: (input: codebuild.UpdateProjectInput) => codebuild.UpdateProjectOutput) {

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

@@ -5,16 +5,30 @@ 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) => {};
 
 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,
+  }, {
+    // these are needed for the waiter API that the Lambda service hotswap uses
+    api: {
+      waiters: {},
+    },
+    makeRequest: mockMakeRequest,
   });
 });
 
@@ -65,3 +79,54 @@ 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
+  const waiters = (hotswapMockSdkProvider.mockSdkProvider.sdk.lambda() as any).api.waiters;
+  expect(mockMakeRequest).toHaveBeenCalledWith('getFunction', { FunctionName: 'my-function' });
+  expect(waiters).toEqual(expect.objectContaining({
+    updateFunctionCodeToFinish: expect.objectContaining({
+      name: 'UpdateFunctionCodeToFinish',
+      delay: 5,
+    }),
+  }));
+});

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

@@ -4,17 +4,29 @@ import * as setup from './hotswap-test-setup';
 let mockUpdateLambdaCode: (params: Lambda.Types.UpdateFunctionCodeRequest) => Lambda.Types.FunctionConfiguration;
 let mockTagResource: (params: Lambda.Types.TagResourceRequest) => {};
 let mockUntagResource: (params: Lambda.Types.UntagResourceRequest) => {};
+let mockMakeRequest: (operation: string, params: any) => {};
 let hotswapMockSdkProvider: setup.HotswapMockSdkProvider;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({});
   mockTagResource = jest.fn();
   mockUntagResource = jest.fn();
+  mockMakeRequest = jest.fn().mockReturnValue({
+    promise: () => Promise.resolve({}),
+    response: {},
+    addListeners: () => {},
+  });
   hotswapMockSdkProvider.stubLambda({
     updateFunctionCode: mockUpdateLambdaCode,
     tagResource: mockTagResource,
     untagResource: mockUntagResource,
+  }, {
+    // these are needed for the waiter API that the Lambda service hotswap uses
+    api: {
+      waiters: {},
+    },
+    makeRequest: mockMakeRequest,
   });
 });
 
@@ -539,3 +551,56 @@ test('does not call the updateLambdaCode() API when a resource with type that is
   expect(deployStackResult).toBeUndefined();
   expect(mockUpdateLambdaCode).not.toHaveBeenCalled();
 });
+
+test('calls getFunction() after function code is updated', async () => {
+  // GIVEN
+  setup.setCurrentCfnStackTemplate({
+    Resources: {
+      Func: {
+        Type: 'AWS::Lambda::Function',
+        Properties: {
+          Code: {
+            S3Bucket: 'current-bucket',
+            S3Key: 'current-key',
+          },
+          FunctionName: 'my-function',
+        },
+        Metadata: {
+          'aws:asset:path': 'old-path',
+        },
+      },
+    },
+  });
+  const cdkStackArtifact = setup.cdkStackArtifactOf({
+    template: {
+      Resources: {
+        Func: {
+          Type: 'AWS::Lambda::Function',
+          Properties: {
+            Code: {
+              S3Bucket: 'current-bucket',
+              S3Key: 'new-key',
+            },
+            FunctionName: 'my-function',
+          },
+          Metadata: {
+            'aws:asset:path': 'new-path',
+          },
+        },
+      },
+    },
+  });
+
+  // WHEN
+  await hotswapMockSdkProvider.tryHotswapDeployment(cdkStackArtifact);
+  const waiters = (hotswapMockSdkProvider.mockSdkProvider.sdk.lambda() as any).api.waiters;
+
+  // THEN
+  expect(mockMakeRequest).toHaveBeenCalledWith('getFunction', { FunctionName: 'my-function' });
+  expect(waiters).toEqual(expect.objectContaining({
+    updateFunctionCodeToFinish: expect.objectContaining({
+      name: 'UpdateFunctionCodeToFinish',
+      delay: 1,
+    }),
+  }));
+});

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

@@ -8,13 +8,25 @@ let hotswapMockSdkProvider: setup.HotswapMockSdkProvider;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({});
   mockTagResource = jest.fn();
   mockUntagResource = jest.fn();
   hotswapMockSdkProvider.stubLambda({
     updateFunctionCode: mockUpdateLambdaCode,
     tagResource: mockTagResource,
     untagResource: mockUntagResource,
+  }, {
+    // these are needed for the waiter API that the Lambda service hotswap uses
+    api: {
+      waiters: {},
+    },
+    makeRequest() {
+      return {
+        promise: () => Promise.resolve({}),
+        response: {},
+        addListeners: () => {},
+      };
+    },
   });
 });
 

packages/aws-cdk/test/api/hotswap/lambda-versions-aliases-hotswap-deployments.test.ts

@@ -8,14 +8,26 @@ let hotswapMockSdkProvider: setup.HotswapMockSdkProvider;
 
 beforeEach(() => {
   hotswapMockSdkProvider = setup.setupHotswapTests();
-  mockUpdateLambdaCode = jest.fn();
+  mockUpdateLambdaCode = jest.fn().mockReturnValue({});
   mockPublishVersion = jest.fn();
   mockUpdateAlias = jest.fn();
   hotswapMockSdkProvider.stubLambda({
     updateFunctionCode: mockUpdateLambdaCode,
     publishVersion: mockPublishVersion,
     updateAlias: mockUpdateAlias,
     waitFor: jest.fn(),
+  }, {
+    // these are needed for the waiter API that the Lambda service hotswap uses
+    api: {
+      waiters: {},
+    },
+    makeRequest() {
+      return {
+        promise: () => Promise.resolve({}),
+        response: {},
+        addListeners: () => {},
+      };
+    },
   });
 });
 

packages/aws-cdk/test/util/mock-sdk.ts

@@ -102,8 +102,8 @@ export class MockSdkProvider extends SdkProvider {
     (this.sdk as any).ssm = jest.fn().mockReturnValue(partialAwsService<AWS.SSM>(stubs));
   }
 
-  public stubLambda(stubs: SyncHandlerSubsetOf<AWS.Lambda>) {
-    (this.sdk as any).lambda = jest.fn().mockReturnValue(partialAwsService<AWS.Lambda>(stubs));
+  public stubLambda(stubs: SyncHandlerSubsetOf<AWS.Lambda>, additionalProperties: { [key: string]: any } = {}) {
+    (this.sdk as any).lambda = jest.fn().mockReturnValue(partialAwsService<AWS.Lambda>(stubs, additionalProperties));
   }
 
   public stubStepFunctions(stubs: SyncHandlerSubsetOf<AWS.StepFunctions>) {