awslabs / aws-lambda-rust-runtime

A Rust runtime for AWS Lambda

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Mixing buffered and streaming responses

GabrielBianconi opened this issue · comments

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.

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.

Makes sense, thank you.

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.