Expose streaming API

[s0l0ist]

📬 Issue #, if available:

The SDK supports streaming, but lambda_http::Adapter only handles buffered responses from Axum's Router (i.e., a Service). This change allows direct streaming from Axum handlers/services, while specifying a custom runtime (i.e., with OTeL).

For example, lambda_http::run_with_streaming_response(app).await doesn't allow you to specify a custom runtime. So this change helps to abstract out constructing the stream response and allowing a custom runtime.

Related (ish):

• run_with_streaming_response and automatically streaming buffered responses #967
• Mixing buffered and streaming responses #805

✍️ Description of changes:

This PR exposes:

• into_streaming_response: Convert a Service into an AWS Lambda streaming response.

Which was originally internal to the run_with_streaming_response function. There are no functional changes.

There might be a more ergonomic way to expose an API like this, but I'm not aware. Happy to make changes as necessary so we're not exposing internals that may change in the future.

This is how you can use it with a custom runtime supporting OTeL:

use crate::{error::BoxError, service::build_service};
use axum::Router;
use lambda_http::into_streaming_response;
use lambda_runtime::{
    Runtime,
    layers::{OpenTelemetryFaasTrigger, OpenTelemetryLayer},
};
use opentelemetry_sdk::trace as sdktrace;

/// The main function to run the AWS Lambda
pub async fn async_main(tracer_provider: sdktrace::SdkTracerProvider) -> Result<(), BoxError> {
    let (app, _): (Router, _) = build_service().await;

    // For a buffered response:
    // let handler = lambda_http::Adapter::from(app);
    // For a streamed response:
    let handler = into_streaming_response(app);

    let runtime = Runtime::new(handler).layer(
        // Create a tracing span for each Lambda invocation
        OpenTelemetryLayer::new(|| {
            if let Err(err) = tracer_provider.force_flush() {
                eprintln!("Error flushing traces: {:#?}", err);
            }
        })
        .with_trigger(OpenTelemetryFaasTrigger::Http),
    );
    runtime.run().await
}

🔏 By submitting this pull request

• I confirm that I've ran cargo +nightly fmt.
• I confirm that I've ran cargo clippy --fix.
• I confirm that I've made a best effort attempt to update all relevant documentation.
• I confirm that my contribution is made under the terms of the Apache 2.0 license.

[jlizen]

Hi @s0l0ist , thanks for cutting this PR. This change seems very reasonable to me.

I do have concerns about leaking internals. This is only naming return types that are already public, but we might want to tweak their composition or otherwise shift bounds in ways that break callers. I think it would be best to type erase the returned service stack.

The easiest way to do that would be via a tower::util::BoxService / BoxLayer (or the sync/clone-bounded variants, though I don't think we need them here - but we could certainly add alternate APIs that do include the sync/clone bounds if it would be useful for how the returned struct is used by a runtime).

That adds a small amount of performance overhead due to an extra allocation and layer of dynamic dispatch. But, I think the ergonomics would be much better compared to a more complex, composable builder-style API using generics and sealed inner layer types. Note that we would probably want some sort of into_streaming_streaming_response_inner() that is what into_streaming_response() is currently, so that our run_with_streaming_response can skip the minor type erasure overhead.

How does that sound to you?

It would also be great to get a small example showing usage of this with a non-tokio runtime, if you'd be open to it! That would both let us validate the API, and make it more discoverable for users. I'd probably be ok with this landing without that, though, if you don't have cycles.

[jlizen]

This formatting could use cleaning up also, some extra newlines.

[s0l0ist]

@jlizen, I think that's the right direction. Currently busy right now, but can revisit this PR in a couple of weeks from now.

[s0l0ist]

@jlizen - does this API work? I'll add an example or two if I have time, but wanted to make sure this was aligned beforehand.

[jlizen]

General approach looks good, but left a note around newtyping the BoxService so we don't lock ourselves into extra Service: Send + 'static bounds for all time. (For now is fine though, just want to leave the door open to removing them).

I'm seeing that we don't have ANY tests currently of the run_with_streaming_response() API... that is unfortunate... An example would be really great if you have the time (which would double as a test that things compile at least). And then you could trivially validate e2e with cargo-lambda at that point. (Or feel free to throw in an integration test if you want our CI/CD to probe the e2e, up to you).

[jlizen]

That's a bit awkward that we have the new Send + 'static bounds on the service since we don't really need them... I see that is required for boxed and BoxService. That is unfortunate.

But, the machinery to spin our own is quite complicated, i'm not sure it's worth it unless we have use cases for it.

I guess one thing I would suggest would be, we could look at newtyping the BoxService that this responds with so that we have the option to swap it out with something with less restrictive bounds down the road?

Today it would be a breaking change to rip out BoxService for a handrolled impl that didn't require a Send + 'static Service.

[jlizen]

Ie something like:

pub struct TypeErasedService<T, U, E> {
    inner: BoxService<T, U, E>
}

impl <T, U, E> Service<T> for TypeErasedService<T, U, E> {
     type Response = U;

    type Error = E;

    type Future = BoxFuture<U, E>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        self.inner.get_mut().poll_ready(cx)
    }


    fn call(&mut self, request: T) -> BoxFuture<U, E> {
        self.inner.get_mut().call(request)
    }
}
impl<T, U, E> fmt::Debug for TypeErasedService<T, U, E> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        fmt.debug_struct("TypeErasedService").finish()
    }
}

That way we could today require the Send + 'static input bound, which will allow us to construct the inner BoxService, but we could later loose that restriction and use a different inner.