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

こんにちは、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
            }
          }
        }
      ]
    }
  }
}

作成したStateMachineFragmentがちゃんと再利用されており、またStateMachineFragment内のすべてのStateのIDにプレフィクスが付いていますね！

{
    "branch": "left"
}

もちろん、State Machineも問題なく実行できました。

おわりに

AWS Step FunctionsのState Machine Fragmentsで再利用可能な一連のStateのクラスをAWS CDKで作ってみました。

State Machine Fragmentsで処理をクラス化しているので、今回のように一つのStack内で使うだけでなく、複数のStack間で共有することももちろん可能です。Step Functionsを積極的に活用しているプロジェクトでは使い所があるのではないでしょうか。

以上