feat(runtime,http): add concurrent entrypoints without changing run

Author: alessandrobologna

lambda-http/src/lib.rs

@@ -165,9 +165,9 @@ where
 
 impl<'a, R, S, E> From<S> for Adapter<'a, R, S>
 where
-    S: Service<Request, Response = R, Error = E> + Clone + Send + 'static,
-    S::Future: Send + 'static,
-    R: IntoResponse + Send + Sync + 'static,
+    S: Service<Request, Response = R, Error = E>,
+    S::Future: Send + 'a,
+    R: IntoResponse,
 {
     fn from(service: S) -> Self {
         Adapter {
@@ -205,14 +205,32 @@ where
 ///
 /// This takes care of transforming the LambdaEvent into a [`Request`] and then
 /// converting the result into a `LambdaResponse`.
-pub async fn run<R, S, E>(handler: S) -> Result<(), Error>
+pub async fn run<'a, R, S, E>(handler: S) -> Result<(), Error>
+where
+    S: Service<Request, Response = R, Error = E>,
+    S::Future: Send + 'a,
+    R: IntoResponse,
+    E: std::fmt::Debug + Into<Diagnostic>,
+{
+    lambda_runtime::run(Adapter::from(handler)).await
+}
+
+/// Starts the Lambda Rust runtime in a mode that is compatible with
+/// Lambda Managed Instances (concurrent invocations).
+///
+/// When `AWS_LAMBDA_MAX_CONCURRENCY` is set to a value greater than 1, this
+/// will use a concurrent `/next` polling loop with a bounded number of
+/// in-flight handler tasks. 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.
+pub async fn run_concurrent<R, S, E>(handler: S) -> Result<(), Error>
 where
     S: Service<Request, Response = R, Error = E> + Clone + Send + 'static,
     S::Future: Send + 'static,
     R: IntoResponse + Send + Sync + 'static,
     E: std::fmt::Debug + Into<Diagnostic> + Send + 'static,
 {
-    lambda_runtime::run(Adapter::from(handler)).await
+    lambda_runtime::run_concurrent(Adapter::from(handler)).await
 }
 
 #[cfg(test)]

lambda-http/src/streaming.rs

@@ -9,7 +9,10 @@ use futures_util::{Stream, TryFutureExt};
 pub use http::{self, Response};
 use http_body::Body;
 use lambda_runtime::{
-    tower::{util::BoxCloneService, ServiceBuilder, ServiceExt},
+    tower::{
+        util::{BoxCloneService, MapRequest, MapResponse},
+        ServiceBuilder, ServiceExt,
+    },
     Diagnostic,
 };
 pub use lambda_runtime::{Error, LambdaEvent, MetadataPrelude, Service, StreamResponse};
@@ -24,22 +27,10 @@ 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> + Clone + Send + 'static,
-    S::Future: Send + 'static,
+    S: Service<Request, Response = Response<B>, Error = E>,
+    S::Future: Send + 'a,
     B: Body + Unpin + Send + 'static,
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
@@ -54,15 +45,15 @@ where
 
 impl<'a, S, B, E> Service<LambdaEvent<LambdaRequest>> for StreamAdapter<'a, S, B>
 where
-    S: Service<Request, Response = Response<B>, Error = E> + Clone + Send + 'static,
-    S::Future: Send + 'static,
+    S: Service<Request, Response = Response<B>, Error = E>,
+    S::Future: Send + 'a,
     B: Body + Unpin + Send + 'static,
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
 {
     type Response = StreamResponse<BodyStream<B>>;
     type Error = E;
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'static>>;
+    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send + 'a>>;
 
     fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         self.service.poll_ready(cx)
@@ -87,7 +78,31 @@ where
 /// Used internally by [`run_with_streaming_response`]; not part of the public
 /// API.
 #[allow(clippy::type_complexity)]
-fn into_stream_service<S, B, E>(
+fn into_stream_service<'a, S, B, E>(
+    handler: S,
+) -> MapResponse<
+    MapRequest<S, impl FnMut(LambdaEvent<LambdaRequest>) -> Request>,
+    impl FnOnce(Response<B>) -> StreamResponse<BodyStream<B>> + Clone,
+>
+where
+    S: Service<Request, Response = Response<B>, Error = E>,
+    S::Future: Send + 'a,
+    E: Debug + Into<Diagnostic>,
+    B: Body + Unpin + Send + 'static,
+    B::Data: Into<Bytes> + Send,
+    B::Error: Into<Error> + Send + Debug,
+{
+    ServiceBuilder::new()
+        .map_request(event_to_request as fn(LambdaEvent<LambdaRequest>) -> Request)
+        .service(handler)
+        .map_response(into_stream_response)
+}
+
+/// 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.
+#[allow(dead_code)]
+#[allow(clippy::type_complexity)]
+fn into_stream_service_boxed<S, B, E>(
     handler: S,
 ) -> BoxCloneService<LambdaEvent<LambdaRequest>, StreamResponse<BodyStream<B>>, E>
 where
@@ -145,7 +160,26 @@ fn event_to_request(req: LambdaEvent<LambdaRequest>) -> Request {
 ///
 /// [AWS docs for response streaming]:
 ///     https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html
-pub async fn run_with_streaming_response<S, B, E>(handler: S) -> Result<(), Error>
+pub async fn run_with_streaming_response<'a, S, B, E>(handler: S) -> Result<(), Error>
+where
+    S: Service<Request, Response = Response<B>, Error = E>,
+    S::Future: Send + 'a,
+    E: Debug + Into<Diagnostic>,
+    B: Body + Unpin + Send + 'static,
+    B::Data: Into<Bytes> + Send,
+    B::Error: Into<Error> + Send + Debug,
+{
+    lambda_runtime::run(into_stream_service(handler)).await
+}
+
+/// Runs the Lambda runtime with a handler that returns **streaming** HTTP
+/// responses, in a mode that is compatible with Lambda Managed Instances.
+///
+/// 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`].
+#[allow(dead_code)]
+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,
     S::Future: Send + 'static,
@@ -154,7 +188,7 @@ where
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
 {
-    lambda_runtime::run(into_stream_service(handler)).await
+    lambda_runtime::run_concurrent(into_stream_service_boxed(handler)).await
 }
 
 pin_project_lite::pin_project! {

lambda-runtime/src/lib.rs

@@ -126,6 +126,30 @@ where
 /// }
 /// ```
 pub async fn run<A, F, R, B, S, D, E>(handler: F) -> Result<(), Error>
+where
+    F: Service<LambdaEvent<A>, Response = R>,
+    F::Future: Future<Output = Result<R, F::Error>>,
+    F::Error: Into<Diagnostic> + fmt::Debug,
+    A: for<'de> Deserialize<'de>,
+    R: IntoFunctionResponse<B, S>,
+    B: Serialize,
+    S: Stream<Item = Result<D, E>> + Unpin + Send + 'static,
+    D: Into<bytes::Bytes> + Send,
+    E: Into<Error> + Send + Debug,
+{
+    let runtime = Runtime::new(handler).layer(layers::TracingLayer::new());
+    runtime.run().await
+}
+
+/// Starts the Lambda Rust runtime in a mode that is compatible with
+/// Lambda Managed Instances (concurrent invocations).
+///
+/// When `AWS_LAMBDA_MAX_CONCURRENCY` is set to a value greater than 1, this
+/// will use a concurrent `/next` polling loop with a bounded number of
+/// in-flight handler tasks. 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.
+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,
     F::Future: Future<Output = Result<R, F::Error>> + Send + 'static,
@@ -138,7 +162,7 @@ where
     E: Into<Error> + Send + Debug + 'static,
 {
     let runtime = Runtime::new(handler).layer(layers::TracingLayer::new());
-    runtime.run().await
+    runtime.run_concurrent().await
 }
 
 /// Spawns a task that will be execute a provided async closure when the process

lambda-runtime/src/runtime.rs

@@ -82,12 +82,12 @@ impl<F, EventPayload, Response, BufferedResponse, StreamingResponse, StreamItem,
         >,
     >
 where
-    F: Service<LambdaEvent<EventPayload>, Response = Response> + Send + Clone + 'static,
-    F::Future: Future<Output = Result<Response, F::Error>> + Send + 'static,
+    F: Service<LambdaEvent<EventPayload>, Response = Response>,
+    F::Future: Future<Output = Result<Response, F::Error>>,
     F::Error: Into<Diagnostic> + Debug,
-    EventPayload: for<'de> Deserialize<'de> + Send + 'static,
-    Response: IntoFunctionResponse<BufferedResponse, StreamingResponse> + Send + 'static,
-    BufferedResponse: Serialize + Send + 'static,
+    EventPayload: for<'de> Deserialize<'de>,
+    Response: IntoFunctionResponse<BufferedResponse, StreamingResponse>,
+    BufferedResponse: Serialize,
     StreamingResponse: Stream<Item = Result<StreamItem, StreamError>> + Unpin + Send + 'static,
     StreamItem: Into<bytes::Bytes> + Send,
     StreamError: Into<BoxError> + Send + Debug,
@@ -160,19 +160,23 @@ where
     S: Service<LambdaInvocation, Response = (), Error = BoxError> + Clone + Send + 'static,
     S::Future: Send,
 {
-    /// Start the runtime and begin polling for events on the Lambda Runtime API.
-    pub async fn run(self) -> Result<(), BoxError> {
+    /// Start the runtime in concurrent mode when configured for Lambda managed-concurrency.
+    ///
+    /// 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.
+    pub async fn run_concurrent(self) -> Result<(), BoxError> {
         if self.config.is_concurrent() {
             let max_concurrency = self.config.max_concurrency.unwrap_or(1);
-            Self::run_concurrent(self.service, self.config, self.client, max_concurrency).await
+            Self::run_concurrent_inner(self.service, self.config, self.client, max_concurrency).await
         } else {
             let incoming = incoming(&self.client);
             Self::run_with_incoming(self.service, self.config, incoming).await
         }
     }
 
     /// Concurrent processing using windowed long-polls (for Lambda managed-concurrency).
-    async fn run_concurrent(
+    async fn run_concurrent_inner(
         service: S,
         config: Arc<Config>,
         client: Arc<ApiClient>,
@@ -301,32 +305,46 @@ impl<S> Runtime<S>
 where
     S: Service<LambdaInvocation, Response = (), Error = BoxError>,
 {
+    /// Start the runtime and begin polling for events on the Lambda Runtime API.
+    pub async fn run(self) -> Result<(), BoxError> {
+        let incoming = incoming(&self.client);
+        Self::run_with_incoming(self.service, self.config, incoming).await
+    }
+
     /// Internal utility function to start the runtime with a customized incoming stream.
-    /// This implements the sequential mode of the runtime.
+    /// This implements the core of the [Runtime::run] method.
     pub(crate) async fn run_with_incoming(
         mut service: S,
         config: Arc<Config>,
         incoming: impl Stream<Item = Result<http::Response<hyper::body::Incoming>, BoxError>> + Send,
     ) -> Result<(), BoxError> {
         tokio::pin!(incoming);
-        while let Some(next_event_response) = tokio_stream::StreamExt::next(&mut incoming).await {
+        while let Some(next_event_response) = incoming.next().await {
             trace!("New event arrived (run loop)");
             let event = next_event_response?;
             let (parts, incoming) = event.into_parts();
 
             #[cfg(debug_assertions)]
             if parts.status == http::StatusCode::NO_CONTENT {
+                // Ignore the event if the status code is 204.
+                // This is a way to keep the runtime alive when
+                // there are no events pending to be processed.
                 continue;
             }
 
+            // Build the invocation such that it can be sent to the service right away
+            // when it is ready
             let body = incoming.collect().await?.to_bytes();
             let context = Context::new(invoke_request_id(&parts.headers)?, config.clone(), &parts.headers)?;
             let invocation = LambdaInvocation { parts, body, context };
 
             // Setup Amazon's default tracing data
             amzn_trace_env(&invocation.context);
 
+            // Wait for service to be ready
             let ready = service.ready().await?;
+
+            // Once ready, call the service which will respond to the Lambda runtime API
             ready.call(invocation).await?;
         }
         Ok(())