Step Functionsで条件を満たすまでループする

以下のページで、Step Functionsの条件分岐(Choice)を取り上げました。

今回はこのChoiceを使用して、条件を満たすまで、繰り返し処理を行うステートマシンを作成します。

例えばStep FunctionsとLambda関数を使用して、RDSからスナップショットを作成し、コピーすることを考えます。
スナップショット作成を実行した直後は、スナップショットのステータスは「creating」となり、すぐにはコピーできません。
つまりステータスが「available」となるまで待ってからコピーを実行しなければなりません。

今回、Choiceを使用して実現したいポイントは「ステータスが変更するまで待つ」という部分です。

今回作成するStep Functionsステートマシンは、ステータスが変更するまでは次のステップに進まず、一時的に停止します。
停止後、改めてステータスチェックを行い、ステータスが条件を満たした場合のみ、次のステップに進みます。

構築する環境

Diagram of Loop in Step Functions until the condition is satisfied.

Step Functionsステートマシン内に、4つのステートを定義します。

2つのステートは、以下のLambda関数を実行するタスクです。

Lambda関数1：現在日時の秒数を取得し、これが偶数か奇数か判定する。
Lambda関数2：ステートマシンが完了した旨の文字列を返す。

関数のランタイム環境はPython3.8とします。

残り2つのステートの働きは以下の通りです。

Choice：ステートを条件分岐させる。秒数が偶数の場合は、Lambda関数2のステートに進む。奇数の場合はWaitステートに進む。
Wait：3秒待機する。待機後はLambda関数1のステートに戻る。

CloudFormationテンプレートファイル

上記の構成をCloudFormationで構築します。
以下のURLにCloudFormationテンプレートを配置しています。

https://github.com/awstut-an-r/awstut-fa/tree/main/140

テンプレートファイルのポイント解説

Lambda関数

Resources:
  Function1:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          import datetime

          def lambda_handler(event, context):
            now = datetime.datetime.now()

            if now.second % 2 == 0:
              return True
            return False
      FunctionName: !Sub "${Prefix}-function-01"
      Handler: !Ref Handler
      Runtime: !Ref Runtime
      Role: !GetAtt FunctionRole.Arn

  Function2:
    Type: AWS::Lambda::Function
    Properties:
      Code:
        ZipFile: |
          def lambda_handler(event, context):
            return 'finish!'
      FunctionName: !Sub "${Prefix}-function-02"
      Handler: !Ref Handler
      Runtime: !Ref Runtime
      Role: !GetAtt FunctionRole.Arn
Code language: YAML (yaml)

2つの関数を作成します。

1つ目はステータス判定をイメージした関数です。
本来は冒頭のRDSの例にように、スナップショットのステータスを確認する内容です。
今回は動作検証ということで、現在日時の秒数に注目します。
この値が偶数か、つまり2で割り切れるかどうかを確認します。
偶数の場合はTrueを返し、奇数の場合はFalseを返します。

2つ目はステートマシン完了を知らせる関数です。
シンプルに文字列「finish!」を返します。

以下が両関数用のIAMロールです。

Resources:
  FunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action: sts:AssumeRole
            Principal:
              Service:
                - lambda.amazonaws.com
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
Code language: YAML (yaml)

Lambda用のAWS管理ポリシーをアタッチしただけのシンプルなIAMロールです。

Step Functions

Resources:
  StateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      Definition:
        Comment: !Sub "${Prefix}-StateMachine"
        StartAt: FirstState
        States:
          FirstState:
            Type: Task
            Resource: !Ref FunctionArn1
            Next: ChoiceState
          ChoiceState:
            Type: Choice
            Choices:
              - Variable: $
                BooleanEquals: false
                Next: WaitState
              - Variable: $
                BooleanEquals: true
                Next: LastState
          WaitState:
            Type: Wait
            Seconds: !Ref WaitSeconds
            Next: FirstState
          LastState:
            Type: Task
            Resource: !Ref FunctionArn2
            End: true
      LoggingConfiguration:
        Destinations:
          - CloudWatchLogsLogGroup:
              LogGroupArn: !GetAtt LogGroup.Arn
        IncludeExecutionData: true
        Level: ALL
      RoleArn: !GetAtt StateMachineRole.Arn
      StateMachineName: !Ref Prefix
      StateMachineType: STANDARD
Code language: YAML (yaml)

ステートマシン全体の構造を確認します。

StartAtプロパティを見ると、FirstStateから処理が開始されることがわかります。
FirstStateでは、関数1を実行します。

その後ChoiceStateに移行し、FirstStateから返ってきた値を参照して、条件分岐の判定が行われます。
Variablesプロパティに「$」と指定することで、関数1の結果を参照できます。
関数1から返された値がTrueの場合は、LastStateに移動し、関数2が実行されて、ステートマシンが終了します。
値がFalseの場合は、WaitStateに移動し、3秒待機した後、再度FirstStateが実行されます。

ポイントはこのWaitStateの後の移動先です。
ここでFirstStateに戻ることで、改めてステータスをチェックするという挙動を再現できます。
つまり特定のステータスに変更するまで、繰り返し処理を行うことができます。

以下が本ステートマシン用のIAMロールです。

Resources:
  StateMachineRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: 2012-10-17
        Statement:
          - Effect: Allow
            Action: sts:AssumeRole
            Principal:
              Service:
                - states.amazonaws.com
      Policies:
        - PolicyName: !Sub "${Prefix}-InvokeTaskFunctions"
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource:
                  - !Ref FunctionArn1
                  - !Ref FunctionArn2
        - PolicyName: !Sub "${Prefix}-DeliverToCloudWatchLogPolicy"
          PolicyDocument:
            Version: 2012-10-17
            Statement:
              - Effect: Allow
                Action:
                  - logs:CreateLogDelivery
                  - logs:GetLogDelivery
                  - logs:UpdateLogDelivery
                  - logs:DeleteLogDelivery
                  - logs:ListLogDeliveries
                  - logs:PutLogEvents
                  - logs:PutResourcePolicy
                  - logs:DescribeResourcePolicies
                  - logs:DescribeLogGroups
                Resource: "*"
Code language: YAML (yaml)