Zero-Friction AWS Lambda Instrumentation with external extensions

If you’ve been in the software business for some time, you’ve probably noticed that creating software isn’t only about adding features. There are usually many different tasks involved. You have to test your system, fix bugs, and ensure it keeps working over its lifetime. 

Take logging, for example. It doesn’t bring direct value to your users, but it helps keep things running by giving you insights into the system. If things go wrong, you can consult the logs and check out what your users did that led to the issue. Logging is notorious for touching all parts of your system while not directly related to its functionality. Another example might be cost optimization. You want to know what each of the features actually costs and if a change in implementation or configuration could save you money. That requires you to measure execution times all over your app. 

This is where AWS Lambda External Extensions come into play, a feature of Lambda that allows you to add code to your Lambda handler code that runs parallel to it in its own process. This way, you can implement features like logging without writing it inside your Lambda handler code.

The Downsides of Putting Code Inside Your Handler Functions

When you write code into the handler function that isn’t directly related to the function’s scope, you can run into a few problems. 

If you build a logging system and sprinkle logging code all over your handler code, you’ll have to update, test, and redeploy your handler code when the logging code changes. They’re now tightly coupled, and you’ve multiplied your maintenance work.

Code isn’t perfect; it can break, and the more code you have in your Lambda function, the higher the risk of a crash. If handler-unrelated code runs in the same process, it can take down your entire Lambda function.

This also brings us to the next point. While the code might be OK and doesn’t break, it still adds additional complexity inside your handler code. When searching for issues, you’ll have to comb through everything you wrote in the handler, not just the crucial parts. This makes maintenance unnecessarily complicated.

Finally, the external extension code will be executed in a different process. While this might not give you a huge performance boost with single CPU configurations, you can use the post-execution-time for additional processing after the function has returned its value.

Separation of Concerns to Improve Maintainability

Software engineering is all about the separation of concerns. It’s touted all around the industry: Don’t put parts of your code together that aren’t related. 

The textbook example of separation of concerns is usually achieved with modularization. You write a Lambda function with multiple modules so the code responsible for logging ends up in its own file. The same goes for your feature code, which will get its own files, too. You call from one file to the other to execute that code, but the main implementations are separated and can now be maintained mostly independently from each other.

There are additional ways to separate your system. For example, executing it in multiple processes. This is a runtime separation independent of the location of the code. You can start various processes from code inside one module. This separation also improves maintainability because if one process crashes, you know where to start looking for a bug. 

But the main gain here is performance. Running code in multiple processes allows multiple CPUs to handle the workload and ensure that unrelated parts of code don’t impact each other. Remember that this is only the case for Lambda functions with more than 1.769 GB of memory configured since those will also get multiple CPUs assigned. 

Separation with External Extensions

AWS helps you with separation of concerns in the form of AWS Lambda Extensions. You can mark an extension as external, prompting the Lambda service to run it in a different process from your handler code. This way, you can benefit from both separations via modularization and runtime separation, where your code is executed by multiple CPUs and has fewer chances of affecting each other’s performance or stability.

Extensions can hook into all three steps of the Lambda lifecycle:

1. During initialization, the Lambda service sets up the runtime and function. It creates network connections and loads the code from storage.
2. Upon invocation, the Lambda service executes the function code.
3. The shutdown step runs when the function finishes executing. Here, the Lambda service turns down the runtime, and network connections are closed.

External extensions can keep running even if the handler has already finished executing and delivered its result to a client. In the logging example, this would allow you to clean everything up and send it to the logging service without making the user wait for it.

How to Create an Extension?

An external Lambda Extension is a file in /opt/extensionsthat starts its own process. So, to create your external extension, you can create a layer with a script that holds your logic. If the script is in the correct directory, the AWS Lambda Service will execute it in the initialization phase of the lambda..

Let’s create an extension with the help of the AWS CDK. 

Prerequisites

To complete this how-to, you need a current installation of Node.js and the AWS CDK. 

Setting Up a New Project

To set up a new CDK project, run the following commands:

$ mkdir lambda-extension && cd lambda-extension
$ cdk init app --language=javascript

Creating the Lambda Extension

The extension will consist of two files:

• A script that starts the extension
• The actual code of the extension

Let’s start with the actual extension code. Create a file at lib/extension/example-extension/index.js with the following content:

#!/usr/bin/env node
const extensionsApiUrl = `http://${process.env.AWS_LAMBDA_RUNTIME_API}/2020-01-01/extension`;

async function registerToExtensionsApi() {
  const res = await fetch(`${extensionsApiUrl}/register`, {
    method: "post",
    headers: {
      "Content-Type": "application/json",
      "Lambda-Extension-Name": "logger-extension",
    },
    body: JSON.stringify({ events: ["INVOKE", "SHUTDOWN"] }),
  });

  if (res.ok) return res.headers.get("lambda-extension-identifier");

  const error = await res.text();
  console.log("Error: " + error);
}

async function getNextEvent(extensionId) {
  const res = await fetch(`${extensionsApiUrl}/event/next`, {
    method: "get",
    headers: {
      "Content-Type": "application/json",
      "Lambda-Extension-Identifier": extensionId,
    },
  });

  if (res.ok) return await res.json();

  const error = await res.text();
  console.log("Error: " + error);
}

async function main() {
  process.on("SIGINT", () => console.log("SHUTDOWN"));
  process.on("SIGTERM", () => console.log("SHUTDOWN"));

  const extensionId = await registerToExtensionsApi();

  while (true) {
    const { eventType } = await getNextEvent(extensionId);
    switch (eventType) {
      case "SHUTDOWN":
        console.log("SHUTDOWN");
        break;
      case "INVOKE":
        console.log("INVOKE");
        break;
      default:
        throw new Error("Unknown: " + eventType);
    }
  }
}

main();

Let’s go through this step by step. 

Since the extension will run separately from your function handler code, it needs a way to get notified about the state of that other process. That’s where the Extension API comes into play.

The first function is to register the extension to the Extensions API. This way, it can get Lambda events later, in this case, the INVOKE and SHUTDOWN events.

The second function is to fetch the next event with the event ID received from the registerToExtensionsApi function.

The main function will listen to process termination events of its own process and also start an infinite loop that fetches events from the Extensions API. If you built a logging extension, the console.log statements would be an excellent place to put your logging code.

Now, we need the script that starts the extension. It needs to have the same name as the extension, but we have to put it in the extensions directory, not in the same directory as the extension. The Lambda service extracts our files later into /opt. The script has to end up in /opt/extensions because the service looks for all extension startup scripts there.

So, let’s create the startup script at lib/extension/example-extension/extensions/example-extension with this code:

#!/bin/bash
set -euo pipefail
OWN_FILENAME="$(basename $0)"
LAMBDA_EXTENSION_NAME="$OWN_FILENAME"
exec "/opt/${LAMBDA_EXTENSION_NAME}/index.js"

In the last line, the script expects our extension inside the /opt/example-extension.

This script is the crucial part of this how-to!

If you just deployed it without the script, you had to require it from your Lambda function code via require("/opt/example-extension"), and it would be internal, i.e., run in the same process as your Lambda function. But since the script starts another process, the extension will run in its own process, separate from your Lambda function code.

Creating a Lambda Function

To invoke something, let’s also create a Lambda function at lib/function/index.js with the following code: 

exports.handler = async (event, context) => {
  console.log("HANDLER");
  return { statusCode: 200, body: "HELLO" };
};

This function only responds via HTTP with a HELLO string. 

Creating the CDK Stack

We’ll fill our empty CDK stack with life to connect everything together. To do so, replace the content of lib/lambda-extension-stack.js with this code:

const path = require("path");
const cdk = require("aws-cdk-lib");
const lambda = require("aws-cdk-lib/aws-lambda");

class LambdaExtensionStack extends cdk.Stack {
  constructor(scope, id, props) {
    super(scope, id, props);

    const myExtension = new lambda.LayerVersion(this, "MyLayer", {
      code: lambda.Code.fromAsset(path.join(__dirname, "extension")),
      compatibleRuntimes: [lambda.Runtime.NODEJS_18_X],
      removalPolicy: cdk.RemovalPolicy.DESTROY,
    });

    const myFunction = new lambda.Function(this, "MyFunction", {
      handler: "index.handler",
      code: lambda.Code.fromAsset(path.join(__dirname, "function")),
      runtime: lambda.Runtime.NODEJS_18_X,
      layers: [myExtension],
    });

    const myFunctionUrl = myFunction.addFunctionUrl({
      authType: lambda.FunctionUrlAuthType.NONE,
    });

    new cdk.CfnOutput(this, "functionUrl", {
      value: myFunctionUrl.url,
    });
  }
}

module.exports = { LambdaExtensionStack };

First, we create the extension. Again, extensions are just layers with a special startup script. We need to point to the extensions directory and tell the CDK what Node.js versions it’s compatible with.

Second, we create the Lambda function. The exciting part is that myExtension is passed to the function’s layers configuration. This tells the Lambda service to start the extension before the Lambda is invoked.

Finally, we create a function URL to make the Lambda function accessible from the internet and export that URL so we see it right after the deployment.

Deploying and Testing the Stack

To deploy the whole stack, run the following command:

$ cdk deploy

Summary

Extensions are a powerful addition to AWS Lambda. They allow you to separate your concerns at implementation and at runtime. This essentially means writing them once and running them alongside all your Lambda functions, which improves maintainability. 

Since they’re running in their own process, they are more resilient, and errors have a lower chance of taking out your entire function. Extensions also allow you to execute code after your function returns, so you don’t impact user experience. And if you configured your Lambda functions with more than 1.8 GB of memory, you will also benefit from parallelism of the CPUs.