🎃

StepFunctionsの変数をリトライカウンターにできるか試してみた

2025/01/27に公開

はじめに

StepFunctionsで変数が使用できることを知り、この変数をリトライカウンターにできるのでは?と思い試してみました。この記事では、その検証結果を紹介します。

検証用のコード

このテンプレートは、AWS SAMを使用してStep Functionsの状態遷移を構築するサンプルです。2つのLambda関数を順番に呼び出し、その結果に基づいて再試行または成功・失敗状態を決定するシンプルなワークフローを構成しています。

template.yamlのコード全体
template.yaml
AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: "Example of step functions"

Mappings:
  ConstValues:
    lambda1:
      name: "hello1"
    lambda2:
      name: "hello2"

Globals:
  Function:
    Timeout: 5

Resources:

  Lambda1:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !FindInMap [ConstValues, lambda1, name]
      CodeUri: lambda1/
      Handler: app.lambda_handler
      Runtime: python3.12

  LogGroup1:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Join
        - ""
        - - /aws/lambda/
          - !FindInMap [ConstValues, lambda1, name]
      RetentionInDays: 1

  Lambda2:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: !FindInMap [ConstValues, lambda2, name]
      CodeUri: lambda2/
      Handler: app.lambda_handler
      Runtime: python3.12

  LogGroup2:
    Type: AWS::Logs::LogGroup
    Properties:
      LogGroupName: !Join
        - ""
        - - /aws/lambda/
          - !FindInMap [ConstValues, lambda2, name]
      RetentionInDays: 1

  HelloStateMachine:
    Type: AWS::Serverless::StateMachine
    Properties:
      Name: !Sub "hello-${AWS::StackName}"
      Definition:
        TimeoutSeconds: 60
        StartAt: Initialize
        States:
          # IN: { "code": "xxx"}
          Initialize:
            Type: Pass
            Assign:
              counter: 0
            Next: InvokeLambda1
          # IN: { "code": "xxx"}
          InvokeLambda1:
            Type: Task
            Resource: !GetAtt Lambda1.Arn
            ResultPath: "$.lambda1result"
            Next: InvokeLambda2
          # IN: { "code": "xxx", "lambda1result": { "status": "xxx" } }
          InvokeLambda2:
            Type: Task
            Resource: !GetAtt Lambda2.Arn
            ResultPath: "$.lambda2result"
            Next: CheckStatus
          # IN: { "code": "xxx", "lambda1result": { "status": "xxx" }, "lambda2result": { "status": "xxx" }}
          CheckStatus:
            Type: Choice
            Choices:
              - Variable: "$.lambda2result.status"
                StringEquals: "retry"
                Next: CheckRetryCount
            Default: SuccessState
          CheckRetryCount:
            Type: Choice
            Choices:
              - Variable: "$counter"
                NumericLessThan: 1
                Next: IncrementRetryCount
            Assign:
              counter: "{% counter %}"
            Default: FailureState
          IncrementRetryCount:
            Type: Pass
            Assign:
              counter: "{% counter + 1 %}"
            Next: WaitState
          WaitState:
            Type: Wait
            Seconds: 3
            Next: InvokeLambda2
          SuccessState:
            Type: Succeed
          FailureState:
            Type: Fail
            Error: "RetriesExhausted"
            Cause: "The maximum number of retries has been reached."
      Policies:
        - LambdaInvokePolicy:
            FunctionName: !Ref Lambda1
        - LambdaInvokePolicy:
            FunctionName: !Ref Lambda2

ワークフローは以下のような状態遷移を持っています。Lambda2の結果(status)が"retry"の場合に、リトライが発動します。リトライ回数の上限は1回です。

  • Initialize: 初期化(カウンタ変数の設定)
  • InvokeLambda1: Lambda1の呼び出し
  • InvokeLambda2: Lambda2の呼び出し
  • CheckStatus: Lambda2の結果に基づき条件分岐
  • CheckRetryCount: 再試行回数の確認(リトライ回数上限の設定)
  • IncrementRetryCount: カウンタ変数を増加
  • WaitState: 一定時間待機
  • SuccessState / FailureState: 成功または失敗の状態に遷移

Lambda1,Lamba2のコードは以下のとおりです。
ワークフロー実行時に入力したJSONデータ(code)がLambda関数に渡されるので、その受け取った値をstatusで返します。

Lambdaのコード全体
import logging

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

def lambda_handler(event, context):
    logger.info(f"event: {event}")
    return {"status": event["code"]}

これをデプロイして、AWSコンソールで見ると下図のフローになります。

ステートマシンの実行

ステートマシンはAWSコンソールから実行します。

リトライしないパターン

まずは、リトライしないパターンで試してみます。
ステートマシン実行時には{"code":"pass"}を入力します。

期待通り、リトライは発動しませんでした。

リトライするパターン

次に、リトライを発動させてみます。
ステートマシン実行時には{"code":"retry"}を入力します。

期待どおり、リトライが発動しました。リトライ回数上限を1回に設定しているので、2回目でタイムアウトして異常終了しています。

まとめ

今回の検証で、StepFunctionsの変数をリトライカウンターとして使用できることが確認できました。これにより、より柔軟なフロー制御が可能になりそうです。

Discussion