Deploying Axum to Lambda and ECS, using Lambda Web Adapter

Bing Image Creator, with prompt: “The Rust crab running code inside a container”

AWS Lambda Web Adapter is “A tool to run web applications on AWS Lambda”. It allows you to build a REST API with tooling that you are probably already familiar with, deploying it on Lambda without any changes to your application code. In this blog post, we will use this tool with Axum, deploying the application to both Lambda and ECS.

Let’s start by looking at the Rust code. Our application exposes two GET endpoints. The root endpoint returns a simple ‘hello world’, while /pet requires the owner as a path parameter and retrieves the pet from our database, based on its owner. (If you’ve read my blog posts on Lambda and Rust macros, you’ll notice that I’m being lazy here. I’m reusing some code — and the database — from those blogs.)

For the API itself, we need Axum and Tokio as dependencies, with the tower dependency used for testing. Because our response will be JSON, we also need Serde. AWS config and the DynamoDB SDK are required for talking to our database.

[dependencies]
axum = {  version = "0.7.4" }
tokio = { version = "1.20.4", features = ["full"] }
aws-config = { version = "1.1.1" }
aws-sdk-dynamodb = "1.9.0"
serde = { version = "1.0.195", features = ["derive"] }
serde_json = "1.0.111"

[dev-dependencies]
tower = { version = "0.4", features = ["util"] }

Our code is some forty lines long. In app, we set up our DynamoDB client and pass it as state to our Axum Router. We also configure our two endpoints in that Router. The first endpoint is defined inline, but the second one is a separate function, get_pet, which takes the owner from Path and uses the database from the state to retrieve the pet. In main, we create the app and expose it on port 8080. For simplicity’s sake, there are a lot of unwraps in the example code. Please use proper error handling for production code.

// imports

#[derive(Serialize)]
struct Response {
    pet: String,
}

struct AppState {
    db_client: Client
}

async fn get_pet(Path(payload): Path<String>, State(state): State<Arc<AppState>>) -> (StatusCode, Json<Response>) {
    let response = &state.db_client.get_item()
        .table_name("example-with-names")
        .key("name".to_string(), AttributeValue::S(payload))
        .send()
        .await
        .unwrap();
    let pet = response.item.as_ref().unwrap().get("pet").unwrap().as_s().unwrap().to_string();

    (StatusCode::OK, Json(Response { pet }))
}

async fn app() -> Router {
    let config = aws_config::load_defaults(aws_config::BehaviorVersion {}).await;
    let client = Client::new(&config);
    let shared_state = Arc::new(AppState { db_client: client });

    Router::new()
        .route("/", get(|| async { "Hello, world!" }))
        .route("/pet/:owner", get(get_pet))
        .with_state(shared_state)
}

#[tokio::main]
async fn main() {
    let app = app().await;
    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

Because we do not initialize the Router within main, this code is fairly easy to unit test.

#[cfg(test)]
mod tests {
    use axum::body::Body;
    // more imports...

    #[tokio::test]
    async fn test_get_root() {
        let app = app().await;

        let response = app
            .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
            .await
            .unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }
}

Testing the pet endpoint would require more work, but is very doable. Depending on our preferences we could hide the concrete database behind a trait, allowing us to ‘mock’ the client, or we could run DynamoDB locally.

What is cool, is that we can run the application with a simple cargo run and everything will work, though you do need valid AWS credentials for calling the pet endpoint. We can also use cargo lambda build — release — arm64 (or similar) to build our project as an executable, and put it in a Dockerfile:

FROM rust:1.75.0-buster

EXPOSE 8080

COPY output/bootstrap .

ENTRYPOINT [ "./bootstrap" ]

We start from a rust-buster image (smaller images would work as well, but might require more setup), expose our port, copy the cargo lambda output, and make it our entry point. After building the image, we can run docker run -p 8080:8080 name-of-the-image and see the output of our API:

Root endpoint response from cargo or docker

Ok, so how about deploying our solution to AWS? We can use the CDK for that, putting its code in an infra directory. We create a ‘LambdaStack’ to take care of the details for us. This Stack builds a Rust Lambda based on our zipped bootstrap. We allow Lambda to call DynamoDB and create an API Gateway to put in front of our compute.

import * as cdk from 'aws-cdk-lib';
// other imports

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

        const webAdapterFunction = createRustFunction(this)({
            zip: '../output/web-adapter.zip',
        }) // create the function
        webAdapterFunction.addToRolePolicy(createDynamoDbGetPolicy()) // allow calls to Dynamo
        buildGateway(this)(webAdapterFunction) // put an API Gateway in front of it
    }
}

Creating the function is easy. The only special thing we have to do is add the web adapter as a layer. In this example code, I have set the region of the layer to eu-west-1, which will only work if the corresponding Lambda is also deployed in that region.

// imports and props type

export const createRustFunction = (scope: Construct) => ({
                                                             zip,
                                                             duration = 3,
                                                             environment = {}
                                                         }: RustFunctionProps): Function => {
    const baseName = zip.split('/')?.pop()?.replace('.zip', '')

    if (!baseName) {
        throw new Error(`Invalid zip: ${zip}`)
    }

    const name = baseName.split('_').map(part => part.charAt(0).toUpperCase() + part.slice(1)).join('')
    // create the layer with the Web Adapter code
    const layer = LayerVersion.fromLayerVersionArn(scope, 'AdapterLayer', 'arn:aws:lambda:eu-west-1:753240598075:layer:LambdaAdapterLayerArm64:18') // note: set to eu-west-1...

    return new Function(scope, name, {
        handler: 'bootstrap',
        code: Code.fromAsset(zip),
        runtime: Runtime.PROVIDED_AL2,
        timeout: Duration.seconds(duration),
        memorySize: 1024,
        environment,
        logRetention: RetentionDays.ONE_MONTH,
        architecture: Architecture.ARM_64, // ARM is simpler for me because I'm building locally from a M1 MacBook
        layers: [
            layer
        ]
    })
}

Our API Gateway is also simple. Our Lambda will respond to every incoming request, so a default integration is plenty.

export const buildGateway = (scope: Construct) => (lambda: Function) => {
    const defaultIntegration = new HttpLambdaIntegration('WebAdapterIntegration', lambda);
    new HttpApi(scope, 'WebAdapterApi',
        {
            createDefaultStage: true,
            defaultIntegration,
        });
}

Deploying this will result in an API Gateway with two endpoints. Nice. But because the adapter has made our app independent from implementation details like ‘this runs on Lambda’, we can deploy to ECS as well. All we need is an ‘ECSStack’ to create a role for our Fargate Task and a cluster.

export class ECSStack extends cdk.Stack {
    constructor(scope: Construct, id: string, props?: cdk.StackProps) {
        super(scope, id, props);
        const ecsTaskRole = createEcsTaskRole(this)
        createRustCluster(this)(ecsTaskRole)
    }
}

The heavy lifting is done by createRustCluster, which uses aws-ecs-patterns to create a full-blown ECS service with a load balancer, running on Fargate, in just a few lines of code. What I was not aware of, is that your instances will be placed in a private subnet and will communicate to the outside world via a NAT Gateway, a famously pricy AWS service, so cost-conscious-reader beware, deploying this version of the project might cost you several dollars.

We are using the Dockerfile we showed earlier in this post, which is located at the root of the project, for creating our Docker Image. Hence why the asset is located one directory up.

export const createRustCluster = (scope: Construct) => (taskRole: IRole) => {
    const cluster = new Cluster(scope, 'WebAdapterCluster');
    new ApplicationLoadBalancedFargateService(scope, 'WebAdapterService', {
        cluster,
        memoryLimitMiB: 1024,
        cpu: 512,
        runtimePlatform: {
            operatingSystemFamily: OperatingSystemFamily.LINUX,
            cpuArchitecture: CpuArchitecture.ARM64,
        },
        taskImageOptions: {
            taskRole,
            containerPort: 8080,
            image: ContainerImage.fromAsset(".."), // get image from root of project
        },
    });
}

When deployed, this setup creates a cluster, service, and DNS endpoint that we can call to get the expected API responses. To make it easy to switch between infrastructures, bin/infra.ts deploys either the LambdaStack or ECSStack based on an environment variable. So if we want to deploy to Lambda, we run STACK_TYPE=lambda cdk deploy and the CDK will take care of everything for us.

#!/usr/bin/env node
import 'source-map-support/register';
import * as cdk from 'aws-cdk-lib';
// ...

const app = new cdk.App();

if(process.env.STACK_TYPE?.toLowerCase() === 'lambda') {
    new LambdaStack(app, 'LambdaWebAdapterStack', {});
} else if(process.env.STACK_TYPE?.toLowerCase() === 'ecs') {
    new ECSStack(app, 'EcsWebAdapterStack', {});
} else {
    throw Error('Expected STACK_TYPE env var to be "lambda" or "ecs"')
}

I also wanted to see how much faster my tasks would start if I used SOCI for the +- 130 MB sized images, so I deployed the ‘AWS Partner Solution’ to my account. In my limited testing, starting tasks went from 3 to a maximum of 40 seconds to a maximum of 20 seconds. Which is a nice improvement. But I was actually hoping that SOCI would make it possible to run an application without noticeable downtime because its single Task would restart really fast. Alas, not only did startup take longer than I thought, but I also failed to take into account that deprovisioning and provisioning can take several minutes. Guess you still have to deploy multiple Tasks for availability. Or, you know, just go for Lambda.

Lambda Web Adapter not only makes it easier to develop and test code: it abstracts away the infrastructure from our application. And being able to seamlessly switch between Lambda and ECS could prove useful when requirements change (we need more time or memory than Lambda can provide), or when price is a major concern. An application that runs 24/7 on prod can be costly on Lambda, but cheap on ECS. On the other hand, if you have one or more infrequently used test environments, Lambda is the better choice. So deploying test environments to Lambda, and production to ECS, might be an idea. It’s a tradeoff because there is now a meaningful difference between the infrastructure your code runs on in test versus production. But in some scenarios, the cost savings might be worth it.

Thanks for reading.