AWS X-RayのAPIを使ってみる

# はじめに
最近AWS X-Rayについて触りながら色々試しています。AWS X-RayにもAPIがあったので少し触ってみたので備忘録として残します。

https://docs.aws.amazon.com/ja_jp/xray/latest/devguide/xray-api-gettingdata.html

:::note warn
実際に試す場合は自己責任でお願いします。
:::

## 今回試したこと
AWS CLIを使ってトレースの情報を取得できます。以下のドキュメントにあるようにセグメントの情報が載ったJSONを取得できるようです。そのJSONを読むことで、AWS Lambda関数の処理時間やコールドスタートにかかった時間など色々取得できそうでした。

今回は、ランダムな時間スリープするAWS Lambda関数を作成して何回か実行したリクエストを、APIを使ってJSONを取得・解析して処理時間を確認してみようと思います。

https://docs.aws.amazon.com/ja_jp/xray/latest/devguide/xray-api-gettingdata.html

https://docs.aws.amazon.com/ja_jp/xray/latest/devguide/xray-api-segmentdocuments.html

```python:AWS Lambda
import json
import random
import time
from aws_xray_sdk.core import xray_recorder

def lambda_handler(event, context):
    
    xray_recorder.begin_subsegment("Test")
    
    # generated by ChatGPT
    min_time = 1  
    max_time = 5  

    sleep_time = random.randint(min_time, max_time)
    time.sleep(sleep_time)
    
    xray_recorder.end_subsegment()
    
    return {
        'statusCode': 200,
        'body': json.dumps('Hello from Lambda!')
    }
```


ドキュメントにあるように、指定した時間範囲のTraceIdを取得して得られたIDのトレース全体を取得するスクリプトを作成しました。また、得られたJSONを解析するスクリプトも簡単に作成しました。今回は、AWS Lambda側で作成した名前がTestのサブセグメントの処理時間（end_time - start_time）をCSVに出力してみます。

```bash:get_segment.sh
EPOCH=$(date +%s)
basedir="trace_$(($EPOCH-$2))_$(($EPOCH-$1))"
TRACEIDS=$(aws xray get-trace-summaries --start-time $(($EPOCH-$2)) --end-time $(($EPOCH-$1)) --query 'TraceSummaries[*].Id' --output text)
echo $TRACEIDS
mkdir -p $basedir
for trace_id in $TRACEIDS
do
    echo $trace_id
    aws xray batch-get-traces --trace-ids $trace_id > $basedir'/'$trace_id.json
    python read_trace.py --traceid $trace_id --basedir $basedir
done
```

```python:read_trace.py
import json
import argparse
import subprocess

parser = argparse.ArgumentParser()
parser.add_argument('--basedir', '-b', help='base dir')
parser.add_argument('--traceid', '-i', help='traceid')
args = parser.parse_args()

basedir = args.basedir + '/'
input_file = basedir + args.traceid + '.json'
output_file = basedir + 'result.csv' 

with open(input_file) as f:
    trace_json = json.load(f)

for trace_obj in trace_json['Traces']:
    trace_id = trace_obj['Id']
    print (trace_id)
    for segment in trace_obj['Segments']:
        document_obj = json.loads(segment['Document'])
        if 'origin' in document_obj and document_obj['origin'] == "AWS::Lambda::Function":
            for subsegment in document_obj['subsegments']:
                if subsegment['name'] == 'Invocation':
                    for subsegment2 in subsegment['subsegments']:
                        if subsegment2['name'] == 'Test':
                            start_time = subsegment['start_time']
                            end_time = subsegment['end_time']
                            with open(output_file, 'a') as f:
                                text = trace_id + ',' + document_obj['name'] + ',' + str(end_time - start_time)
                                f.write(text)
                                f.write('\n')
```


例えば、600秒（10分）以内のトレースを取得するには以下のように実行すればよいです。

```bash
$ bash get_segment.sh 0 600 # 600秒前までのトレースを取得
```

AWS LambdaはAPI Gatewayから叩けるように設定して、Pythonで10回叩いてみました。

```python:lambda_request.py
import requests

url = "api-gateway-execute-url"
api_key ="api-key"

headers = {'x-api-key' : api_key}

for i in range(10):
    requests.get(url, headers=headers)
```

![Trace.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118168/d0d40f1a-2ae5-1422-ddd6-1c3b650363cf.png)

![スクリーンショット 2023-05-28 19.12.12.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/118168/3e5b396b-530b-9da1-c71c-55b1d117cf03.png)


得られるJSONは以下のような形です。


```json:Trace.json
{
    "Traces": [
        {
            "Id": "1-647208a0-1daefbe244cb264546375b6b",
            "Duration": 1.028,
            "LimitExceeded": false,
            "Segments": [
                {
                    "Id": "588258985084eff2",
                    "Document": "{\"id\":\"588258985084eff2\",\"name\":\"testLambda1\",\"start_time\":1.6851949128601005E9,\"trace_id\":\"1-647208a0-1daefbe244cb264546375b6b\",\"end_time\":1.6851949138627546E9,\"parent_id\":\"00dad9c947a091fd\",\"aws\":{\"account_id\":\"xxxx\",\"function_arn\":\"xxxx\",\"resource_names\":[\"testLambda1\"]},\"origin\":\"AWS::Lambda::Function\",\"subsegments\":[{\"id\":\"2676a9ee21eb86fe\",\"name\":\"Overhead\",\"start_time\":1.6851949138624983E9,\"end_time\":1.6851949138626897E9,\"aws\":{\"function_arn\":\"xxxx\"}},{\"id\":\"77150bd07b269df3\",\"name\":\"Invocation\",\"start_time\":1.6851949128601909E9,\"end_time\":1.6851949138624642E9,\"aws\":{\"function_arn\":\"xxxx\"},\"subsegments\":[{\"id\":\"90a1036a97a7a592\",\"name\":\"Test\",\"start_time\":1.6851949128606284E9,\"end_time\":1.685194913861826E9,\"in_progress\":false,\"namespace\":\"local\"}]}]}"
                },
                {
                    "Id": "00dad9c947a091fd",
                    "Document": "{\"id\":\"00dad9c947a091fd\",\"name\":\"testLambda1\",\"start_time\":1.685194912851E9,\"trace_id\":\"1-647208a0-1daefbe244cb264546375b6b\",\"end_time\":1.685194913863E9,\"parent_id\":\"4f31e15d8ffd5546\",\"http\":{\"response\":{\"status\":200}},\"aws\":{\"request_id\":\"7ca49352-e68e-4969-9319-c935dd75886d\"},\"origin\":\"AWS::Lambda\",\"resource_arn\":\"xxxx\"}"
                },
                {
                    "Id": "5a4b683f41923a69",
                    "Document": "{\"id\":\"5a4b683f41923a69\",\"name\":\"testLambda1-API/default\",\"start_time\":1.685194912837E9,\"trace_id\":\"1-647208a0-1daefbe244cb264546375b6b\",\"end_time\":1.685194913865E9,\"http\":{\"request\":{\"url\":\"xxxx\",\"method\":\"GET\",\"user_agent\":\"python-requests/2.26.0\",\"client_ip\":\"xxxx\",\"x_forwarded_for\":true},\"response\":{\"status\":200,\"content_length\":0}},\"aws\":{\"api_gateway\":{\"account_id\":\"xxxx\",\"rest_api_id\":\"xxxx\",\"stage\":\"default\",\"request_id\":\"45ff8314-691e-4f68-a39e-3e4396da79bd\"}},\"annotations\":{\"aws:api_id\":\"xxxx\",\"aws:api_stage\":\"default\"},\"metadata\":{\"default\":{\"extended_request_id\":\"FlZJLH-oNjMFWLg=\",\"request_id\":\"45ff8314-691e-4f68-a39e-3e4396da79bd\"}},\"origin\":\"AWS::ApiGateway::Stage\",\"resource_arn\":\"xxxx\",\"subsegments\":[{\"id\":\"4f31e15d8ffd5546\",\"name\":\"Lambda\",\"start_time\":1.685194912842E9,\"end_time\":1.685194913865E9,\"http\":{\"request\":{\"url\":\"xxxx\",\"method\":\"ANY\"},\"response\":{\"status\":200,\"content_length\":20}},\"aws\":{\"function_name\":\"testLambda1\",\"region\":\"ap-northeast-1\",\"operation\":\"Invoke\",\"resource_names\":[\"testLambda1\"]},\"namespace\":\"aws\"}]}"
                }
            ]
        }
    ],
    "UnprocessedTraceIds": []
}
```

Subsegmentに関する情報を抜き出したJSONは以下のような形になります。

```json
{
    "id": "7fb33877422529f6",
    "name": "testLambda1",
    "start_time": 1685196597.317368,
    "trace_id": "1-64720f35-7eed683d28d035345dde57af",
    "end_time": 1685196599.3208942,
    "parent_id": "1b300d0d9257ab37",
    "aws": {
      "account_id": "xxxx",
      "function_arn": "xxxx",
      "resource_names": [
        "testLambda1"
      ]
    },
    "origin": "AWS::Lambda::Function",
    "subsegments": [
      {
        "id": "2a4933c558cdaa93",
        "name": "Invocation",
        "start_time": 1685196597.3174572,
        "end_time": 1685196599.3206277,
        "aws": {
          "function_arn": "xxxx"
        },
        "subsegments": [
          {
            "id": "a21dd67efcc46849",
            "name": "Test",
            "start_time": 1685196597.3178673,
            "end_time": 1685196599.320048,
            "in_progress": false,
            "namespace": "local"
          }
        ]
      },
      {
        "id": "af02ee737492a3a1",
        "name": "Overhead",
        "start_time": 1685196599.3206565,
        "end_time": 1685196599.3208442,
        "aws": {
          "function_arn": "xxxx"
        }
      }
    ]
  }
```

この中で、「"name": "Invocation"」の中にあるSubsegment「"name": "Test"」内のstart_timeとend_timeを計算すると処理時間になると思います。CSVに出力してみましたがランダムな時間スリープした処理になっていそうです。

```csvs
1-64732786-3f8223292aeed41d74033bb0,testLambda1,5.006308317184448
1-64732778-6b9267382602169a43edead8,testLambda1,5.006357908248901
1-64732773-13a477724f2561d60710ac84,testLambda1,5.006493806838989
1-64732784-4264c9405cb227dc154bd80a,testLambda1,1.0023705959320068
1-64732782-7ac9aabc0b9206ac439c8cdf,testLambda1,2.003343105316162
1-6473278b-480be45142af544536a667e4,testLambda1,1.0014824867248535
1-6473277d-03afa6f52111fb3639849285,testLambda1,2.003434658050537
1-6473277f-3666e7494e5fe6a65ca1075f,testLambda1,1.002347707748413
1-6473276f-7858828d3739229d402c68ce,testLambda1,3.004973888397217
1-64732780-31501b272d8dcb0e3d53cf8f,testLambda1,2.0033411979675293
```
## 参考

コールドスタート時の時間を測った例がありました。

https://dev.classmethod.jp/articles/lambda-measure/

## おわりに

AWS Lambdaの実行結果の情報などが取れるので一例として簡単に取得してみました。SDKを使って自分で実装したSubsegmentの情報も取得できるので、なにか使えそう（使えたら良いな）と思います。