fix: aspect in the middle

Author: rix0rrr

packages/aws-cdk-lib/aws-s3/test/bucket.test.ts

@@ -6,6 +6,9 @@ import * as kms from '../../aws-kms';
 import * as cdk from '../../core';
 import * as s3 from '../lib';
 import { ReplicationTimeValue } from '../lib/bucket';
+import { IConstruct } from 'constructs';
+import * as lambda from '../../aws-lambda';
+import * as lambdaev from '../../aws-lambda-event-sources';
 
 // to make it easy to copy & paste from output:
 /* eslint-disable quote-props */
@@ -4795,3 +4798,36 @@ describe('bucket', () => {
     });
   });
 });
+
+test('aspect prio test', () => {
+  class TagNameAspect implements cdk.IAspect {
+    public visit(node: IConstruct): void {
+      if (node instanceof s3.Bucket) {
+        cdk.Tags.of(node).add('Name', 'asdf', { priority: 10 });
+      }
+    }
+  }
+  const app = new cdk.App({
+    postCliContext: {
+      '@aws-cdk/core:aspectStabilization': true,
+    },
+  });
+  const stack = new cdk.Stack(app, 'stack');
+  var bucket = new s3.Bucket(stack, 'test-bucket', {
+    bucketName: 'test-bucket',
+  });
+  const fn = new lambda.Function(stack, 'test-function', {
+    functionName: 'test-function',
+    handler: 'handler',
+    runtime: lambda.Runtime.NODEJS_22_X,
+    code: lambda.Code.fromInline('function(){ }'),
+  });
+
+  fn.addEventSource(new lambdaev.S3EventSourceV2(bucket, {
+    events: [s3.EventType.OBJECT_CREATED],
+  }));
+
+  cdk.Aspects.of(stack).add(new TagNameAspect(), { priority: 10 });
+
+  app.synth();
+});

packages/aws-cdk-lib/core/lib/aspect.ts

@@ -90,6 +90,7 @@ export class Aspects {
       return;
     }
     this._appliedAspects.push(newApplication);
+    bumpAspectTreeRevision(this._scope);
   }
 
   /**
@@ -109,6 +110,33 @@ export class Aspects {
   }
 }
 
+/**
+ * Bump the tree revision for the tree containing the given construct.
+ *
+ * Module-local function since I don't want to open up the door for *anyone* calling this.
+ */
+function bumpAspectTreeRevision(construct: IConstruct) {
+  const root = construct.node.root;
+  const rev = (root as any)[TREE_REVISION_SYM] ?? 0;
+  (root as any)[TREE_REVISION_SYM] = rev + 1;
+}
+
+/**
+ * Return the aspect revision of the tree that the given construct is part of
+ *
+ * The revision number is changed every time an aspect is added to the tree.
+ *
+ * Instead of returning the version, returns a function that returns the
+ * version: we can cache the root lookup inside the function; this saves
+ * crawling the tree on every read, which is frequent.
+ *
+ * @internal
+ */
+export function _aspectTreeRevisionReader(construct: IConstruct): () => number {
+  const root = construct.node.root;
+  return () => (root as any)[TREE_REVISION_SYM] ?? 0;
+}
+
 /**
  * Object respresenting an Aspect application. Stores the Aspect, the pointer to the construct it was applied
  * to, and the priority value of that Aspect.
@@ -155,3 +183,5 @@ export class AspectApplication {
     this._priority = priority;
   }
 }
+
+const TREE_REVISION_SYM = Symbol.for('@aws-cdk/core.Aspects.treeRevision');

packages/aws-cdk-lib/core/lib/private/synthesis.ts

@@ -8,7 +8,7 @@ import { CloudAssembly } from '../../../cx-api';
 import * as cxapi from '../../../cx-api';
 import { Annotations } from '../annotations';
 import { App } from '../app';
-import { AspectApplication, AspectPriority, Aspects } from '../aspect';
+import { _aspectTreeRevisionReader, AspectApplication, AspectPriority, Aspects } from '../aspect';
 import { FileSystem } from '../fs';
 import { Stack } from '../stack';
 import { ISynthesisSession } from '../stack-synthesizers/types';
@@ -37,7 +37,7 @@ export function synthesize(root: IConstruct, options: SynthesisOptions = { }): c
   // we start by calling "synth" on all nested assemblies (which will take care of all their children)
   synthNestedAssemblies(root, options);
 
-  if (options.aspectStabilization == true) {
+  if (options.aspectStabilization) {
     invokeAspectsV2(root);
   } else {
     invokeAspects(root);
@@ -274,32 +274,33 @@ function invokeAspects(root: IConstruct) {
  * than the most recently invoked Aspect on that node.
  */
 function invokeAspectsV2(root: IConstruct) {
-  const invokedByPath: { [nodePath: string]: AspectApplication[] } = { };
+  const invokedByPath = new Map<string, AspectApplication[]>();
 
   recurse(root, []);
 
   for (let i = 0; i <= 100; i++) {
     const didAnythingToTree = recurse(root, []);
 
-    if (!didAnythingToTree) {
+    // Keep on invoking until nothing gets invoked anymore
+    if (didAnythingToTree === 'nothing') {
       return;
     }
   }
 
   throw new Error('We have detected a possible infinite loop while invoking Aspects. Please check your Aspects and verify there is no configuration that would cause infinite Aspect or Node creation.');
 
-  function recurse(construct: IConstruct, inheritedAspects: AspectApplication[]): boolean {
+  function recurse(construct: IConstruct, inheritedAspects: AspectApplication[]): 'invoked' | 'aspect-added' | 'nothing' {
     const node = construct.node;
 
-    let didSomething = false;
-
-    let localAspects = getAspectApplications(construct);
-    const allAspectsHere = sortAspectsByPriority(inheritedAspects, localAspects);
+    let ret: ReturnType<typeof recurse> = 'nothing';
+    const currentAspectTreeRevision = _aspectTreeRevisionReader(construct);
+    const versionAtStart = currentAspectTreeRevision();
 
+    const allAspectsHere = sortAspectsByPriority(inheritedAspects, getAspectApplications(construct));
     for (const aspectApplication of allAspectsHere) {
-      let invoked = invokedByPath[node.path];
+      let invoked = invokedByPath.get(node.path);
       if (!invoked) {
-        invoked = invokedByPath[node.path] = [];
+        invokedByPath.set(node.path, invoked = []);
       }
 
       if (invoked.some(invokedApp => invokedApp.aspect === aspectApplication.aspect)) {
@@ -310,25 +311,44 @@ function invokeAspectsV2(root: IConstruct) {
       const lastInvokedAspect = invoked[invoked.length - 1];
       if (lastInvokedAspect && lastInvokedAspect.priority > aspectApplication.priority) {
         throw new Error(
-          `Cannot invoke Aspect ${aspectApplication.aspect.constructor.name} with priority ${aspectApplication.priority} on node ${node.path}: an Aspect ${lastInvokedAspect.aspect.constructor.name} with a lower priority (${lastInvokedAspect.priority}) was already invoked on this node.`,
+          `Cannot invoke Aspect ${aspectApplication.aspect.constructor.name} with priority ${aspectApplication.priority} on node ${node.path}: an Aspect ${lastInvokedAspect.aspect.constructor.name} with a lower priority (added at ${lastInvokedAspect.construct.node.path} with priority ${lastInvokedAspect.priority}) was already invoked on this node.`,
         );
       }
 
       aspectApplication.aspect.visit(construct);
 
-      didSomething = true;
+      ret = 'invoked';
 
       // mark as invoked for this node
       invoked.push(aspectApplication);
+
+      // If this aspect added another aspect, we need to reconsider everything;
+      // it might have added an aspect above us and we need to restart the
+      // entire tree. This could probably be made more efficient, but restarting
+      // the tree from the top currently does it as well.
+      if (currentAspectTreeRevision() !== versionAtStart) {
+        return 'aspect-added';
+      }
     }
 
-    let childDidSomething = false;
     for (const child of construct.node.children) {
       if (!Stage.isStage(child)) {
-        childDidSomething = recurse(child, allAspectsHere) || childDidSomething;
+        const childDidSomething = recurse(child, allAspectsHere);
+        switch (childDidSomething) {
+          case 'aspect-added':
+            // Must abort immediately if we see this.
+            return 'aspect-added';
+          case 'invoked':
+            // We have to keep on invoking until there's nothing left to do.
+            ret = 'invoked';
+            break;
+          default:
+            break;
+        }
       }
     }
-    return didSomething || childDidSomething;
+
+    return ret;
   }
 }
 

packages/aws-cdk-lib/core/test/aspect.test.ts

@@ -1,5 +1,4 @@
 import { Construct, IConstruct } from 'constructs';
-import { getWarnings } from './util';
 import { Template } from '../../assertions';
 import { Bucket, CfnBucket } from '../../aws-s3';
 import * as cxschema from '../../cloud-assembly-schema';
@@ -298,7 +297,7 @@ describe('aspect', () => {
 
     expect(() => {
       app.synth();
-    }).toThrow('Cannot invoke Aspect Tag with priority 0 on node My-Stack: an Aspect Tag with a lower priority (10) was already invoked on this node.');
+    }).toThrow('an Aspect Tag with a lower priority');
   });
 
   test('Infinitely recursing Aspect is caught with error', () => {
@@ -338,14 +337,13 @@ describe('aspect', () => {
     Aspects.of(root).add(new Tag('AspectA', 'Visited'));
 
     // "Monkey patching" - Override `applied` to simulate its absence
-    Object.defineProperty(Aspects.prototype, 'applied', {
-      value: undefined,
-      configurable: true,
-    });
+    const appliedGetter = jest.spyOn(Aspects.prototype, 'applied', 'get').mockReturnValue(undefined as any);
 
     expect(() => {
       app.synth();
     }).not.toThrow();
+
+    appliedGetter.mockRestore();
   });
 
   test('RemovalPolicy: higher priority wins', () => {
@@ -456,4 +454,54 @@ describe('aspect', () => {
       DeletionPolicy: 'Retain',
     });
   });
+
+  test.each([
+    'on self',
+    'on parent',
+  ] as const)('aspect can insert another aspect %s in between other ones', (addWhere) => {
+    // Test that between two pre-existing aspects with LOW and HIGH priorities,
+    // the LOW aspect can add one in the middle, when stabilization is enabled.
+    //
+    // With stabilization (Aspects V2), we are stricter about the ordering and we would
+    // throw if something doesn't work.
+
+    // GIVEN
+    const app = new App();
+    const root = new MyConstruct(app, 'Root');
+    const aspectTarget = new MyConstruct(root, 'MyConstruct');
+
+    // WHEN
+    // - Aspect with prio 100 adds an Aspect with prio 500
+    // - An Aspect with prio 700 also exists
+    // We should execute all and in the right order.
+    const executed = new Array<string>();
+
+    class TestAspect implements IAspect {
+      constructor(private readonly name: string, private readonly block?: () => void) {
+      }
+
+      visit(node: IConstruct): void {
+        // Only do something for the target node
+        if (node.node.path === 'Root/MyConstruct') {
+          executed.push(this.name);
+          this.block?.();
+        }
+      }
+    }
+
+    Aspects.of(aspectTarget).add(new TestAspect('one', () => {
+      // We either add the new aspect on ourselves, or on an ancestor.
+      //
+      // In either case, it should execute next, before the 700 executes.
+      const addHere = addWhere === 'on self' ? aspectTarget : root;
+
+      Aspects.of(addHere).add(new TestAspect('two'), { priority: 500 });
+    }), { priority: 100 });
+    Aspects.of(aspectTarget).add(new TestAspect('three'), { priority: 700 });
+
+    // THEN: should not throw and execute in the right order
+    app.synth({ aspectStabilization: true });
+
+    expect(executed).toEqual(['one', 'two', 'three']);
+  });
 });