Amazon SNS でサポートされた AWS X-Ray のアクティブトレースを試してみる

今回は、つい先日にアナウンスがあった Amazon SNS の下記の新しい機能を試してみます。

aws.amazon.com

なお、この記事の内容は 2023 年 2 月 12 日時点の検証結果に基づいて記載しています。

何ができるようになったのか

これまでも、Amazon SNS トピックは、呼び出し先が AWS Lambda 関数または HTTP/HTTPS のエンドポイントの場合は、AWS X-Ray のトレースを取得するために必要なコンテキストを伝播することはできました。

そのため、下図のようにトピックを経由して AWS Lambda 関数を呼び出すまでの一連のフローを AWS X-Ray で可視化することはできました。

ただし、この場合、Amazon SNS のトピックはトレースのコンテキストを AWS Lambda 関数や HTTP/S のエンドポイントに伝播しているだけで、Amazon SNS のトピック自体は AWS X-Ray のトレースを出力することはできませんでした。

そのため、例えばトピックのサブスクライバーが Amazon SQS のキューの場合は、トピックからキューに対してメッセージを送信したときのトレースは取得できませんでした。

今回の Update で Amazon SNS トピックとしてどのような動作を行ったかを AWS X-Ray のトレースとして出力できるようになりました。

この機能を適用できるのは標準トピックだけで、FIFO トピックには適用できません。

また、トレース取得の対象にできるのは、サブスクライバーが下記の場合のみです。

AWS Lambda 関数
Amazon SQS のキュー
Amazon Kinesis Dara Firehosse
HTTP/S エンドポイント

上記以外のEmail や SMS ( Short Message Service ) には適用できないことは留意しておきましょう。

それでは試してみましょう！

サブスクライバーが AWS Lambda 関数や Amazon SQS キュー、Email の構成で試してみる

Amazon SNS のトピックで X-Ray トレースの取得を有効化して、下図のような構成でトピックにメッセージを発行してみます。

AWS マネジメントコンソールでの設定

X-Ray トレースの有効化は、AWS マネジメントコンソールの場合はトピックの設定ページの [編集] ボタンから行えます。

また、トピックの設定ページの下部にある [統合] タブから有効化を確認できます。

取得したトレース

AWS マネジメントコンソールからトピックに対してメッセージを発行してトレースを取得して表示してみましょう。

今回は、Amazon CloudWatch のコンソールの左側のメニューで [ X-Ray トレース ] - [ トレース ] を選択して表示してみます。

トレースマップでは、Amazon SNS トピックとサブスクライバーである Amazon SQS キューや、AWS Lambda 関数が表示されています。

セグメントのタイムラインでは、Amazon SNS が Amazon SQS キューに SendMessage API を発行、また AWS Lambda 関数を invoke する API を発行したことも表示され、その処理にかかった時間も表示されています。

実際はサブスクライバーとして Email も設定しメールも送信されているのですが、トレースには含まれていません。ただ Email はもともとトレース取得の対象外なので、これは想定通りですね。

次に、敢えてサブスクライバーとして設定している AWS Lambda 関数を削除してからトピックにメッセージを送信したときのトレースをみてみましょう。

AWS Lambda 関数の invoke がエラーコード 403 で失敗していることがわかりますね。またそのため、サブスクリプションに設定したデッドレターキューにメッセージが送信されたこともわかります。

シンプルなファンアウトの構成で試してみる

次に、Amazon SNS と Amazon SQS、AWS Lambda を組み合わせて下図のようなシンプルなファンアウトを実装した場合、AWS X-Ray ではどのようにトレースを可視化できるのかを試してみます。

使用する AWS Lambda 関数のコード

今回は、Python の Lambda 関数を使用します。下記は、Amazon SNS の標準トピックにメッセージを発行する Lambda関数です。

メッセージとして、ハンドラ関数の引数であるイベントオブジェクトを文字列化したものを送信する前提とします。

import json
import boto3
import datetime
import os
from botocore.config import Config
from aws_xray_sdk.core import patch
patch(['boto3'])

sns = boto3.client('sns')

def lambda_handler(event, context):
    topic_arn = os.getenv('SNS_TOPIC_ARN')
    dt = datetime.datetime.now().strftime('%Y/%m/%d %H:%M:%S %Z')
   
    # Topicへメッセージ送信
    response = sns.publish(
        TopicArn=topic_arn,
        Message=json.dumps(event),
        Subject='Message from Lambda Function ' + dt
    )
    
    #
    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "Done:"+ dt
        }),
    }

次は、トピックをサブスクライブしている Amazon SQS のキューからメッセージを取得する Lambda 関数です。

関数名とイベントオブジェクトをそのまま print 関数で出力するだけのシンプルな内容です。

ファンアウトではありますが、今回はキューからメッセージを取得する 2 つの Lambda 関数は同じコードを使用します。

import json
import boto3
import datetime
import os
from botocore.config import Config
from aws_xray_sdk.core import patch
patch(['boto3'])

def lambda_handler(event, context):
    print("Lambda function name:", context.function_name)
    print("-----")
    print(event)
    #
    return {
        "statusCode": 200,
        "body": json.dumps({
            "message": "Done"
        }),
    }

なお、AWS X-Ray の SDK を使用していますが、AWS X-Ray SDK のパッケージは AWS Lambda レイヤーとして用意し、今回使用する全てのLambda 関数から共用する前提にしています。

シンプルなファンアウトを構築するための AWS SAM テンプレート

もちろん、AWS マネジメントコンソールから構築しても良かったんですが、再利用性を考え今回は AWS SAM テンプレートで構築することにしました。

Amazon API Gateway のAPI 、AWS Lambda 関数、そして今回のテーマである Amazon SNS トピックと全て AWS X-Ray のトレース取得を有効化する前提です。

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  sam-demo-xray-fanout

  Sample SAM Template for sam-demo-xray-fanout

Parameters:
  SNSTOPIC:
    Type: String
    Default: 'DemoXRayFanoutTopic'
  SQSQUEUE1:
    Type: String
    Default: 'DemoXRayFanoutQueue1'
  SQSQUEUE2:
    Type: String
    Default: 'DemoXRayFanoutQueue2'  
Globals:
  Function:
    Timeout: 15
    Tracing: Active
    Layers:
      - !Ref DemoXRayLambdaLayer
  Api:
    TracingEnabled: True
Resources:
  DemoXRayFanoutSNSPublishFunction:
    Type: AWS::Serverless::Function 
    Properties:
      FunctionName: 'DemoXRayFanoutSNSPublishFunction'
      CodeUri: demo_xray_fanout_sns_publish_function/
      Handler: app.lambda_handler
      Runtime: python3.8
      Role: !GetAtt DemoXRayFunctionRole.Arn
      Environment:
        Variables:
          SNS_TOPIC_ARN: !GetAtt DemoXRayTopic.TopicArn
      Events:
        DemoXRayEvent:
          Type: Api 
          Properties:
            Path: /xray
            Method: get
  DemoXRayFanoutSQSReceiveFunction1:
    Type: AWS::Serverless::Function
    Properties:
      FunctionName: 'DemoXRayFanoutSQSReceiveFunction1'
      CodeUri: demo_xray_fanout_sqs_receive_function/
      Handler: app.lambda_handler
      Runtime: python3.8
      Role: !GetAtt DemoXRayFunctionRole.Arn
      Events:
        DemoSQSEvent:
          Type: SQS
          Properties:
            Queue: !GetAtt DemoXRayQueue1.Arn
            BatchSize: 1
  DemoXRayFanoutSQSReceiveFunction2:
    Type: AWS::Serverless::Function 
    Properties:
      FunctionName: 'DemoXRayFanoutSQSReceiveFunction2'
      CodeUri: demo_xray_fanout_sqs_receive_function/
      Handler: app.lambda_handler
      Runtime: python3.8
      Role: !GetAtt DemoXRayFunctionRole.Arn
      Events:
        DemoSQSEvent:
          Type: SQS
          Properties:
            Queue: !GetAtt DemoXRayQueue2.Arn
            BatchSize: 1
  DemoXRayLambdaLayer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      LayerName: 'DemoXRayLambdaLayer'
      CompatibleRuntimes:
        - python3.8
      ContentUri: demo_xray_layer/xray-python.zip
  DemoXRayFunctionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: "Allow"
            Action: "sts:AssumeRole"
            Principal: 
              Service: "lambda.amazonaws.com"
      Policies:
        - PolicyName: "my-demo-xray-policy"
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:*"
                  - "sns:*"
                  - "sqs:*"
                  - "xray:PutTraceSegments"
                  - "xray:PutTelemetryRecords"
                Resource: "*"
  DemoXRayTopic:
    Type: AWS::SNS::Topic
    Properties:
      TopicName: !Ref SNSTOPIC
      TracingConfig: Active
  DemoXRayQueue1:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: !Ref SQSQUEUE1
  DemoXRayQueue2:
    Type: AWS::SQS::Queue
    Properties:
      QueueName: !Ref SQSQUEUE2  
  DemoXRayTopicSubscription1:
    Type: 'AWS::SNS::Subscription'
    Properties:
      TopicArn: !Ref DemoXRayTopic
      Endpoint: !GetAtt 
        - DemoXRayQueue1
        - Arn
      Protocol: sqs
      RawMessageDelivery: 'true'   
  DemoXRayTopicSubscription2:
    Type: 'AWS::SNS::Subscription'
    Properties:
      TopicArn: !Ref DemoXRayTopic
      Endpoint: !GetAtt 
        - DemoXRayQueue2
        - Arn
      Protocol: sqs
      RawMessageDelivery: 'true'
  DemoXRayQueuePolicy: 
    Type: AWS::SQS::QueuePolicy
    Properties: 
      Queues: 
        - !Ref DemoXRayQueue1
        - !Ref DemoXRayQueue2
      PolicyDocument: 
        Statement: 
          - 
            Action: 
              - "SQS:SendMessage" 
            Effect: "Allow"
            Resource: 
              - !GetAtt DemoXRayQueue1.Arn
              - !GetAtt DemoXRayQueue2.Arn
            Principal:  
              Service: 
                -  "sns.amazonaws.com"
            Condition:
              ArnLike:
                aws:SourceArn:
                  - !GetAtt DemoXRayTopic.TopicArn
  DemoXRayResoucePolicy:
    Type: AWS::XRay::ResourcePolicy
    Properties: 
      BypassPolicyLockoutCheck: True
      PolicyName: "PolicyforDemoXRayFanoutTopic" 
      PolicyDocument: !Sub | 
        {
          "Version": "2012-10-17",
          "Statement": [
          {
            "Effect": "Allow",
            "Principal": {
              "Service": "sns.amazonaws.com"
            },
            "Action": [
              "xray:PutTraceSegments",
              "xray:GetSamplingRules",
              "xray:GetSamplingTargets"
            ],
            "Resource": "*",
            "Condition": {
              "StringEquals": {
                "aws:SourceAccount": "${AWS::AccountId}"
              },
              "StringLike": {
                "aws:SourceArn": "${DemoXRayTopic.TopicArn}"
              }
            }
          }
        ]
        }
Outputs:
  DemoXRayFunctionApi:
    Description: "API Gateway endpoint URL for Prod stage "
    Value: !Sub "https://${ServerlessRestApi}.execute-api.${AWS::Region}.amazonaws.com/Prod/xray/"

この AWS SAM テンプレートでポイントになる部分は 2つです。

1 つ目は、104 行目 (下記）です。

TracingConfig: Active

この Amazon SNS トピックの TracingConfig プロパティで Active を指定することで、AWS X-Ray トレースの取得を有効化します。

2 つ目は、153 行目から 183 行目で、AWS X-Ray のリソースベースのポリシーを作成している部分です。

AWS マネジメントコンソールから Amazon SNS トピックに対してトレースを有効化したときには自動的に AWS X-Ray のリソースベースのポリシーも設定してくれたのですが、AWS SAM などを使用する場合は、明示的に作成する必要があります。

また、このポリシーの内容については 158行目からの PolicyDocument で指定するのですが、このプロパティは Type (型）が String になっています。

docs.aws.amazon.com

そのため、例えば 89 行目にある IAM ロールのポリシードキュメントの記述方法が異なるので注意が必要です。

89行目からの IAM ロールのポリシードキュメントは、Type が Json なので、下記のように記述できます。

PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action:
                  - "logs:*"
                  - "sns:*"
                  - "sqs:*"
                  - "xray:PutTraceSegments"
                  - "xray:PutTelemetryRecords"
                Resource: "*"

しかし 153行目からの AWS X-Ray のリソースポリシーのポリシードキュメントは、Type が String なので、下記のように記述する必要があります。

 PolicyDocument: !Sub | 
        {
          "Version": "2012-10-17",
          "Statement": [
          {
            "Effect": "Allow",
            "Principal": {
              "Service": "sns.amazonaws.com"
            },
            "Action": [
              "xray:PutTraceSegments",
              "xray:GetSamplingRules",
              "xray:GetSamplingTargets"
            ],
            "Resource": "*",
            "Condition": {
              "StringEquals": {
                "aws:SourceAccount": "${AWS::AccountId}"
              },
              "StringLike": {
                "aws:SourceArn": "${DemoXRayTopic.TopicArn}"
              }
            }
          }
        ]
        }