refactor(toolkit-lib): batch asset-active lookups in cdk gc with a multi-pattern search

packages/@aws-cdk/toolkit-lib/lib/api/garbage-collection/garbage-collector.ts

@@ -274,7 +274,8 @@ export class GarbageCollector {
       for await (const batch of this.readRepoInBatches(ecr, repo, batchSize, currentTime)) {
         await backgroundStackRefresh.noOlderThan(600_000); // 10 mins
 
-        const { included: isolated, excluded: notIsolated } = partition(batch, asset => !asset.tags.some(t => activeAssets.contains(t)));
+        const foundTags = activeAssets.containsAny(batch.flatMap(asset => asset.tags));
+        const { included: isolated, excluded: notIsolated } = partition(batch, asset => !asset.tags.some(t => foundTags.has(t)));
 
         await this.ioHelper.defaults.debug(`${isolated.length} isolated images`);
         await this.ioHelper.defaults.debug(`${notIsolated.length} not isolated images`);
@@ -350,7 +351,8 @@ export class GarbageCollector {
       for await (const batch of this.readBucketInBatches(s3, bucket, batchSize, currentTime)) {
         await backgroundStackRefresh.noOlderThan(600_000); // 10 mins
 
-        const { included: isolated, excluded: notIsolated } = partition(batch, asset => !activeAssets.contains(asset.fileName()));
+        const foundFileNames = activeAssets.containsAny(batch.map(asset => asset.fileName()));
+        const { included: isolated, excluded: notIsolated } = partition(batch, asset => !foundFileNames.has(asset.fileName()));
 
         await this.ioHelper.defaults.debug(`${isolated.length} isolated assets`);
         await this.ioHelper.defaults.debug(`${notIsolated.length} not isolated assets`);

packages/@aws-cdk/toolkit-lib/lib/api/garbage-collection/stack-refresh.ts

@@ -3,20 +3,128 @@ import { ToolkitError } from '../../toolkit/toolkit-error';
 import type { ICloudFormationClient } from '../aws-auth/private';
 import type { IoHelper } from '../io/private';
 
+/**
+ * A multi-pattern substring search index (Aho-Corasick).
+ *
+ * Finds every pattern (from a fixed set) that occurs anywhere in a text, in a
+ * single pass over that text -- O(text.length + sum(pattern.length)) total,
+ * regardless of how many patterns are being searched for. This is what makes
+ * `ActiveAssetCache.containsAny()` scale: checking N asset hashes against M
+ * stack templates costs O(M templates scanned once + N pattern lengths), not
+ * O(N asset hashes x M stacks x template size).
+ */
+class AhoCorasickIndex {
+  private readonly children: Array<Map<string, number>> = [new Map()];
+  private readonly fail: number[] = [0];
+  // Pattern indices whose match ends at this node, INCLUDING those inherited
+  // via the fail-link chain -- so a single lookup at match time is enough.
+  private readonly output: Array<number[]> = [[]];
+
+  constructor(private readonly patterns: string[]) {
+    patterns.forEach((pattern, id) => this.insert(pattern, id));
+    this.buildFailureLinks();
+  }
+
+  private insert(pattern: string, id: number) {
+    let node = 0;
+    for (const ch of pattern) {
+      let next = this.children[node].get(ch);
+      if (next === undefined) {
+        next = this.children.length;
+        this.children.push(new Map());
+        this.fail.push(0);
+        this.output.push([]);
+        this.children[node].set(ch, next);
+      }
+      node = next;
+    }
+    if (pattern.length > 0) {
+      this.output[node].push(id);
+    }
+  }
+
+  private buildFailureLinks() {
+    const queue: number[] = [];
+    for (const child of this.children[0].values()) {
+      this.fail[child] = 0;
+      queue.push(child);
+    }
+
+    let head = 0;
+    while (head < queue.length) {
+      const node = queue[head++];
+      for (const [ch, child] of this.children[node]) {
+        queue.push(child);
+
+        let f = this.fail[node];
+        while (f !== 0 && !this.children[f].has(ch)) {
+          f = this.fail[f];
+        }
+        const candidate = this.children[f].get(ch);
+        this.fail[child] = candidate !== undefined && candidate !== child ? candidate : 0;
+
+        if (this.output[this.fail[child]].length > 0) {
+          this.output[child] = this.output[child].concat(this.output[this.fail[child]]);
+        }
+      }
+    }
+  }
+
+  /**
+   * Scans `text` once and adds the (pattern, not id) of every pattern found to `into`.
+   */
+  public search(text: string, into: Set<string>) {
+    let node = 0;
+    for (let i = 0; i < text.length; i++) {
+      const ch = text[i];
+      while (node !== 0 && !this.children[node].has(ch)) {
+        node = this.fail[node];
+      }
+      node = this.children[node].get(ch) ?? 0;
+
+      for (const id of this.output[node]) {
+        into.add(this.patterns[id]);
+      }
+    }
+  }
+}
+
 export class ActiveAssetCache {
-  private readonly stacks: Set<string> = new Set();
+  private readonly stacks: string[] = [];
 
   public rememberStack(stackTemplate: string) {
-    this.stacks.add(stackTemplate);
+    this.stacks.push(stackTemplate);
   }
 
+  /**
+   * Whether `asset` occurs anywhere in any remembered stack template.
+   */
   public contains(asset: string): boolean {
+    return this.containsAny([asset]).has(asset);
+  }
+
+  /**
+   * For a batch of candidate asset identifiers, returns the subset that occur
+   * anywhere in any remembered stack template. Scans every stack template exactly
+   * once no matter how many candidates are passed in -- callers that need to check
+   * many assets (as `cdk gc` does, in batches of up to 1000) should always prefer
+   * this over calling `contains()` in a loop.
+   */
+  public containsAny(assets: string[]): Set<string> {
+    const found = new Set<string>();
+    if (assets.length === 0) {
+      return found;
+    }
+
+    const uniqueAssetCount = new Set(assets).size;
+    const index = new AhoCorasickIndex(assets);
     for (const stack of this.stacks) {
-      if (stack.includes(asset)) {
-        return true;
+      if (found.size === uniqueAssetCount) {
+        break;
       }
+      index.search(stack, found);
     }
-    return false;
+    return found;
   }
 }
 

packages/@aws-cdk/toolkit-lib/test/api/garbage-collection/stack-refresh.test.ts

@@ -0,0 +1,115 @@
+import { ActiveAssetCache } from '../../../lib/api/garbage-collection/stack-refresh';
+
+describe(ActiveAssetCache, () => {
+  test('contains returns false when no stacks are remembered', () => {
+    const cache = new ActiveAssetCache();
+    expect(cache.contains('some-asset-hash')).toBe(false);
+  });
+
+  test('contains finds an asset referenced anywhere within a remembered template', () => {
+    const cache = new ActiveAssetCache();
+    cache.rememberStack('{"Resources":{"Bucket":{"Properties":{"Key":"prefix-abc123def-suffix.zip"}}}}');
+    expect(cache.contains('abc123def')).toBe(true);
+    expect(cache.contains('not-present')).toBe(false);
+  });
+
+  test('contains searches across all remembered stacks, not just the first', () => {
+    const cache = new ActiveAssetCache();
+    cache.rememberStack('{"Resources":{"A":"nothing-relevant-here"}}');
+    cache.rememberStack('{"Resources":{"B":"has-the-hash-xyz789-in-it"}}');
+    cache.rememberStack('{"Resources":{"C":"also-irrelevant"}}');
+    expect(cache.contains('xyz789')).toBe(true);
+  });
+
+  test('an asset hash cannot false-positive-match across a template boundary', () => {
+    const cache = new ActiveAssetCache();
+    // Concatenating these naively (without a separator) would form "foobar" at the boundary
+    cache.rememberStack('...foo');
+    cache.rememberStack('bar...');
+    expect(cache.contains('foobar')).toBe(false);
+  });
+
+  test('remembering a new stack after a contains() call is still picked up (no stale cache)', () => {
+    const cache = new ActiveAssetCache();
+    cache.rememberStack('template-one');
+    expect(cache.contains('late-hash')).toBe(false);
+
+    cache.rememberStack('template-two-with-late-hash');
+    expect(cache.contains('late-hash')).toBe(true);
+  });
+
+  describe('containsAny', () => {
+    test('returns exactly the subset of candidates that are present, across many stacks', () => {
+      const cache = new ActiveAssetCache();
+      cache.rememberStack('irrelevant-stack-one');
+      cache.rememberStack('stack-with-hashA-and-hashB');
+      cache.rememberStack('irrelevant-stack-two');
+      cache.rememberStack('another-stack-with-hashC');
+
+      const result = cache.containsAny(['hashA', 'hashB', 'hashC', 'hashD-not-present']);
+      expect(result).toEqual(new Set(['hashA', 'hashB', 'hashC']));
+    });
+
+    test('returns an empty set for an empty candidate list', () => {
+      const cache = new ActiveAssetCache();
+      cache.rememberStack('anything');
+      expect(cache.containsAny([])).toEqual(new Set());
+    });
+
+    test('handles duplicate candidates', () => {
+      const cache = new ActiveAssetCache();
+      cache.rememberStack('has-dupe-hash');
+      expect(cache.containsAny(['dupe-hash', 'dupe-hash', 'missing'])).toEqual(new Set(['dupe-hash']));
+    });
+
+    test('one pattern being a substring of another does not cause a miss', () => {
+      const cache = new ActiveAssetCache();
+      cache.rememberStack('template-contains-abcdef-only');
+      // 'abc' is a prefix of 'abcdef' -- both must independently register as found/not-found correctly
+      const result = cache.containsAny(['abcdef', 'abc', 'xyz']);
+      expect(result).toEqual(new Set(['abcdef', 'abc']));
+    });
+
+    test('does not false-positive-match across a template boundary', () => {
+      const cache = new ActiveAssetCache();
+      cache.rememberStack('...foo');
+      cache.rememberStack('bar...');
+      expect(cache.containsAny(['foobar'])).toEqual(new Set());
+    });
+
+    test('agrees with the naive per-candidate contains() on a randomized workload (no false negatives)', () => {
+      function randHash(rng: () => number) {
+        return Array.from({ length: 12 }, () => Math.floor(rng() * 16).toString(16)).join('');
+      }
+      // Simple deterministic PRNG so failures are reproducible. Math.imul keeps
+      // the multiply within a 32-bit-safe integer (no bitwise operators, no
+      // precision loss from JS doubles on the full product).
+      let seed = 42;
+      const rng = () => {
+        seed = Math.imul(seed, 1103515245) + 12345;
+        seed = seed % 0x7fffffff;
+        if (seed < 0) {
+          seed += 0x7fffffff;
+        }
+        return seed / 0x7fffffff;
+      };
+
+      const cache = new ActiveAssetCache();
+      const embeddedHashes: string[] = [];
+      for (let i = 0; i < 20; i++) {
+        const h = randHash(rng);
+        embeddedHashes.push(h);
+        cache.rememberStack(`{"Resources":{"R${i}":{"Key":"prefix-${h}-suffix"}}}`);
+      }
+
+      const candidates: string[] = [];
+      for (let i = 0; i < 100; i++) {
+        candidates.push(rng() < 0.4 ? embeddedHashes[Math.floor(rng() * embeddedHashes.length)] : randHash(rng));
+      }
+
+      const expected = new Set(candidates.filter((c) => cache.contains(c)));
+      const actual = cache.containsAny(candidates);
+      expect(actual).toEqual(expected);
+    });
+  });
+});