fix(stepfunctions): prefixes not appended to states in parallel branches (#17806)

There is an order mismatch issue when calling `Parallel.branch()`. `branch` used to immediately create the state graph of the branch which locks in the `stateIds` of the branch states. This fails to take into account the following scenario: 

```ts
class ParallelMachineFragment extends stepfunctions.StateMachineFragment {
  public readonly startState: stepfunctions.State;
  public readonly endStates: stepfunctions.INextable[];

  constructor(scope: Construct, id: string) {
    super(scope, id);

    const step1 = new stepfunctions.Pass(this, 'Step 1');
    const parallelState = new stepfunctions.Parallel(this, 'Parallel State');
    const chain = parallelState.branch(step1);
    this.startState = parallelState;
    this.endStates = [chain];
  }
}

class MyStack extends Stack {
  constructor(scope: Construct, id: string) {
    const fragment1 = new ParallelMachineFragment(this, 'Fragment 1').prefixStates();
    const fragment2 = new ParallelMachineFragment(this, 'Fragment 2').prefixStates();
    
    new stepfunctions.StateMachine(stack, 'State Machine', {
      definition: fragment1.next(fragment2),
    });
  }
}
```

Here, the branches in the `parallelState` do not get prefixes before they are evaluated, causing an error since they have the same `stateId`.

The fix is to make the state graph creation late-bound; it now happens when we call `Parallel.bindToGraph()`. At this point, it will know of any prefixes. 

Fixes #17354. 

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

Author: kaizencc

packages/@aws-cdk/aws-stepfunctions/lib/states/parallel.ts

@@ -71,6 +71,7 @@ export interface ParallelProps {
  */
 export class Parallel extends State implements INextable {
   public readonly endStates: INextable[];
+  private readonly _branches: IChainable[] = [];
 
   constructor(scope: Construct, id: string, props: ParallelProps = {}) {
     super(scope, id, props);
@@ -112,11 +113,23 @@ export class Parallel extends State implements INextable {
    * Define one or more branches to run in parallel
    */
   public branch(...branches: IChainable[]): Parallel {
-    for (const branch of branches) {
+    // Store branches for late-bound stategraph creation when we call bindToGraph.
+    this._branches.push(...branches);
+    return this;
+  }
+
+  /**
+   * Overwrites State.bindToGraph. Adds branches to
+   * the Parallel state here so that any necessary
+   * prefixes are appended first.
+   */
+  public bindToGraph(graph: StateGraph) {
+    for (const branch of this._branches) {
       const name = `Parallel '${this.stateId}' branch ${this.branches.length + 1}`;
       super.addBranch(new StateGraph(branch.startState, name));
     }
-    return this;
+    this._branches.splice(0, this._branches.length);
+    return super.bindToGraph(graph);
   }
 
   /**

packages/@aws-cdk/aws-stepfunctions/test/state-machine-fragment.test.ts

@@ -0,0 +1,57 @@
+import { Match, Template } from '@aws-cdk/assertions';
+import * as cdk from '@aws-cdk/core';
+import { Construct } from 'constructs';
+import * as stepfunctions from '../lib';
+
+describe('State Machine Fragment', () => {
+  test('Prefix applied correctly on Fragments with Parallel states', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    const fragment1 = new ParallelMachineFragment(stack, 'Fragment 1').prefixStates();
+    const fragment2 = new ParallelMachineFragment(stack, 'Fragment 2').prefixStates();
+
+    new stepfunctions.StateMachine(stack, 'State Machine', {
+      definition: fragment1.next(fragment2),
+    });
+
+    // THEN
+    Template.fromStack(stack).hasResourceProperties('AWS::StepFunctions::StateMachine', {
+      DefinitionString: Match.serializedJson({
+        StartAt: 'Fragment 1: Parallel State',
+        States: {
+          'Fragment 1: Parallel State': Match.objectLike({
+            Branches: [Match.objectLike({
+              States: {
+                'Fragment 1: Step 1': Match.anyValue(),
+              },
+            })],
+          }),
+          'Fragment 2: Parallel State': Match.objectLike({
+            Branches: [Match.objectLike({
+              States: {
+                'Fragment 2: Step 1': Match.anyValue(),
+              },
+            })],
+          }),
+        },
+      }),
+    });
+  });
+});
+
+class ParallelMachineFragment extends stepfunctions.StateMachineFragment {
+  public readonly startState: stepfunctions.State;
+  public readonly endStates: stepfunctions.INextable[];
+
+  constructor(scope: Construct, id: string) {
+    super(scope, id);
+
+    const step1 = new stepfunctions.Pass(this, 'Step 1');
+    const parallelState = new stepfunctions.Parallel(this, 'Parallel State');
+    const chain = parallelState.branch(step1);
+    this.startState = parallelState;
+    this.endStates = [chain];
+  }
+}

packages/@aws-cdk/aws-stepfunctions/test/states-language.test.ts

@@ -613,9 +613,12 @@ describe('States Language', () => {
     const state1 = new stepfunctions.Pass(stack, 'State1');
     const state2 = new stepfunctions.Pass(stack, 'State2');
 
-    expect(() => new stepfunctions.Parallel(stack, 'Parallel')
+    const parallel = new stepfunctions.Parallel(stack, 'Parallel')
       .branch(state1.next(state2))
-      .branch(state2)).toThrow();
+      .branch(state2);
+
+    // WHEN
+    expect(() => render(parallel)).toThrow();
   }),
 
   describe('findReachableStates', () => {