[AWS Step Functions] State Machine Fragmentsで再利用可能な一連のStateのクラスを作ってみた（AWS CDK）

若槻龍太

2022.09.18

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

こんにちは、CX事業本部 IoT事業部の若槻です。

AWS Step Functions State MachineをAWS CDKで実装する際に、1つまたは複数のStateから成る似通った処理の記述を1つにまとめたい場合があります。

しかしTaskやStateなどIChainableなオブジェクトは、処理をループさせる場合などを除いてCDKコード内で再利用することはできません。（しようとするとDeploy時にエラーとなる）

しかし今回、ドキュメントを読んでいると、再利用可能なStateを作ることができるState Machine Fragmentsなるクラスを見つけたので、試してみました。

試してみた

ドキュメントにCDK（TypeScript）の良さそうなサンプルが載っていたのでほぼそのまま利用します。

import { Stack } from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as sfn from 'aws-cdk-lib/aws-stepfunctions';

interface MyJobProps {
  jobFlavor: string;
}

class MyJob extends sfn.StateMachineFragment {
  public readonly startState: sfn.State;
  public readonly endStates: sfn.INextable[];

  constructor(parent: Construct, id: string, props: MyJobProps) {
    super(parent, id);

    const choice = new sfn.Choice(this, 'Choice')
      .when(
        sfn.Condition.stringEquals('$.branch', 'left'),
        new sfn.Pass(this, 'Left Branch')
      )
      .when(
        sfn.Condition.stringEquals('$.branch', 'right'),
        new sfn.Pass(this, 'Right Branch')
      );

    this.startState = choice;
    this.endStates = choice.afterwards().endStates;
  }
}

export class MyStack extends Stack {
  constructor(scope: Construct, id: string) {
    super(scope, id);
    // Do 3 different variants of MyJob in parallel
    const parallel = new sfn.Parallel(this, 'All jobs')
      .branch(new MyJob(this, 'Quick', { jobFlavor: 'quick' }).prefixStates())
      .branch(new MyJob(this, 'Medium', { jobFlavor: 'medium' }).prefixStates())
      .branch(new MyJob(this, 'Slow', { jobFlavor: 'slow' }).prefixStates());

    new sfn.StateMachine(this, 'MyStateMachine', {
      definition: parallel,
    });
  }
}

StateMachineFragmentにより1連のStateから成る処理をMyJobとしてクラス化しています。
- startStateでは処理の冒頭となるState、endStatesでは処理の終端となるStateを指定します。
- afterwards().endStatesと指定することにより、choiceのような終端が複数ある可能性がある場合に対応させています。
MyStackでMyJobをインスタンス化して、parallel stateの中で複数回再利用しています。
- prefixStates()を使用することにより、StateMachineFragment内の全てのStateのIDに同じプレフィクス（Quick、MediumおよびSlow）を持たせることができます。

上記をCDK Deployしてスタックをデプロイします。これにより次の定義のState Machineが作成されました。

{
  "StartAt": "All jobs",
  "States": {
    "All jobs": {
      "Type": "Parallel",
      "End": true,
      "Branches": [
        {
          "StartAt": "Quick: Choice",
          "States": {
            "Quick: Choice": {
              "Type": "Choice",
              "Choices": [
                {
                  "Variable": "$.branch",
                  "StringEquals": "left",
                  "Next": "Quick: Left Branch"
                },
                {
                  "Variable": "$.branch",
                  "StringEquals": "right",
                  "Next": "Quick: Right Branch"
                }
              ]
            },
            "Quick: Left Branch": {
              "Type": "Pass",
              "End": true
            },
            "Quick: Right Branch": {
              "Type": "Pass",
              "End": true
            }
          }
        },
        {
          "StartAt": "Medium: Choice",
          "States": {
            "Medium: Choice": {
              "Type": "Choice",
              "Choices": [
                {
                  "Variable": "$.branch",
                  "StringEquals": "left",
                  "Next": "Medium: Left Branch"
                },
                {
                  "Variable": "$.branch",
                  "StringEquals": "right",
                  "Next": "Medium: Right Branch"
                }
              ]
            },
            "Medium: Left Branch": {
              "Type": "Pass",
              "End": true
            },
            "Medium: Right Branch": {
              "Type": "Pass",
              "End": true
            }
          }
        },
        {
          "StartAt": "Slow: Choice",
          "States": {
            "Slow: Choice": {
              "Type": "Choice",
              "Choices": [
                {
                  "Variable": "$.branch",
                  "StringEquals": "left",
                  "Next": "Slow: Left Branch"
                },
                {
                  "Variable": "$.branch",
                  "StringEquals": "right",
                  "Next": "Slow: Right Branch"
                }
              ]
            },
            "Slow: Left Branch": {
              "Type": "Pass",
              "End": true
            },
            "Slow: Right Branch": {
              "Type": "Pass",
              "End": true
            }
          }
        }
      ]
    }
  }
}