feat: add SnapStart support via SnapStartResource trait#1150
Merged
Conversation
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
darklight3it
approved these changes
Jul 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds AWS Lambda SnapStart support to the Rust runtime. SnapStart reduces cold-start latency by snapshotting the initialized execution environment and restoring from it, but resources created during init (connections, credentials, unique values) may be invalid after restore. This PR gives functions a place to release/rebuild them, and makes the runtime itself restore-safe.
What's added
SnapStartResourcetrait (lambda-runtime/src/snapstart.rs) — a CRaC-style trait for custom snapshot/restore logic:BoxFuture(notasync fn) so the trait stays object-safe asdyn SnapStartResourcewith noasync-traitdependency.Runtime::register_snapstart_resource(Arc<...>).Lifecycle the runtime drives (from both
run()andrun_concurrent()):Register foundations first; teardown runs reverse (LIFO), rebuild forward (FIFO). When SnapStart is not enabled (
AWS_LAMBDA_INITIALIZATION_TYPE != "snap-start"), hooks are never invoked.Error handling — a
before_snapshot//restore/nextfailure is reported toPOST /init/error; anafter_restorefailure toPOST /restore/error. Both propagate out via the normalResultpath (noprocess::exit, so graceful-shutdown handlers still run).Non-breaking connection-pool reset — new
PooledClient(lock-freeOnceLock-basedreset_pool()) and aRuntimeApiClienttrait in the api-client crate. The service layer takes a defaulted genericC = Clientbounded by the trait and usesPooledClientinternally. The oldClientis frozen at its 1.0.3 shape and superseded.cargo-semver-checksconfirms all three crates stay non-breaking (api-client is a 1.1.0 minor bump).Changes
PooledClient+RuntimeApiClienttrait +PooledClient::builder; supersedeClient/build; bump to 1.1.0SnapStartResourcetrait, Init/Restore error requests,Runtime::is_snapstart()andregister_snapstart_resource(); run the restore lifecycle; re-exportSnapStartResourceandBoxFutureruntime()/runtime_concurrent()andstreaming_runtime()helpers; re-exportSnapStartResourceandBoxFuturePooledClientbasic-snapstartandhttp-snapstart(with Dockerfiles onprovided:al2023)Testing
Unit tests cover request serialization, the full restore lifecycle, LIFO/FIFO ordering, no-op when not SnapStart, and the before/after/restore-next failure paths.