awslabs / aws-lambda-rust-runtime

A Rust runtime for AWS Lambda
Apache License 2.0
3.31k stars 338 forks source link

Mixing buffered and streaming responses #805

Closed GabrielBianconi closed 7 months ago

GabrielBianconi commented 7 months ago

I'm trying to build a Lambda with routes for both streaming (Axum SSE) and regular (buffered) handlers.

I've tried using run_with_streaming_response, which works as expected for the streaming handlers but sends an empty response for regular handlers.

With run, the regular handler works as expected but it also buffers the streaming handler before responding.

Here's a minimal example:

use axum::{
    response::{
        sse::{Event, Sse},
        IntoResponse,
    },
    routing::get,
    Router,
};
use futures::stream::{self, Stream};
use std::{convert::Infallible, time::Duration};
use tokio_stream::StreamExt as _;

async fn regular_handler() -> impl IntoResponse {
    "Buffered hello world!"
}

async fn sse_handler() -> Sse<impl Stream<Item = Result<Event, Infallible>>> {
    let stream = stream::repeat_with(|| Event::default().data("Streamed hello world!"))
        .map(Ok)
        .throttle(Duration::from_secs(1))
        .take(5);

    Sse::new(stream).keep_alive(
        axum::response::sse::KeepAlive::new()
            .interval(Duration::from_secs(1))
            .text("keep-alive-text"),
    )
}

#[tokio::main]
async fn main() -> Result<(), lambda_http::Error> {
    let app = Router::new()
        .route("/a", get(regular_handler))
        .route("/b", get(sse_handler));

    // lambda_http::run(app).await // --> /a works fine, but also buffers /b before responding

    lambda_http::run_with_streaming_response(app).await // --> /a returns an empty reply from server, /b streams as expected
}

Is there a way to set up a Lambda that accepts both types of handlers?

Thank you.

bnusunny commented 7 months ago

Lambda response streaming is decided by the invoke mode of function URLs. When you configure a function URL to response streaming, all the invokes are streamed. It is not possible to switch between buffered and streamed responses per invoke, unless you build a custom frontend to invoke the lambda function using Lambda API directly.

GabrielBianconi commented 7 months ago

Makes sense, thank you.

github-actions[bot] commented 7 months ago

This issue is now closed. Comments on closed issues are hard for our team to see. If you need more assistance, please either tag a team member or open a new issue that references this one.