wip

Author: s0l0ist

lambda-http/src/streaming.rs

@@ -2,16 +2,21 @@ use crate::{http::header::SET_COOKIE, request::LambdaRequest, Request, RequestEx
 use bytes::Bytes;
 use core::{
     fmt::Debug,
-    future::Future,
     pin::Pin,
     task::{Context, Poll},
 };
-use futures_util::Stream;
+use futures_util::{Stream, TryFutureExt};
 pub use http::{self, Response};
 use http_body::Body;
-use lambda_runtime::Diagnostic;
+use lambda_runtime::{
+    tower::{
+        util::{MapRequest, MapResponse},
+        ServiceBuilder, ServiceExt,
+    },
+    Diagnostic,
+};
 pub use lambda_runtime::{Error, LambdaEvent, MetadataPrelude, Service, StreamResponse};
-use std::marker::PhantomData;
+use std::{future::Future, marker::PhantomData};
 
 /// An adapter that lifts a standard [`Service<Request>`] into a
 /// [`Service<LambdaEvent<LambdaRequest>>`] which produces streaming Lambda HTTP
@@ -25,7 +30,6 @@ impl<'a, S, B, E> From<S> for StreamAdapter<'a, S, B>
 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,
@@ -42,46 +46,87 @@ impl<'a, S, B, E> Service<LambdaEvent<LambdaRequest>> for StreamAdapter<'a, S, B
 where
     S: Service<Request, Response = Response<B>, Error = E>,
     S::Future: Send + 'a,
-    B: Body + Send + 'static,
+    B: Body + Unpin + Send + 'static,
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
-    E: Debug + Into<Diagnostic>,
 {
     type Response = StreamResponse<BodyStream<B>>;
     type Error = E;
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, E>> + Send + 'a>>;
+    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)
     }
 
     fn call(&mut self, req: LambdaEvent<LambdaRequest>) -> Self::Future {
         let event: Request = req.payload.into();
-        let fut = self.service.call(event.with_lambda_context(req.context));
-        Box::pin(async move {
-            let res = fut.await?;
-            let (parts, body) = res.into_parts();
-
-            let mut headers = parts.headers;
-            let cookies = headers
-                .get_all(SET_COOKIE)
-                .iter()
-                .map(|c| String::from_utf8_lossy(c.as_bytes()).to_string())
-                .collect::<Vec<_>>();
-            headers.remove(SET_COOKIE);
-
-            Ok(StreamResponse {
-                metadata_prelude: MetadataPrelude {
-                    headers,
-                    status_code: parts.status,
-                    cookies,
-                },
-                stream: BodyStream { body },
-            })
-        })
+        Box::pin(
+            self.service
+                .call(event.with_lambda_context(req.context))
+                .map_ok(into_stream_response),
+        )
     }
 }
 
+/// Converts an `http::Response<B>` into a streaming Lambda response.
+fn into_stream_response<B>(res: Response<B>) -> StreamResponse<BodyStream<B>>
+where
+    B: Body + Unpin + Send + 'static,
+    B::Data: Into<Bytes> + Send,
+    B::Error: Into<Error> + Send + Debug,
+{
+    let (parts, body) = res.into_parts();
+
+    let mut headers = parts.headers;
+    let cookies = headers
+        .get_all(SET_COOKIE)
+        .iter()
+        .map(|c| String::from_utf8_lossy(c.as_bytes()).to_string())
+        .collect::<Vec<_>>();
+    headers.remove(SET_COOKIE);
+
+    StreamResponse {
+        metadata_prelude: MetadataPrelude {
+            headers,
+            status_code: parts.status,
+            cookies,
+        },
+        stream: BodyStream { body },
+    }
+}
+
+/// Builds a streaming-aware Tower service from a `Service<Request>` **without**
+/// boxing its future (no heap allocation / vtable).
+///
+/// Transforms `LambdaEvent<LambdaRequest>` into `Request` with Lambda context
+/// and wraps `Response<B>` into `StreamResponse<BodyStream<B>>`.
+///
+/// Used internally by [`run_with_streaming_response`]; not part of the public
+/// API.
+#[allow(clippy::type_complexity)]
+fn into_streaming_response<'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(|req: LambdaEvent<LambdaRequest>| {
+            let event: Request = req.payload.into();
+            event.with_lambda_context(req.context)
+        })
+        .service(handler)
+        .map_response(into_stream_response)
+}
+
 /// Runs the Lambda runtime with a handler that returns **streaming** HTTP
 /// responses.
 ///
@@ -98,7 +143,7 @@ where
     B::Data: Into<Bytes> + Send,
     B::Error: Into<Error> + Send + Debug,
 {
-    lambda_runtime::run(StreamAdapter::from(handler)).await
+    lambda_runtime::run(into_streaming_response(handler)).await
 }
 
 pin_project_lite::pin_project! {
@@ -127,3 +172,59 @@ where
         }
     }
 }
+
+#[cfg(test)]
+mod test_stream_adapter {
+    use super::*;
+
+    use crate::{
+        tower::{ServiceBuilder, ServiceExt},
+        Body, Request, StreamAdapter,
+    };
+    use http::StatusCode;
+
+    // A middleware that logs requests before forwarding them to another service
+    struct LogService<S> {
+        inner: S,
+    }
+
+    impl<S> Service<LambdaEvent<LambdaRequest>> for LogService<S>
+    where
+        S: Service<LambdaEvent<LambdaRequest>>,
+    {
+        type Response = S::Response;
+        type Error = S::Error;
+        type Future = S::Future;
+
+        fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            self.inner.poll_ready(cx)
+        }
+
+        fn call(&mut self, event: LambdaEvent<LambdaRequest>) -> Self::Future {
+            // Log the request
+            println!("Lambda event: {event:#?}");
+
+            self.inner.call(event)
+        }
+    }
+
+    /// This tests that `StreamAdapter` can be used in a `tower::Service` where
+    /// the user may require additional middleware between `lambda_runtime::run`
+    /// and where the `LambdaEvent` is converted into a `Request`.
+    #[test]
+    fn stream_adapter_is_boxable() {
+        let _svc = ServiceBuilder::new()
+            .layer_fn(|service| {
+                // This could be any middleware that logs, inspects, or
+                // manipulates the `LambdaEvent` before it's converted to a
+                // `Request` by `Adapter`.
+
+                LogService { inner: service }
+            })
+            .layer_fn(StreamAdapter::from)
+            .service_fn(
+                |_req: Request| async move { http::Response::builder().status(StatusCode::OK).body(Body::Empty) },
+            )
+            .boxed();
+    }
+}