Rust macros taking care of some Lambda boilerplate
When you have a shiny golden hammer, you go looking for nice nails to hit.
While I was writing my Rust macros book, I was thinking of how to use those macros to really nail setting up an AWS Lambda. Why? Well, look at this example from the AWS documentation on how you can write a simple Lambda function with tracing in Rust.
use lambda_runtime::{service_fn, Error, LambdaEvent};
use serde_json::{json, Value};
#[tracing::instrument(skip(event), fields(req_id = %event.context.request_id))]
async fn handler(event: LambdaEvent<Value>) -> Result<Value, Error> {
// do stuff
}
#[tokio::main]
async fn main() -> Result<(), Error> {
tracing_subscriber::fmt().json()
.with_max_level(tracing::Level::INFO)
.with_current_span(false)
.with_ansi(false)
.without_time()
.with_target(false)
.init();
lambda_runtime::run(service_fn(handler)).await
}There is definitely boilerplate here. Basically, we create a main function, set up tracing, and call the function that will serve as our ‘handler’. The reason for separating our code into at least two functions is the way Lambda works under the hood. In all its implementations, your code will have a handler function that runs every time the Lambda is invoked and a bit outside that handler that will run once every time AWS has to create a ‘container’ for your code to run in.
This is useful, because you may have to do some expensive setup that could serve multiple invocations. If you have to connect to a database, that could take a relatively long time and there is little reason to assume that connections would suddenly become unusable when your code is invoked a second or third time.
In Javascript or Python, simply putting initialization logic outside of your handler function will make it run once per ‘container creation’. In Java, you typically put that code inside the handler class, but outside the handler method. In Rust’s case, the obligatory main was a logical choice for expensive ‘one-time’ operations. In that scenario, the ability to customize both handler and main function is definitely valuable. For example, the following code fragment shows how to create an S3 client, passing it to our handler.
async fn handler(client: &Client, event: LambdaEvent<Request>) -> Result<Response, Error> {
// do stuff like getting keys from the bucket
Ok(Response { keys })
}
#[tokio::main]
async fn main() -> Result<(), Error> {
let shared_config = aws_config::from_env().load().await;
// create an S3 client
let client = Client::new(&shared_config);
// and pass a reference to the handler
let shared_client = &client;
lambda_runtime::run(service_fn(move |event: LambdaEvent<Request>| async move {
handler(&shared_client, event).await
}))
.await
}But we are getting ahead of ourselves. In this blog post, we focus on the simplest case: using a procedural macro to get rid of all boilerplate from the first code example. In that particular setup everything inside main looks like boilerplate. Only the code within the handler is interesting.
We will need to pick a suitable procedural macro. A derive macro would be great in theory since it only adds things, e.g. a main, without manipulating existing code. But unfortunately, we can’t use derive on a function. (And while we could create a struct with a handler method to circumvent this restriction, that raises other problems.) So we’ll go with an attribute macro.
Before continuing: would you ever want to use a macro like this? For real applications: default to no. In most cases, either the boilerplate won’t be that much of a bother, or you will want to do some custom things like initializations in your main function (see the next blog post, though!). On the other hand, if you have a lot of very simple Lambdas, if you have a very low tolerance for boilerplate, or if you like to experiment a bit with Rust macros and serverless, this might be just the thing for you.
To get started, we need to create a new Rust library, which we can do with cargo init --lib. The generated Cargo.toml has to declare that this is a procedural macro. We will also want to add two dependencies, syn and quote to make our lives easier.
[dependencies]
quote = "1.0.35"
syn = "2.0.48"
[lib]
proc-macro = trueNext, lib.rs, which contains the implementation, has only 29 lines for this naive implementation. You can see that this specific Lambda example had a high ratio of boilerplate to business code, as the macro needs very little input to spit out quite a few lines of code.
use quote::quote;
use syn::{ItemFn};
use proc_macro::TokenStream;
#[proc_macro_attribute]
pub fn lambda_setup(_: TokenStream, input: TokenStream) -> TokenStream {
let item: ItemFn = syn::parse(input.clone()).unwrap();
let function_name = &item.sig.ident;
quote!(
#[tokio::main]
async fn main() -> Result<(), lambda_runtime::Error> {
tracing_subscriber::fmt()
.with_max_level(tracing::Level::INFO)
.with_current_span(false)
.with_ansi(false)
.without_time()
.with_target(false)
.init();
lambda_runtime::run(lambda_runtime::service_fn(#function_name))
.await?;
Ok(())
}
#item
).into()
}The annotation tells Rust that this is an attribute macro. That means it takes two parameters, both of type TokenStream. We only need the second argument, which contains the function we annotated as a stream of tokens. Using the parse function from syn, we change this stream into an ItemFn struct, perfect for representing functions. We could also have gone for custom parsing, because as you can see we only need one thing from that entire input: the name of the function, which can be found in its signature (sig). We require this function_name because our main needs to know which function is the handler, as we saw in the example.
#[tokio::main]
async fn main() -> Result<(), Error> {
// tracing setup...
lambda_runtime::run(service_fn(handler)).await
}Back to the macro implementation. With quote we return the generated main, invoke the function by its name, and re-add the original input. Since this is an attribute macro, forgetting the last step erases the original function, causing compilation to fail. Plus, erasing the function is probably not what the user wants…
That’s it. Now we need to add the macro to our application. In my case, the library is in a nested directory of the example, so my application’s Cargo.toml needs a relative path:
[dependencies]
# path to the macro
rust-lambda-macro = { path = "./rust-lambda-macro" }
# other dependencies
aws-config = { version = "1.1.1" }
lambda_runtime = "0.9.0"
serde = "1.0.195"
serde_json = "1.0.111"
tokio = { version = "1", features = ["macros"] }
tracing = { version = "0.1", features = ["log"] }
tracing-subscriber = { version = "0.3", default-features = false, features = ["fmt"] }Next, add the attribute to the handler function.
use serde::{Deserialize, Serialize};
use lambda_runtime::{Error, LambdaEvent};
use rust_lambda_macro::lambda_setup;
// not using any of the input
#[derive(Deserialize)]
struct Request {}
// return something simple in our response
#[derive(Serialize)]
struct Response {
pet: String,
}
// put the macro on our handler
#[lambda_setup]
async fn func(_: LambdaEvent<Request>) -> Result<Response, Error> {
Ok(Response {
pet: "cat".to_string()
})
}
// no main, no cry!You can now build and deploy the Lambda to your AWS account, e.g. by using cargo lambda: cargo lambda build — release && cargo lambda deploy — iam-role arn:aws:iam::{account}:role/{role}.
When invoked, your Lambda should complete successfully. But it will only keep working as long as your code stays simple and you don’t need additional dependencies, like, for example, a client for calling DynamoDB.
#[lambda_setup]
async fn func(client: aws_sdk_dynamodb::Client, _: LambdaEvent<Request>) -> Result<Response, Error> {
Ok(Response {
pet: "cat".to_string()
})
}Because you will get the following error:
error[E0593]: function is expected to take 1 argument, but it takes 2 arguments
--> src/main.rs:13:1
|
13 | #[lambda_setup]
| ^^^^^^^^^^^^^^^ expected function that takes 1 argument
14 | async fn func(client: aws_sdk_dynamodb::Client, _: LambdaEvent<Request>) -> Result<Response, Error> {
| --------------------------------------------------------------------------------------------------- takes 2 argumentsFixing that is for next time. :)
