runtime: improve concurrent worker supervision

alessandrobologna · alessandrobologna · commit 3ddc948242f2 · 2025-12-25T02:38:34.000-05:00
- aggregate worker exits with task IDs and surface clean exits in errors
- add per-worker spans and log unexpected exits at error level
- guard run() when AWS_LAMBDA_MAX_CONCURRENCY is set; docs prefer run_concurrent
- make streaming adapter cloneable for concurrent HTTP path
- document pool size guidance and LMI test target; log example via tracing
diff --git a/README.md b/README.md
@@ -405,6 +405,12 @@ make test-rie EXAMPLE=basic-sqs
 make test-rie EXAMPLE=http-basic-lambda
 ```
 
+To test Lambda Managed Instances (concurrent polling), use:
+
+```bash
+make test-rie-lmi EXAMPLE=basic-lambda-concurrent
+```
+
 This command will:
 1. Build a Docker image with Rust toolchain and RIE
 2. Compile the specified example inside the Linux container
diff --git a/examples/basic-lambda-concurrent/src/main.rs b/examples/basic-lambda-concurrent/src/main.rs
@@ -29,7 +29,7 @@ async fn main() -> Result<(), Error> {
 
     let func = service_fn(my_handler);
     if let Err(err) = lambda_runtime::run_concurrent(func).await {
-        eprintln!("run error: {:?}", err);
+        tracing::error!(error = %err, "run error");
         return Err(err);
     }
     Ok(())
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -207,6 +207,12 @@ where
 ///
 /// This takes care of transforming the LambdaEvent into a [`Request`] and then
 /// converting the result into a `LambdaResponse`.
+///
+/// # Managed concurrency
+/// If `AWS_LAMBDA_MAX_CONCURRENCY` is set, this function returns an error because
+/// it does not enable concurrent polling. If your handler can satisfy `Clone`,
+/// prefer [`run_concurrent`], which honors managed concurrency and falls back to
+/// sequential behavior when unset.
 pub async fn run<'a, R, S, E>(handler: S) -> Result<(), Error>
 where
     S: Service<Request, Response = R, Error = E>,
diff --git a/lambda-http/src/streaming.rs b/lambda-http/src/streaming.rs
@@ -12,7 +12,7 @@ pub use http::{self, Response};
 use http_body::Body;
 use lambda_runtime::{
     tower::{
-        util::{BoxCloneService, MapRequest, MapResponse},
+        util::{MapRequest, MapResponse},
         ServiceBuilder, ServiceExt,
     },
     Diagnostic,
@@ -29,6 +29,18 @@ pub struct StreamAdapter<'a, S, B> {
     _phantom_data: PhantomData<&'a B>,
 }
 
+impl<'a, S, B> Clone for StreamAdapter<'a, S, B>
+where
+    S: Clone,
+{
+    fn clone(&self) -> Self {
+        Self {
+            service: self.service.clone(),
+            _phantom_data: PhantomData,
+        }
+    }
+}
+
 impl<'a, S, B, E> From<S> for StreamAdapter<'a, S, B>
 where
     S: Service<Request, Response = Response<B>, Error = E>,
@@ -104,10 +116,12 @@ where
 
 /// Builds a streaming-aware Tower service from a `Service<Request>` that can be
 /// cloned and sent across tasks. This is used by the concurrent HTTP entrypoint.
+type EventToRequest = fn(LambdaEvent<LambdaRequest>) -> Request;
+
 #[allow(clippy::type_complexity)]
-fn into_stream_service_boxed<S, B, E>(
+fn into_stream_service_cloneable<S, B, E>(
     handler: S,
-) -> BoxCloneService<LambdaEvent<LambdaRequest>, StreamResponse<BodyStream<B>>, E>
+) -> MapResponse<MapRequest<S, EventToRequest>, impl FnOnce(Response<B>) -> StreamResponse<BodyStream<B>> + Clone>
 where
     S: Service<Request, Response = Response<B>, Error = E> + Clone + Send + 'static,
     S::Future: Send + 'static,
@@ -116,12 +130,10 @@ where
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
 {
-    let svc = ServiceBuilder::new()
-        .map_request(event_to_request as fn(LambdaEvent<LambdaRequest>) -> Request)
+    ServiceBuilder::new()
+        .map_request(event_to_request as EventToRequest)
         .service(handler)
-        .map_response(into_stream_response);
-
-    BoxCloneService::new(svc)
+        .map_response(into_stream_response)
 }
 
 /// Converts an `http::Response<B>` into a streaming Lambda response.
@@ -163,6 +175,11 @@ fn event_to_request(req: LambdaEvent<LambdaRequest>) -> Request {
 ///
 /// See the [AWS docs for response streaming].
 ///
+/// # Managed concurrency
+/// If `AWS_LAMBDA_MAX_CONCURRENCY` is set, this function returns an error because
+/// it does not enable concurrent polling. Use [`run_with_streaming_response_concurrent`]
+/// instead.
+///
 /// [AWS docs for response streaming]:
 ///     https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html
 pub async fn run_with_streaming_response<'a, S, B, E>(handler: S) -> Result<(), Error>
@@ -192,7 +209,7 @@ where
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
 {
-    lambda_runtime::run_concurrent(into_stream_service_boxed(handler)).await
+    lambda_runtime::run_concurrent(into_stream_service_cloneable(handler)).await
 }
 
 pin_project_lite::pin_project! {
diff --git a/lambda-runtime-api-client/src/lib.rs b/lambda-runtime-api-client/src/lib.rs
@@ -120,6 +120,10 @@ impl ClientBuilder {
     }
 
     /// Provide a pool size hint for the underlying Hyper client.
+    ///
+    /// When using concurrent polling, this should be at least the maximum
+    /// concurrency (e.g., `AWS_LAMBDA_MAX_CONCURRENCY`) to avoid connection
+    /// starvation.
     pub fn with_pool_size(self, pool_size: usize) -> Self {
         Self {
             pool_size: Some(pool_size),
diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs
@@ -94,6 +94,12 @@ where
 /// If you need more control over the runtime and add custom middleware, use the
 /// [Runtime] type directly.
 ///
+/// # Managed concurrency
+/// If `AWS_LAMBDA_MAX_CONCURRENCY` is set, this function returns an error because
+/// it does not enable concurrent polling. If your handler can satisfy `Clone`,
+/// prefer [`run_concurrent`], which honors managed concurrency and falls back to
+/// sequential behavior when unset.
+///
 /// # Example
 /// ```no_run
 /// use lambda_runtime::{Error, service_fn, LambdaEvent};
diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs
@@ -4,20 +4,14 @@ use crate::{
     types::{invoke_request_id, IntoFunctionResponse, LambdaEvent},
     Config, Context, Diagnostic,
 };
-use futures::stream::FuturesUnordered;
+use futures::{future::BoxFuture, stream::FuturesUnordered};
 use http_body_util::BodyExt;
 use lambda_runtime_api_client::{BoxError, Client as ApiClient};
 use serde::{Deserialize, Serialize};
-use std::{
-    env,
-    fmt::Debug,
-    future::Future,
-    io,
-    sync::{Arc, OnceLock},
-};
+use std::{env, fmt, fmt::Debug, future::Future, io, sync::Arc};
 use tokio_stream::{Stream, StreamExt};
 use tower::{Layer, Service, ServiceExt};
-use tracing::{error, trace, warn};
+use tracing::{error, info_span, trace, warn, Instrument};
 
 /* ----------------------------------------- INVOCATION ---------------------------------------- */
 
@@ -63,9 +57,6 @@ pub struct Runtime<S> {
     concurrency_limit: u32,
 }
 
-/// One-time marker to log X-Ray behavior in concurrent mode.
-static XRAY_LOGGED: OnceLock<()> = OnceLock::new();
-
 impl<F, EventPayload, Response, BufferedResponse, StreamingResponse, StreamItem, StreamError>
     Runtime<
         RuntimeApiClientService<
@@ -170,6 +161,7 @@ where
     /// classic Lambda and Lambda Managed Instances.
     pub async fn run_concurrent(self) -> Result<(), BoxError> {
         if self.concurrency_limit > 1 {
+            trace!("Concurrent mode: _X_AMZN_TRACE_ID is not set; use context.xray_trace_id");
             Self::run_concurrent_inner(self.service, self.config, self.client, self.concurrency_limit).await
         } else {
             let incoming = incoming(&self.client);
@@ -186,72 +178,144 @@ where
     ) -> Result<(), BoxError> {
         let limit = concurrency_limit as usize;
 
-        let mut workers = FuturesUnordered::new();
+        // Use FuturesUnordered so we can observe worker exits as they happen,
+        // rather than waiting for all workers to finish (join_all).
+        let mut workers: FuturesUnordered<WorkerJoinFuture> = FuturesUnordered::new();
+        let spawn_worker = |service: S, config: Arc<Config>, client: Arc<ApiClient>| -> WorkerJoinFuture {
+            let handle = tokio::spawn(concurrent_worker_loop(service, config, client));
+            let task_id = handle.id();
+            Box::pin(async move { (task_id, handle.await) })
+        };
+        // Spawn one worker per concurrency slot; the last uses the owned service to avoid an extra clone.
         for _ in 1..limit {
-            workers.push(tokio::spawn(concurrent_worker_loop(
-                service.clone(),
-                config.clone(),
-                client.clone(),
-            )));
+            workers.push(spawn_worker(service.clone(), config.clone(), client.clone()));
         }
-        workers.push(tokio::spawn(concurrent_worker_loop(service, config, client)));
+        workers.push(spawn_worker(service, config, client));
 
-        // Track the first infrastructure error. A single worker failing should
-        // not terminate the whole runtime (LMI keeps running with the remaining
+        // Track worker exits across tasks. A single worker failing should not
+        // terminate the whole runtime (LMI keeps running with the remaining
         // healthy workers). We only return an error once there are no workers
         // left (i.e., we cannot keep at least 1 worker alive).
         //
         // Note: Handler errors (Err returned from user code) do NOT trigger this.
         // They are reported to Lambda via /invocation/{id}/error and the worker
         // continues. This only captures unrecoverable runtime failures like
         // API client failures, runtime panics, etc.
-        let mut first_error: Option<BoxError> = None;
+        let mut errors: Vec<WorkerError> = Vec::new();
         let mut remaining_workers = limit;
-        while let Some(result) = futures::StreamExt::next(&mut workers).await {
+        while let Some((task_id, result)) = futures::StreamExt::next(&mut workers).await {
             remaining_workers = remaining_workers.saturating_sub(1);
             match result {
                 Ok(Ok(())) => {
                     // `concurrent_worker_loop` runs indefinitely, so an Ok return indicates
                     // an unexpected worker exit; we still decrement because the task is gone.
-                    warn!(
+                    let clean_exit_msg = "Concurrent worker exited cleanly (unexpected - loop should run forever)";
+                    error!(
+                        task_id = %task_id,
                         remaining_workers,
-                        "Concurrent worker exited cleanly (unexpected - loop should run forever)"
+                        "{}",
+                        clean_exit_msg
                     );
-                    if first_error.is_none() {
-                        first_error = Some(Box::new(io::Error::other(
-                            "all concurrent workers exited cleanly (unexpected - loop should run forever)",
-                        )));
-                    }
+                    errors.push(WorkerError::CleanExit(task_id));
                 }
                 Ok(Err(err)) => {
-                    error!(error = %err, remaining_workers, "Concurrent worker exited with error");
-                    if first_error.is_none() {
-                        first_error = Some(err);
-                    }
+                    error!(
+                        task_id = %task_id,
+                        error = %err,
+                        remaining_workers,
+                        "Concurrent worker exited with error"
+                    );
+                    errors.push(WorkerError::Failure(task_id, err));
                 }
                 Err(join_err) => {
                     let err: BoxError = Box::new(join_err);
-                    error!(error = %err, remaining_workers, "Concurrent worker panicked");
-                    if first_error.is_none() {
-                        first_error = Some(err);
-                    }
+                    error!(
+                        task_id = %task_id,
+                        error = %err,
+                        remaining_workers,
+                        "Concurrent worker panicked"
+                    );
+                    errors.push(WorkerError::Failure(task_id, err));
                 }
             }
         }
 
-        match first_error {
-            Some(err) => Err(err),
-            None => Ok(()),
+        match errors.len() {
+            0 => Ok(()),
+            _ => Err(Box::new(ConcurrentWorkerErrors { errors })),
+        }
+    }
+}
+
+#[derive(Debug)]
+enum WorkerError {
+    CleanExit(tokio::task::Id),
+    Failure(tokio::task::Id, BoxError),
+}
+
+type WorkerJoinResult = (tokio::task::Id, Result<Result<(), BoxError>, tokio::task::JoinError>);
+type WorkerJoinFuture = BoxFuture<'static, WorkerJoinResult>;
+
+#[derive(Debug)]
+struct ConcurrentWorkerErrors {
+    errors: Vec<WorkerError>,
+}
+
+impl fmt::Display for ConcurrentWorkerErrors {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let mut clean = Vec::new();
+        let mut failures = Vec::new();
+        for error in &self.errors {
+            match error {
+                WorkerError::CleanExit(task_id) => clean.push(task_id),
+                WorkerError::Failure(task_id, err) => failures.push((task_id, err)),
+            }
+        }
+
+        if failures.is_empty() && !clean.is_empty() {
+            write!(
+                f,
+                "all concurrent workers exited cleanly (unexpected - loop should run forever)"
+            )?;
+            for task_id in clean {
+                write!(f, " [task {task_id}]")?;
+            }
+            return Ok(());
+        }
+
+        write!(f, "concurrent workers exited unexpectedly")?;
+        if !clean.is_empty() {
+            write!(f, "; clean exits:")?;
+            for task_id in clean {
+                write!(f, " [task {task_id}]")?;
+            }
+        }
+        if !failures.is_empty() {
+            write!(f, "; failures:")?;
+            for (task_id, err) in failures {
+                write!(f, " [task {task_id}] {err}")?;
+            }
         }
+        Ok(())
     }
 }
 
+impl std::error::Error for ConcurrentWorkerErrors {}
+
 impl<S> Runtime<S>
 where
     S: Service<LambdaInvocation, Response = (), Error = BoxError>,
 {
     /// Start the runtime and begin polling for events on the Lambda Runtime API.
+    ///
+    /// If `AWS_LAMBDA_MAX_CONCURRENCY` is set, this returns an error because it does not enable
+    /// concurrent polling. Use [`Runtime::run_concurrent`] instead.
     pub async fn run(self) -> Result<(), BoxError> {
+        if let Some(raw) = concurrency_env_value() {
+            return Err(Box::new(io::Error::other(format!(
+                "AWS_LAMBDA_MAX_CONCURRENCY is set to '{raw}', but Runtime::run does not support concurrent polling; use Runtime::run_concurrent instead"
+            ))));
+        }
         let incoming = incoming(&self.client);
         Self::run_with_incoming(self.service, self.config, incoming).await
     }
@@ -332,21 +396,29 @@ fn max_concurrency_from_env() -> Option<u32> {
         .filter(|&c| c > 0)
 }
 
+fn concurrency_env_value() -> Option<String> {
+    env::var("AWS_LAMBDA_MAX_CONCURRENCY").ok()
+}
+
 async fn concurrent_worker_loop<S>(mut service: S, config: Arc<Config>, client: Arc<ApiClient>) -> Result<(), BoxError>
 where
     S: Service<LambdaInvocation, Response = (), Error = BoxError>,
     S::Future: Send,
 {
+    let task_id = tokio::task::id();
+    let span = info_span!("concurrent_worker_loop", task_id = %task_id);
     loop {
         let event = match next_event_future(client.clone()).await {
             Ok(event) => event,
             Err(e) => {
-                warn!(error = %e, "Error polling /next, retrying");
+                warn!(task_id = %task_id, error = %e, "Error polling /next, retrying");
                 continue;
             }
         };
 
-        process_invocation(&mut service, &config, event, false).await?;
+        process_invocation(&mut service, &config, event, false)
+            .instrument(span.clone())
+            .await?;
     }
 }
 
@@ -378,11 +450,6 @@ where
     if set_amzn_trace_env {
         // Setup Amazon's default tracing data
         amzn_trace_env(&invocation.context);
-    } else {
-        // Inform users that X-Ray is available via context, not env var, in concurrent mode.
-        XRAY_LOGGED.get_or_init(|| {
-            trace!("Concurrent mode: _X_AMZN_TRACE_ID is not set; use context.xray_trace_id");
-        });
     }
 
     // Wait for service to be ready

Original file line number	Diff line number	Diff line change
`@@ -29,7 +29,7 @@ async fn main() -> Result<(), Error> {`
`29`	`29`
`30`	`30`	`let func = service_fn(my_handler);`
`31`	`31`	`if let Err(err) = lambda_runtime::run_concurrent(func).await {`
`32`		`- eprintln!("run error: {:?}", err);`
	`32`	`+ tracing::error!(error = %err, "run error");`
`33`	`33`	`return Err(err);`
`34`	`34`	`}`
`35`	`35`	`Ok(())`