☁️

[AWS CDK] Step Functions の Lambda タスクを同期呼び出しする場合の ASL の書き方について紹介・比較する

2023/05/01に公開

TypeScript

AWS CDK

tech

Step Functions から Lambda を呼び出す場合、ASL の表現としてはいくつかのパターンがあります。

その中でも同期的な呼び出しを行う場合、以下のような2パターンが使えます。

-	Type	Resource	備考
方法(1)	"Task"	Function ARN	ステートの入出力はハンドラの入出力と一致する
方法(2)	"Task"	"arn:aws:states:::lambda:invoke"	タスクに与える入力フォーマットに決まりがある。また、ハンドラの戻り値はタスクの出力と一致しない。 "SDK service integrations" で Lambda Invoke API を呼び出す実装パターン

この2つについて、ASL 上での違いと、対応する CDK での実装を紹介します。

方法(1) に寄せた方がいいんじゃないか？というのがこの記事で述べる私の主張になります。

ASL の実装スタイルとして、私は「できることなら ASL 側で大きなデータ加工は行わないようにしておきたいし、記述をシンプルに保っておきたい」と思っています。

このスタンスでいくと、Lambda ハンドラのコード内で扱う event オブジェクトの中身が冗長になる可能性はあります。しかし、Step Functions で実装されたシステムを読み解くには ASL とハンドラの2つを往復しながらコードを見ていく必要があるので、ASL の実装が入り組んでいると ASL とハンドラ実装の双方をがっつり読まないと全容の理解がしづらく、それはちょっと嬉しくないな...という感覚を持っています

なので、これ以降出てくる見解的なものは上記のような思想を持った人が書いてるものと思って受け取ってください

ASL の比較

方法(1) の場合、ASL に書かれている内容と Lambda ハンドラの実装で記述する　event, return の型は基本的に一致します。

InputPath, Parameters, ResultPath, OutputPath を使って入出力を加工せず、ステートと Lambda 実行タスクの入出力を一致させる場合の ASL は以下のような書き方になります。

{
  "Type": "Task",
  "Resource": ":aws:lambda:<region>:<aws-account>:function:<function-name>"
}

前提として、私はこの書き方が好きです。シンプルですね。「ステートの入出力」と「タスクの入出力」が一致しているので、この実装の中身を理解するのに必要なコンテキストは少ないです。

一方、方法(2) の場合は少々工夫が必要です。まず、タスクに渡す入力はトップレベルに FunctionName, Payload を含む必要があります。ASL で表現すると、例えば Parameters を使って以下のように書く必要があります。

{
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
        "FunctionName": "arn:aws:lambda:<region>:<aws-account>:function:<function-name>",
        "Payload.$": "$"
      }
}

上記の例は、このステートが手前のステートから受け取った入力値をそのまま Lambda ハンドラに渡す場合の書き方になります。この書き方をした場合、Lambda 実行タスクからステートが受け取る戻り値は次のような形式になります。（一部の値は例）

{
  "ExecutedVersion": "$LATEST", 
  // Lambda ハンドラのロジックが返した値
  "Payload": {
    // ...
  },
  // AWS SDK で　Invoke API を実行した際の戻り値が含まれる
  "SdkHttpMetadata": {
    // ...
  },
  "SdkResponseMetadata": {
    // ...
  }
}

見ての通り、方法(1) と比べて情報量が増えています。ほとんどの場合、この戻り値のうち次のステートに渡したいデータはハンドラロジックが返すペイロード ($.Payload) 内の構造のみではないでしょうか？そうすると、ResultSelecttor や ResultPath などを駆使して不要なタスク出力をカットする整形を実装することが必須になります。

方法(1) よりも記述が冗長になっていますし、ASL の字面から明示的に読み取れない暗黙のデータ構造が戻り値側に出現しています。この点がイマイチに感じます。また、Step Functions が提供するステートの入出力加工・整形の仕組み（InputPath など）がすでに少々ややこしいので、その前提からさらに余計なコンテキストが増えてしまっているのがまたイマイチだなぁと思います。後から引き継ぐ人に優しくない感。

Lambda の実行タスクなら方法(1)がよりシンプルなので、わざわざ方法(2)の書き方をする必要はないだろうと考えています。

CDK での実装方法

サンプルコードを書きましたので、そちらもご覧ください。

2023/05 時点の L2 Constrauct では、Lambda 実行タスクは aws-cdk-lib » aws_stepfunctions_tasks » LambdaInvoke を使って記述可能です。

この Construct は、デフォルトで上記の方法(2)を使った ASL 実装を提供します。

さきほど方法(2) はあまり直感的でなく余分な情報が増えている、としてあまり好みではないと言及しましたが、CDK のコードで見ると更に非直感的さが顕著になります（主観）。

幸い、2つの表現方法を切り替える方法は提供されています。payloadResponseOnly という　boolean 型パラメータがそれです。リファレンスには次のように書かれてますが、私の読解力では初見で何しているかわかりませんでした。あと気づけませんでした。

payloadResponseOnly?
Type: boolean (optional, default: false)

Invoke the Lambda in a way that only returns the payload response without additional metadata.

The payloadResponseOnly property cannot be used if integrationPattern, invocationType, clientContext, or qualifier are specified. It always uses the REQUEST_RESPONSE behavior.

上記で書かれているように、デフォルトは　false です。false の場合は方法(1)による ASL 実装になります。

(CDK) payloadResponseOnly	生成される ASL の表現
true	方法(1)の形式
false(default)	方法(2)の形式（AWS SDK service integrations）

CDK (TypeScript) でそれぞれのやり方のタスクを実装してみます。サンプルコードから関係する部分のみ抜粋すると以下です。

import * as aws_sfn_tasks from "aws-cdk-lib/aws-stepfunctions-tasks";

/**
 * 方法（１） の実装
 */
new aws_sfn_tasks.LambdaInvoke(this, 'LambdaFunctionTask1', {
  lambdaFunction: lambdaFunctionConstruct,
  payloadResponseOnly: true,
})

/**
 * 方法（2） の実装
 * 
 * 暗黙的に SDK service integrations に適合した入力ペイロードとなるような ASL が生成される。
 * ハンドラの戻り値「のみ」後続に渡したい場合は ResultSelector や ResultPath, 
 * OutputPath などを用いて `$.Payload` を取り出すような記述が追加で必要
 * 
 */
new aws_sfn_tasks.LambdaInvoke(this, 'LambdaFunctionTask2', {
  lambdaFunction: lambdaFunctionConstruct,
  payloadResponseOnly: false,　　// 無指定の場合と同値
})

デプロイして検証

実際にデプロイして、生成された ASL を見てみます。関係する部分以外は端折ります。

// 方法（１） で実装したタスクの ASL
{
  "Type": "Task",
  "Resource": "arn:aws:lambda:<region>:<aws-account>:function:FirstTask"
}

// 方法（2） で実装したタスクの ASL
{
  "Type": "Task",
  "Resource": "arn:aws:states:::lambda:invoke",
  "Parameters": {
    "FunctionName": "arn:aws:lambda:<region>:<aws-account>:function:SecondTask",
    "Payload.$": "$"
  }
}

ステートマシンを実行してみます。出力はこんな感じです。

方法(1) で実装したタスクの出力

このタスクに対する入力は以下です。

{
  "execution": { /* 省略... */},
  "data": {
    "name": "foo"
  }
}

ハンドラの実装はペイロードの $.data.message に値を追加したデータ構造を返すだけです。

方法(2) で実装したタスクの出力

このタスクに対する入力は以下です。

{
  "execution": { /* 省略... */},
  "data": {
    "name": "foo",
    "message": "hello, foo"
  }
}

ハンドラの実装はペイロードの $.data.success に値を追加したデータ構造を返すだけです。

Lambda ハンドラが返した値は、ステート側から見ると $.Payload 以下にある様子がわかります。

$.SdkHttpMetadata などの構造が後続処理で要らない場合は、次のステートに渡す前に ResultSelector などを使って不要部分をよしなにカットするような定義を追加する必要があります。
ステートの入力を破棄して出力をそのまま利用する場合（ResultPath をデフォルト値の $ に指定した場合の挙動）は割とシンプルで、OutputPath を用いてタスク出力の一部 ($.Paylaod) のみ抜き取ってあげれば良いです。
CDK のコードでは次のような実装になります。

new aws_sfn_tasks.LambdaInvoke(this, 'LambdaFunctionTask2', {
  lambdaFunction: lambdaFunctionConstruct,
  payloadResponseOnly: false,
  // タスクの出力のうち `$.Payload` キーの値だけをステートの出力として使用する
  outputPath: aws_sfn.JsonPath.stringAt('$.Payload'),
})

ResultPath に　$ や null 以外を指定する場合...つまりステートの入力と出力を組み合わせたいケースではもっとややこしくなります。正直私は考えたくないなと思いましたし、Step Functions や CDK の経験値が少ない状態でパッと理解できる気がしません。

さいごに

ここまでの検証を経て、AWS SDK service integrations のやり方で Lambda タスクを実装するよりは payloadResponseOnly: true で実装しておいた方が素直に読み書きできそうだなぁ、と思いました。

Workflow Studio でステートマシンを実装した場合は AWS SDK service integrations の作法に準拠した ASL が出来上がるらしいので（私は未確認ですが）、AWS 的にはそっちが標準になっていくのでしょうかね。個人的には嬉しくない...。

ちなみに、Serverless Framework (の serverless-step-functions プラグイン）では基本的に ASL の実装は CloudFormation の記法ほぼそのままです。

yaml なので IDE の定義ジャンプなどが使えない点が非常にネックであるものの、ASL の読み書きに関しては　CloudFormation と ASL の仕様だけ理解していれば素直に読み解けるので私的にはそちらの方がわかりやすく感じます。

※私は普段は Serverless Framework ユーザーで CDK はまだまだライトユーザーです。なので、慣れによる贔屓目線もあるかもしれません

definition に関しては、もうちょい抽象化のレイヤーが薄く、かつ型定義の恩恵が受けられるような L2 があったら良かったのになぁ、と思いました。OSS なんだから自分で作れ、って話ですな。がんばりたい（がんばるとは言っていない）