feat(aspects): priority-ordered aspect invocation

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

@@ -24,7 +24,7 @@ export class Aspects {
   public static of(scope: IConstruct): Aspects {
     let aspects = (scope as any)[ASPECTS_SYMBOL];
     if (!aspects) {
-      aspects = new Aspects();
+      aspects = new Aspects(scope);
 
       Object.defineProperty(scope, ASPECTS_SYMBOL, {
         value: aspects,
@@ -35,18 +35,23 @@ export class Aspects {
     return aspects;
   }
 
+  private readonly _scope: IConstruct;
   private readonly _aspects: IAspect[];
+  private readonly _appliedAspects: AspectApplication[];
 
-  private constructor() {
+  private constructor(scope: IConstruct) {
     this._aspects = [];
+    this._appliedAspects = [];
+    this._scope = scope;
   }
 
   /**
    * Adds an aspect to apply this scope before synthesis.
    * @param aspect The aspect to add.
    */
-  public add(aspect: IAspect) {
+  public add(aspect: IAspect, priority?: number) {
     this._aspects.push(aspect);
+    this._appliedAspects.push(new AspectApplication(this._scope, aspect, priority ?? 600));
   }
 
   /**
@@ -55,4 +60,66 @@ export class Aspects {
   public get all(): IAspect[] {
     return [...this._aspects];
   }
-}
+
+  /**
+   * The list of aspects with priority which were directly applied on this scope.
+   *
+   * Also returns inherited Aspects of this node.
+   */
+  public get list(): AspectApplication[] {
+    return this._appliedAspects;
+  }
+}
+
+/**
+ * Object respresenting an Aspect application. Stores the Aspect, the pointer to the construct it was applied
+ * to, and the priority value of that Aspect.
+ */
+export class AspectApplication {
+  /**
+   * The construct that the Aspect was applied to.
+   */
+  public readonly construct: IConstruct;
+
+  /**
+   * The Aspect that was applied.
+   */
+  public readonly aspect: IAspect;
+
+  /**
+   * The priority value of this Aspect. Must be non-negative integer.
+   */
+  private _priority: number;
+
+  /**
+   * Initializes AspectApplication object
+   */
+  public constructor(construct: IConstruct, aspect: IAspect, priority: number) {
+    this.construct = construct;
+    this.aspect = aspect;
+    this._priority = priority;
+  }
+
+  /**
+   * Gets the priority value.
+   */
+  public get priority(): number {
+    return this._priority;
+  }
+
+  // Setter for priority
+  private set priority(priority: number) {
+    if (priority < 0) {
+      throw new Error('Priority must be a non-negative number');
+    }
+    this._priority = priority;
+  }
+
+  /**
+   * Method that will change the priority of this AspectApplication
+   * @param priority - non-negative integer value
+   */
+  public changePriority(priority: number): void {
+    this.priority = priority;
+  }
+}

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

@@ -6,9 +6,8 @@ import { prepareApp } from './prepare-app';
 import { TreeMetadata } from './tree-metadata';
 import { CloudAssembly } from '../../../cx-api';
 import * as cxapi from '../../../cx-api';
-import { Annotations } from '../annotations';
 import { App } from '../app';
-import { Aspects, IAspect } from '../aspect';
+import { AspectApplication, Aspects, IAspect } from '../aspect';
 import { FileSystem } from '../fs';
 import { Stack } from '../stack';
 import { ISynthesisSession } from '../stack-synthesizers/types';
@@ -214,44 +213,96 @@ function synthNestedAssemblies(root: IConstruct, options: StageSynthesisOptions)
  * Aspects are not propagated across Assembly boundaries. The same Aspect will not be invoked
  * twice for the same construct.
  */
+// function invokeAspects(root: IConstruct) {
+//   const invokedByPath: { [nodePath: string]: IAspect[] } = { };
+
+//   let nestedAspectWarning = false;
+//   recurse(root, []);
+
+//   function recurse(construct: IConstruct, inheritedAspects: IAspect[]) {
+//     const node = construct.node;
+//     const aspects = Aspects.of(construct);
+//     const allAspectsHere = [...inheritedAspects ?? [], ...aspects.all];
+//     const nodeAspectsCount = aspects.all.length;
+//     for (const aspect of allAspectsHere) {
+//       let invoked = invokedByPath[node.path];
+//       if (!invoked) {
+//         invoked = invokedByPath[node.path] = [];
+//       }
+
+//       if (invoked.includes(aspect)) { continue; }
+
+//       aspect.visit(construct);
+
+//       // if an aspect was added to the node while invoking another aspect it will not be invoked, emit a warning
+//       // the `nestedAspectWarning` flag is used to prevent the warning from being emitted for every child
+//       if (!nestedAspectWarning && nodeAspectsCount !== aspects.all.length) {
+//         Annotations.of(construct).addWarningV2('@aws-cdk/core:ignoredAspect', 'We detected an Aspect was added via another Aspect, and will not be applied');
+//         nestedAspectWarning = true;
+//       }
+
+//       // mark as invoked for this node
+//       invoked.push(aspect);
+//     }
+
+//     for (const child of construct.node.children) {
+//       if (!Stage.isStage(child)) {
+//         recurse(child, allAspectsHere);
+//       }
+//     }
+//   }
+// }
 function invokeAspects(root: IConstruct) {
-  const invokedByPath: { [nodePath: string]: IAspect[] } = { };
-
-  let nestedAspectWarning = false;
-  recurse(root, []);
-
-  function recurse(construct: IConstruct, inheritedAspects: IAspect[]) {
-    const node = construct.node;
-    const aspects = Aspects.of(construct);
-    const allAspectsHere = [...inheritedAspects ?? [], ...aspects.all];
-    const nodeAspectsCount = aspects.all.length;
-    for (const aspect of allAspectsHere) {
-      let invoked = invokedByPath[node.path];
-      if (!invoked) {
-        invoked = invokedByPath[node.path] = [];
-      }
+  const aspectsMap = collectAllAspectApplications(root);
+
+  const orderedPriorities = Array.from(aspectsMap.keys()).sort();
 
-      if (invoked.includes(aspect)) { continue; }
+  for (const priority of orderedPriorities) {
+    for (const aspectApplication of aspectsMap.get(priority)!) {
+      invokeAspect(aspectApplication.construct, aspectApplication.aspect);
+    }
+  }
+}
 
-      aspect.visit(construct);
+/**
+ * Invokes an individual Aspect on a construct and all of its children in the construct tree.
+ */
+function invokeAspect(construct: IConstruct, aspect: IAspect) {
+  aspect.visit(construct);
+  for (const child of construct.node.children) {
+    invokeAspect(child, aspect);
+  }
+}
 
-      // if an aspect was added to the node while invoking another aspect it will not be invoked, emit a warning
-      // the `nestedAspectWarning` flag is used to prevent the warning from being emitted for every child
-      if (!nestedAspectWarning && nodeAspectsCount !== aspects.all.length) {
-        Annotations.of(construct).addWarningV2('@aws-cdk/core:ignoredAspect', 'We detected an Aspect was added via another Aspect, and will not be applied');
-        nestedAspectWarning = true;
+/**
+ * Returns a Map of AspectApplications organized by Priority value.
+ *
+ * Returs a Map where the keys are Priority (number) and the values are Set<AspectApplication>
+ */
+function collectAllAspectApplications(root: IConstruct): Map<number, Set<AspectApplication>> {
+  let aspectsMap = new Map<number, Set<AspectApplication>>();
+
+  function aspectsFromNode(node: IConstruct) {
+    for (const aspectApplication of Aspects.of(node).list) {
+      const curPriority = aspectApplication.priority;
+      if (!aspectsMap.has(curPriority)) {
+        aspectsMap.set(curPriority, new Set());
       }
 
-      // mark as invoked for this node
-      invoked.push(aspect);
+      let aspectsSet = aspectsMap.get(curPriority)!;
+      aspectsSet.add(aspectApplication);
+
+      aspectsMap.set(curPriority, aspectsSet);
     }
 
-    for (const child of construct.node.children) {
-      if (!Stage.isStage(child)) {
-        recurse(child, allAspectsHere);
-      }
+    for (const child of node.node.children) {
+      aspectsFromNode(child);
     }
   }
+
+  aspectsFromNode(root);
+
+  return aspectsMap;
 }
 
 /**

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

@@ -1,7 +1,9 @@
 import { Construct, IConstruct } from 'constructs';
 import * as cxschema from '../../cloud-assembly-schema';
-import { App } from '../lib';
+import { App, CfnResource, Stack } from '../lib';
 import { IAspect, Aspects } from '../lib/aspect';
+import { Bucket, CfnBucket } from '../../aws-s3';
+import { Template } from '../../assertions';
 
 class MyConstruct extends Construct {
   public static IsMyConstruct(x: any): x is MyConstruct {
@@ -24,6 +26,29 @@ class MyAspect implements IAspect {
   }
 }
 
+class EnableBucketVersioningAspect implements IAspect {
+  public visit(node: IConstruct): void {
+    if (node instanceof CfnBucket) {
+      node.versioningConfiguration = {
+        status: 'Enabled'
+      };
+    }
+  }
+}
+
+class AddLoggingBucketAspect implements IAspect {
+  public visit(node: IConstruct): void {
+    // Check if the node is an S3 Bucket
+    if (node instanceof CfnBucket) {
+      // Add a new logging bucket Bucket to the stack this bucket belongs to.
+      const stack = Stack.of(node);
+      new Bucket(stack, 'my-new-logging-bucket-from-aspect', {
+        bucketName: 'my-new-logging-bucket-from-aspect'
+      });
+    }
+  }
+}
+
 describe('aspect', () => {
   test('Aspects are invoked only once', () => {
     const app = new App();
@@ -69,4 +94,93 @@ describe('aspect', () => {
     expect(root.node.metadata.length).toEqual(1);
     expect(child.node.metadata.length).toEqual(1);
   });
+
+  test('Aspects applied without priority get the default priority value', () => {
+    const app = new App();
+    const root = new MyConstruct(app, 'Construct');
+    const child = new MyConstruct(root, 'ChildConstruct');
+
+    // WHEN - adding an Aspect without priority specified
+    Aspects.of(root).add(new MyAspect());
+
+    // THEN - the priority is set to default (600)
+    let aspectApplication = Aspects.of(root).list[0];
+    expect(aspectApplication.priority).toEqual(600);
+  });
+
+  test('Can override Aspect priority with changePriority() function', () => {
+    const app = new App();
+    const root = new MyConstruct(app, 'Construct');
+    const child = new MyConstruct(root, 'ChildConstruct');
+
+    // WHEN - adding an Aspect without priority specified and resetting it.
+    Aspects.of(root).add(new MyAspect());
+    let aspectApplication = Aspects.of(root).list[0];
+
+    // THEN - we can reset the priority of an Aspect
+    aspectApplication.changePriority(0);
+    expect(Aspects.of(root).list[0].priority).toEqual(0);
+  });
+
+  test('In-place mutating Aspect gets applied', () => {
+    const app = new App();
+    const stack = new Stack(app, 'My-Stack');
+
+    // GIVEN - Bucket with versioning disabled
+    const bucket = new Bucket(stack, 'my-bucket', {
+      versioned: false,
+    });
+
+    // WHEN - adding the Aspect to enable bucket versioning gets applied:
+    Aspects.of(stack).add(new EnableBucketVersioningAspect());
+
+    // THEN - Aspect is successfully applied
+    Template.fromStack(stack).hasResourceProperties('AWS::S3::Bucket', {
+      'VersioningConfiguration': {
+        'Status': 'Enabled',
+      }
+    });
+  });
+
+  test('mutating Aspect that creates a node gets applied', () => {
+    const app = new App();
+    const stack = new Stack(app, 'My-Stack');
+
+    // GIVEN - Bucket with versioning disabled
+    const bucket = new Bucket(stack, 'my-bucket', {
+      bucketName: 'my-original-bucket',
+      versioned: false,
+    });
+
+    // WHEN - adding the Aspect to add a logging bucket:
+    Aspects.of(stack).add(new AddLoggingBucketAspect());
+
+    // THEN - Aspect is successfully applied, new logging bucket is added
+    Template.fromStack(stack).hasResourceProperties('AWS::S3::Bucket', {
+      'BucketName': 'my-new-logging-bucket-from-aspect'
+    });
+  });
+
+  test('can set mutating Aspects in specified orcder and visit newly created node', () => {
+    const app = new App();
+    const stack = new Stack(app, 'My-Stack');
+
+    // GIVEN - Bucket with versioning disabled
+    const bucket = new Bucket(stack, 'my-bucket', {
+      bucketName: 'my-original-bucket',
+      versioned: false,
+    });
+
+    // WHEN - adding both Aspects but making LoggingBucket Aspect run first
+    Aspects.of(stack).add(new AddLoggingBucketAspect(), 0);
+    Aspects.of(stack).add(new EnableBucketVersioningAspect());
+
+    // THEN - both Aspects are successfully applied, new logging bucket is added with versioning enabled
+    Template.fromStack(stack).hasResourceProperties('AWS::S3::Bucket', {
+      'BucketName': 'my-new-logging-bucket-from-aspect',
+      'VersioningConfiguration': {
+        'Status': 'Enabled',
+      }
+    });
+  });
 });