2017年頃からサーバレスに浸っている、TISの小西啓介です。

概要

SAM(Serverless Application Model)を使ったnode.jsのLambdaプログラム等に対して、
パフォーマンス計測等を行いX-Ray機能の組み込みについて説明します。

また、X-Ray機能を組み込むと、aws-sdk-mockを使用したユニットテストでは
通らなくなるので、aws-xray-sdkのモック化についても説明します。

SAMへのX-Rayの追加

API Gateway - Lambda - DynamoDBのような構成だとして、
SAM(Serverless Application Model)へのX-Rayの組み込みは、
以下のように、SAMのテンプレーレートのGlobalセクションの
Function に対して、Tracing: Activeと
Api に対して、TracingEnabled: Trueを追加するだけです。

もちろん、Globalではなく、AWS::Serverless::Functionや
AWS::Serverless::ApiのPropertiesに追加しても良いです。

変更前SAM-Template


Globals:
  Function:
    Timeout: 15
  Api:
    # https://github.com/awslabs/serverless-application-model/issues/191
    OpenApiVersion: 3.0.0

変更後SAM-Template

Globals:
  Function:
    Timeout: 15
    Tracing: Active # 追加
  Api:
    # https://github.com/awslabs/serverless-application-model/issues/191
    OpenApiVersion: 3.0.0
    TracingEnabled: True # 追加

Lambda内のAWSサービス呼出へのX-Rayの組み込み

API GatewayからLambdaへの呼出のみであれば、上記のSAMへの組み込みのみで十分なのですが、
Lambda内で、AWS-SDKを使ってAWSのサービスを呼び出している場合、その呼出のトレースを行う場合、コードの変更が必要です。
ただ、コードの変更と言っても、 npm i -S aws-xray-sdk-core でaws-xray-sdk を組み込んで以下のコードの変更を行うだけです。

変更前js

const AWS = require('aws-sdk');

変更後js

const AWSXRay = require('aws-xray-sdk-core');
const AWS = AWSXRay.captureAWS(require('aws-sdk'));

なお、aws-xray-sdkには、以下の４つのモジュールが含まれます。

aws-xray-sdk-core
aws-xray-sdk-express
aws-xray-sdk-mysql
aws-xray-sdk-postgres

Lambdaの場合、mysqlやpostgresに接続しないのであれば、aws-xray-sdk-coreで十分だと思います。
Lambda内で、AWS Serverless Expressを使用している場合、aws-xray-sdk-express が気になりますが、これはEC2用の為、使用できません（Lambdaでは、呼出時にX-Rayのセグメントが既に作られているので、新たにセグメントを設定することは出来ないため）

ユニットテスト(jest)でのモック化

上記のようなAWS-SDKを使用しているプログラムのユニットテストで便利な AWS-SDK-Mockがあるのですが、AWS-SDKをラップしている関係で、そのままのコードでは動かなくなってしまします。

AWSXRayが使用する部分AWS-SDKの処理をAWS-SDK-Mockで書けば良いのもしれないのですが、情報が無くよくわかりませんでした。
(AWS_XRAY_CONTEXT_MISSING などの環境変数を設定してもだめでした)

また、Jest用にaws-xray-Mockというのが有るのですが、上記のような呼び出し方(captureAWS)には対応しておらず、個別の(AWSサービス毎の)クライアントにしか対応していないようです。プルリクも出ていますが、５月から、放置されてしまっています。
https://github.com/MechanicalRock/aws-xray-mock/pull/4

そこで、上記のコードを参考に、captureAWSのみに対応する場合、
以下の追加で、aws-sdk-mockを使ったテストが全て通るようになりました。

beforeAll(async () => {
    jest.mock('aws-xray-sdk-core', () => {
        return { captureAWS: aws => aws };
    });
});

サブセグメントの追加やAWSクライアント(サービス)単位の設定等を行う場合、
以下のように、利用するメソッドを追加しましょう。
(出典：aws-xray-mock 及び aws-xray-mockのプルリク)

let xRaySegment;
class MockXraySegment {
    constructor(name) {
        this.name = name;
        xRaySegment = this;
    }
    addNewSubsegment(name) {
        return new MockXraySegment(name);
    }
    addError(error) { }
    close() { }
}
beforeAll(async () => {
    // process.env.AWS_XRAY_CONTEXT_MISSING = 'LOG_ERROR';
    jest.mock('aws-xray-sdk-core', () => {
        return {
            captureAWS: aws => aws,
            captureAWSClient: client => client,
            captureHTTPs: http => http,
            captureHTTPsGlobal: httpGlobal => httpGlobal,
            getSegment: () => {
                if (!xRaySegment) {
                    xRaySegment = new MockXraySegment('Parent Segment');
                    return xRaySegment;
                }
                return xRaySegment;
            }
        };
    });
});
beforeEach(async () => {
    xRaySegment = null;
});

やはり、 aws-xray-mock のプルリクマージして欲しいですね。
(aws-xray-sdkだけでなくて、aws-xray-sdk-core もモック化してほしいですが。。)

おわり

参考