Distributed map with AWS CDK using custom state

distributed-map-custom.ts

import { Map, CustomState } from "aws-cdk-lib/aws-stepfunctions";

export class MyStack extends Stack {
  constructor(scope: App, id: string, props: StackProps) {

    const mapJob = new Map(this, "Map", {
      itemsPath: "$.Items",
      parameters: {
        "Item.$": "$$.Map.Item.Value",
        ...props.mapProps.parameters,
      },
      ...props.mapProps,
    });
    mapJob.iterator(props.mapIterator);

    const iteratorJson = mapJob.toStateJson();

    // custom state to make distributed maps work
    // https://docs.aws.amazon.com/cdk/api/v1/docs/@aws-cdk_aws-stepfunctions.CustomState.html
    const customMapJob = new CustomState(this, "CustomMap", {
      stateJson: {
        ...iteratorJson,
        // End: undefined, // you may need this if you want to have a "next" state
        Iterator: {
          // @ts-ignore
          ...iteratorJson.Iterator,
          ProcessorConfig: {
            Mode: "DISTRIBUTED",
            ExecutionType: "STANDARD",
          },
        },
      },
    });

    const passState = new Pass(...);
    
    // you will need a role so we can allow the state machine to invoke itself
    const stateMachineRole = new Role(this, "StateMachineRole", {
      assumedBy: new ServicePrincipal("states"),
    });

    const statemachine = new StateMachine(this, "Statemachine", {
      role: stateMachineRole,
      definition: passState.next(customMapJob).next(passState), // you can use the custom job like so
      timeout: Duration.hours(72),
    });
    
    // The state machine will also need permission to invoke itself
    stateMachineRole.attachInlinePolicy(
      new Policy(this, "ExecuteStateMachine", {
        statements: [
          new PolicyStatement({
            effect: Effect.ALLOW,
            actions: ["states:StartExecution"],
            resources: [statemachine.stateMachineArn],
          }),
        ],
      }),
    );

    // the custom state won't handle permissions for you like the built-in states, so add your own policies as needed

    // if you need to invoke another state machine
    stateMachineRole.addToPolicy(
      new PolicyStatement({
        effect: Effect.ALLOW,
        actions: ["states:StartExecution", "states:DescribeExecution", "states:StopExecution"],
        resources: [subExpressMachine.stateMachineArn, `${subExpressMachine.stateMachineArn}*`],
      }),
    );

    stateMachineRole.addToPolicy(
      new PolicyStatement({
        effect: Effect.ALLOW,
        actions: ["events:PutTargets", "events:PutRule", "events:DescribeRule"],
        resources: [
          `arn:${this.partition}:events:${this.region}:${this.account}:rule/StepFunctionsGetEventsForStepFunctionsExecutionRule`,
        ],
      }),
    );
}