feat: support Lambda Managed Instances (draft)

[alessandrobologna]

📬 Issue #: #1065

✍️ Description of changes:

Support Lambda Managed Instances / concurrent runtime (refs #1065)

Note

I originally built this to experiment with Lambda Managed Instances in a test workload and thought it might be a useful starting point for upstream support. Parts of the implementation were drafted with the help of an AI assistant and have only been exercised in my own workloads and the existing test suite so far, so please treat this as a starting point for discussion rather than a final design. A super simple demo/benchmark of this implementation is at https://github.com/alessandrobologna/rust-lmi-demo

This draft PR implements a concurrent runtime mode for Lambda Managed Instances, plus the API changes required in lambda-runtime and lambda-http. It preserves existing single-invocation behavior for classic Lambda. The implementation is based on the AWS guidance for building runtimes that support Lambda Managed Instances, as described in Building custom runtimes for Lambda Managed Instances.

Architecture (high level)

The runtime chooses between a classic "one invocation at a time" loop and a concurrent mode controlled by AWS_LAMBDA_MAX_CONCURRENCY:

• Sequential mode (default): a single /next long-poll loop fetches one event at a time and invokes the handler synchronously, preserving existing behavior.
• Concurrent mode (AWS_LAMBDA_MAX_CONCURRENCY >= 2): the runtime spawns AWS_LAMBDA_MAX_CONCURRENCY worker tasks. Each worker runs its own /next long-poll loop and processes invocations sequentially within that task, so overall concurrency is bounded by the number of workers. If a worker crashes, the runtime continues operating with the remaining workers and only terminates if it can't keep at least 1 worker alive.

sequenceDiagram
    participant Lambda as Lambda orchestrator
    participant Runtime as lambda-runtime
    participant Worker as Runtime worker task
    participant API as Runtime API
    participant Handler as User handler

    Lambda->>Runtime: Start process (env AWS_LAMBDA_MAX_CONCURRENCY = N)
    Runtime->>Worker: Spawn N worker tasks
    loop each worker (identical)
        loop each invocation
            Worker->>API: GET /runtime/invocation/next
            API-->>Worker: Invocation event + context headers
            Worker->>Handler: Call handler(event, context)
            Handler-->>Worker: Result or error
            Worker->>API: POST /response or /error
        end
    end

Loading

Breaking changes & compatibility

• Existing entrypoints (lambda_runtime::run, lambda_http::run, and lambda_http::run_with_streaming_response) keep their original signatures and sequential behavior. Handlers that compiled against the current release should continue to compile unchanged.
• New concurrent entrypoints are added:
  • lambda_runtime::run_concurrent
  • lambda_http::run_concurrent
  • lambda_http::run_with_streaming_response_concurrent
    These require handler services to implement Clone + Send + 'static (with responses/stream bodies Send/Sync + 'static) so they can be safely cloned and driven by the concurrent runtime.
• AWS_LAMBDA_MAX_CONCURRENCY is read directly by the runtime to decide concurrency and size the HTTP pool, without adding new public fields to Config (avoids semver breakage for struct literals).
• In concurrent mode (when the new entrypoints are used and AWS_LAMBDA_MAX_CONCURRENCY > 1), the runtime no longer sets _X_AMZN_TRACE_ID in the process environment. The per-invocation X-Ray trace ID is available via Context::xray_trace_id and tracing spans instead. Sequential mode behavior is unchanged.
• For the concurrent entrypoints, the behavior when AWS_LAMBDA_MAX_CONCURRENCY is set changes from "always sequential per environment" to "per-environment concurrency up to that value". Code that continues to call the existing run functions will remain strictly sequential even if the env var is set.

In other words: the earlier versions of this branch tightened the bounds on the existing run functions, but after maintainer feedback those entrypoints are left as-is and concurrency is opt-in via the new *_concurrent APIs.

Below is a concise summary of the changes (unfortunately many) by area.

Runtime & Config (lambda-runtime)

• Read AWS_LAMBDA_MAX_CONCURRENCY directly inside the runtime to decide whether concurrent mode should be enabled and to size the HTTP client pool, without extending Config.
• Runtime::new now sizes the lambda_runtime_api_client HTTP pool from max_concurrency so the number of idle connections matches expected concurrency.
• Runtime::run remains the original sequential /next loop via run_with_incoming, preserving existing behavior.
• New Runtime::run_concurrent spawns AWS_LAMBDA_MAX_CONCURRENCY worker tasks, each running its own /next loop. Sequential and concurrent modes share the same invocation processing helper, so there are not two separate event loop implementations to maintain. When Config::is_concurrent() is false, it falls back to the same sequential loop as Runtime::run.
• Worker failures are isolated: if a worker exits/panics, the runtime logs the failure and continues with fewer workers. The runtime terminates only if it can't keep at least 1 worker alive (i.e., all workers have exited).
• Unexpected clean worker exits are treated as infrastructure errors so the runtime returns an error once all workers are gone.
• Handler errors (user code returning Err) are reported to Lambda via /invocation/{id}/error and do not terminate the runtime.
• The existing graceful shutdown handler stays aligned with upstream. We do not attempt to preempt/coordinate in-flight invocations on SIGTERM (the expectation is that Lambda only shuts down environments when they are idle; local emulators may behave differently).
• Internal layers (api_client, api_response, trace) now implement/derive Clone so they can be composed into a cloneable service stack for the concurrent entrypoints.
• Context::new is more robust when lambda-runtime-client-context / lambda-runtime-cognito-identity headers are present but empty (treated as None instead of failing JSON parse).
• Added a regression test (concurrent_worker_crash_does_not_stop_other_workers) to verify worker isolation behavior.

HTTP & streaming (lambda-http, lambda-runtime-api-client)

• lambda_runtime_api_client::Client:
  • Added with_pool_size(usize) on the builder and threads a pool_size: Option<usize> into the Hyper client to set pool_max_idle_per_host.
  • Still works as before when pool_size is not provided.
• lambda_http::run and run_with_streaming_response keep their existing signatures and sequential behavior, delegating to lambda_runtime::run.
• New lambda_http::run_concurrent and lambda_http::run_with_streaming_response_concurrent wrap the same handler types but require them to be Clone + Send + 'static (with response/stream bounds aligned to lambda_runtime::run_concurrent) so they can be driven by the concurrent runtime.
• HTTP adapters (Adapter, StreamAdapter) are now Clone when the inner service is Clone, and the streaming path uses BoxCloneService internally for the concurrent entrypoint so the composed service stack can be cloned.
  • Note: run_with_streaming_response_concurrent is exported from lambda_http, similarly to run_with_streaming_response.

Tooling & examples

• Makefile and scripts/test-rie.sh:
  • Add RIE_MAX_CONCURRENCY and a test-rie-lmi target that runs RIE with AWS_LAMBDA_MAX_CONCURRENCY set, making it easy to exercise managed-instances behavior locally.
• examples/basic-lambda-concurrent:
  • New minimal example that calls lambda_runtime::run_concurrent so the concurrent code path can be exercised under RIE/LMI.
• examples/basic-lambda/src/main.rs:
  • Wraps lambda_runtime::run(func).await in an if let Err(err) block to log and propagate runtime errors when testing under RIE.

---

Validation

On this branch, I ran:

• cargo +nightly fmt --all -- --check
• cargo clippy --workspace --all-features -- -D warnings
• cargo +stable test --all-features -p lambda_runtime_api_client -p lambda_runtime -p lambda_http
• make test-rie EXAMPLE=basic-lambda
• make test-rie-lmi EXAMPLE=basic-lambda-concurrent (sets AWS_LAMBDA_MAX_CONCURRENCY inside the RIE container)

---

If maintainers prefer, this could be split into smaller PRs (e.g., builder/Config prep, handler Clone changes, and finally the concurrent runtime), but this branch shows the full "end-to-end" implementation so that it can be tested with Lambda Managed Instances.

---

🔏 By submitting this pull request

• I confirm that I've ran cargo +nightly fmt --all -- --check.
• I confirm that I've ran cargo clippy --workspace --all-features -- -D warnings.
• I confirm that I've made a best effort attempt to update all relevant documentation.
• I confirm that my contribution is made under the terms of the Apache 2.0 license.

[alessandrobologna]

Thanks @bnusunny for the approval!
Currently CI is failing on the semver check because Config has a new public field. I have a small follow‑up that keeps Config unchanged by moving the concurrency limit into Runtime (derived from AWS_LAMBDA_MAX_CONCURRENCY), which avoids a major version bump.

Patch: alessandrobologna@a8af68a

If you’re OK with that approach, I can push it directly to this PR branch so the checks go green.

[bnusunny]

Yes, please.

[alessandrobologna]

Ok, I've pushed the update to keep concurrency settings internal and to avoid the semver bump (so no Config public field needed). I also rebased the branch on current main.

[jlizen]

Overall implementation looks great!

To summarize my feedback:

• I think naming things based on _concurrent is overly limiting limiting if we expect future features to require Clone, i'd suggest _extended instead maybe? And we should probably nudge callers to use _extended, _concurrent, _maybe_concurrent, etc, unless they specifically need !Clone bounds?
• The non-*_concurrent should have doc comments warning that it ignores concurrency variables and probably an error log or panic if the caller tries to invoke them with concurrency ENV set
• I'd like to make it easier to trace per-worker issues by injecting task IDs into both worker loop spans and worker loop results
• I'd prefer to preserve all errors across workers rather than only the first. Also it seems like we don't really need the full-batteries FuturesUnordered instead of join_all?
• I don't like the Oncelock that allows showing a single trace message, it seems like it could be simplified?

[jlizen]

(Question, not sure): If the fallback case if AWS_LAMBDA_MAX_CONCURRENCY if unset is just lambda_runtime::run(), should this be our new suggested entrypoint, with fn run() marked as deprecated? I guess the main issue would be the Clone bounds are more restrictive, so maybe deprecation is a bit aggressive, but we could add that in the run() doc comments?

The regular run() should definitely also add a warning that it WON'T respect AWS_LAMBDA_MAX_CONCURRENCY.

Context: My suspicion is, most consumers don't mind the Clone bound, and the 'principal of least surprise' would be for it to respect the ENV. So we probably should be pointing people to prefer this method where possible.

I would suggest at minimum the following:

• probably this should be named run_maybe_concurrent? Or maybe run_extended to avoid a combinatorial explosion of new entrypoint names as featureset grows that needs Clone?
• We should have a guard clause in lambda_runtime::run() that throws either either crashes the lambda runtime (preferred), or emits error logs (more conservative), if AWS_LAMBDA_MAX_CONCURRENCY is set, since that is clearly a misconfiguration that would prove confusing to debug.

[alessandrobologna]

Thanks! I added a doc warning on run() and a fail‑fast guard when AWS_LAMBDA_MAX_CONCURRENCY is set. I kept the *_concurrent naming for now to limit "code churn", and also because _extended seems to me a little vague. But I am open to change it if there's consensus on renaming all of them to what you suggested.
I have also added a note to the run() docstring to make it more clean that run_concurrent is preferred if the handler can satisfy Clone

[jlizen]

_concurrent is probably ok... ultimately future features that require Clone will likely involve concurrency, so the naming probably makes sense...

[jlizen]

In the current impl, None should essentially never be reachable, right? Since even if we have no task-specific errors, we instead create an error out of a clean exit?

If the desired API is to return Ok(()) on all workers cleanly exiting, then we could use the enum mentioned above to gather that info, and emit an error log line if all exited cleanly, but then drop the error.

[alessandrobologna]

Yes, I am now using the WorkerError and aggregate everything.

For the other point, I currently treat all workers exiting (clean or not) as an error because then the runtime has no workers left and can't serve any more invocations. Returning Ok(()) here would make the runtime exit quietly and hide the failure. But I can change it if you prefer.

[jlizen]

Would it be worth instrumenting this future with a span containing the taskId? Seems pretty useful for debug

[alessandrobologna]

Yes, I have added a span on the worker loop with the Tokio task id and instrumented the per-invocation processing with it.

[jlizen]

The OnceLock + get_or_init every invocation seems like unnecessary overhead / complexity. Why can't we just display this message once, at the start of run_concurrent?

[alessandrobologna]

Agree.

[jlizen]

Thanks for all the tweaks! New error handling, removal of XRAY_LOGGED, into_stream_service_cloneable(), etc, look great!

There are a few nits that are totally up to you.

Main load bearing feedback has to do with the usage of Box::pin() for the worker future. If it it turns out to be tricky to remove, that's fine / just say so, we can punt that as it is something we can change later.

I also think having a debug log indicating when sequential fallback is used in run_concurrent, would be pretty handy (ref)

[jlizen]

Why the extra Box::pin() here? The tokio::task::Id is Unpin? Is there away to avoid this extra allocation / indirection?

Can we just get away with using an unnameable type like we did previously?

[alessandrobologna]

Good point. I used Box::pin originally just to erase the unnameable async move type for FuturesUnordered, but, as you pointed out, that adds an allocation. I switched to spawning a task that returns (task_id, Result<(), BoxError>), so the set holds a concrete JoinHandle and there’s no Box::pin. For panics I use JoinError::id() a little below.

[jlizen]

Thanks for the detailed comment on this passage

[jlizen]

nit: i would probably just put this string slice inline since we only use it once

[jlizen]

nit: this is a bit odd, i would have expected something like [task_1, task_2, ..]

Might be easier with write!(f, " {}", clean.join(", ")) (and more efficient anyway, fewer write calls, though probably we aren't perf sensitive in this code path)

[alessandrobologna]

Agreed. I switched the aggregate error formatting to just be valid JSON so it's easy to query in Log Insights/analytics tools (if the caller logs/prints it). The per‑worker logs are unchanged.

[jlizen]

Same nit.

Also, I wonder, would it be worth something like: [{"id": "foo", "err": "foo" },..] to make it easier to query against in log insights or similar? Even if you need to regex match only the [..] part of the message, having some signpost that there is an error message is pretty nice.

[alessandrobologna]

Yes, same reasoning here. I ended up switching the aggregate error formatting to valid JSON to keep it simple and queryable, rather than bespoke formatting

[jlizen]

So perf impact here for non-concurrent path is, exactly one environment variable check? That seems fine, good tradeoff for footgun avoidance.

[jlizen]

nit: concurrent_worker_loop is a bit verbose to see in all inner log lines, i'd suggest worker probably

[alessandrobologna]

I have changed it just "worker"

[jlizen]

Any reason not to instrument the hyper future as well? I seem to recall hyper has some sort of tracing feature support that might benefit from task id context?

[jlizen]

Side question - do we really need to clone the client here? Could we just pass it by reference and save the atomic increment? call() just requires &self right?

[alessandrobologna]

Agreed, I wrapped the /next request future in the worker span so it inherits the task_id context. I didn't enable hyper's tracing feature (to keep deps minimal, and it seems that tracing in 1.0 is still unstable), but this gives us span context for the request.

Also agreed, I have removed the cloning and passing &ApiClient into next_event_future now.

[jlizen]

The CI failure related to MSRV seems unrelated... I'll fix separately

Ref: #1077