Skip to content

Commit 65b7b73

Browse files
committed
fix: communicate EOF properly in lambda streaming mode
1 parent 2dc9581 commit 65b7b73

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

handler/functionurl.go

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"context"
88
"encoding/base64"
99
"github.com/aws/aws-lambda-go/events"
10+
"io"
1011
"net/http"
1112
"strconv"
1213
"strings"
@@ -139,7 +140,7 @@ func NewFunctionURLHandler(adapter AdapterFunc) func(context.Context, events.Lam
139140
// region streaming
140141
type functionURLStreamingResponseWriter struct {
141142
headers http.Header
142-
body bytes.Buffer
143+
body io.WriteCloser
143144
res *events.LambdaFunctionURLStreamingResponse
144145
resCh chan<- *events.LambdaFunctionURLStreamingResponse
145146
}
@@ -155,10 +156,12 @@ func (w *functionURLStreamingResponseWriter) Write(p []byte) (int, error) {
155156

156157
func (w *functionURLStreamingResponseWriter) WriteHeader(statusCode int) {
157158
if w.res == nil {
159+
pr, pw := io.Pipe()
160+
w.body = pw
158161
w.res = &events.LambdaFunctionURLStreamingResponse{
159162
StatusCode: statusCode,
160163
Headers: make(map[string]string),
161-
Body: &w.body,
164+
Body: pr,
162165
Cookies: make([]string, 0),
163166
}
164167

@@ -195,6 +198,8 @@ func handleFunctionURLStreaming(ctx context.Context, event events.LambdaFunction
195198
}
196199

197200
go func() {
201+
defer w.body.Close()
202+
198203
if err := adapter(ctx, req, &w); err != nil {
199204
errCh <- err
200205
}

0 commit comments

Comments
 (0)