capsule.adrianhesketh.com

Using AWS X-Ray with a TypeScript Lambda

A while back, I set up a repository to demonstrate using CloudWatch Embedded Metric Log format [0]. Before this was introduced, I used to write log entries in JSON format, and then create a metric extraction filter to turn the log files into metrics. This was a bit time consuming (just another thing to set up), so it's great that in AWS Lambda, if you use the CloudWatch Embedded Metric Log format, metrics are automatically extracted with no work required. There's a node.js library for this called aws-embedded-metrics [1].

[0]

[1]

Once I had an example of using the metrics, I added some more features along the way to demonstrate how it's possible to use AWS X-Ray to trace requests all the way from AWS API Gateway, through Lambda, and on to 3rd party HTTP APIs and other AWS services. However, the code was pretty ugly, due to some complexities in how Lambda, X-Ray and Node.js work together, failure to "close" X-Ray segments when a request completes can mean that the data is not displayed correctly. It's a requirement to "flush" the logs, and "close" the X-Ray segment.

This makes the actual logic code dwarfed by a barrage of X-Ray and metric related stuff.

Before

export const hello: APIGatewayProxyHandler = async (event, context) => {
  const metrics = createMetricsLogger();

  // We need to create a subsegment, because in AWS Lambda, the top-level X-Ray segment is readonly.
  return await AWSXRay.captureAsyncFunc(
    "handler",
    async (segment) => {
      // Annotate the segment with metadata to allow it to be searched.
      segment.addAnnotation("userId", "user123");
      try {
        return await helloFunction(event, context, metrics, segment);
      } catch (err) {
        segment.close(err);
        throw err;
      } finally {
        // Metrics and segments MUST be closed.
        metrics.flush();
        if (!segment.isClosed()) {
          segment.close();
        }
      }
    },
    AWSXRay.getSegment()
  );
};

There was a lot of complexity in the setup of capturing AWS SDK requests, and outbound HTTP requests using X-Ray. It turns out that you can't follow the examples in the documentation when you're in Lambda, you have to do it differently. [2]

[2]

The resolution states that you have to have your X-Ray stuff "inside" the Lambda handler. This makes it sound like we're stuck with having lots of code within the handlers of our functions, and we can't use helpers.

X-Ray in Lambda also has the issues in that you can't just attach annotations to the default X-Ray segment, you have to create a subsegment against the segment, and then you can modify it. All in all, there's a few traps to suprise you and waste your time.

During a rewrite of a service, I decided to visit this, to see if there's been any improvements since the last time I looked (nope!), but I came up with a better solution this time around.

After

In the new design, I wanted to use the `metricsScope` function that's provided by the embedded metrics node.js library, and to create a version of that idea that works for the X-Ray segment too.

Adopting this design allowed me to get the handler down to size. The Lambda handler is now uses two helper functions `xrayScope` and `metricScope`. These functions are nested so that the Lambda handler has access to both the X-Ray `segment` and the `metrics` object.

Lambda handler

export const handler = xrayScope((segment) =>
  metricScope((metrics) =>
    async (_event: APIGatewayProxyEvent, _context: Context): Promise<APIGatewayProxyResult> => {
      log.info("hello/get starting")
      segment.addAnnotation("source", "apiHandler");
      metrics.putMetric("count", 1, Unit.Count);
      try {
        await axios.get("http://httpstat.us/500");
      } catch (err) {
        log.warn("httpstat.us gave error, as expected", { status: 500 });
        // {"status":500,"level":"warn","message":"httpstat.us gave error, as expected","timestamp":"2020-09-01T18:33:30.296Z"}
      }
      await axios.get("https://jsonplaceholder.typicode.com/todos/1");
      metrics.putMetric("exampleApiCallsMade", 2);
      await putEvent();
      return {
        statusCode: 200,
        body: "World!"
      }
    }
  )
);

xray.ts

The `xrayScope` function uses some rather unusual looking features of TypeScript to keep hold of the strong typing that would otherwise be lost.

First, I had to redefine the `Handler` type - the function signature of every Lambda function - to remove the ancient callback parameter to stop it getting in the way. I've been using async/await for years, and I haven't seen any callback based Lambda functions in years either, so it doesn't make sense for it to be kept around.

import * as AWSXRay from "aws-xray-sdk";
import {  Context } from "aws-lambda";
import { Subsegment } from "aws-xray-sdk";

export type Handler<TEvent = any, TResult = any> = (
    event: TEvent,
    context: Context,
) => Promise<void | TResult>;

With that defined, I could create the `xrayScope` function. The `xrayScope` function has one argument (`fn`), a function of generic type `F`, and it returns a Lambda `Handler`.

`F` is defined as being a function that receives a `Subsegment` and returns a Lambda `Handler` function.

It's hard to see what's going on with all of that syntax (it took me a while to get it right...), but the `xrayScope` function wraps the `fn` function argument with the required AWS Lambda X-Ray setup and teardown.

export const xrayScope = <TEvent, TResult, F extends (segment: Subsegment) => Handler<TEvent, TResult>>(fn: F): Handler<TEvent, TResult> => async (e, c) => {
  AWSXRay.captureAWS(require("aws-sdk"));
  AWSXRay.captureHTTPsGlobal(require("http"), true);
  AWSXRay.captureHTTPsGlobal(require("https"), true);
  AWSXRay.capturePromise();
  const segment = AWSXRay.getSegment().addNewSubsegment("handler");
  try {
    return await fn(segment)(e, c)
  } finally {
    if (!segment.isClosed()) {
      segment.close();
    }
  }
};

With that in place, it's much simpler to apply X-Ray to each Lambda function.

export const handler = xrayScope((segment) =>
    async (_event: APIGatewayProxyEvent, _context: Context): Promise<APIGatewayProxyResult> => {
      // Any code here can use the segment (e.g. appending data to it), and it will use X-Ray for everything.
    }
)

CDK setup

I included a CDK setup that creates the Lambda function as an API Gateway, so it's easy to try out. I use the `@aws-cdk/aws-lambda-nodejs` package to package up the TypeScript and all its dependencies. Note the use of the `tracing` and `deployOptions` parameters to enable X-Ray.

import * as cdk from "@aws-cdk/core";
import * as apigw from "@aws-cdk/aws-apigateway";
import * as lambda from "@aws-cdk/aws-lambda";
import * as path from "path";
import * as lambdaNode from "@aws-cdk/aws-lambda-nodejs";
import { EventBus } from "@aws-cdk/aws-events";

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

    const helloGetHandler = new lambdaNode.NodejsFunction(
      this,
      "helloGetHandler",
      {
        runtime: lambda.Runtime.NODEJS_14_X,
        entry: path.join(__dirname, "../../api/handlers/http/hello/get.ts"),
        handler: "handler",
        memorySize: 1024,
        description: `Build time: ${new Date().toISOString()}`,
        tracing: lambda.Tracing.ACTIVE,
        timeout: cdk.Duration.seconds(15),
      }
    );
    EventBus.grantPutEvents(helloGetHandler);

    const api = new apigw.LambdaRestApi(this, "lambdaXRayExample", {
      handler: helloGetHandler,
      proxy: false,
      deployOptions: {
        dataTraceEnabled: true,
        tracingEnabled: true,
      },
    });

    api.root
      .addResource("hello")
      .addMethod("GET", new apigw.LambdaIntegration(helloGetHandler));
  }
}

Results

It's possible to validate that the results are there in the X-Ray console. You can see the call out to send events to EventBridge, and two calls to 3rd party APIs.

x-ray-console.png

You can download the full example code at [3]

[3]

Trying out npm and yarn workspaces

Launch a Gemini capsule on AWS with the CDK