feat: gate concurrent APIs behind experimental feature

alessandrobologna · alessandrobologna · commit 8ed50a3ffba6 · 2025-12-30T00:39:55.000-05:00
- add experimental-concurrency feature in lambda-runtime and lambda-http
- cfg-gate concurrent entrypoints and streaming helpers
- keep run() guard and X-Ray header source on default path (partial gating)
- fix default clippy by gating concurrent-only imports
- update docs and concurrent example dependency
diff --git a/examples/basic-lambda-concurrent/Cargo.toml b/examples/basic-lambda-concurrent/Cargo.toml
@@ -4,6 +4,6 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
-lambda_runtime = { path = "../../lambda-runtime" }
+lambda_runtime = { path = "../../lambda-runtime", features = ["experimental-concurrency"] }
 serde = "1.0.219"
 tokio = { version = "1", features = ["macros"] }
diff --git a/lambda-http/Cargo.toml b/lambda-http/Cargo.toml
@@ -29,6 +29,7 @@ opentelemetry = ["lambda_runtime/opentelemetry"] # enables access to the OpenTel
 anyhow = ["lambda_runtime/anyhow"] # enables From<T> for Diagnostic for anyhow error types, see README.md for more info
 eyre = ["lambda_runtime/eyre"] # enables From<T> for Diagnostic for eyre error types, see README.md for more info
 miette = ["lambda_runtime/miette"] # enables From<T> for Diagnostic for miette error types, see README.md for more info
+experimental-concurrency = ["lambda_runtime/experimental-concurrency"]
 
 [dependencies]
 bytes = { workspace = true }
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -102,7 +102,9 @@ use std::{
 };
 
 mod streaming;
-pub use streaming::{run_with_streaming_response, run_with_streaming_response_concurrent, StreamAdapter};
+#[cfg(feature = "experimental-concurrency")]
+pub use streaming::run_with_streaming_response_concurrent;
+pub use streaming::{run_with_streaming_response, StreamAdapter};
 
 /// Type alias for `http::Request`s with a fixed [`Body`](enum.Body.html) type
 pub type Request = http::Request<Body>;
@@ -195,7 +197,7 @@ where
         let LambdaEvent { payload, context } = req;
         let request_origin = payload.request_origin();
         let mut event: Request = payload.into();
-        update_xray_trace_id_header_from_context(event.headers_mut(), &context);
+        update_xray_trace_id_header(event.headers_mut(), &context);
         let fut = Box::pin(self.service.call(event.with_lambda_context(context)));
 
         TransformResponse::Request(request_origin, fut)
@@ -211,8 +213,9 @@ where
 /// # 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.
+/// prefer [`run_concurrent`] (requires the `experimental-concurrency` feature),
+/// 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>,
@@ -226,11 +229,15 @@ where
 /// Starts the Lambda Rust runtime in a mode that is compatible with
 /// Lambda Managed Instances (concurrent invocations).
 ///
+/// Requires the `experimental-concurrency` feature.
+///
 /// When `AWS_LAMBDA_MAX_CONCURRENCY` is set to a value greater than 1, this
 /// will spawn `AWS_LAMBDA_MAX_CONCURRENCY` worker tasks, each running its own
 /// `/next` polling loop. When the environment variable is unset or `<= 1`,
 /// it falls back to the same sequential behavior as [`run`], so the same
 /// handler can run on both classic Lambda and Lambda Managed Instances.
+#[cfg(feature = "experimental-concurrency")]
+#[cfg_attr(docsrs, doc(cfg(feature = "experimental-concurrency")))]
 pub async fn run_concurrent<R, S, E>(handler: S) -> Result<(), Error>
 where
     S: Service<Request, Response = R, Error = E> + Clone + Send + 'static,
@@ -241,8 +248,8 @@ where
     lambda_runtime::run_concurrent(Adapter::from(handler)).await
 }
 
-// Replaces update_xray_trace_id_header (env var), now set from Context
-fn update_xray_trace_id_header_from_context(headers: &mut http::HeaderMap, context: &Context) {
+// In concurrent mode we must use the per-request context.
+fn update_xray_trace_id_header(headers: &mut http::HeaderMap, context: &Context) {
     if let Some(trace_id) = context.xray_trace_id.as_deref() {
         if let Ok(header_value) = http::HeaderValue::from_str(trace_id) {
             headers.insert(http::header::HeaderName::from_static("x-amzn-trace-id"), header_value);
diff --git a/lambda-http/src/streaming.rs b/lambda-http/src/streaming.rs
@@ -1,6 +1,4 @@
-use crate::{
-    http::header::SET_COOKIE, request::LambdaRequest, update_xray_trace_id_header_from_context, Request, RequestExt,
-};
+use crate::{http::header::SET_COOKIE, request::LambdaRequest, update_xray_trace_id_header, Request, RequestExt};
 use bytes::Bytes;
 use core::{
     fmt::Debug,
@@ -76,7 +74,7 @@ where
     fn call(&mut self, req: LambdaEvent<LambdaRequest>) -> Self::Future {
         let LambdaEvent { payload, context } = req;
         let mut event: Request = payload.into();
-        update_xray_trace_id_header_from_context(event.headers_mut(), &context);
+        update_xray_trace_id_header(event.headers_mut(), &context);
         Box::pin(
             self.service
                 .call(event.with_lambda_context(context))
@@ -116,8 +114,10 @@ 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.
+#[cfg(feature = "experimental-concurrency")]
 type EventToRequest = fn(LambdaEvent<LambdaRequest>) -> Request;
 
+#[cfg(feature = "experimental-concurrency")]
 #[allow(clippy::type_complexity)]
 fn into_stream_service_cloneable<S, B, E>(
     handler: S,
@@ -166,7 +166,7 @@ where
 fn event_to_request(req: LambdaEvent<LambdaRequest>) -> Request {
     let LambdaEvent { payload, context } = req;
     let mut event: Request = payload.into();
-    update_xray_trace_id_header_from_context(event.headers_mut(), &context);
+    update_xray_trace_id_header(event.headers_mut(), &context);
     event.with_lambda_context(context)
 }
 
@@ -178,7 +178,7 @@ fn event_to_request(req: LambdaEvent<LambdaRequest>) -> Request {
 /// # 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.
+/// (requires the `experimental-concurrency` feature) instead.
 ///
 /// [AWS docs for response streaming]:
 ///     https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html
@@ -197,9 +197,13 @@ where
 /// Runs the Lambda runtime with a handler that returns **streaming** HTTP
 /// responses, in a mode that is compatible with Lambda Managed Instances.
 ///
+/// Requires the `experimental-concurrency` feature.
+///
 /// This uses a cloneable, boxed service internally so it can be driven by the
 /// concurrent runtime. When `AWS_LAMBDA_MAX_CONCURRENCY` is not set or `<= 1`,
 /// it falls back to the same sequential behavior as [`run_with_streaming_response`].
+#[cfg(feature = "experimental-concurrency")]
+#[cfg_attr(docsrs, doc(cfg(feature = "experimental-concurrency")))]
 pub async fn run_with_streaming_response_concurrent<S, B, E>(handler: S) -> Result<(), Error>
 where
     S: Service<Request, Response = Response<B>, Error = E> + Clone + Send + 'static,
diff --git a/lambda-runtime/Cargo.toml b/lambda-runtime/Cargo.toml
@@ -25,6 +25,7 @@ miette = ["dep:miette"] # enables From<T> for Diagnostic for miette error types,
 # as well as default features
 # https://github.com/aws/aws-lambda-rust-runtime/issues/984
 graceful-shutdown = ["tokio/rt", "tokio/signal", "dep:lambda-extension"]
+experimental-concurrency = []
 
 [dependencies]
 anyhow = { version = "1.0.86", optional = true }
diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs
@@ -97,8 +97,9 @@ where
 /// # 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.
+/// prefer [`run_concurrent`] (requires the `experimental-concurrency` feature),
+/// which honors managed concurrency and falls back to sequential behavior when
+/// unset.
 ///
 /// # Example
 /// ```no_run
@@ -135,6 +136,8 @@ where
 /// Starts the Lambda Rust runtime in a mode that is compatible with
 /// Lambda Managed Instances (concurrent invocations).
 ///
+/// Requires the `experimental-concurrency` feature.
+///
 /// When `AWS_LAMBDA_MAX_CONCURRENCY` is set to a value greater than 1, this
 /// will spawn `AWS_LAMBDA_MAX_CONCURRENCY` worker tasks, each running its own
 /// `/next` polling loop. When the environment variable is unset or `<= 1`, it
@@ -160,6 +163,8 @@ where
 ///     Ok(event.payload)
 /// }
 /// ```
+#[cfg(feature = "experimental-concurrency")]
+#[cfg_attr(docsrs, doc(cfg(feature = "experimental-concurrency")))]
 pub async fn run_concurrent<A, F, R, B, S, D, E>(handler: F) -> Result<(), Error>
 where
     F: Service<LambdaEvent<A>, Response = R> + Clone + Send + 'static,
diff --git a/lambda-runtime/src/runtime.rs b/lambda-runtime/src/runtime.rs
@@ -4,14 +4,19 @@ use crate::{
     types::{invoke_request_id, IntoFunctionResponse, LambdaEvent},
     Config, Context, Diagnostic,
 };
+#[cfg(feature = "experimental-concurrency")]
 use futures::stream::FuturesUnordered;
 use http_body_util::BodyExt;
 use lambda_runtime_api_client::{BoxError, Client as ApiClient};
 use serde::{Deserialize, Serialize};
-use std::{env, fmt, fmt::Debug, future::Future, io, sync::Arc};
+#[cfg(feature = "experimental-concurrency")]
+use std::fmt;
+use std::{env, fmt::Debug, future::Future, io, sync::Arc};
 use tokio_stream::{Stream, StreamExt};
 use tower::{Layer, Service, ServiceExt};
-use tracing::{debug, error, info_span, trace, warn, Instrument};
+use tracing::trace;
+#[cfg(feature = "experimental-concurrency")]
+use tracing::{debug, error, info_span, warn, Instrument};
 
 /* ----------------------------------------- INVOCATION ---------------------------------------- */
 
@@ -149,6 +154,7 @@ impl<S> Runtime<S> {
     }
 }
 
+#[cfg(feature = "experimental-concurrency")]
 impl<S> Runtime<S>
 where
     S: Service<LambdaInvocation, Response = (), Error = BoxError> + Clone + Send + 'static,
@@ -159,6 +165,7 @@ where
     /// If `AWS_LAMBDA_MAX_CONCURRENCY` is not set or is `<= 1`, this falls back to the
     /// sequential `run_with_incoming` loop so that the same handler can run on both
     /// classic Lambda and Lambda Managed Instances.
+    #[cfg_attr(docsrs, doc(cfg(feature = "experimental-concurrency")))]
     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");
@@ -252,17 +259,20 @@ where
     }
 }
 
+#[cfg(feature = "experimental-concurrency")]
 #[derive(Debug)]
 enum WorkerError {
     CleanExit(tokio::task::Id),
     Failure(tokio::task::Id, BoxError),
 }
 
+#[cfg(feature = "experimental-concurrency")]
 #[derive(Debug)]
 struct ConcurrentWorkerErrors {
     errors: Vec<WorkerError>,
 }
 
+#[cfg(feature = "experimental-concurrency")]
 #[derive(Serialize)]
 struct ConcurrentWorkerErrorsPayload<'a> {
     message: &'a str,
@@ -272,12 +282,14 @@ struct ConcurrentWorkerErrorsPayload<'a> {
     failures: Vec<WorkerFailurePayload>,
 }
 
+#[cfg(feature = "experimental-concurrency")]
 #[derive(Serialize)]
 struct WorkerFailurePayload {
     id: String,
     err: String,
 }
 
+#[cfg(feature = "experimental-concurrency")]
 impl fmt::Display for ConcurrentWorkerErrors {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut clean = Vec::new();
@@ -314,6 +326,7 @@ impl fmt::Display for ConcurrentWorkerErrors {
     }
 }
 
+#[cfg(feature = "experimental-concurrency")]
 impl std::error::Error for ConcurrentWorkerErrors {}
 
 impl<S> Runtime<S>
@@ -323,11 +336,12 @@ where
     /// 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.
+    /// concurrent polling. Enable the `experimental-concurrency` feature and 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"
+                "AWS_LAMBDA_MAX_CONCURRENCY is set to '{raw}', but Runtime::run does not support concurrent polling; enable the experimental-concurrency feature and use Runtime::run_concurrent instead"
             ))));
         }
         let incoming = incoming(&self.client);
@@ -398,6 +412,7 @@ fn incoming(
 }
 
 /// Creates a future that polls the `/next` endpoint.
+#[cfg(feature = "experimental-concurrency")]
 async fn next_event_future(client: &ApiClient) -> Result<http::Response<hyper::body::Incoming>, BoxError> {
     let req = NextEventRequest.into_req()?;
     client.call(req).await
@@ -414,6 +429,7 @@ fn concurrency_env_value() -> Option<String> {
     env::var("AWS_LAMBDA_MAX_CONCURRENCY").ok()
 }
 
+#[cfg(feature = "experimental-concurrency")]
 async fn concurrent_worker_loop<S>(mut service: S, config: Arc<Config>, client: Arc<ApiClient>) -> Result<(), BoxError>
 where
     S: Service<LambdaInvocation, Response = (), Error = BoxError>,
@@ -744,6 +760,7 @@ mod endpoint_tests {
         .await
     }
 
+    #[cfg(feature = "experimental-concurrency")]
     #[tokio::test]
     async fn concurrent_worker_crash_does_not_stop_other_workers() -> Result<(), Error> {
         let next_calls = Arc::new(AtomicUsize::new(0));
