Skip to content

feat: add SnapStart support via SnapStartResource trait#1150

Merged
bnusunny merged 1 commit into
mainfrom
feat/snapstart-resource-trait
Jul 9, 2026
Merged

feat: add SnapStart support via SnapStartResource trait#1150
bnusunny merged 1 commit into
mainfrom
feat/snapstart-resource-trait

Conversation

@bnusunny

@bnusunny bnusunny commented Jul 9, 2026

Copy link
Copy Markdown
Collaborator

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

SnapStartResource trait (lambda-runtime/src/snapstart.rs) — a CRaC-style trait for custom snapshot/restore logic:

pub trait SnapStartResource: Send + Sync {
    fn before_snapshot(&self) -> BoxFuture<'_, Result<(), Error>> { /* no-op default */ }
    fn after_restore(&self)   -> BoxFuture<'_, Result<(), Error>> { /* no-op default */ }
}
  • Returns BoxFuture (not async fn) so the trait stays object-safe as dyn SnapStartResource with no async-trait dependency.
  • Both methods default to no-ops; register instances via Runtime::register_snapstart_resource(Arc<...>).

Lifecycle the runtime drives (from both run() and run_concurrent()):

before_snapshot hooks (LIFO) -> GET /restore/next (blocks until restore)
  -> rebuild RAPID connection pool -> after_restore hooks (FIFO)

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/next failure is reported to POST /init/error; an after_restore failure to POST /restore/error. Both propagate out via the normal Result path (no process::exit, so graceful-shutdown handlers still run).

Non-breaking connection-pool reset — new PooledClient (lock-free OnceLock-based reset_pool()) and a RuntimeApiClient trait in the api-client crate. The 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 superseded. cargo-semver-checks confirms all three crates stay non-breaking (api-client is a 1.1.0 minor bump).

Changes

  • lambda-runtime-api-client: add PooledClient + RuntimeApiClient trait + PooledClient::builder; supersede 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; 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 (with Dockerfiles on provided: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.

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
@bnusunny bnusunny requested review from Copilot, darklight3it and jlizen and removed request for Copilot July 9, 2026 07:14
@bnusunny bnusunny merged commit 094dbf9 into main Jul 9, 2026
32 checks passed
@bnusunny bnusunny deleted the feat/snapstart-resource-trait branch July 9, 2026 14:28
This was referenced Jul 9, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants