Handler design for streaming responses

AWS recently released support for streaming responses from Lambda functions. I can see that work is currently underway to add this to the Go SDK in #494, #495. It might be too late, but I'd like to raise an alternative API design for your consideration.

Is your feature request related to a problem? Please describe. Right now it looks like streaming responses will be supported by having a handler function return an io.Reader. I think that this could potentially make the implementation of streaming-response Lambda functions more error-prone and provide less control over the streamed output than using an io.Writer.

My understanding is that a streaming response handler ~today should be implemented like this:

func handleReturningReader(ctx context.Context, input json.RawMessage) (io.Reader, error) {
    pr, pw := io.Pipe()

    go func() {
        // if this function panics, the entire process crashes. it can't be caught by
        // callBytesHandlerFunc like a panic in handleReturningReader would be. i suspect
        // a lot of users will be affected by this.

        for idx := 0; idx < 10; idx++ {
            fmt.Fprintln(pw, "hello world!")
            time.Sleep(time.Second)
        }

        if err := doSomething(); err != nil {
            // all mid-stream errors need to be funneled through the pipe. that's not
            // the end of the world, but how many users will know how to do this? it's
            // not as convenient as returning an error
            pw.CloseWithError(err)
        }

        // finally, if the user forgets to close the pipe (e.g. maybe they returned early)
        // then the invocation will just hang until it times out. the docs can advise
        // that the pipe must always be closed, but it's another opportunity for error
        pw.Close()
    }()

    return pr, nil
}

Describe the solution you'd like Consider a handler signature like this:

func handleReceivingWriter(ctx context.Context, output io.Writer, input json.RawMessage) error {
    // any panics in this function will be caught by the invoke loop and turned into an error
    // response like usual

    for idx := 0; idx < 10; idx++ {
        // the user has control over the size of each chunk in the chunked transfer-encoding.
        // this can simplify the usage of InvokeWithResponseStream on the client side
        fmt.Fprintln(output, "hello world!")
        time.Sleep(time.Second)
    }

    if err := doSomething(); err != nil {
        // error handling is the normal if err != nil { return nil } convention
        return err
    }

    return nil
}

This handler signature avoids the necessity of spawning a goroutine and recovering from panics in it. It allows for standard error handling conventions. It also allows the user to have control over the size of the response chunks.

This potential design doesn't address how Lambda function URL response streaming would work. The application/vnd.awslambda.http-integration-response content-type doesn't appear to documented yet, but I've figured out how it works from #494. I can think of two options.

The first option is that there can be a lambda.WriteResponseHeaders(writer io.Writer, statusCode int, headers map[string]string, cookies []string) helper function that the user calls to write the prelude before streaming their response body. This would suffer from the issue of being non-obvious and users would have to read the docs to know that they need to use it.

The second option is that we can copy the design of stdlib http.Handler and instead of a plain io.Writer, we pass in a lambda.StreamingResponseWriter, e.g. like so:

// this just so happens to be identical to the http.ResponseWriter interface
type StreamingResponseWriter interface {
    Header() http.Header
    Write([]byte) (int, error)
    WriteHeader(statusCode int)
}

func handleReceivingResponseWriter(ctx context.Context, w lambda.StreamingResponseWriter, input json.RawMessage) error {
    w.WriteHeader(200)
    w.Header().Set("Content-Type", "text/plain")
    w.Header().Set("My-Header", "my value")

    // this works because lambda.StreamingResponseWriter matches http.ResponseWriter
    http.SetCookie(w, &http.Cookie{Name: "cookie", Value: "val", Path: "/my-path"})

    for idx := 0; idx < 10; idx++ {
        fmt.Fprintln(w, "hello world!")
        time.Sleep(time.Second)
    }

    if err := doSomething(); err != nil {
        // error handling is the normal if err != nil { return nil } convention
        return err
    }

    return nil
}

I haven't fully thought through this second option yet. I think there's definitely value in the familiarity aspect - most Go developers will recognise the http.Handler-like signature and know how to use it. An outstanding question on my mind is "do we need to support non-function URL streaming responses?" i.e. does there need to be a way to set the "top-level" content-type and not just a content-type within the function URL response "prelude". But it seems like the ability to set that content-type doesn't yet exist today anyway so maybe it's not necessary 🤔

Additional context I noticed that the Lambda-Runtime-Function-Response-Mode: streaming request header isn't set anywhere yet. Is it not needed?

I'm also aware that supporting a second handler type might be a bigger undertaking than squeezing support in via an io.Reader return type. I just wanted to mention the alternatives to see if they had been considered yet.

Supporting two different handler signatures might seem to over-complicate the package, but I also suppose it would match the Node.js library introducing the awslambda.streamifyResponse decorator and inner function with three parameters.

To add to this, using io.Writer instead would make JSON streams easier too, which I believe are common use cases here: json.NewEncoder(lambdaStream).Encode(data)

Note that the streams implementation should respect the error returned in those streaming cases and respond accordingly (error behavior should be explicitly defined then by the stream response handler).

It might be too late

It's not :)

The plumbing in https://github.com/aws/aws-lambda-go/pull/494 is the bare-minimum addition so that current users of events.LambdaFunctionURLResponse can make a tiny code change and start using InvokeMode: RESPONSE_STREAMING to unlock jumbo response sizes. I totally expected proposals like yours for an interface for doing things the io.Writer/http.Handler way too, so thanks for getting to conversation started!

i.e. does there need to be a way to set the "top-level" content-type and not just a content-type within the function URL response "prelude". But it seems like the ability to set that content-type doesn't yet exist today anyway so maybe it's not necessary

This is where I've had trouble in designing an singular io.Writer centric interface - today the only use case for setting the top-level content type is to switch on the new mode for Function URLs, which then also requires the prelude! So exposing a SetContentType in the primary writer interface everyone sees could cause confusion when it doesn't behave like the stdlib's w.Header().Set("Content-Type", "foo-bar")

There can of course be an additional Function URL specific interface that stacks ontop of a lower level lambda.ResponseWriter interface - the tradeoff being more bloated documentation.

The second option is that we can copy the design of stdlib http.Handler and instead of a plain io.Writer, we pass in a lambda.StreamingResponseWriter, e.g. like so:

Do you have an opinion on mimicking http.Handler vs directly using http.HandlerFunc? Like a subset of what gets done in https://github.com/awslabs/aws-lambda-go-api-proxy - I've thought about this as an alternative to designing a new interface that's only usable for streaming Function URL responses.

I'd be OK with a pairing like

// StartWriter is similar to lambda.StartHandlerFunc, but passes in the io.Writer that is directly conected to the Runtime API response stream
func StartWriter[EventT any](handler func(context.Context, io.Writer, EventT) error, options ...lambda.Option) {}

// StartHTTPWriter is similar to StartWriter, but handles the construction of a response that conforms to "application/vnd.awslambda.http-integration-response"
func StartHTTPWriter[EventT any](handler func(context.Context, http.ResponseWriter, EventT) error, options ...lambda.Option) {}

My lingering doubt is, if we do something that ties into the http stdlib, should we just go all way way to doing something like the following:

// StartHTTP transform an `http.HandlerFunc` into an `lambda.HandlerFunc[*events.LambdaFunctionURLRequest, *events.LambdaFunctionURLStreamingResponse]`
func StartHTTP(handler http.HandlerFunc, options ...lambda.Option) {}

Ah, that makes perfect sense now. I couldn't understand the purpose of #494 in isolation, but allowing larger payloads explains it.

HTTP responses I personally like the pattern enabled by the aws-lambda-go-api-proxy package. I always write my web Lambda functions using (wrapped) http.Handler because then I can use http.ListenAndServe during local development, not to mention frameworks that can extract path parameters, etc. The potential cons are:

Function URL response payloads don't support multi-valued response headers. Would anyone be led astray by expecting multi-valued response headers to work because http.ResponseWriter permits it?
Cookie handling. Would you just pass through Set-Cookie headers as-is, or would you convert them to the Cookies []string response expected by Lambda? (What even happens if you have Set-Cookie headers?)

I still err on the side of selfishly answering "yes, it would be great to have that functionality in aws-lambda-go", but I'm not responsible for maintaining it 😅

My lingering doubt is, if we do something that ties into the http stdlib, should we just go all [the] way

I feel it would make sense to go all the way. I would probably also expect the underlying *events.LambdaFunctionURLRequest to be extractible from the *http.Request.Context(). That said, if StartHTTP existed I would expect it to work in all scenarios:

API Gateway v1 (REST)
API Gateway v2 (HTTP)
Application Load Balancer
Function URLs (buffered)
Function URLs (streaming)

If it didn't work in all those scenarios, it could potentially be a liability: I could imagine people trying to use it with APIGW, it failing, and issues being opened on GitHub.

Non-HTTP responses I suppose the existence of the HTTP-specific functionality would mean that a lambda.ResponseWriter having a SetContentType(string) method is less confusing? The method's docs should probably have a note advising the user to refer to StartHTTP if they're trying to do HTTP.

Yeah there's a lot to get wrong in trying to support all the existing http proxy integrations. I think I wanna namespace any new http.Handler stuff into a new sub-package specific to Streaming Lambda Function URLs, like what was done for CloudFormation custom resources. https://github.com/aws/aws-lambda-go/pull/503 has where I've gotten to so far on this.

Closing this. https://github.com/aws/aws-lambda-go/releases/tag/v1.41.0 added the lambdaurl package, which in addition to the existing low-level option to return an io.Reader, I believe covers things.

If Lambda ever opens up the ability to do something clever with the Runtime API Response Content-Type, then I belive something like a handler entry with a lambda.ResponseWriter input can make sense.

aws / aws-lambda-go

Handler design for streaming responses #500