StepFunctionsで実行したECS Runtaskをトレースする。

AWS Java 基盤グループ

マネックス証券 システム管理部のHです。

はじめに

Step Functionsはワークフローのサービスで、連続する複数のジョブ実行として現在私が関わっているプロジェクトでは使用しています。 Step FunctionsはX-Rayと連携されるため、ステートマシン内で実行されるジョブのトレースも取得できるのですが、ややはまったところがあったため書いていきます。 X-Rayについては下記ブログでも触れています。

blog.tech-monex.com

やりたいこと、困っていること

上記に書いた通り、Step FunctionsはX-Rayとも連携されており、Step Function自体は設定でチェックボックスをオンにするだけで簡単にトレースを取ることができます。 Step Functionのステートマシン内で実行されるAWSリソースについても連携できるものがあり、Lambdaについてはこれも設定でチェックボックスをオンにするだけでX-Rayの画面でStep FunctionsとLambda合わせてトレースを確認することができます。
一方でECS RuntaskについてもStep Functionsと連携を取りたいのですが、Step FunctionsとECS Runtaskの間でそれぞれX-Rayの設定をしても、連携されません・・・。

Step Functions＋Lambdaの場合

上記の通り、それぞれ設定でチェックボックスをオンにするだけでトレースデータが連携されます。

Step Functionsステートマシーン

Lambda関数

単純なレスポンスを返す関数です。

export const handler = async (event) => {
  // TODO implement
  const response = {
    statusCode: 200,
    body: JSON.stringify('こんにちはマネックス'),
  };
  return response;
};
X-Rayトレースマップ

Step FunctionsステートマシンとLambdaのトレースが表示されています。

Step Functions＋ECS Runtaskの場合

Step Functionsステートマシーン

Step Functionsでは設定でチェックボックスをオンにしています。一方ECS Runtaskについては特に設定しておりません。

ECS Runtask

Runtask内で実行するSpring Batchアプリケーションについて設定します。ここを参考に設定します。

  • バッチ設定クラス
    @XRayEnabledのアノテーションを設定することで、X-Rayのトレースマップに「job」として表示されます。
@Slf4j
@XRayEnabled
@Configuration
public class BatchConfiguration {

    @Bean
    public Job job(JobRepository jobRepository, PlatformTransactionManager txManager) {
        Step step = new StepBuilder("Step1", jobRepository)
                .tasklet(new TestTasklet(), txManager)
                .build();
        
        log.info("Execute Step");
        Job job = new JobBuilder("TestJob", jobRepository)
                .start(step)
                .build();
        log.info("Executed Step: " + job.getName());
        return job;
    }
}
  • タスクレット
    タスクレットは自動ではトレースが出力されないので、明示的にサブセグメントを設定します。
@Slf4j
public class TestTasklet implements Tasklet {

    @Override
    public RepeatStatus execute(StepContribution contribution, ChunkContext chunkContext) throws Exception {
        Subsegment subsegment = AWSXRay.beginSubsegment("TestTasklet::execute");
        log.info("TestTasklet::execute");
        Thread.sleep(EnvUtil.getSleepTime()); // トレースを見やすくするため数秒間停止させる
        subsegment.end();
        return RepeatStatus.FINISHED;
    }
}

  • X-Rayトレースマップ
    Step FunctionsとECS Runtaskは別々のトレースIDとなるため、1画面ですべての情報は表示できませんでした。1画面で見れないとそれぞれで画面を開かないといけないため、不便です。⇐ここが今回困っているところ。

    • Step Functionsステートマシン
      ECS Runtaskのトレースは表示されていますが、ECS内のトレースが表示されません。

    • ECS Runtask
      Spring Batchが実行しているジョブ、タスクレットの情報が表示されます。

どうやって実現するか考える

X-RayとStep Functionsのデベロッパーガイドを見ても、Lambdaのように簡単にはいかないようです。ECS Runtaskの場合にStep Functionsとトレースの画面が分かれてしまうのは、トレースIDがそれぞれで違うためです。ならばStep Functionsで実行しているステートマシーンのトレースIDをECS Runtaskに渡せればいけるのではないかと思いました。しかしどこからトレースIDが取得できるのだろうか。

実装してみる

トレースIDを取得する方法として、やや強引ですが、ECS Runtaskの前にLambdaを追加しました。
LambdaはHTTPヘッダーの中にトレースIDを持っており、それをECS Runtaskの入力で設定します。
Spring Batchでは、X-RayのSDKでトレースIDに環境変数から取得した値を設定します。

Step Functionsステートマシン

Lambda関数

出力の中にトレースIDが含まれています。

ECS Runtask

環境変数にトレースIDを設定します。

  • パラメータ設定

  • 実際の入力値

  • Spring Batchメインクラス
    メインクラス内で環境変数からトレースIDを取得します。

@Slf4j
@SpringBootApplication
public class EblogMainApplication {
    public static final String TRACE_NAME = "Eblog-202308";
    
    public static void main(String[] args) {
        List<String> ids = getTraceIds();
        if (ids.size() > 0) {
            TraceID traceId = TraceID.fromString(ids.get(0)); // rootの値がトレースID
            AWSXRay.getGlobalRecorder().beginSegment(TRACE_NAME, traceId, null);
        } else {
            AWSXRay.getGlobalRecorder().beginSegment(TRACE_NAME);
        }
        SpringApplication.run(EblogMainApplication.class, args);
        AWSXRay.endSegment();
    }

    // 環境変数から値を取得し、「;」「=」を区切り文字として分割してリストに代入する。
    // root=1-xxxxxxxx-xxxxxxxx;parent=xxxxxxxx;sampled=1;lineage=xxxxxxxx:0
    private static List<String> getTraceIds() {
        List<String> list = new ArrayList<>();
        String envTraceIds = System.getenv("TRACE_ID");
        if (StringUtils.isEmpty(envTraceIds)) {
            return list;
        }
        log.info("trace data: " + envTraceIds);
        String[] arr = envTraceIds.split(";");
        if (arr.length >= 2) {
            for (String tmpStr: arr) {
                String[] val = tmpStr.split("=");
                if (val.length >= 2) {
                    list.add(val[1]);
                }
            }
        } else {
            return list;
        }

        if (list.size() < 2) {
            return new ArrayList<>();
        }

        return list;
    }
}

確認してみる

Step Functionsのステートマシンを実行したトレースです。1画面内にECS Runtaskのトレースも含まれています。

一見するとうまくいっているように見えます。

ですがトレースマップを見てみるとStep Functionsの部分とECS Runtaskの部分が分かれました。ECSの下にEblog-202308がついてほしかったのですが、そうなりませんでした・・・。
理由はSpring Batchアプリケーション内で親のID(この場合ECSのID)を指定していないためです。 下記でnullを指定している部分に本来なら親のIDを指定するのですが、親のIDを取得することができなかったため仕方なくnullを指定しました。Step FunctionsとECS Runtaskの連携についてはAWSの今後に期待です。

AWSXRay.getGlobalRecorder().beginSegment(TRACE_NAME, traceId, null);

おわりに

課題はあるものの、ステートマシン実行のトレースについて、1画面内でStep FunctionsとECS Runtaskの確認をすることができました。ステートマシン内でLambdaとECS Runtask両方のジョブがある場合には利用できる方法かと思います。

人材募集