Support Lambda Managed Instances

[bnusunny]

AWS just launched Lambda Managed Instances What's New. It supports multiple concurrent requests on a single execution environment. We should support it to get the benefit to the Rust Runtime users.

To support multi-concurrency, we can spawn a new async task for each invoke. Both lambda-runtime and lambda-http crates need new run methods because the tower service instances need implement Clone trait. Adding Clone trait to the exiting run methods will be a breaking change.

X-ray trace header propagation should not use the environment variable. We need use a thread-safe way to store the x-tray trace header for each request.

[jlizen]

Re:

we can spawn a new async task for each invoke

(This is assuming this feature is only going to be implemented for tokio, initially).

Just as an implementation suggestion, we probably should make sure that we aren't dispatching work from the tokio block_on thread (aka non-worker thread), but rather from a worker thread (ie inside a task).

We don't control whether the caller awaits the lambda_runtime future on a worker thread or not, so it seems prudent to always spawn a task to drive the receive loop.

This is a performance optimization to avoid ping ponging the socket between the I/O driver on a worker thread, and the non-worker thread.

[jlizen]

Another thing to think about would be, how does this fit into turnkey graceful shutdown handling. The simple answer is, just fire the graceful shutdown once, globally, but I imagine the nicer thing might be some signaling to allow each accept loop to clean up its own state.

[jlizen]

new async task for each invoke

Correction (discussed offline): new task per event loop based on max concurrency, not per invoke

[garysassano]

Thanks for opening this; now I understand why my containerized function that uses the Lambda Rust Runtime breaks as soon as I attach a Capacity Provider to it.

I also noticed that AWS recently updated the Lambda docs with a section on building custom runtimes for Lambda Managed Instances, which helps clarify how custom runtimes need to be implemented in that setup..

[alessandrobologna]

I have pushed a (very) draft PR for review to implement concurrent requests and published a super simple demo/benchmark comparing a tiny 128MB rust function with a 2GB one on managed instances. Currently the lambda is just a hello world so it's not very useful as a real world example, but it's a start...

[bnusunny]

@alessandrobologna Thanks for the PR. I took a very brief look. One thing I would recommend to change: don't change the existing run methods in the runtimes, create additional run_concurrent() methods instead. This is because adding Clone + Send + 'static to run methods is a big breaking change. Most of the existing functions won't even compile without adding those traits.

[alessandrobologna]

Thanks @bnusunny, that makes sense. I did (just briefly) consider either separate crates (maybe lambda_runtime_concurrent and lambda_http_concurrent) or a concurrent::run namespace, but I agree that your suggestion is the most practical.
I've updated the PR to add new *_concurrent entry points (lambda_runtime::run_concurrent, lambda_http::run_concurrent, and lambda_http::run_with_streaming_response_concurrent) and left the existing run functions (and their signatures/behavior) untouched. Concurrency is now opt‑in via the new APIs.

Thanks for the PR!

I am curious if it is at all possible to relax the trait requirements here.

[alessandrobologna]

Hi, (I guess the user has been deleted?), I couldn't really think of a way to make the current approach work without Clone + Send + 'static.
In the concurrent code we spawn tasks per invocation, with tokio::spawn(async move { ... }), so the handler (and anything it captures) needs to be Send + 'static on Tokio’s multithread runtime.
And it also needs to be Clone because Tower Service::call takes &mut self. In concurrent mode, once you have multiple invocations in flight, a single service instance can’t be called concurrently from multiple tasks.

But it's also very possible that there are other possible approaches that I am not thinking of.

[bnusunny]

@alessandrobologna Thanks for the updates. Two more comments from me:

1. Ideally we should create AWS_LAMBDA_MAX_CONCURRENCY numbers of tasks and within each task to run an event loop pulling from Runtime API. The event loop could reuse the existing one with small updates. We don't really need two event loop implementations in the runtime.
2. Regarding the graceful shutdown handler, Lambda will only shutdown the sandboxes when they are idle. The existing graceful shutdown handler should work for LMI as well.

[alessandrobologna]

@bnusunny very good points, thank you.

1. You are right, in the implementation as it is now, the concurrent code is reimplementing some of the same logic that already exists in the sequential loop, so it's a bit duplicative. Spawning N (the max concurrency) worker tasks and running one /next loop per worker is a cleaner design, and should let us reuse more of the existing event-loop code. I'm going to refactor in that direction.

2. I was thinking that in case of running with RIE or other local emulator we could have situations where there are in flight invocations, but you are right, it's probably over-engineered to handle scenarios that are unlikely to happen in the real world.

[alessandrobologna]

Hi,
I’ve updated the draft PR #1067, now it should be better aligned with the suggested new design:

1. The concurrent runtime now follows the recommended LMI model: it spawns AWS_LAMBDA_MAX_CONCURRENCY worker tasks, and each worker runs its own /next loop (no task-per-invocation approach).
2. I refactored the runtime so sequential and concurrent paths share the same invocation processing logic, rather than maintaining two separate event loop implementations.
3. I kept the graceful shutdown handler aligned with upstream (run the hook then exit), and removed the extra "drain in-flight invocations" coordination. In concurrent mode, the runtime still broadcasts shutdown on first worker infra error/panic so other workers stop polling and we join them before returning.
4. Added a small basic-lambda-concurrent example and the test-rie-lmi make target to make it easier to test locally.

Overall there are less changes now compared to main, which is good I guess. If you have a chance to give it another look, please let me know if it's going in the right direction.

[bnusunny]

@alessandrobologna Thanks for the update! It looks great. Only one problem, LMI handles worker failures differently: If a runtime worker crashes, the execution environment continues operating with the remaining healthy workers.

We don't really need the shutdown channel to broadcast the shutdown signal. With It, a single worker failure shuts down the whole runtime. We should just remove that to maintain the error isolation of workers.

We probably need a few tests to cover this behavior.