From 476bef37c1e8e1a9471aeac88476446356f88ead Mon Sep 17 00:00:00 2001 From: Harold Sun Date: Thu, 18 Jun 2026 16:15:43 +0000 Subject: [PATCH] feat: add SnapStart support via SnapStartResource trait Implement the SnapStart restore lifecycle so Rust Lambda functions can participate in SnapStart optimizations when deployed as container images with custom base images. Custom logic around the snapshot/restore boundary is expressed by implementing the CRaC-style `SnapStartResource` trait and registering instances on the runtime. The trait methods return `BoxFuture` (no `async-trait` dependency) so the trait stays object-safe as `dyn SnapStartResource`. The runtime drives the lifecycle internally: before_snapshot hooks (LIFO) -> GET /restore/next (blocks until restore) -> rebuild RAPID connection pool -> after_restore hooks (FIFO) Resources use stack ordering: register foundations first; before_snapshot runs in reverse registration order and after_restore in registration order. Error handling: a before-snapshot/`/restore/next` failure is reported to Lambda via POST /init/error; an after_restore failure via POST /restore/error. Both are propagated outward so the runtime exits via the normal Result path (no process::exit, so graceful-shutdown handlers still run). Failures to send the report are logged. When SnapStart is not enabled the hooks are never invoked. Connection-pool reset after restore is non-breaking: rather than mutating the existing `Client`, a new `PooledClient` (with lock-free `OnceLock`-based `reset_pool()`) and a `RuntimeApiClient` trait are added to the api-client crate. `lambda_runtime`'s service layer takes a defaulted generic `C = Client` bounded by the trait and uses `PooledClient` internally; the old `Client` is frozen at its 1.0.3 shape and deprecated. cargo-semver-checks confirms all three crates stay non-breaking (api-client is a 1.1.0 minor bump for the deprecation + additive API). Changes: - lambda-runtime-api-client: add `PooledClient` + `RuntimeApiClient` trait + `ClientBuilder::build_pooled`; deprecate `Client`/`build`; bump to 1.1.0 - lambda-runtime: add `SnapStartResource` trait, Init/Restore error requests, `Runtime::is_snapstart()` and `register_snapstart_resource()`; run the restore lifecycle from run()/run_concurrent(); defaulted-generic service layer using `PooledClient`; re-export `SnapStartResource` and `BoxFuture` - lambda-http: add runtime()/runtime_concurrent() and streaming_runtime() helpers; re-export `SnapStartResource` and `BoxFuture` - lambda-extension: use `PooledClient` - examples: add basic-snapstart and http-snapstart - tests: request serialization, full restore lifecycle, LIFO/FIFO ordering, no-op when not SnapStart, and before/after/restore-next failure paths --- examples/basic-snapstart/Cargo.toml | 10 + examples/basic-snapstart/Dockerfile | 38 ++ examples/basic-snapstart/README.md | 58 +++ examples/basic-snapstart/src/main.rs | 85 ++++ examples/http-snapstart/Cargo.toml | 8 + examples/http-snapstart/Dockerfile | 41 ++ examples/http-snapstart/README.md | 59 +++ examples/http-snapstart/src/main.rs | 77 ++++ lambda-extension/Cargo.toml | 2 +- lambda-extension/src/extension.rs | 6 +- lambda-http/Cargo.toml | 2 +- lambda-http/src/lib.rs | 127 +++++- lambda-http/src/streaming.rs | 58 +++ lambda-runtime-api-client/CHANGELOG.md | 19 + lambda-runtime-api-client/Cargo.toml | 2 +- lambda-runtime-api-client/src/lib.rs | 268 ++++++++++--- lambda-runtime/Cargo.toml | 2 +- lambda-runtime/src/layers/api_client.rs | 27 +- lambda-runtime/src/lib.rs | 9 + lambda-runtime/src/requests.rs | 105 +++++ lambda-runtime/src/runtime.rs | 489 +++++++++++++++++++++++- lambda-runtime/src/snapstart.rs | 104 +++++ 22 files changed, 1505 insertions(+), 91 deletions(-) create mode 100644 examples/basic-snapstart/Cargo.toml create mode 100644 examples/basic-snapstart/Dockerfile create mode 100644 examples/basic-snapstart/README.md create mode 100644 examples/basic-snapstart/src/main.rs create mode 100644 examples/http-snapstart/Cargo.toml create mode 100644 examples/http-snapstart/Dockerfile create mode 100644 examples/http-snapstart/README.md create mode 100644 examples/http-snapstart/src/main.rs create mode 100644 lambda-runtime/src/snapstart.rs diff --git a/examples/basic-snapstart/Cargo.toml b/examples/basic-snapstart/Cargo.toml new file mode 100644 index 000000000..77ec94e58 --- /dev/null +++ b/examples/basic-snapstart/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "basic-snapstart" +version = "0.1.0" +edition = "2021" + +[dependencies] +lambda_runtime = { path = "../../lambda-runtime" } +serde = "1.0.219" +serde_json = "1.0" +tokio = { version = "1", features = ["macros"] } diff --git a/examples/basic-snapstart/Dockerfile b/examples/basic-snapstart/Dockerfile new file mode 100644 index 000000000..50d8068d2 --- /dev/null +++ b/examples/basic-snapstart/Dockerfile @@ -0,0 +1,38 @@ +# Container image for the basic-snapstart example. +# +# SnapStart for functions packaged as OCI images requires a custom base image +# whose runtime implements the restore lifecycle — which `lambda_runtime` now +# does. This image builds the example as the Lambda `bootstrap` on top of the +# AWS-provided base (`provided:al2023`), which supports SnapStart. +# +# Build from the REPOSITORY ROOT (the example uses path dependencies on the +# workspace crates, so the build context must include them): +# +# docker build -f examples/basic-snapstart/Dockerfile -t basic-snapstart . +# +# Then push to ECR and create the function with SnapStart enabled +# (PublishedVersions) and the architecture matching your build host. + +# ---- build stage ---- +FROM public.ecr.aws/lambda/provided:al2023 AS builder + +RUN dnf install -y gcc && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +ENV PATH="/root/.cargo/bin:${PATH}" + +# Copy only the workspace crates this example depends on via path, then the +# example itself: lambda_runtime -> lambda_runtime_api_client. +COPY Cargo.* /build/ +COPY lambda-runtime /build/lambda-runtime +COPY lambda-runtime-api-client /build/lambda-runtime-api-client +COPY examples/basic-snapstart /build/examples/basic-snapstart + +WORKDIR /build/examples/basic-snapstart +RUN cargo build --release + +# ---- runtime stage ---- +FROM public.ecr.aws/lambda/provided:al2023 + +COPY --from=builder /build/examples/basic-snapstart/target/release/basic-snapstart ${LAMBDA_RUNTIME_DIR}/bootstrap + +ENTRYPOINT [ "/var/runtime/bootstrap" ] diff --git a/examples/basic-snapstart/README.md b/examples/basic-snapstart/README.md new file mode 100644 index 000000000..c6a9ee40f --- /dev/null +++ b/examples/basic-snapstart/README.md @@ -0,0 +1,58 @@ +# AWS Lambda SnapStart example (lambda_runtime) + +This example shows how to use AWS Lambda **SnapStart** with the `lambda_runtime` +crate. It implements the [`SnapStartResource`] trait on an `AppState` and +registers it on the runtime, so the runtime runs the resource's `before_snapshot` +hook before the VM snapshot and its `after_restore` hook after each restore. + +SnapStart reduces cold-start latency by snapshotting the initialized execution +environment and restoring from it. Resources created during init (connections, +credentials, unique values) may be invalid after restore — `before_snapshot` / +`after_restore` are where you release and re-establish them. When SnapStart is +not enabled, the hooks are never called and the function behaves normally. + +## Why a container image? + +SnapStart for functions packaged as OCI images requires a custom base image +whose runtime implements the restore lifecycle — which `lambda_runtime` now does. +This example ships a `Dockerfile` that builds the binary as the Lambda +`bootstrap` on top of `public.ecr.aws/lambda/provided:al2023`, which supports +SnapStart. + +## Build & Deploy + +Build the image from the **repository root** (the example depends on the +workspace crates via path, so the build context must include them): + +```sh +docker build -f examples/basic-snapstart/Dockerfile -t basic-snapstart . +``` + +Push it to ECR and create the function from the image: + +```sh +aws ecr create-repository --repository-name basic-snapstart +docker tag basic-snapstart:latest .dkr.ecr..amazonaws.com/basic-snapstart:latest +docker push .dkr.ecr..amazonaws.com/basic-snapstart:latest + +aws lambda create-function \ + --function-name basic-snapstart \ + --package-type Image \ + --code ImageUri=.dkr.ecr..amazonaws.com/basic-snapstart:latest \ + --role \ + --snap-start ApplyOn=PublishedVersions +``` + +SnapStart applies to **published versions**, so publish a version to trigger +snapshot creation, then invoke that version (or an alias pointing at it): + +```sh +aws lambda publish-version --function-name basic-snapstart +aws lambda invoke --function-name basic-snapstart:1 --payload '{"name":"world"}' out.json +``` + +## Architecture + +The compiled `bootstrap` is architecture-specific. Build on (or for) the same +architecture as the target function — `provided:al2023` is available for both +`x86_64` and `arm64`. diff --git a/examples/basic-snapstart/src/main.rs b/examples/basic-snapstart/src/main.rs new file mode 100644 index 000000000..ed9f39f6c --- /dev/null +++ b/examples/basic-snapstart/src/main.rs @@ -0,0 +1,85 @@ +// This example demonstrates using SnapStart with the lambda_runtime crate. +// +// Implement the `SnapStartResource` trait on the types that hold +// snapshot-sensitive state (connections, credentials, cached values) and +// register them on the runtime. When deployed with SnapStart enabled, the +// runtime will: +// 1. Initialize and create resources (e.g., database connections) +// 2. Run each resource's `before_snapshot` hook (reverse registration order) +// before the VM snapshot +// 3. Call `/restore/next`, which blocks until the VM is restored +// 4. Run each resource's `after_restore` hook (registration order) after restore +// 5. Enter the normal invocation loop +// +// When SnapStart is NOT enabled, the hooks are never called and the runtime +// behaves exactly as before. + +use lambda_runtime::{service_fn, tracing, BoxFuture, Error, LambdaEvent, Runtime, SnapStartResource}; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +#[derive(Deserialize)] +struct Request { + name: String, +} + +#[derive(Serialize)] +struct Response { + message: String, + invocation_count: u64, +} + +struct AppState { + counter: AtomicU64, +} + +impl AppState { + fn new() -> Self { + Self { + counter: AtomicU64::new(0), + } + } +} + +impl SnapStartResource for AppState { + fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async move { + tracing::info!("Releasing resources before snapshot"); + self.counter.store(0, Ordering::Relaxed); + Ok(()) + }) + } + + fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async move { + tracing::info!("Re-establishing resources after restore"); + Ok(()) + }) + } +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + tracing::init_default_subscriber(); + + let state = Arc::new(AppState::new()); + + let state_ref = state.clone(); + let handler = service_fn(move |event: LambdaEvent| { + let state = state_ref.clone(); + async move { + let count = state.counter.fetch_add(1, Ordering::Relaxed) + 1; + Ok::<_, Error>(Response { + message: format!("Hello, {}!", event.payload.name), + invocation_count: count, + }) + } + }); + + Runtime::new(handler) + .register_snapstart_resource(state.clone()) + .run() + .await?; + Ok(()) +} diff --git a/examples/http-snapstart/Cargo.toml b/examples/http-snapstart/Cargo.toml new file mode 100644 index 000000000..a2bad8b88 --- /dev/null +++ b/examples/http-snapstart/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "http-snapstart" +version = "0.1.0" +edition = "2021" + +[dependencies] +lambda_http = { path = "../../lambda-http" } +tokio = { version = "1", features = ["macros"] } diff --git a/examples/http-snapstart/Dockerfile b/examples/http-snapstart/Dockerfile new file mode 100644 index 000000000..0a72a338e --- /dev/null +++ b/examples/http-snapstart/Dockerfile @@ -0,0 +1,41 @@ +# Container image for the http-snapstart example. +# +# SnapStart for functions packaged as OCI images requires a custom base image +# whose runtime implements the restore lifecycle — which `lambda_runtime` now +# does. This image builds the example as the Lambda `bootstrap` on top of the +# AWS-provided base (`provided:al2023`), which supports SnapStart. +# +# Build from the REPOSITORY ROOT (the example uses path dependencies on the +# workspace crates, so the build context must include them): +# +# docker build -f examples/http-snapstart/Dockerfile -t http-snapstart . +# +# Then push to ECR and create the function with SnapStart enabled +# (PublishedVersions) and the architecture matching your build host. + +# ---- build stage ---- +FROM public.ecr.aws/lambda/provided:al2023 AS builder + +RUN dnf install -y gcc && \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +ENV PATH="/root/.cargo/bin:${PATH}" + +# Copy only the workspace crates this example depends on via path, then the +# example itself: lambda_http -> lambda_runtime -> lambda_runtime_api_client, +# plus lambda-events (the aws_lambda_events crate used by lambda_http). +COPY Cargo.* /build/ +COPY lambda-http /build/lambda-http +COPY lambda-runtime /build/lambda-runtime +COPY lambda-runtime-api-client /build/lambda-runtime-api-client +COPY lambda-events /build/lambda-events +COPY examples/http-snapstart /build/examples/http-snapstart + +WORKDIR /build/examples/http-snapstart +RUN cargo build --release + +# ---- runtime stage ---- +FROM public.ecr.aws/lambda/provided:al2023 + +COPY --from=builder /build/examples/http-snapstart/target/release/http-snapstart ${LAMBDA_RUNTIME_DIR}/bootstrap + +ENTRYPOINT [ "/var/runtime/bootstrap" ] diff --git a/examples/http-snapstart/README.md b/examples/http-snapstart/README.md new file mode 100644 index 000000000..ecf459993 --- /dev/null +++ b/examples/http-snapstart/README.md @@ -0,0 +1,59 @@ +# AWS Lambda SnapStart example (lambda_http) + +This example shows how to use AWS Lambda **SnapStart** with the `lambda_http` +crate. It wraps a shared connection pool in a `SnapStartResource` and registers +it via the [`lambda_http::runtime`] helper, so the runtime drains the pool before +the VM snapshot and reconnects it after each restore. The HTTP handler keeps the +usual `lambda_http` request/response ergonomics. + +SnapStart reduces cold-start latency by snapshotting the initialized execution +environment and restoring from it. Resources created during init (connections, +credentials, unique values) may be invalid after restore — `before_snapshot` / +`after_restore` are where you release and re-establish them. When SnapStart is +not enabled, the hooks are never called and the function behaves normally. + +If you don't need custom hooks, plain `lambda_http::run(...)` already gets +SnapStart support for free: the runtime calls `/restore/next` and rebuilds its +internal RAPID connection pool on restore without any extra code. + +## Why a container image? + +SnapStart for functions packaged as OCI images requires a custom base image +whose runtime implements the restore lifecycle — which `lambda_runtime` (used by +`lambda_http`) now does. This example ships a `Dockerfile` that builds the binary +as the Lambda `bootstrap` on top of `public.ecr.aws/lambda/provided:al2023`, +which supports SnapStart. + +## Build & Deploy + +Build the image from the **repository root** (the example depends on the +workspace crates via path, so the build context must include them): + +```sh +docker build -f examples/http-snapstart/Dockerfile -t http-snapstart . +``` + +Push it to ECR and create the function from the image: + +```sh +aws ecr create-repository --repository-name http-snapstart +docker tag http-snapstart:latest .dkr.ecr..amazonaws.com/http-snapstart:latest +docker push .dkr.ecr..amazonaws.com/http-snapstart:latest + +aws lambda create-function \ + --function-name http-snapstart \ + --package-type Image \ + --code ImageUri=.dkr.ecr..amazonaws.com/http-snapstart:latest \ + --role \ + --snap-start ApplyOn=PublishedVersions +``` + +SnapStart applies to **published versions**, so publish a version to trigger +snapshot creation, then expose it behind a function URL or API Gateway and invoke +it (e.g. `?name=world`). + +## Architecture + +The compiled `bootstrap` is architecture-specific. Build on (or for) the same +architecture as the target function — `provided:al2023` is available for both +`x86_64` and `arm64`. diff --git a/examples/http-snapstart/src/main.rs b/examples/http-snapstart/src/main.rs new file mode 100644 index 000000000..bcfffff61 --- /dev/null +++ b/examples/http-snapstart/src/main.rs @@ -0,0 +1,77 @@ +// This example demonstrates using SnapStart with the lambda_http crate. +// +// Use lambda_http::runtime() to get a Runtime on which you can register +// `SnapStartResource`s. This gives you access to the SnapStart lifecycle while +// keeping the lambda_http request/response ergonomics. +// +// For the simple case where no custom hooks are needed, just use +// lambda_http::run() — SnapStart works automatically. + +use lambda_http::{ + service_fn, tracing, BoxFuture, Body, Error, IntoResponse, Request, RequestExt, Response, SnapStartResource, +}; +use std::sync::Arc; +use tokio::sync::RwLock; + +struct DbPool { + connected: bool, +} + +impl DbPool { + async fn connect() -> Self { + tracing::info!("Establishing database connection pool"); + Self { connected: true } + } +} + +/// A `SnapStartResource` wrapper around the shared connection pool. The handler +/// and the resource share the same `Arc>`, so draining before the +/// snapshot and reconnecting after restore are visible to invocations. +struct PoolResource(Arc>); + +impl SnapStartResource for PoolResource { + fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async move { + tracing::info!("Draining database connections before snapshot"); + self.0.write().await.connected = false; + Ok(()) + }) + } + + fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async move { + tracing::info!("Re-establishing database connections after restore"); + self.0.write().await.connected = true; + Ok(()) + }) + } +} + +#[tokio::main] +async fn main() -> Result<(), Error> { + tracing::init_default_subscriber(); + + let pool = Arc::new(RwLock::new(DbPool::connect().await)); + + let pool_ref = pool.clone(); + let handler = service_fn(move |req: Request| { + let pool = pool_ref.clone(); + async move { + let pool = pool.read().await; + assert!(pool.connected, "Pool should be connected during invocation"); + + let name = req + .query_string_parameters_ref() + .and_then(|params| params.first("name")) + .unwrap_or("world"); + + Ok::, Error>(format!("Hello, {name}!").into_response().await) + } + }); + + lambda_http::runtime(handler) + .register_snapstart_resource(Arc::new(PoolResource(pool.clone()))) + .run() + .await?; + Ok(()) +} diff --git a/lambda-extension/Cargo.toml b/lambda-extension/Cargo.toml index c319bd8f5..6c8a4bd10 100644 --- a/lambda-extension/Cargo.toml +++ b/lambda-extension/Cargo.toml @@ -25,7 +25,7 @@ http = { workspace = true } http-body-util = { workspace = true } hyper = { workspace = true, features = ["http1", "client", "server"] } hyper-util = { workspace = true } -lambda_runtime_api_client = { version = "1.0.3", path = "../lambda-runtime-api-client" } +lambda_runtime_api_client = { version = "1.1.0", path = "../lambda-runtime-api-client" } serde = { version = "1", features = ["derive"] } serde_json = "^1" tokio = { version = "1.0", features = [ diff --git a/lambda-extension/src/extension.rs b/lambda-extension/src/extension.rs index 36ffe5913..16cfd1f66 100644 --- a/lambda-extension/src/extension.rs +++ b/lambda-extension/src/extension.rs @@ -3,7 +3,7 @@ use http_body_util::BodyExt; use hyper::{body::Incoming, server::conn::http1, service::service_fn}; use hyper_util::rt::tokio::TokioIo; -use lambda_runtime_api_client::Client; +use lambda_runtime_api_client::PooledClient as Client; use serde::{de::DeserializeOwned, Deserialize}; use std::{ convert::Infallible, @@ -244,7 +244,7 @@ where /// extension, it is safe to call `lambda_runtime::run` once the future returned by this /// function resolves. pub async fn register(self) -> Result, Error> { - let client = &Client::builder().build()?; + let client = &Client::builder().build(); let register_res = register(client, self.extension_name, self.events).await?; @@ -428,7 +428,7 @@ where /// [shutdown](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtime-environment.html#runtimes-lifecycle-shutdown) /// Lambda lifecycle phases. pub async fn run(self) -> Result<(), Error> { - let client = &Client::builder().build()?; + let client = &Client::builder().build(); let mut ep = self.events_processor; let extension_id = &self.extension_id; diff --git a/lambda-http/Cargo.toml b/lambda-http/Cargo.toml index 10600be53..183398188 100644 --- a/lambda-http/Cargo.toml +++ b/lambda-http/Cargo.toml @@ -61,7 +61,7 @@ features = ["alb", "apigw", "vpc_lattice"] axum-core = "0.5.4" axum-extra = { version = "0.12.5", features = ["query"] } -lambda_runtime_api_client = { version = "1.0.3", path = "../lambda-runtime-api-client" } +lambda_runtime_api_client = { version = "1.1.0", path = "../lambda-runtime-api-client" } log = "^0.4" maplit = "1.0" diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs index 496110db9..caea86181 100644 --- a/lambda-http/src/lib.rs +++ b/lambda-http/src/lib.rs @@ -74,7 +74,7 @@ pub use http::{self, Response}; #[cfg_attr(docsrs, doc(cfg(feature = "tracing")))] pub use lambda_runtime::tracing; use lambda_runtime::Diagnostic; -pub use lambda_runtime::{self, service_fn, tower, Context, Error, LambdaEvent, Service}; +pub use lambda_runtime::{self, service_fn, tower, BoxFuture, Context, Error, LambdaEvent, Service, SnapStartResource}; use request::RequestFuture; use response::ResponseFuture; @@ -105,9 +105,9 @@ use std::{ }; mod streaming; +pub use streaming::{run_with_streaming_response, streaming_runtime, StreamAdapter}; #[cfg(feature = "concurrency-tokio")] -pub use streaming::run_with_streaming_response_concurrent; -pub use streaming::{run_with_streaming_response, StreamAdapter}; +pub use streaming::{run_with_streaming_response_concurrent, streaming_runtime_concurrent}; /// Type alias for `http::Request`s with a fixed [`Body`](enum.Body.html) type pub type Request = http::Request; @@ -268,6 +268,96 @@ where lambda_runtime::run_concurrent(Adapter::from(handler)).await } +/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given +/// handler, without starting the event loop. +/// +/// Use this when you need to register +/// [`SnapStartResource`]s for the SnapStart +/// snapshot/restore lifecycle. For the common case where no custom hooks are +/// needed, use [`run()`] instead — SnapStart support works automatically. +/// +/// # Example +/// +/// ```no_run +/// use lambda_http::{service_fn, BoxFuture, Error, Request, SnapStartResource}; +/// use std::sync::Arc; +/// +/// struct Pool; +/// impl SnapStartResource for Pool { +/// fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { +/// Box::pin(async move { Ok(()) }) +/// } +/// fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { +/// Box::pin(async move { Ok(()) }) +/// } +/// } +/// +/// #[tokio::main] +/// async fn main() -> Result<(), Error> { +/// let pool = Arc::new(Pool); +/// let runtime = lambda_http::runtime(service_fn(handler)) +/// .register_snapstart_resource(pool.clone()); +/// +/// runtime.run().await?; +/// Ok(()) +/// } +/// +/// async fn handler(_req: Request) -> Result<&'static str, std::convert::Infallible> { +/// Ok("hello") +/// } +/// ``` +/// +/// # Panics +/// +/// This function panics if required Lambda environment variables are missing +/// (`AWS_LAMBDA_FUNCTION_NAME`, `AWS_LAMBDA_FUNCTION_MEMORY_SIZE`, +/// `AWS_LAMBDA_FUNCTION_VERSION`, `AWS_LAMBDA_RUNTIME_API`). +pub fn runtime( + handler: S, +) -> lambda_runtime::Runtime< + impl lambda_runtime::Service, +> +where + S: Service + Send + 'static, + S::Future: Send + 'static, + R: IntoResponse + Send + Sync + 'static, + E: std::fmt::Debug + Into + Send + 'static, +{ + lambda_runtime::Runtime::new(Adapter::from(handler)) +} + +/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given +/// handler for concurrent execution, without starting the event loop. +/// +/// This is the concurrent variant of [`runtime()`]. Use it when you need SnapStart +/// hooks AND your handler supports concurrent invocations. +/// +/// # Panics +/// +/// This function panics if required Lambda environment variables are missing. +#[cfg(feature = "concurrency-tokio")] +#[cfg_attr(docsrs, doc(cfg(feature = "concurrency-tokio")))] +pub fn runtime_concurrent( + handler: S, +) -> lambda_runtime::Runtime< + impl lambda_runtime::Service< + lambda_runtime::LambdaInvocation, + Response = (), + Error = lambda_runtime::Error, + Future: Send, + > + Clone + + Send + + 'static, +> +where + S: Service + Clone + Send + 'static, + S::Future: Send + 'static, + R: IntoResponse + Send + Sync + 'static, + E: std::fmt::Debug + Into + Send + 'static, +{ + lambda_runtime::Runtime::new(Adapter::from(handler)) +} + // In concurrent mode we must use the per-request context. fn update_xray_trace_id_header(headers: &mut http::HeaderMap, context: &Context) { if let Some(trace_id) = context.xray_trace_id.as_deref() { @@ -331,4 +421,35 @@ mod test_adapter { .service_fn(|_event: Request| async move { Response::builder().status(StatusCode::OK).body(Body::Empty) }) .boxed(); } + + async fn http_handler(_req: Request) -> Result<&'static str, std::convert::Infallible> { + Ok("hello") + } + + /// `runtime()` accepts a (non-`Clone`) handler for the sequential path. This + /// is a compile-time check that the helper's bounds are satisfiable. + #[test] + fn runtime_helper_accepts_handler() { + let _ = || { + let _runtime = crate::runtime(crate::service_fn(http_handler)); + }; + } + + /// `runtime_concurrent()` requires a `Clone` handler (Lambda Managed + /// Instances). `service_fn` over an `async fn` is `Clone`, so this compiles; + /// it guards against the `Clone` bound regressing on the concurrent helper. + /// + /// Crucially, this also calls `.run_concurrent()` on the result: that method + /// requires the wrapped service to be `Clone + Send + 'static` with a `Send` + /// future, so this compiling proves the helper's return type exposes those + /// bounds (without them, the chained call would not type-check). + #[cfg(feature = "concurrency-tokio")] + #[test] + fn runtime_concurrent_helper_exposes_lmi_bounds() { + let _ = || async { + let _ = crate::runtime_concurrent(crate::service_fn(http_handler)) + .run_concurrent() + .await; + }; + } } diff --git a/lambda-http/src/streaming.rs b/lambda-http/src/streaming.rs index 2c0a99cea..361a925e1 100644 --- a/lambda-http/src/streaming.rs +++ b/lambda-http/src/streaming.rs @@ -230,6 +230,64 @@ where lambda_runtime::run_concurrent(into_stream_service_cloneable(handler)).await } +/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given +/// streaming handler, without starting the event loop. +/// +/// Use this when you need access to the [`Runtime`](lambda_runtime::Runtime) for +/// SnapStart lifecycle hooks with a streaming response handler. +/// +/// # Panics +/// +/// This function panics if required Lambda environment variables are missing. +pub fn streaming_runtime( + handler: S, +) -> lambda_runtime::Runtime< + impl lambda_runtime::Service, +> +where + S: Service, Error = E> + Send + 'static, + S::Future: Send + 'static, + E: Debug + Into + Send + 'static, + B: Body + Unpin + Send + 'static, + B::Data: Into + Send, + B::Error: Into + Send + Debug, +{ + lambda_runtime::Runtime::new(into_stream_service(handler)) +} + +/// Returns a configured [`Runtime`](lambda_runtime::Runtime) wrapping the given +/// streaming handler for concurrent execution, without starting the event loop. +/// +/// This is the concurrent variant of [`streaming_runtime()`]. +/// +/// # Panics +/// +/// This function panics if required Lambda environment variables are missing. +#[cfg(feature = "concurrency-tokio")] +#[cfg_attr(docsrs, doc(cfg(feature = "concurrency-tokio")))] +pub fn streaming_runtime_concurrent( + handler: S, +) -> lambda_runtime::Runtime< + impl lambda_runtime::Service< + lambda_runtime::LambdaInvocation, + Response = (), + Error = lambda_runtime::Error, + Future: Send, + > + Clone + + Send + + 'static, +> +where + S: Service, Error = E> + Clone + Send + 'static, + S::Future: Send + 'static, + E: Debug + Into + Send + 'static, + B: Body + Unpin + Send + 'static, + B::Data: Into + Send, + B::Error: Into + Send + Debug, +{ + lambda_runtime::Runtime::new(into_stream_service_cloneable(handler)) +} + pin_project_lite::pin_project! { #[non_exhaustive] pub struct BodyStream { diff --git a/lambda-runtime-api-client/CHANGELOG.md b/lambda-runtime-api-client/CHANGELOG.md index c56572383..0899561fd 100644 --- a/lambda-runtime-api-client/CHANGELOG.md +++ b/lambda-runtime-api-client/CHANGELOG.md @@ -7,6 +7,25 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- `PooledClient`: a Runtime API client that can rebuild its connection pool after a + SnapStart VM restore via `reset_pool()`, so stale connections to the Runtime API + (RAPID) are discarded. Implemented lock-free with `OnceLock`; `call()` has no + hot-path lock. Construct via `PooledClient::builder().build()`. +- `PooledClientBuilder`, returned by `PooledClient::builder()`. +- `RuntimeApiClient` trait abstracting `call`, implemented by both `Client` and + `PooledClient`. + +### Deprecated + +- `Client` and `ClientBuilder::build` are superseded by `PooledClient` / + `PooledClient::builder().build()`. `Client` is unchanged and fully functional; it + simply cannot reset its connection pool after a SnapStart restore. They are + documented as superseded but do **not** carry a `#[deprecated]` attribute, so no + compiler warnings are emitted; they are expected to be removed in the next major + version (2.0.0). This keeps the crate non-breaking (no major version bump). + ## [1.0.3](https://github.com/aws/aws-lambda-rust-runtime/compare/lambda_runtime_api_client-v1.0.2...lambda_runtime_api_client-v1.0.3) - 2026-03-19 ### Other diff --git a/lambda-runtime-api-client/Cargo.toml b/lambda-runtime-api-client/Cargo.toml index ec4ee2d67..ef8a78ce3 100644 --- a/lambda-runtime-api-client/Cargo.toml +++ b/lambda-runtime-api-client/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lambda_runtime_api_client" -version = "1.0.3" +version = "1.1.0" edition = "2021" rust-version = "1.84.0" authors = [ diff --git a/lambda-runtime-api-client/src/lib.rs b/lambda-runtime-api-client/src/lib.rs index adbd5e1bc..ead8abe75 100644 --- a/lambda-runtime-api-client/src/lib.rs +++ b/lambda-runtime-api-client/src/lib.rs @@ -12,7 +12,7 @@ use http::{ }; use hyper::body::Incoming; use hyper_util::client::legacy::connect::HttpConnector; -use std::{convert::TryInto, fmt::Debug, future}; +use std::{convert::TryInto, fmt::Debug, future, sync::OnceLock}; const USER_AGENT_HEADER: &str = "User-Agent"; const DEFAULT_USER_AGENT: &str = concat!("aws-lambda-rust/", env!("CARGO_PKG_VERSION")); @@ -26,7 +26,20 @@ pub mod body; #[cfg_attr(docsrs, doc(cfg(feature = "tracing")))] pub mod tracing; +/// The single method the runtime needs from any client: send a request to the +/// Lambda Runtime API (RAPID). Implemented by both [`Client`] and [`PooledClient`]. +pub trait RuntimeApiClient { + /// Send a given request to the Runtime API. + fn call(&self, req: Request) -> BoxFuture<'static, Result, BoxError>>; +} + /// API client to interact with the AWS Lambda Runtime API. +/// +/// **Superseded by [`PooledClient`]**, which additionally supports rebuilding its +/// connection pool after a SnapStart VM restore via [`PooledClient::reset_pool`]. +/// `Client` remains fully functional; it simply cannot reset its pool after a +/// restore. Prefer [`PooledClient`] for new code; `Client` is expected to be +/// removed in the next major version (2.0.0). #[derive(Debug)] pub struct Client { /// The runtime API URI @@ -38,62 +51,143 @@ pub struct Client { impl Client { /// Create a builder struct to configure the client. pub fn builder() -> ClientBuilder { - ClientBuilder { - connector: HttpConnector::new(), - uri: None, - pool_size: None, - } + ClientBuilder::new() } -} -impl Client { /// Send a given request to the Runtime API. /// Use the client's base URI to ensure the API endpoint is correct. pub fn call(&self, req: Request) -> BoxFuture<'static, Result, BoxError>> { + ::call(self, req) + } + + /// Create a new client with a given base URI, HTTP connector, and optional pool size hint. + fn with(base: Uri, connector: HttpConnector, pool_size: Option) -> Self { + Self { + base, + client: build_hyper_client(&connector, pool_size), + } + } +} + +impl RuntimeApiClient for Client { + fn call(&self, req: Request) -> BoxFuture<'static, Result, BoxError>> { // NOTE: This method returns a boxed future such that the future has a static lifetime. // Due to limitations around the Rust async implementation as of Mar 2024, this is // required to minimize constraints on the handler passed to [lambda_runtime::run]. - let req = match self.set_origin(req) { + let req = match set_origin(&self.base, req) { Ok(req) => req, Err(err) => return future::ready(Err(err)).boxed(), }; self.client.request(req).map_err(Into::into).boxed() } +} + +/// API client to interact with the AWS Lambda Runtime API, with support for +/// rebuilding its connection pool after a SnapStart VM restore. +#[derive(Debug)] +pub struct PooledClient { + /// The runtime API URI + pub base: Uri, + /// The client that manages the API connections + client: hyper_util::client::legacy::Client, + /// The HTTP connector used to rebuild the client after a SnapStart restore + connector: HttpConnector, + /// Optional pool size hint for the hyper client + pool_size: Option, + /// Holds a freshly built client after a SnapStart restore. When set, [`PooledClient::call`] + /// uses it instead of [`PooledClient::client`]; otherwise the original client is used. + /// + /// This is populated exactly once by [`PooledClient::reset_pool`] during the restore + /// lifecycle, before any concurrent polling starts, so the `OnceLock` write never + /// races with reads. Using `OnceLock` keeps `call` lock-free on the hot path. + restored: OnceLock>, +} + +impl PooledClient { + /// Create a builder struct to configure the client. + pub fn builder() -> PooledClientBuilder { + PooledClientBuilder::new() + } + + /// Send a given request to the Runtime API. + /// Use the client's base URI to ensure the API endpoint is correct. + pub fn call(&self, req: Request) -> BoxFuture<'static, Result, BoxError>> { + ::call(self, req) + } /// Create a new client with a given base URI, HTTP connector, and optional pool size hint. fn with(base: Uri, connector: HttpConnector, pool_size: Option) -> Self { - let mut builder = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()); - builder.http1_max_buf_size(1024 * 1024); - - if let Some(size) = pool_size { - builder.pool_max_idle_per_host(size); + let client = build_hyper_client(&connector, pool_size); + Self { + base, + client, + connector, + pool_size, + restored: OnceLock::new(), } + } - let client = builder.build(connector); - Self { base, client } + /// Reset the internal connection pool by building a fresh hyper client that + /// subsequent calls will use. + /// + /// This is useful after a SnapStart VM restore where existing connections + /// to the Lambda Runtime API (RAPID) are stale and must be discarded. It is + /// called exactly once during the restore lifecycle, before concurrent + /// polling starts, so the underlying `OnceLock` write never races. + pub fn reset_pool(&self) { + let _ = self.restored.set(build_hyper_client(&self.connector, self.pool_size)); } +} - fn set_origin(&self, req: Request) -> Result, BoxError> { - let (mut parts, body) = req.into_parts(); - let (scheme, authority, base_path) = { - let scheme = self.base.scheme().unwrap_or(&Scheme::HTTP); - let authority = self.base.authority().expect("Authority not found"); - let base_path = self.base.path().trim_end_matches('/'); - (scheme, authority, base_path) +impl RuntimeApiClient for PooledClient { + fn call(&self, req: Request) -> BoxFuture<'static, Result, BoxError>> { + let req = match set_origin(&self.base, req) { + Ok(req) => req, + Err(err) => return future::ready(Err(err)).boxed(), }; - let path = parts.uri.path_and_query().expect("PathAndQuery not found"); - let pq: PathAndQuery = format!("{base_path}{path}").parse().expect("PathAndQuery invalid"); + // After a SnapStart restore, `restored` holds a fresh client with no stale + // connections; otherwise fall back to the original client. Lock-free read. + let hyper_client = self.restored.get().unwrap_or(&self.client); + hyper_client.request(req).map_err(Into::into).boxed() + } +} - let uri = Uri::builder() - .scheme(scheme.as_ref()) - .authority(authority.as_ref()) - .path_and_query(pq) - .build() - .map_err(Box::new)?; +/// Build a fresh hyper client from the given connector and pool size settings. +fn build_hyper_client( + connector: &HttpConnector, + pool_size: Option, +) -> hyper_util::client::legacy::Client { + let mut builder = hyper_util::client::legacy::Client::builder(hyper_util::rt::TokioExecutor::new()); + builder.http1_max_buf_size(1024 * 1024); - parts.uri = uri; - Ok(Request::from_parts(parts, body)) + if let Some(size) = pool_size { + builder.pool_max_idle_per_host(size); } + + builder.build(connector.clone()) +} + +/// Rewrite a request's URI against the client's base (scheme/authority/base path). +fn set_origin(base: &Uri, req: Request) -> Result, BoxError> { + let (mut parts, body) = req.into_parts(); + let (scheme, authority, base_path) = { + let scheme = base.scheme().unwrap_or(&Scheme::HTTP); + let authority = base.authority().expect("Authority not found"); + let base_path = base.path().trim_end_matches('/'); + (scheme, authority, base_path) + }; + let path = parts.uri.path_and_query().expect("PathAndQuery not found"); + let pq: PathAndQuery = format!("{base_path}{path}").parse().expect("PathAndQuery invalid"); + + let uri = Uri::builder() + .scheme(scheme.as_ref()) + .authority(authority.as_ref()) + .path_and_query(pq) + .build() + .map_err(Box::new)?; + + parts.uri = uri; + Ok(Request::from_parts(parts, body)) } /// Builder implementation to construct any Runtime API clients. @@ -104,6 +198,14 @@ pub struct ClientBuilder { } impl ClientBuilder { + fn new() -> ClientBuilder { + ClientBuilder { + connector: HttpConnector::new(), + uri: None, + pool_size: None, + } + } + /// Create a new builder with a given HTTP connector. pub fn with_connector(self, connector: HttpConnector) -> ClientBuilder { ClientBuilder { @@ -132,18 +234,73 @@ impl ClientBuilder { } /// Create the new client to interact with the Runtime API. + /// + /// **Superseded by [`PooledClient::builder`]**, which produces a + /// [`PooledClient`] that supports SnapStart connection-pool reset. Prefer + /// `PooledClient::builder().build()` for new code; this method is expected to + /// be removed in the next major version (2.0.0). pub fn build(self) -> Result { - let uri = match self.uri { - Some(uri) => uri, - None => { - let uri = std::env::var("AWS_LAMBDA_RUNTIME_API").expect("Missing AWS_LAMBDA_RUNTIME_API env var"); - uri.try_into().expect("Unable to convert to URL") - } - }; + let uri = resolve_uri(self.uri); Ok(Client::with(uri, self.connector, self.pool_size)) } } +/// Builder to construct a [`PooledClient`]. +pub struct PooledClientBuilder { + connector: HttpConnector, + uri: Option, + pool_size: Option, +} + +impl PooledClientBuilder { + fn new() -> PooledClientBuilder { + PooledClientBuilder { + connector: HttpConnector::new(), + uri: None, + pool_size: None, + } + } + + /// Use a given HTTP connector. + pub fn with_connector(self, connector: HttpConnector) -> PooledClientBuilder { + PooledClientBuilder { connector, ..self } + } + + /// Use a given base URI. Inherits all other attributes from the existing builder. + pub fn with_endpoint(self, uri: http::Uri) -> Self { + Self { uri: Some(uri), ..self } + } + + /// Provide a pool size hint for the underlying Hyper client. + /// + /// When using concurrent polling, this should be at least the maximum + /// concurrency (e.g., `AWS_LAMBDA_MAX_CONCURRENCY`) to avoid connection + /// starvation. + pub fn with_pool_size(self, pool_size: usize) -> Self { + Self { + pool_size: Some(pool_size), + ..self + } + } + + /// Create the new [`PooledClient`] to interact with the Runtime API. + pub fn build(self) -> PooledClient { + let uri = resolve_uri(self.uri); + PooledClient::with(uri, self.connector, self.pool_size) + } +} + +/// Resolve the base URI from an explicit value or the `AWS_LAMBDA_RUNTIME_API` env var. +fn resolve_uri(uri: Option) -> Uri { + match uri { + Some(uri) => uri, + None => { + let uri = std::env::var("AWS_LAMBDA_RUNTIME_API").expect("Missing AWS_LAMBDA_RUNTIME_API env var"); + uri.try_into().expect("Unable to convert to URL") + } + } +} + /// Create a request builder. /// This builder uses `aws-lambda-rust/CRATE_VERSION` as /// the default User-Agent. @@ -163,13 +320,12 @@ mod tests { #[test] fn test_set_origin() { - let base = "http://localhost:9001"; - let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap(); + let base: Uri = "http://localhost:9001".parse().unwrap(); let req = build_request() .uri("/2018-06-01/runtime/invocation/next") .body(()) .unwrap(); - let req = client.set_origin(req).unwrap(); + let req = set_origin(&base, req).unwrap(); assert_eq!( "http://localhost:9001/2018-06-01/runtime/invocation/next", &req.uri().to_string() @@ -178,25 +334,23 @@ mod tests { #[test] fn test_set_origin_with_base_path() { - let base = "http://localhost:9001/foo"; - let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap(); + let base: Uri = "http://localhost:9001/foo".parse().unwrap(); let req = build_request() .uri("/2018-06-01/runtime/invocation/next") .body(()) .unwrap(); - let req = client.set_origin(req).unwrap(); + let req = set_origin(&base, req).unwrap(); assert_eq!( "http://localhost:9001/foo/2018-06-01/runtime/invocation/next", &req.uri().to_string() ); - let base = "http://localhost:9001/foo/"; - let client = Client::builder().with_endpoint(base.parse().unwrap()).build().unwrap(); + let base: Uri = "http://localhost:9001/foo/".parse().unwrap(); let req = build_request() .uri("/2018-06-01/runtime/invocation/next") .body(()) .unwrap(); - let req = client.set_origin(req).unwrap(); + let req = set_origin(&base, req).unwrap(); assert_eq!( "http://localhost:9001/foo/2018-06-01/runtime/invocation/next", &req.uri().to_string() @@ -207,12 +361,22 @@ mod tests { fn builder_accepts_pool_size() { let base = "http://localhost:9001"; let expected: Uri = base.parse().unwrap(); - let client = Client::builder() + let client = PooledClient::builder() .with_pool_size(4) .with_endpoint(base.parse().unwrap()) - .build() - .unwrap(); + .build(); assert_eq!(client.base, expected); } + + #[test] + fn test_reset_pool() { + let base = "http://localhost:9001"; + let client = PooledClient::builder() + .with_pool_size(4) + .with_endpoint(base.parse().unwrap()) + .build(); + client.reset_pool(); + assert_eq!(client.base, base.parse::().unwrap()); + } } diff --git a/lambda-runtime/Cargo.toml b/lambda-runtime/Cargo.toml index f283aa92e..68ca9039f 100644 --- a/lambda-runtime/Cargo.toml +++ b/lambda-runtime/Cargo.toml @@ -42,7 +42,7 @@ http-body-util = { workspace = true } http-serde = { workspace = true } hyper = { workspace = true, features = ["http1", "client"] } lambda-extension = { version = "1.0", path = "../lambda-extension", default-features = false, optional = true } -lambda_runtime_api_client = { version = "1.0.3", path = "../lambda-runtime-api-client", default-features = false } +lambda_runtime_api_client = { version = "1.1.0", path = "../lambda-runtime-api-client", default-features = false } miette = { version = "7.2.0", optional = true } opentelemetry-semantic-conventions = { version = "0.32", optional = true, features = ["semconv_experimental"] } pin-project = "1" diff --git a/lambda-runtime/src/layers/api_client.rs b/lambda-runtime/src/layers/api_client.rs index df9b864db..bc2bcd0f3 100644 --- a/lambda-runtime/src/layers/api_client.rs +++ b/lambda-runtime/src/layers/api_client.rs @@ -1,7 +1,10 @@ +// `Client` is referenced only as the default for the `C` generic parameter, kept +// for backward compatibility of the public `RuntimeApiClientService`/future types. + use crate::LambdaInvocation; use futures::{future::BoxFuture, ready, FutureExt, TryFutureExt}; use hyper::body::Incoming; -use lambda_runtime_api_client::{body::Body, BoxError, Client}; +use lambda_runtime_api_client::{body::Body, BoxError, Client, RuntimeApiClient}; use pin_project::pin_project; use std::{future::Future, pin::Pin, sync::Arc, task}; use tower::Service; @@ -13,25 +16,26 @@ use tracing::error; /// This type is only meant for internal use in the Lambda runtime crate. It neither augments the /// inner service's request type nor its error type. However, this service returns an empty /// response `()` as the Lambda request has been completed. -pub struct RuntimeApiClientService { +pub struct RuntimeApiClientService { inner: S, - client: Arc, + client: Arc, } -impl RuntimeApiClientService { - pub fn new(inner: S, client: Arc) -> Self { +impl RuntimeApiClientService { + pub fn new(inner: S, client: Arc) -> Self { Self { inner, client } } } -impl Service for RuntimeApiClientService +impl Service for RuntimeApiClientService where S: Service, S::Future: Future, BoxError>>, + C: RuntimeApiClient, { type Response = (); type Error = S::Error; - type Future = RuntimeApiClientFuture; + type Future = RuntimeApiClientFuture; fn poll_ready(&mut self, cx: &mut task::Context<'_>) -> task::Poll> { self.inner.poll_ready(cx) @@ -44,7 +48,7 @@ where } } -impl Clone for RuntimeApiClientService +impl Clone for RuntimeApiClientService where S: Clone, { @@ -57,14 +61,15 @@ where } #[pin_project(project = RuntimeApiClientFutureProj)] -pub enum RuntimeApiClientFuture { - First(#[pin] F, Arc), +pub enum RuntimeApiClientFuture { + First(#[pin] F, Arc), Second(#[pin] BoxFuture<'static, Result, BoxError>>), } -impl Future for RuntimeApiClientFuture +impl Future for RuntimeApiClientFuture where F: Future, BoxError>>, + C: RuntimeApiClient, { type Output = Result<(), BoxError>; diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs index 8755a130e..69c04ebc6 100644 --- a/lambda-runtime/src/lib.rs +++ b/lambda-runtime/src/lib.rs @@ -13,6 +13,7 @@ use std::{ env, fmt::{self, Debug}, future::Future, + pin::Pin, sync::Arc, }; use tokio_stream::Stream; @@ -31,6 +32,8 @@ mod deserializer; pub mod layers; mod requests; mod runtime; +/// SnapStart snapshot/restore lifecycle support. +pub mod snapstart; /// Utilities for Lambda Streaming functions. pub mod streaming; @@ -44,8 +47,14 @@ mod types; use requests::EventErrorRequest; pub use runtime::{LambdaInvocation, Runtime}; +pub use snapstart::SnapStartResource; pub use types::{Context, FunctionResponse, IntoFunctionResponse, LambdaEvent, MetadataPrelude, StreamResponse}; +/// A boxed, `Send` future with a bound lifetime — the return type of +/// [`SnapStartResource`] hooks. Inlined (rather than re-exported from `futures`) +/// to avoid leaking the `futures` crate into this crate's public type surface. +pub type BoxFuture<'a, T> = Pin + Send + 'a>>; + /// Error type that lambdas may result in pub type Error = lambda_runtime_api_client::BoxError; diff --git a/lambda-runtime/src/requests.rs b/lambda-runtime/src/requests.rs index 5384af1fd..b03f14c79 100644 --- a/lambda-runtime/src/requests.rs +++ b/lambda-runtime/src/requests.rs @@ -24,6 +24,60 @@ impl IntoRequest for NextEventRequest { } } +// /runtime/restore/next +#[derive(Debug, Eq, PartialEq)] +pub(crate) struct RestoreNextRequest; + +impl IntoRequest for RestoreNextRequest { + fn into_req(self) -> Result, Error> { + let req = build_request() + .method(Method::GET) + .uri(Uri::from_static("/2018-06-01/runtime/restore/next")) + .body(Default::default())?; + Ok(req) + } +} + +// /runtime/init/error +pub(crate) struct InitErrorRequest { + pub(crate) diagnostic: Diagnostic, +} + +impl IntoRequest for InitErrorRequest { + fn into_req(self) -> Result, Error> { + let uri = Uri::from_static("/2018-06-01/runtime/init/error"); + let body = serde_json::to_vec(&self.diagnostic)?; + let body = Body::from(body); + + let req = build_request() + .method(Method::POST) + .uri(uri) + .header("lambda-runtime-function-error-type", "Runtime.InitError") + .body(body)?; + Ok(req) + } +} + +// /runtime/restore/error +pub(crate) struct RestoreErrorRequest { + pub(crate) diagnostic: Diagnostic, +} + +impl IntoRequest for RestoreErrorRequest { + fn into_req(self) -> Result, Error> { + let uri = Uri::from_static("/2018-06-01/runtime/restore/error"); + let body = serde_json::to_vec(&self.diagnostic)?; + let body = Body::from(body); + + let req = build_request() + .method(Method::POST) + .uri(uri) + .header("lambda-runtime-function-error-type", "Runtime.AfterRestoreError") + .body(body)?; + Ok(req) + } +} + // /runtime/invocation/{AwsRequestId}/response pub(crate) struct EventCompletionRequest<'a, R, B, S, D, E> where @@ -218,6 +272,57 @@ mod tests { }); } + #[test] + fn test_restore_next_request() { + let req = RestoreNextRequest; + let req = req.into_req().unwrap(); + assert_eq!(req.method(), Method::GET); + assert_eq!(req.uri(), &Uri::from_static("/2018-06-01/runtime/restore/next")); + assert!(req + .headers() + .get("User-Agent") + .unwrap() + .to_str() + .unwrap() + .starts_with("aws-lambda-rust/")); + } + + #[test] + fn test_restore_error_request() { + let req = RestoreErrorRequest { + diagnostic: Diagnostic { + error_type: "Runtime.AfterRestoreError".into(), + error_message: "Failed to reconnect to database".into(), + }, + }; + let req = req.into_req().unwrap(); + let expected = Uri::from_static("/2018-06-01/runtime/restore/error"); + assert_eq!(req.method(), Method::POST); + assert_eq!(req.uri(), &expected); + assert_eq!( + req.headers().get("lambda-runtime-function-error-type").unwrap(), + "Runtime.AfterRestoreError" + ); + } + + #[test] + fn test_init_error_request() { + let req = InitErrorRequest { + diagnostic: Diagnostic { + error_type: "Runtime.InitError".into(), + error_message: "Failed to release resources before snapshot".into(), + }, + }; + let req = req.into_req().unwrap(); + let expected = Uri::from_static("/2018-06-01/runtime/init/error"); + assert_eq!(req.method(), Method::POST); + assert_eq!(req.uri(), &expected); + assert_eq!( + req.headers().get("lambda-runtime-function-error-type").unwrap(), + "Runtime.InitError" + ); + } + #[test] #[cfg(tokio_unstable)] fn streaming_send_data_error_is_ignored() { diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs index 94efeca42..ae00bc20b 100644 --- a/lambda-runtime/src/runtime.rs +++ b/lambda-runtime/src/runtime.rs @@ -1,13 +1,14 @@ use crate::{ layers::{CatchPanicService, RuntimeApiClientService, RuntimeApiResponseService}, - requests::{IntoRequest, NextEventRequest}, + requests::{InitErrorRequest, IntoRequest, NextEventRequest, RestoreErrorRequest, RestoreNextRequest}, + snapstart::SnapStartResource, types::{invoke_request_id, IntoFunctionResponse, LambdaEvent}, Config, Context, Diagnostic, }; #[cfg(feature = "concurrency-tokio")] use futures::stream::FuturesUnordered; use http_body_util::BodyExt; -use lambda_runtime_api_client::{BoxError, Client as ApiClient}; +use lambda_runtime_api_client::{BoxError, PooledClient as ApiClient}; use serde::{Deserialize, Serialize}; #[cfg(feature = "concurrency-tokio")] use std::fmt; @@ -60,6 +61,7 @@ pub struct Runtime { config: Arc, client: Arc, concurrency_limit: u32, + snapstart_resources: Vec>, } impl @@ -74,6 +76,7 @@ impl, + ApiClient, >, > where @@ -109,17 +112,13 @@ where let concurrency_limit = max_concurrency_from_env().unwrap_or(1).max(1); // Strategy: allocate all worker tasks up-front, so size the client pool to match. let pool_size = concurrency_limit as usize; - let client = Arc::new( - ApiClient::builder() - .with_pool_size(pool_size) - .build() - .expect("Unable to create a runtime client"), - ); + let client = Arc::new(ApiClient::builder().with_pool_size(pool_size).build()); Self { service: wrap_handler(handler, client.clone()), config, client, concurrency_limit, + snapstart_resources: Vec::new(), } } } @@ -157,10 +156,59 @@ impl Runtime { config: self.config, service: layer.layer(self.service), concurrency_limit: self.concurrency_limit, + snapstart_resources: self.snapstart_resources, } } } +impl Runtime { + /// Returns `true` if the current execution environment has SnapStart enabled. + pub fn is_snapstart(&self) -> bool { + is_snapstart_env() + } + + /// Register a [`SnapStartResource`] whose `before_snapshot`/`after_restore` + /// hooks should run around the SnapStart snapshot/restore boundary. + /// + /// Register resources in dependency order — **foundations first** (e.g. + /// credentials before the pool that depends on them). The runtime runs + /// `before_snapshot` in reverse registration order (LIFO) and `after_restore` + /// in registration order (FIFO), so teardown and rebuild both happen in the + /// correct relative order. See the [`snapstart`](crate::snapstart) module + /// docs for details. + /// + /// When SnapStart is not enabled, registered resources are never invoked and + /// add no runtime overhead beyond the cost of holding the `Arc`. + /// + /// # Example + /// ```no_run + /// use lambda_runtime::{Error, LambdaEvent, Runtime, SnapStartResource}; + /// use std::sync::Arc; + /// use serde_json::Value; + /// use tower::service_fn; + /// + /// // Uses the default no-op hooks; override before_snapshot/after_restore as needed. + /// struct Pool; + /// impl SnapStartResource for Pool {} + /// + /// #[tokio::main] + /// async fn main() -> Result<(), Error> { + /// let pool = Arc::new(Pool); + /// let runtime = Runtime::new(service_fn(handler)) + /// .register_snapstart_resource(pool.clone()); + /// runtime.run().await + /// } + /// + /// async fn handler(event: LambdaEvent) -> Result { + /// Ok(event.payload) + /// } + /// ``` + pub fn register_snapstart_resource(mut self, resource: Arc) -> Self { + self.snapstart_resources.push(resource); + self + } +} + #[cfg(feature = "concurrency-tokio")] impl Runtime where @@ -185,6 +233,8 @@ where panic!("`run_concurrent` must be called from within a Tokio runtime"); } + snapstart_lifecycle(&self.client, &self.snapstart_resources).await?; + if self.concurrency_limit > 1 { trace!("Concurrent mode: _X_AMZN_TRACE_ID is not set; use context.xray_trace_id"); Self::run_concurrent_inner(self.service, self.config, self.client, self.concurrency_limit).await @@ -370,6 +420,7 @@ where fallback: eprintln!("AWS_LAMBDA_MAX_CONCURRENCY is set to '{raw}', but the concurrency-tokio feature is not enabled; running sequentially") ); } + snapstart_lifecycle(&self.client, &self.snapstart_resources).await?; let incoming = incoming(&self.client); Self::run_with_incoming(self.service, self.config, incoming).await } @@ -393,6 +444,136 @@ where /* ------------------------------------------- UTILS ------------------------------------------- */ +/// Returns `true` if the current execution environment has SnapStart enabled. +fn is_snapstart_env() -> bool { + env::var("AWS_LAMBDA_INITIALIZATION_TYPE").as_deref() == Ok("snap-start") +} + +/// Runs the SnapStart restore lifecycle when SnapStart is enabled, and is a +/// no-op otherwise. +/// +/// The sequence is: run `before_snapshot` hooks (LIFO) → call `/restore/next` +/// (blocks until the VM is restored) → reset the internal RAPID connection pool +/// → run `after_restore` hooks (FIFO). A failure in the before-snapshot phase +/// (a hook or `/restore/next`) is reported to Lambda via `/init/error`; an +/// `after_restore` failure is reported via `/restore/error`. Either way the +/// error is returned so the runtime exits cleanly (no `process::exit`, so +/// graceful-shutdown handlers still run). +/// +/// This is a free function taking only the fields it needs (rather than `&self`) +/// so it does not borrow the handler `S` across an `.await`, which would require +/// `Runtime: Sync` in the by-value `run`/`run_concurrent` methods. +async fn snapstart_lifecycle( + client: &Arc, + resources: &[Arc], +) -> Result<(), BoxError> { + if !is_snapstart_env() { + return Ok(()); + } + run_restore_lifecycle(client, resources).await +} + +/// Performs the restore lifecycle unconditionally (the env gate lives in the +/// caller, [`snapstart_lifecycle`]). Split out so it can be tested directly +/// without mutating the process-global `AWS_LAMBDA_INITIALIZATION_TYPE`. +async fn run_restore_lifecycle( + client: &Arc, + resources: &[Arc], +) -> Result<(), BoxError> { + // before_snapshot runs in REVERSE registration order (LIFO): tear down + // dependents before their foundations. A failure here happens during init, + // so report it to /init/error. + for resource in resources.iter().rev() { + if let Err(e) = resource.before_snapshot().await { + return Err(report_init_error(client, e).await); + } + } + + // Signal init complete and block until the VM is restored from snapshot. + let req = RestoreNextRequest.into_req()?; + let resp = match client.call(req).await { + Ok(resp) => resp, + // A transport failure calling /restore/next is still an init-phase error. + Err(e) => return Err(report_init_error(client, e).await), + }; + // `call()` only surfaces transport errors; a non-2xx status from RAPID still + // resolves to `Ok`. Treat a failed `/restore/next` as fatal rather than + // silently proceeding into the invocation loop with an un-restored VM. This + // still happens during init, so report it to /init/error. + let status = resp.status(); + if !status.is_success() { + let err: BoxError = format!("/restore/next returned a non-success status: {status}").into(); + return Err(report_init_error(client, err).await); + } + // Stale connections to RAPID won't survive the snapshot; rebuild the pool. + client.reset_pool(); + + // after_restore runs in registration order (FIFO): rebuild foundations + // before the dependents that need them. A failure here is reported to + // /restore/error and propagated so the runtime exits. + for resource in resources { + if let Err(e) = resource.after_restore().await { + return Err(report_restore_error(client, e).await); + } + } + + Ok(()) +} + +/// Reports an init-phase error (a `before_snapshot` hook or `/restore/next`) to +/// Lambda via `/init/error`, then returns the original error so the caller can +/// propagate it and exit the runtime. +async fn report_init_error(client: &Arc, err: BoxError) -> BoxError { + report_diagnostic(client, err, ReportKind::Init).await +} + +/// Reports an after-restore error to Lambda via `/restore/error`, then returns +/// the original error so the caller can propagate it and exit the runtime. +async fn report_restore_error(client: &Arc, err: BoxError) -> BoxError { + report_diagnostic(client, err, ReportKind::Restore).await +} + +enum ReportKind { + Init, + Restore, +} + +/// Posts a `Diagnostic` (built from the error) to the appropriate RAPID error +/// endpoint, then returns the **original** `BoxError` so the caller propagates it +/// with its concrete type and `source()` chain intact. Failures to build or send +/// the report are logged (never swallowed silently) but do not mask the error. +async fn report_diagnostic(client: &Arc, err: BoxError, kind: ReportKind) -> BoxError { + // Build the diagnostic from a reference so we can return `err` untouched. + let diagnostic = Diagnostic { + error_type: crate::diagnostic::type_name_of_val(&err), + error_message: err.to_string(), + }; + + let (endpoint, req) = match kind { + ReportKind::Init => ("/init/error", InitErrorRequest { diagnostic }.into_req()), + ReportKind::Restore => ("/restore/error", RestoreErrorRequest { diagnostic }.into_req()), + }; + + match req { + Ok(req) => { + if let Err(e) = client.call(req).await { + log_or_print!( + tracing: tracing::error!(error = ?e, "failed to report SnapStart error to {endpoint}"), + fallback: eprintln!("failed to report SnapStart error to {endpoint}: {e:?}") + ); + } + } + Err(e) => { + log_or_print!( + tracing: tracing::error!(error = ?e, "failed to build SnapStart {endpoint} request"), + fallback: eprintln!("failed to build SnapStart {endpoint} request: {e:?}") + ); + } + } + + err +} + #[allow(clippy::type_complexity)] fn wrap_handler<'a, F, EventPayload, Response, BufferedResponse, StreamingResponse, StreamItem, StreamError>( handler: F, @@ -407,6 +588,7 @@ fn wrap_handler<'a, F, EventPayload, Response, BufferedResponse, StreamingRespon StreamItem, StreamError, >, + ApiClient, > where F: Service, Response = Response>, @@ -532,10 +714,9 @@ mod endpoint_tests { use super::{incoming, wrap_handler}; use crate::{ requests::{EventCompletionRequest, EventErrorRequest, IntoRequest, NextEventRequest}, - Config, Diagnostic, Error, Runtime, + BoxFuture, Config, Diagnostic, Error, Runtime, }; use bytes::Bytes; - use futures::future::BoxFuture; use http::{HeaderValue, Method, Request, Response, StatusCode}; use http_body_util::{BodyExt, Full}; use httpmock::prelude::*; @@ -545,7 +726,7 @@ mod endpoint_tests { rt::{tokio::TokioIo, TokioExecutor}, server::conn::auto::Builder as ServerBuilder, }; - use lambda_runtime_api_client::Client; + use lambda_runtime_api_client::PooledClient as Client; use std::{ convert::Infallible, env, @@ -574,7 +755,7 @@ mod endpoint_tests { }); let base = server.base_url().parse().expect("Invalid mock server Uri"); - let client = Client::builder().with_endpoint(base).build()?; + let client = Client::builder().with_endpoint(base).build(); let req = NextEventRequest.into_req()?; let rsp = client.call(req).await.expect("Unable to send request"); @@ -607,7 +788,7 @@ mod endpoint_tests { }); let base = server.base_url().parse().expect("Invalid mock server Uri"); - let client = Client::builder().with_endpoint(base).build()?; + let client = Client::builder().with_endpoint(base).build(); let req = EventCompletionRequest::new("156cb537-e2d4-11e8-9b34-d36013741fb9", "{}"); let req = req.into_req()?; @@ -637,7 +818,7 @@ mod endpoint_tests { }); let base = server.base_url().parse().expect("Invalid mock server Uri"); - let client = Client::builder().with_endpoint(base).build()?; + let client = Client::builder().with_endpoint(base).build(); let req = EventErrorRequest { request_id: "156cb537-e2d4-11e8-9b34-d36013741fb9", @@ -673,7 +854,7 @@ mod endpoint_tests { }); let base = server.base_url().parse().expect("Invalid mock server Uri"); - let client = Client::builder().with_endpoint(base).build()?; + let client = Client::builder().with_endpoint(base).build(); async fn func(event: crate::LambdaEvent) -> Result { let (event, _) = event.into_parts(); @@ -708,6 +889,7 @@ mod endpoint_tests { config: Arc::new(config), service: wrap_handler(f, client), concurrency_limit: 1, + snapstart_resources: Vec::new(), }; let client = &runtime.client; let incoming = incoming(client).take(1); @@ -745,7 +927,7 @@ mod endpoint_tests { }); let base = server.base_url().parse().expect("Invalid mock server Uri"); - let client = Client::builder().with_endpoint(base).build()?; + let client = Client::builder().with_endpoint(base).build(); let f = crate::service_fn(func); @@ -763,6 +945,7 @@ mod endpoint_tests { config, service: wrap_handler(f, client), concurrency_limit: 1, + snapstart_resources: Vec::new(), }; let client = &runtime.client; let incoming = incoming(client).take(1); @@ -904,7 +1087,7 @@ mod endpoint_tests { } let handler = crate::service_fn(func); - let client = Arc::new(Client::builder().with_endpoint(base).build()?); + let client = Arc::new(Client::builder().with_endpoint(base).build()); let runtime = Runtime { client: client.clone(), config: Arc::new(Config { @@ -916,6 +1099,7 @@ mod endpoint_tests { }), service: wrap_handler(handler, client), concurrency_limit: 2, + snapstart_resources: Vec::new(), }; let res = tokio::time::timeout(Duration::from_secs(2), runtime.run_concurrent()).await; @@ -1032,7 +1216,7 @@ mod endpoint_tests { } let handler = crate::service_fn(test_handler); - let client = Arc::new(Client::builder().with_endpoint(base).build()?); + let client = Arc::new(Client::builder().with_endpoint(base).build()); // Add tracing layer to capture span fields use crate::layers::trace::TracingLayer; @@ -1052,6 +1236,7 @@ mod endpoint_tests { }), service, concurrency_limit: 3, + snapstart_resources: Vec::new(), }; let runtime_handle = tokio::spawn(async move { runtime.run_concurrent().await }); @@ -1108,4 +1293,272 @@ mod endpoint_tests { Ok(()) } + + /// Records the order in which `before_snapshot`/`after_restore` fire across + /// resources, so tests can assert LIFO/FIFO ordering. + struct OrderRecorder { + label: &'static str, + log: Arc>>, + } + + impl crate::SnapStartResource for OrderRecorder { + fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { + let log = self.log.clone(); + let label = self.label; + Box::pin(async move { + log.lock().unwrap().push(format!("before:{label}")); + Ok(()) + }) + } + + fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { + let log = self.log.clone(); + let label = self.label; + Box::pin(async move { + log.lock().unwrap().push(format!("after:{label}")); + Ok(()) + }) + } + } + + /// Which lifecycle phase a [`FailingResource`] should fail in. + #[derive(Clone, Copy)] + enum Phase { + BeforeSnapshot, + AfterRestore, + } + + /// A resource that returns an error from the selected phase (and is a no-op + /// in the other), used to exercise the error-reporting paths. + struct FailingResource { + phase: Phase, + } + + impl crate::SnapStartResource for FailingResource { + fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { + let fail = matches!(self.phase, Phase::BeforeSnapshot); + Box::pin(async move { + if fail { + Err("before_snapshot failed".into()) + } else { + Ok(()) + } + }) + } + + fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { + let fail = matches!(self.phase, Phase::AfterRestore); + Box::pin(async move { + if fail { + Err("after_restore failed".into()) + } else { + Ok(()) + } + }) + } + } + + #[tokio::test] + async fn test_snapstart_restore_lifecycle_calls_restore_next() -> Result<(), Error> { + let server = MockServer::start(); + + let restore_mock = server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(200).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let log = Arc::new(std::sync::Mutex::new(Vec::new())); + let resources: Vec> = vec![Arc::new(OrderRecorder { + label: "pool", + log: log.clone(), + })]; + + // Exercise the restore lifecycle directly (no env mutation, no races with + // other tests that call run()/run_concurrent()). + super::run_restore_lifecycle(&client, &resources).await?; + + restore_mock.assert_async().await; + let recorded = log.lock().unwrap().clone(); + assert_eq!(recorded, vec!["before:pool", "after:pool"]); + Ok(()) + } + + #[tokio::test] + async fn test_snapstart_resource_ordering_is_lifo_then_fifo() -> Result<(), Error> { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(200).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let log = Arc::new(std::sync::Mutex::new(Vec::new())); + // Registration order: credentials → pool → cache (foundations first). + let resources: Vec> = vec![ + Arc::new(OrderRecorder { + label: "credentials", + log: log.clone(), + }), + Arc::new(OrderRecorder { + label: "pool", + log: log.clone(), + }), + Arc::new(OrderRecorder { + label: "cache", + log: log.clone(), + }), + ]; + + super::run_restore_lifecycle(&client, &resources).await?; + + let recorded = log.lock().unwrap().clone(); + assert_eq!( + recorded, + vec![ + // before_snapshot: LIFO (reverse registration) + "before:cache", + "before:pool", + "before:credentials", + // after_restore: FIFO (registration order) + "after:credentials", + "after:pool", + "after:cache", + ] + ); + Ok(()) + } + + #[tokio::test] + async fn test_snapstart_restore_next_non_success_is_fatal() -> Result<(), Error> { + let server = MockServer::start(); + server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(500).body(""); + }); + let init_err_mock = server.mock(|when, then| { + when.method(POST).path("/2018-06-01/runtime/init/error"); + then.status(202).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let log = Arc::new(std::sync::Mutex::new(Vec::new())); + let resources: Vec> = vec![Arc::new(OrderRecorder { + label: "pool", + log: log.clone(), + })]; + + let result = super::run_restore_lifecycle(&client, &resources).await; + + // A non-2xx /restore/next must surface as an error rather than silently + // continuing into the invocation loop with an un-restored VM, and it must + // be reported to /init/error. + assert!(result.is_err(), "expected non-success /restore/next to be fatal"); + init_err_mock.assert_async().await; + // before_snapshot ran; after_restore must NOT run when the restore call fails. + let recorded = log.lock().unwrap().clone(); + assert_eq!(recorded, vec!["before:pool"]); + Ok(()) + } + + #[tokio::test] + async fn test_snapstart_lifecycle_is_noop_when_not_snapstart() -> Result<(), Error> { + // Only meaningful when the SnapStart env var is absent. Never mutate the + // global env (would race with parallel tests); skip if it's set. + if env::var("AWS_LAMBDA_INITIALIZATION_TYPE").is_ok() { + return Ok(()); + } + + let server = MockServer::start(); + let restore_mock = server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(200).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let log = Arc::new(std::sync::Mutex::new(Vec::new())); + let resources: Vec> = vec![Arc::new(OrderRecorder { + label: "pool", + log: log.clone(), + })]; + + // The env-gated entry point must do nothing when SnapStart is not enabled. + super::snapstart_lifecycle(&client, &resources).await?; + + assert_eq!( + restore_mock.calls(), + 0, + "/restore/next must not be called when not snap-start" + ); + assert!( + log.lock().unwrap().is_empty(), + "no hooks should fire when not snap-start" + ); + Ok(()) + } + + #[tokio::test] + async fn test_before_snapshot_failure_reports_init_error() -> Result<(), Error> { + let server = MockServer::start(); + let init_err_mock = server.mock(|when, then| { + when.method(POST).path("/2018-06-01/runtime/init/error"); + then.status(202).body(""); + }); + let restore_mock = server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(200).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let resources: Vec> = vec![Arc::new(FailingResource { + phase: Phase::BeforeSnapshot, + })]; + + let result = super::run_restore_lifecycle(&client, &resources).await; + + assert!(result.is_err(), "before_snapshot failure must propagate"); + init_err_mock.assert_async().await; + // /restore/next must NOT be reached when before_snapshot fails. + assert_eq!(restore_mock.calls(), 0); + Ok(()) + } + + #[tokio::test] + async fn test_after_restore_failure_reports_restore_error() -> Result<(), Error> { + let server = MockServer::start(); + let restore_mock = server.mock(|when, then| { + when.method(GET).path("/2018-06-01/runtime/restore/next"); + then.status(200).body(""); + }); + let restore_err_mock = server.mock(|when, then| { + when.method(POST).path("/2018-06-01/runtime/restore/error"); + then.status(202).body(""); + }); + + let base = server.base_url().parse().expect("Invalid mock server Uri"); + let client = Arc::new(Client::builder().with_endpoint(base).build()); + + let resources: Vec> = vec![Arc::new(FailingResource { + phase: Phase::AfterRestore, + })]; + + let result = super::run_restore_lifecycle(&client, &resources).await; + + // Now that the runtime no longer calls process::exit, the after_restore + // failure path returns an error and is unit-testable. + assert!(result.is_err(), "after_restore failure must propagate"); + restore_mock.assert_async().await; + restore_err_mock.assert_async().await; + Ok(()) + } } diff --git a/lambda-runtime/src/snapstart.rs b/lambda-runtime/src/snapstart.rs new file mode 100644 index 000000000..c9d81b11d --- /dev/null +++ b/lambda-runtime/src/snapstart.rs @@ -0,0 +1,104 @@ +//! SnapStart lifecycle support. +//! +//! AWS Lambda SnapStart reduces cold start latency by taking a Firecracker +//! microVM snapshot of the initialized execution environment and restoring from +//! it on subsequent invocations. Resources initialized before the snapshot may +//! become invalid after restore: TCP connections are dropped, credentials may +//! expire, and unique values generated during init would be shared across every +//! restored instance. +//! +//! Implement [`SnapStartResource`] on types that hold such resources and +//! register them with +//! [`Runtime::register_snapstart_resource`](crate::Runtime::register_snapstart_resource). +//! The runtime invokes the registered hooks around the snapshot/restore boundary +//! when SnapStart is enabled (`AWS_LAMBDA_INITIALIZATION_TYPE == "snap-start"`) +//! and does nothing otherwise. +//! +//! # Ordering +//! +//! Resources are orchestrated with stack (LIFO/FIFO) ordering, mirroring how +//! Go's `defer`, middleware unwinding, and Rust destructors work: +//! +//! - [`before_snapshot`](SnapStartResource::before_snapshot) runs in **reverse** +//! registration order (LIFO) — dependents tear down before their foundations. +//! - [`after_restore`](SnapStartResource::after_restore) runs in **registration** +//! order (FIFO) — foundations rebuild before the dependents that need them. +//! +//! The rule for users is one line: **register foundations first** (e.g. +//! credentials before the pool that uses them, the pool before the cache that +//! uses it). Teardown and rebuild then both happen in the correct relative order +//! automatically: +//! +//! ```text +//! register: credentials → pool → cache +//! before_snapshot: cache → pool → credentials (reverse) +//! after_restore: credentials → pool → cache (forward) +//! ``` +//! +//! # Example +//! +//! ```no_run +//! use lambda_runtime::{BoxFuture, Error, SnapStartResource}; +//! +//! struct DbPool { +//! // e.g. a connection pool handle +//! } +//! +//! impl SnapStartResource for DbPool { +//! fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { +//! Box::pin(async move { +//! // release connections before the snapshot +//! Ok(()) +//! }) +//! } +//! +//! fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { +//! Box::pin(async move { +//! // re-establish connections after restore +//! Ok(()) +//! }) +//! } +//! } +//! ``` + +use crate::{BoxFuture, Error}; + +/// A resource that needs custom logic around the SnapStart snapshot/restore +/// boundary. +/// +/// Implement this trait on types holding snapshot-sensitive state (connection +/// pools, credentials, cached DNS, etc.) and register them with +/// [`Runtime::register_snapstart_resource`](crate::Runtime::register_snapstart_resource). +/// +/// Both methods default to a no-op, so implementors only override the hook they +/// need. See the [module docs](crate::snapstart) for ordering semantics. +/// +/// Each method returns a [`BoxFuture`] (rather than being an `async fn`) so the +/// trait stays object-safe and can be stored as `dyn SnapStartResource` without +/// pulling in the `async-trait` crate. Implementations wrap their body in +/// `Box::pin(async move { .. })`. +pub trait SnapStartResource: Send + Sync { + /// Called before the VM snapshot is taken. + /// + /// Use this to release resources that will not survive the snapshot, such as + /// open connections or buffered data. + /// + /// Runs in reverse registration order (LIFO) across all registered + /// resources. + fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } + + /// Called after the VM is restored from a snapshot. + /// + /// Use this to re-establish resources released in + /// [`before_snapshot`](Self::before_snapshot), refresh credentials, or + /// regenerate values that must be unique per execution environment. + /// + /// Runs in registration order (FIFO) across all registered resources. If + /// this returns an error, the runtime reports it to Lambda via + /// `/restore/error` and the runtime exits. + fn after_restore(&self) -> BoxFuture<'_, Result<(), Error>> { + Box::pin(async { Ok(()) }) + } +}