From 841a25540cf4fc43ce04f15a92de2064ea216dae Mon Sep 17 00:00:00 2001
From: Adam Ruka <adamruka@amazon.com>
Date: Mon, 21 Feb 2022 14:54:49 -0800
Subject: [PATCH] fix(cli): hotswapping is slow for many resources deployed at
 once

In the `LazyListStackResources` class,
we were incorrectly caching the results of the ListStackResources CloudFormation API call.
This resulted in executing redundant calls,
and very visibly impacted the performance of hotswapping when the Stack included many resources.

In manual testing, this change brought down deployment time of a large Stack from 80 to 7 seconds.

Closes #19021
---
 .../api/evaluate-cloudformation-template.ts   | 24 +++++++--------
 .../api/lazy-list-stack-resources.test.ts     | 29 +++++++++++++++++++
 2 files changed, 41 insertions(+), 12 deletions(-)
 create mode 100644 packages/aws-cdk/test/api/lazy-list-stack-resources.test.ts

diff --git a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
index 61d6a750b4482..405e51d7c2f9b 100644
--- a/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
+++ b/packages/aws-cdk/lib/api/evaluate-cloudformation-template.ts
@@ -7,30 +7,30 @@ export interface ListStackResources {
 }
 
 export class LazyListStackResources implements ListStackResources {
-  private stackResources: AWS.CloudFormation.StackResourceSummary[] | undefined;
+  private stackResources: Promise<AWS.CloudFormation.StackResourceSummary[]> | undefined;
 
   constructor(private readonly sdk: ISDK, private readonly stackName: string) {
   }
 
   public async listStackResources(): Promise<AWS.CloudFormation.StackResourceSummary[]> {
     if (this.stackResources === undefined) {
-      this.stackResources = await this.getStackResources();
+      this.stackResources = this.getStackResources(undefined);
     }
     return this.stackResources;
   }
 
-  private async getStackResources(): Promise<AWS.CloudFormation.StackResourceSummary[]> {
+  private async getStackResources(nextToken: string | undefined): Promise<AWS.CloudFormation.StackResourceSummary[]> {
     const ret = new Array<AWS.CloudFormation.StackResourceSummary>();
-    let nextToken: string | undefined;
-    do {
-      const stackResourcesResponse = await this.sdk.cloudFormation().listStackResources({
-        StackName: this.stackName,
-        NextToken: nextToken,
-      }).promise();
+    return this.sdk.cloudFormation().listStackResources({
+      StackName: this.stackName,
+      NextToken: nextToken,
+    }).promise().then(async stackResourcesResponse => {
       ret.push(...(stackResourcesResponse.StackResourceSummaries ?? []));
-      nextToken = stackResourcesResponse.NextToken;
-    } while (nextToken);
-    return ret;
+      if (stackResourcesResponse.NextToken) {
+        ret.push(...await this.getStackResources(stackResourcesResponse.NextToken));
+      }
+      return ret;
+    });
   }
 }
 
diff --git a/packages/aws-cdk/test/api/lazy-list-stack-resources.test.ts b/packages/aws-cdk/test/api/lazy-list-stack-resources.test.ts
new file mode 100644
index 0000000000000..64d0b98d62733
--- /dev/null
+++ b/packages/aws-cdk/test/api/lazy-list-stack-resources.test.ts
@@ -0,0 +1,29 @@
+import * as AWS from 'aws-sdk';
+import { LazyListStackResources } from '../../lib/api/evaluate-cloudformation-template';
+import { MockSdk } from '../util/mock-sdk';
+
+describe('Lazy ListStackResources', () => {
+  test('correctly caches calls to the CloudFormation API', async () => {
+    // GIVEN
+    const listStackResMock: jest.Mock<AWS.CloudFormation.ListStackResourcesOutput, AWS.CloudFormation.ListStackResourcesInput[]> = jest.fn();
+    const mockSdk = new MockSdk();
+    mockSdk.stubCloudFormation({
+      listStackResources: listStackResMock,
+    });
+    listStackResMock.mockReturnValue({
+      StackResourceSummaries: [],
+      NextToken: undefined,
+    });
+    const res = new LazyListStackResources(mockSdk, 'StackName');
+
+    // WHEN
+    void res.listStackResources();
+    void res.listStackResources();
+    void res.listStackResources();
+    const result = await res.listStackResources();
+
+    // THEN
+    expect(result.length).toBe(0);
+    expect(listStackResMock).toHaveBeenCalledTimes(1);
+  });
+});