feat: add SnapStart support via SnapStartResource trait

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"] }

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" ]

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 <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/basic-snapstart:latest
+docker push <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/basic-snapstart:latest
+
+aws lambda create-function \
+  --function-name basic-snapstart \
+  --package-type Image \
+  --code ImageUri=<ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/basic-snapstart:latest \
+  --role <YOUR_EXECUTION_ROLE_ARN> \
+  --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`.

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

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"] }

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" ]

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 <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/http-snapstart:latest
+docker push <ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/http-snapstart:latest
+
+aws lambda create-function \
+  --function-name http-snapstart \
+  --package-type Image \
+  --code ImageUri=<ACCOUNT>.dkr.ecr.<REGION>.amazonaws.com/http-snapstart:latest \
+  --role <YOUR_EXECUTION_ROLE_ARN> \
+  --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`.

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<RwLock<DbPool>>`, so draining before the
+/// snapshot and reconnecting after restore are visible to invocations.
+struct PoolResource(Arc<RwLock<DbPool>>);
+
+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::<Response<Body>, Error>(format!("Hello, {name}!").into_response().await)
+        }
+    });
+
+    lambda_http::runtime(handler)
+        .register_snapstart_resource(Arc::new(PoolResource(pool.clone())))
+        .run()
+        .await?;
+    Ok(())
+}

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 = [