Step Functions内でループ処理を作成したら無限ループになっていた

はじめに

タイトルの通りなのだが、AWSから届いたアラートメールで大量にStep Functionsが実行されていたとに気づいた。Step Functionsは基本的に状態遷移の回数によって課金される。アラートが来たということは、大量の状態遷移が実行されているということになる。

図：届いたアラートメールのスクリーンショット

実際に実行履歴を見てみると、一つの実行内で3000以上の遷移が起きていることがわかる。回数か実行時間のどちらかでエラー終了とさせたい。

図：Step Functionsのマネコン。3496回も実行されている･･･

無限実行されていたのは前回記事のステートマシン
https://qiita.com/bd8z/items/13faf2dd985578a8033c

対策を考える

検討用に簡単なStep Functionsを用意する。このChoiceの部分にエラー判定＆エラー終了処理を追加することで、無限ループから抜けられるようにする。

対策1: ループ回数でエラー終了させる

無限ループしているならループ回数が異常に多い時、エラー終了させればよい。なのでループ回数をカウントさせる。Step Functionsは、グローバル変数を持てないのでLambdaの入出力のペイロードにカウント回数を仕込こむ。一定以上のループ回数の時、エラーで終了するようなChoiceのルールを追加する。

図：エラー終了を追加したループ処理をもつステートマシン

lambda_function.py

import json

def lambda_handler(event, context):
    if "loopCount" in list(event.keys()):
        loopCount = event["loopCount"] + 1
    else:
        loopCount = 0
        
    event["loopCount"] = loopCount
    return event

対策2: ステートマシンの実行時間でエラー終了させる。

Step FunctionsのContextを使って、無限ループにならないようなChoiceを作る。Contextにはタイムスタンプの項目があるので、「ループ処理開始時刻」あるいは「ステートマシンの実行開始時刻」を取得し、Lambda内部で経過時間を計算させる。一定の経過時間以上がChoiceにて確認されたらエラー終了させる。

コンテキストオブジェクト
https://docs.aws.amazon.com/ja_jp/step-functions/latest/dg/input-output-contextobject.html

Passを使ってLambdaのペイロードに時刻情報仕込む。今回はステートマシンの実行開始の時刻を利用するので、Passと$$.Execution.StartTimeを使った。

図：エラー終了とPassを追加したループ処理をもつステートマシンとPassの入力パラメータ設定

lambda_function.py

import json
import datetime

def lambda_handler(event, context):
    timeDiff = datetime.datetime.now(datetime.timezone.utc) - datetime.datetime.fromisoformat(event["startTime"])
    event["timeDiff"] = timeDiff.seconds
    return return_json

対策1/2のまとめ

ループ処理、時間経過両方の処理を追加したLambda関数の例。

lambda_function.py

import datetime
import json

def countStepFunctionMetrics(event):
    if "loopCount" in list(event.keys()):
        loopCount = event["loopCount"] + 1
    else:
        loopCount = 0
    timeDiff = datetime.datetime.now(datetime.timezone.utc) - datetime.datetime.fromisoformat(event["startTime"])
    
    metric = {}
    metric["loopCount"] = loopCount
    metric["timeDiffInSecond"] = timeDiff.seconds

    return metric
    

def lambda_handler(event, context):
    metric_ = countStepFunctionMetrics(event)
    
    return_json = event

    return_json["loopCount"] = metric_["loopCount"]
    return_json["timeDiffInSecond"] = metric_["timeDiffInSecond"]
    return return_json

ステートマシンのjson定義の例。この例では、ループ回数が3回あるいは開始から10秒経過したときステートマシンがエラー終了するようにした。

{
  "Comment": "A description of my state machine",
  "StartAt": "Pass",
  "States": {
    "Pass": {
      "Type": "Pass",
      "Next": "Lambda Invoke",
      "Parameters": {
        "startTime.$": "$$.Execution.StartTime"
      }
    },
    "Lambda Invoke": {
      "Type": "Task",
      "Resource": "arn:aws:states:::lambda:invoke",
      "Parameters": {
        "FunctionName": "arn:aws:lambda:ap-northeast-1:{$awsId}:function:{$fuctionName}:$LATEST",
        "Payload.$": "$"
      },
      "Retry": [
        {
          "ErrorEquals": [
            "Lambda.ServiceException",
            "Lambda.AWSLambdaException",
            "Lambda.SdkClientException",
            "Lambda.TooManyRequestsException"
          ],
          "IntervalSeconds": 1,
          "MaxAttempts": 3,
          "BackoffRate": 2
        }
      ],
      "Next": "Choice (1)",
      "OutputPath": "$.Payload"
    },
    "Choice (1)": {
      "Type": "Choice",
      "Choices": [
        {
          "Variable": "$.loopCount",
          "NumericGreaterThan": 10,
          "Next": "Success"
        },
        {
          "Or": [
            {
              "Variable": "$.loopCount",
              "NumericGreaterThan": 3
            },
            {
              "Variable": "$.timeDiffInSecond",
              "NumericGreaterThan": 10
            }
          ],
          "Next": "Fail"
        }
      ],
      "Default": "Wait"
    },
    "Success": {
      "Type": "Succeed"
    },
    "Wait": {
      "Type": "Wait",
      "Next": "Lambda Invoke",
      "Seconds": 3
    },
    "Fail": {
      "Type": "Fail"
    }
  }
}

最後に

状態遷移の無限ループに陥るということはそもそもの設計が悪いだけなのだが、ループ回数や実行時間みたいなメトリクスを利用するにも、そういった変数をペイロードにあえて乗せたり、Lambdaにカウントアップ処理をやらせる必要があったり、少し煩わしさを感じた。グローバル変数的な機能がないのは、SSMパラメータストア使って欲しいからなのかな。Step Functionsには組込み関数もあるみたいだが、時刻の計算には対応していなかったので今回はLambdaで計算させるようにした。コンテキストオブジェクトに対してできる処理は読み取りだけなので、値を出し入れ出来たり、もう少し詳細なメトリクスが取れたりすると便利なんだけど、そのあたりはもうAWSの考えるStep Functionsの思想的なところからずれているよな、仕方ないか･･･。