x-ray custom trace collector sample

cdk.ts

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import { NodejsFunction } from 'aws-cdk-lib/aws-lambda-nodejs';
import * as iam from 'aws-cdk-lib/aws-iam';
import { Effect } from 'aws-cdk-lib/aws-iam';
import * as lambda from 'aws-cdk-lib/aws-lambda';
import * as apigw from 'aws-cdk-lib/aws-apigatewayv2';
import { CorsHttpMethod } from 'aws-cdk-lib/aws-apigatewayv2';
import { RetentionDays } from 'aws-cdk-lib/aws-logs';
import * as integrations from 'aws-cdk-lib/aws-apigatewayv2-integrations';

export class CustomXrayStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const role = new iam.Role(this, 'GptSlackGatewayLambdaRole', {
      roleName: `${this.stackName}-trace-role`,
      assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
      managedPolicies: [iam.ManagedPolicy.fromAwsManagedPolicyName('service-role/AWSLambdaBasicExecutionRole')],
      inlinePolicies: {
        XRayPutCustomSegment: new iam.PolicyDocument({
          statements: [new iam.PolicyStatement({
            effect: Effect.ALLOW,
            resources: ["*"],
            actions: ['xray:PutTraceSegments']
          })]
        }),
      }
    });

    const lambdaFn = new NodejsFunction(this, 'XRayCustomTrace', {
      role: role,
      functionName: `${this.stackName}-custom-xray-trace`,
      description: 'custom xray entry point',
      entry: 'lambda/handler.ts',
      handler: 'handler',
      runtime: lambda.Runtime.NODEJS_20_X,
      logRetention: RetentionDays.ONE_WEEK,
    })

    const httpApi = new apigw.HttpApi(this, 'HttpApi', {
      corsPreflight: {
        allowHeaders: ["*"],
        allowMethods: [CorsHttpMethod.ANY],
        allowOrigins: ["*"]
      }
    });
    httpApi.addRoutes({
      path: '/',
      methods: [apigw.HttpMethod.POST],
      integration: new integrations.HttpLambdaIntegration('Integration', lambdaFn)
    });

    new cdk.CfnOutput(this, 'HttpApiUrl', {
      value: httpApi.url!,
      description: 'The url of the HttpApi',
      exportName: `${this.stackName}-http-api-url`,
    });
  }
}

handler.ts

import { APIGatewayProxyHandlerV2 } from 'aws-lambda';
import { PutTraceSegmentsCommand, XRayClient } from '@aws-sdk/client-xray';
import * as crypto from 'crypto'

type Trace = {
  segments: {
    [key: string]: {
      startTime: number;
      endTime: number;
      metrics?: {
        [key: string]: number
      },
    }
  }
}

const client = new XRayClient();
function generateId(len: number = 8) {
  return crypto.randomBytes(len).toString('hex');
}

export const handler: APIGatewayProxyHandlerV2 = async (event) => {
  const trace: Trace = JSON.parse(event.body ?? '{}');
  const origin = event.headers['origin'] ?? 'unknown';
  const traceId = `1-${Math.floor(Date.now() / 1000).toString(16)}-${generateId(12)}`;
  const min = Math.min(...Object.values(trace.segments).map(v => v.startTime));
  const max = Math.max(...Object.values(trace.segments).map(v => v.endTime));
  const input = {
    name: origin,
    id: generateId(),
    trace_id: traceId,
    start_time: min / 1000,
    end_time: max / 1000,
    subsegments: Object.entries(trace.segments).map(([key, value]) => ({
      id: generateId(),
      name: key,
      start_time: value.startTime / 1000,
      end_time: value.endTime / 1000,
      annotations: value.metrics,
    })),
    annotations: {
      userAgent: event.headers['user-agent'] ?? 'unknown'
    }
  };
  console.log(input)

  const command: PutTraceSegmentsCommand = new PutTraceSegmentsCommand({
    TraceSegmentDocuments: [JSON.stringify(input)]
  })

  const resp = await client.send(command)
  console.log(resp)

  return {
    statusCode: 204,
  }
};

{
  "$metadata": {
    "httpStatusCode": 200,
    "requestId": "96ac8d92-9902-4ee9-a932-1f6d33681b51",
    "attempts": 1,
    "totalRetryDelay": 0
  },
  "Traces": [
    {
      "Duration": 8,
      "Id": "1-66592a23-c94baf8ea8de377eac545fc9",
      "LimitExceeded": false,
      "Segments": [
        {
          "Document": "{\"id\":\"c5cfe1da9b692c3d\",\"name\":\"http://localhost:3000\",\"start_time\":1.717119513474E9,\"trace_id\":\"1-66592a23-c94baf8ea8de377eac545fc9\",\"end_time\":1.717119521474E9,\"annotations\":{\"userAgent\":\"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36\"},\"subsegments\":[{\"id\":\"d5ce13de6b80f459\",\"name\":\"/complete\",\"start_time\":1.717119518474E9,\"end_time\":1.717119521474E9},{\"id\":\"8604cc723d83fbcc\",\"name\":\"/option\",\"start_time\":1.717119515474E9,\"end_time\":1.717119516474E9},{\"id\":\"0580cff5bdfa56c3\",\"name\":\"/\",\"start_time\":1.717119513474E9,\"end_time\":1.717119514474E9,\"annotations\":{\"FID\":2.5,\"FCP\":36.0,\"LCP\":36.0,\"TTFB\":18.2}},{\"id\":\"9e189f3297b153c5\",\"name\":\"/input\",\"start_time\":1.717119517474E9,\"end_time\":1.717119517974E9}]}",
          "Id": "c5cfe1da9b692c3d"
        }
      ]
    }
  ],
  "UnprocessedTraceIds": []
}