Add support for Lambda streaming response (awslabs#628)

bnusunny · web-flow · commit 3ee4389df868 · 2023-04-09T05:06:54.000Z
diff --git a/examples/basic-streaming-response/Cargo.toml b/examples/basic-streaming-response/Cargo.toml
@@ -0,0 +1,18 @@
+[package]
+name = "basic-streaming-response"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+hyper = { version = "0.14", features = [
+    "http1",
+    "client",
+    "stream",
+] }
+lambda_runtime = { path = "../../lambda-runtime" }
+tokio = { version = "1", features = ["macros"] }
+tracing = { version = "0.1", features = ["log"] }
+tracing-subscriber = { version = "0.3", default-features = false, features = ["ansi", "fmt"] }
+serde_json = "1.0"
diff --git a/examples/basic-streaming-response/README.md b/examples/basic-streaming-response/README.md
@@ -0,0 +1,13 @@
+# AWS Lambda Function example
+
+## Build & Deploy
+
+1. Install [cargo-lambda](https://github.com/cargo-lambda/cargo-lambda#installation)
+2. Build the function with `cargo lambda build --release`
+3. Deploy the function to AWS Lambda with `cargo lambda deploy --enable-function-url --iam-role YOUR_ROLE`
+4. Enable Lambda streaming response on Lambda console: change the function url's invoke mode to `RESPONSE_STREAM`
+5. Verify the function works: `curl <function-url>`. The results should be streamed back with 0.5 second pause between each word.
+
+## Build for ARM 64
+
+Build the function with `cargo lambda build --release --arm64`
diff --git a/examples/basic-streaming-response/src/main.rs b/examples/basic-streaming-response/src/main.rs
@@ -0,0 +1,42 @@
+use hyper::{body::Body, Response};
+use lambda_runtime::{service_fn, Error, LambdaEvent};
+use serde_json::Value;
+use std::{thread, time::Duration};
+
+async fn func(_event: LambdaEvent<Value>) -> Result<Response<Body>, Error> {
+    let messages = vec!["Hello", "world", "from", "Lambda!"];
+
+    let (mut tx, rx) = Body::channel();
+
+    tokio::spawn(async move {
+        for message in messages.iter() {
+            tx.send_data((message.to_string() + "\n").into()).await.unwrap();
+            thread::sleep(Duration::from_millis(500));
+        }
+    });
+
+    let resp = Response::builder()
+        .header("content-type", "text/html")
+        .header("CustomHeader", "outerspace")
+        .body(rx)?;
+
+    Ok(resp)
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Error> {
+    // required to enable CloudWatch error logging by the runtime
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        // disable printing the name of the module in every log line.
+        .with_target(false)
+        // this needs to be set to false, otherwise ANSI color codes will
+        // show up in a confusing manner in CloudWatch logs.
+        .with_ansi(false)
+        // disabling time is handy because CloudWatch will add the ingestion time.
+        .without_time()
+        .init();
+
+    lambda_runtime::run_with_streaming_response(service_fn(func)).await?;
+    Ok(())
+}
diff --git a/lambda-http/Cargo.toml b/lambda-http/Cargo.toml
@@ -23,19 +23,20 @@ apigw_websockets = []
 alb = []
 
 [dependencies]
-base64 = "0.13.0"
-bytes = "1"
+base64 = "0.21"
+bytes = "1.4"
+futures = "0.3"
 http = "0.2"
 http-body = "0.4"
-hyper = "0.14.20"
+hyper = "0.14"
 lambda_runtime = { path = "../lambda-runtime", version = "0.7" }
-serde = { version = "^1", features = ["derive"] }
-serde_json = "^1"
-serde_urlencoded = "0.7.0"
-mime = "0.3.16"
-encoding_rs = "0.8.31"
-url = "2.2.2"
-percent-encoding = "2.2.0"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+serde_urlencoded = "0.7"
+mime = "0.3"
+encoding_rs = "0.8"
+url = "2.2"
+percent-encoding = "2.2"
 
 [dependencies.aws_lambda_events]
 version = "^0.7.2"
diff --git a/lambda-http/src/lib.rs b/lambda-http/src/lib.rs
@@ -92,6 +92,9 @@ use std::{
     task::{Context as TaskContext, Poll},
 };
 
+mod streaming;
+pub use streaming::run_with_streaming_response;
+
 /// Type alias for `http::Request`s with a fixed [`Body`](enum.Body.html) type
 pub type Request = http::Request<Body>;
 
diff --git a/lambda-http/src/streaming.rs b/lambda-http/src/streaming.rs
@@ -0,0 +1,34 @@
+use crate::request::LambdaRequest;
+use crate::tower::ServiceBuilder;
+use crate::{Request, RequestExt};
+pub use aws_lambda_events::encodings::Body as LambdaEventBody;
+use bytes::Bytes;
+pub use http::{self, Response};
+use http_body::Body;
+use lambda_runtime::LambdaEvent;
+pub use lambda_runtime::{self, service_fn, tower, Context, Error, Service};
+use std::fmt::{Debug, Display};
+
+/// Starts the Lambda Rust runtime and stream response back [Configure Lambda
+/// Streaming Response](https://docs.aws.amazon.com/lambda/latest/dg/configuration-response-streaming.html).
+///
+/// This takes care of transforming the LambdaEvent into a [`Request`] and
+/// accepts [`http::Response<http_body::Body>`] as response.
+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 + Display,
+    B: Body + Unpin + Send + 'static,
+    B::Data: Into<Bytes> + Send,
+    B::Error: Into<Error> + Send + Debug,
+{
+    let svc = ServiceBuilder::new()
+        .map_request(|req: LambdaEvent<LambdaRequest>| {
+            let event: Request = req.payload.into();
+            event.with_lambda_context(req.context)
+        })
+        .service(handler);
+
+    lambda_runtime::run_with_streaming_response(svc).await
+}
diff --git a/lambda-runtime-api-client/src/lib.rs b/lambda-runtime-api-client/src/lib.rs
@@ -53,7 +53,9 @@ where
 
     /// Create a new client with a given base URI and HTTP connector.
     pub fn with(base: Uri, connector: C) -> Self {
-        let client = hyper::Client::builder().build(connector);
+        let client = hyper::Client::builder()
+            .http1_max_buf_size(1024 * 1024)
+            .build(connector);
         Self { base, client }
     }
 
diff --git a/lambda-runtime/src/lib.rs b/lambda-runtime/src/lib.rs
@@ -34,6 +34,9 @@ mod simulated;
 /// Types available to a Lambda function.
 mod types;
 
+mod streaming;
+pub use streaming::run_with_streaming_response;
+
 use requests::{EventCompletionRequest, EventErrorRequest, IntoRequest, NextEventRequest};
 pub use types::{Context, LambdaEvent};
 
diff --git a/lambda-runtime/src/streaming.rs b/lambda-runtime/src/streaming.rs

Original file line number	Diff line number	Diff line change
`@@ -53,7 +53,9 @@ where`
`53`	`53`
`54`	`54`	`/// Create a new client with a given base URI and HTTP connector.`
`55`	`55`	`pub fn with(base: Uri, connector: C) -> Self {`
`56`		`- let client = hyper::Client::builder().build(connector);`
	`56`	`+ let client = hyper::Client::builder()`
	`57`	`+ .http1_max_buf_size(1024 * 1024)`
	`58`	`+ .build(connector);`
`57`	`59`	`Self { base, client }`
`58`	`60`	`}`
`59`	`61`