X-Ray Traceの分断に対処した話(供養)

2024年7月19日 12:00

はじめに

こんにちは、build サービス部で DevOps エンジニアをやっているt.g.です。
今日は、マイクロサービス構成のアプリケーションなどで威力を発揮するX-Ray Trace でちょっとカスタマイズをした話をしたいと思います。

実は、ここでお話するやり方よりもベターな方法があります。
ですが Trace の再構成は、あまり情報がなく興味深い内容と思いますので、供養がてら共有させてください。

X-Ray Trace とは

X-Ray Traceは、AWSの分散トレースサービスであり、アプリケーションのパフォーマンスを分析し、デバッグするために使用されます。
このサービスを利用することで、クライアントのリクエストがどのようにシステム全体を通過するかを視覚化し、各リソースの実行時間やエラーの発生場所を特定することができます。
例えば、マイクロサービスで構成したアプリケーションなどで威力を発揮します。

主な機能と特徴

分散トレースの可視化
クライアントのリクエストがシステム内の各サービスを通過する際の詳細な経路を追跡できる
パフォーマンス分析
各サービスやリソースの実行時間を計測し、ボトルネックを特定できる
エラー検出
エラーや例外の発生箇所を特定し、迅速なトラブルシューティングができる
統計データの集約
リクエストの遅延、エラーレート、サンプル数などの統計データを収集・表示できる

これを利用することで、開発者にはデバッグが、DevOps エンジニアにはシステム全体の健全性が見られるようになります。

Trace 分断の問題

今回、下記のような構成でトレースも取得するようにしました。

前段は、API-Gateway がリクエストを受けとり、前処理として Lambda を起動、DynamoDB にデータを保存します。
後段として、更新された DynamoDB のアイテムを後処理 Lambda が非同期で実行する形になっています。

以下が Trace Map の結果になります。

API Gateway から前処理 Lambda(FrontLambda)、DynamoDB Table までのトレースは一本でまとまっています。
ただ、後処理Lambda (StreamLambda) は別トレースとして分断されてしまいました。

Trace-Id を追跡

確認のため Log から要所要所の Trace-Id を取得し見比べてみます。
やはり、DDB-Stream でTrace-Id が別物に切り替わっていました。

API-Gateway
Root=1-6697729e-1ce38255066e7abe1830887b
前処理 Lambda
traceId: 1-6697729e-1ce38255066e7abe1830887b
後処理 Lambda
traceId: 1-66976ef6-f250332b983a2a0b121d585f

なぜ分断されるのか?

たいていのAWSサービスは Trace-Id をネイティブサポートしており、サービスをまたいでも自動で Trace-Id を伝播してくれます。
その証拠に、API-Gateway <-> Lambda 間でのトレースIDは自動で伝播しています。

DynamoDB Stream はこのようなサポートをしていないため、Trace-Id の分断が発生しました。

Trace 分断への対処 (注:ベターな方法あり)

今回、以下のような対処を行いました。

前処理 Lambda: DDB への書き込み Item 内に、Trace-Id と Segment-Id を忍ばせる
後段 Lambda: Trace-Id と Segment-Id を取得。Trace として再構成

1. 前処理Lambda

Trace-Id と Segment-Id は以下のコードで取得し、DDB への書き込み Item に付属させました
ここでは、Power Tools for TypeScript を使っています。

// Create a DynamoDB Client and patch it for tracering
const ddbClient = tracer.captureAWSv3Client(new DynamoDBClient({}));
const docClient = DynamoDBDocumentClient.from(ddbClient);

class Lambda implements LambdaInterface {

  @tracer.captureLambdaHandler()
  public async handler(event: APIGatewayProxyEvent, context: Context): Promise<APIGatewayProxyResult> {

    // Create a DDBItem to add
    const item: DdbItem = {
      id: uuidv4(),
      name: 'something new',
      debug_info: {
        trace_id: tracer.getRootXrayTraceId() || 'no-trace-id', // 取得 & 書き込み
        segment_id: tracer.getSegment()?.id || 'no-segment-id', // 取得 & 書き込み
      },
    };

    try {
      await docClient.send(new PutCommand({
        TableName: tableName,
        Item: item,
      }));

  // 以下省略

この結果、DDB には以下のような debug_info 付きの Item が保存されます。

{
  "item": {
    "id": "b9dca9c4-95df-4572-a4c9-90977a75d3c2",
    "name": "something new",
    "debug_info": {
      "trace_id": "1-6697788a-262f9e1c780d1f6575e47657",
      "segment_id": "f137be300c4b46a3"
    }
  }
}

2. 後処理 Lambda

後処理 Lambda は、DDB-Stream から受け取った event から、trace_id とsegment_id を取得。
新しく Segment 作成し、それをDDB-Stream から受け取ったTrace/Segment に紐づけています

PowerTools for Lambda ではここまでマニュアルな方法はみつからなかったため、xray-sdk を利用しています。

ここでセットしている segment の各種メンバ変数はこちらを参照してください。

import { Segment, setSegment } from 'aws-xray-sdk-core';
import { Context, DynamoDBStreamEvent } from 'aws-lambda';

function createLambdaSegment(context: Context, streamRecord: DynamoDBStreamEvent) {

  // DynamoDB Stream から受け取ったイベントから、trace_id と (parent)segment_id を取得
  const ddbRecord = streamRecord.Records[0];
  const dynamodb = ddbRecord.dynamodb;
  const newImage = dynamodb!.NewImage; // ある前提
  const debugInfo = newImage!.debug_info;
  const traceId = debugInfo.M!.trace_id.S;
  const parentSegmentId = debugInfo.M!.segment_id.S;

  // 新しいセグメントを作成
  // 手作業作成のため、各種情報をマニュアル追加する
  const segment = new Segment(
    context.functionName,
    traceId,
    parentSegmentId
  );
  segment.origin = 'AWS::Lambda::Function';
  segment.start_time = new Date().getTime() / 1000;
  segment.addPluginData({
    function_arn: context.invokedFunctionArn,
    region: ddbRecord.awsRegion,
    request_id: context.awsRequestId,
  });
  return segment;
}


/**
 * Creates and returns a Lambda segment for the given context and DynamoDB stream record.
 * @param {Context} context - The Lambda execution context.
 * @param {DynamoDBStreamEvent} streamRecord - The DynamoDB stream record.
 * @returns {Segment} - The created Lambda segment.
 */
export function continueLambdaSegment(context: Context, streamRecord: DynamoDBStreamEvent) {
  const segment = createLambdaSegment(context, streamRecord);
  setSegment(segment);
  return segment;
}

この結果、他 Trace 扱いだった後処理 Lambda が、一個の Trace に紐づきました

さらなる最適化案

いちおう一本の Trace に入りはしました。
が、DDB の後に後処理 Lambda が入ってほしいです。

これは、前処理 Lambda での Segment-Id の取得タイミングが原因です。

以下、前処理 Lambda でのSegment-Id 取得コード(再掲)

// Create a DynamoDB Client and patch it for tracering
const ddbClient = tracer.captureAWSv3Client(new DynamoDBClient({}));
const docClient = DynamoDBDocumentClient.from(ddbClient);

class Lambda implements LambdaInterface {

  @tracer.captureLambdaHandler()
  public async handler(event: APIGatewayProxyEvent, context: Context): Promise<APIGatewayProxyResult> {

    // Create a DDBItem to add
    const item: DdbItem = {
      id: uuidv4(),
      name: 'something new',
      debug_info: {
        trace_id: tracer.getRootXrayTraceId() || 'no-trace-id', // 取得 & 書き込み
        segment_id: tracer.getSegment()?.id || 'no-segment-id', // 取得 & 書き込み
      },
    };

    try {
      await docClient.send(new PutCommand({ // これが作成する segment-Id が本当にほしかったもの
        TableName: tableName,
        Item: item,
      }));

  // 以下省略

本来なら docClient.send() が作成する Segment-Id を使うべきですが、今回は handler で取得した Segment-Id を利用しています。

capture される ddbClient の Segment-Id をあらかじめ取得したり、指定することができれば、DDB-Stream 時の Trace もよりきれいな形になるはずです

もっと良い方法 (PowerTools の batch機能)

上記の方法では
複数 Item のそれぞれの処理に対して Segment で追跡したい
場合に対応できません。

実は、PowerTools for Lambda の batch 機能を使うと

複数 Item に対する処理をコントロール
- 一部だけ fail したときの対処も制御
Trace も対応 (subsegment を作成)

という良いとこずくめな方法がありました。
今度はこちらで試してみたいと思います。

まとめ

今回、Dynamo DB Stream から受け取った trace を再構成することで、Trace-Id の分断化を防ぎ、一本のトレースとして追跡できることができるようになりました。

なお、ここで Trace の再構成は xray-sdk を利用しました。

ですが実は、PowerTools for Lambda (TypeScript) の batch 機能を使うと同様のことがもっと高機能にできるはず

ぜひ次回はこれを使ってみたいと思います