Skip to content

Commit

Permalink
🕳️ http: avoid draining response body when the client disconnects
Browse files Browse the repository at this point in the history
  • Loading branch information
database64128 committed Jun 6, 2024
1 parent 30df8d3 commit dd46d9f
Showing 1 changed file with 33 additions and 1 deletion.
34 changes: 33 additions & 1 deletion http/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ func NewHttpStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, logger *za
plbr := bufio.NewReader(pl)
plbw := bufio.NewWriter(pl)
rwbw := bufio.NewWriter(rw)
rwbwpcw := newPipeClosingWriter(rwbw, pl)

// The current implementation only supports a fixed destination host.
fixedHost := req.Host
Expand Down Expand Up @@ -144,7 +145,15 @@ func NewHttpStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, logger *za
}

// Write response.
if err = resp.Write(rwbw); err != nil {
//
// The Write method always drains the response body, even when the destination writer returns an error.
// This can become a problem when the response body is large or unbounded in size.
// A typical scenario is when the client downloads a large file but aborts the connection midway.
// The Write call will block until the entire file is downloaded, which is a total waste of resources.
// To mitigate this, we wrap the destination writer in a [*pipeClosingWriter] that stops further reads on write error.
//
// When we migrate to using [*http.Client], [*pipeClosingWriter] needs to be updated to cancel the request context instead.
if err = resp.Write(rwbwpcw); err != nil {
err = fmt.Errorf("failed to write HTTP response: %w", err)
break
}
Expand Down Expand Up @@ -234,3 +243,26 @@ func send502(w io.Writer) error {
_, err := w.Write([]byte("HTTP/1.1 502 Bad Gateway\r\n\r\n"))
return err
}

// pipeClosingWriter passes writes to the underlying [io.Writer] and closes the [*pipe.DuplexPipeEnd] on error.
type pipeClosingWriter struct {
w io.Writer
p *pipe.DuplexPipeEnd
}

// newPipeClosingWriter returns a new [pipeClosingWriter].
func newPipeClosingWriter(w io.Writer, p *pipe.DuplexPipeEnd) *pipeClosingWriter {
return &pipeClosingWriter{
w: w,
p: p,
}
}

// Write implements [io.Writer.Write].
func (w *pipeClosingWriter) Write(b []byte) (int, error) {
n, err := w.w.Write(b)
if err != nil {
w.p.CloseWithError(err)
}
return n, err
}

0 comments on commit dd46d9f

Please sign in to comment.