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(()) }) + } +}