Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

fix: propagate body stream error to close function (#1743) #1757

Merged
merged 6 commits into from
Apr 22, 2024
Merged

fix: propagate body stream error to close function (#1743) #1757

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b863648
fix: propagate body stream error to close function (#1743)
Apr 12, 2024
911cf1f
fix: http test
Apr 12, 2024
efce2c7
fix: close body stream with error in encoding functions
Apr 13, 2024
c1f124c
Merge remote-tracking branch 'origin/master'
Apr 22, 2024
0345190
Merge branch 'master' into fix/stream-body-error-propagate
Apr 22, 2024
c1d0020
fix: lint
Apr 22, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2975,12 +2975,12 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret
closeConn := resetConnection || req.ConnectionClose() || resp.ConnectionClose() || isConnRST
if customStreamBody && resp.bodyStream != nil {
rbs := resp.bodyStream
resp.bodyStream = newCloseReader(rbs, func() error {
resp.bodyStream = newCloseReaderWithError(rbs, func(wErr error) error {
hc.releaseReader(br)
if r, ok := rbs.(*requestStream); ok {
releaseRequestStream(r)
}
if closeConn || resp.ConnectionClose() {
if closeConn || resp.ConnectionClose() || wErr != nil {
hc.closeConn(cc)
} else {
hc.releaseConn(cc)
Expand Down
76 changes: 48 additions & 28 deletions http.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -321,26 +321,31 @@ func (resp *Response) BodyStream() io.Reader {
}

func (resp *Response) CloseBodyStream() error {
return resp.closeBodyStream()
return resp.closeBodyStream(nil)
}

type ReadCloserWithError interface {
io.Reader
CloseWithError(err error) error
}

type closeReader struct {
io.Reader
closeFunc func() error
closeFunc func(err error) error
}

func newCloseReader(r io.Reader, closeFunc func() error) io.ReadCloser {
func newCloseReaderWithError(r io.Reader, closeFunc func(err error) error) ReadCloserWithError {
if r == nil {
panic(`BUG: reader is nil`)
}
return &closeReader{Reader: r, closeFunc: closeFunc}
}

func (c *closeReader) Close() error {
func (c *closeReader) CloseWithError(err error) error {
if c.closeFunc == nil {
return nil
}
return c.closeFunc()
return c.closeFunc(err)
}

// BodyWriter returns writer for populating request body.
Expand Down Expand Up @@ -394,7 +399,7 @@ func (resp *Response) Body() []byte {
bodyBuf := resp.bodyBuffer()
bodyBuf.Reset()
_, err := copyZeroAlloc(bodyBuf, resp.bodyStream)
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(err) //nolint:errcheck
if err != nil {
bodyBuf.SetString(err.Error())
}
Expand Down Expand Up @@ -618,7 +623,7 @@ func (req *Request) BodyWriteTo(w io.Writer) error {
func (resp *Response) BodyWriteTo(w io.Writer) error {
if resp.bodyStream != nil {
_, err := copyZeroAlloc(w, resp.bodyStream)
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(err) //nolint:errcheck
return err
}
_, err := w.Write(resp.bodyBytes())
Expand All @@ -629,29 +634,29 @@ func (resp *Response) BodyWriteTo(w io.Writer) error {
//
// It is safe re-using p after the function returns.
func (resp *Response) AppendBody(p []byte) {
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
resp.bodyBuffer().Write(p) //nolint:errcheck
}

// AppendBodyString appends s to response body.
func (resp *Response) AppendBodyString(s string) {
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
resp.bodyBuffer().WriteString(s) //nolint:errcheck
}

// SetBody sets response body.
//
// It is safe re-using body argument after the function returns.
func (resp *Response) SetBody(body []byte) {
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
bodyBuf := resp.bodyBuffer()
bodyBuf.Reset()
bodyBuf.Write(body) //nolint:errcheck
}

// SetBodyString sets response body.
func (resp *Response) SetBodyString(body string) {
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
bodyBuf := resp.bodyBuffer()
bodyBuf.Reset()
bodyBuf.WriteString(body) //nolint:errcheck
Expand All @@ -660,7 +665,7 @@ func (resp *Response) SetBodyString(body string) {
// ResetBody resets response body.
func (resp *Response) ResetBody() {
resp.bodyRaw = nil
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
if resp.body != nil {
if resp.keepBodyBuffer {
resp.body.Reset()
Expand Down Expand Up @@ -700,7 +705,7 @@ func (resp *Response) ReleaseBody(size int) {
return
}
if cap(resp.body.B) > size {
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(nil) //nolint:errcheck
resp.body = nil
}
}
Expand Down Expand Up @@ -734,7 +739,7 @@ func (resp *Response) SwapBody(body []byte) []byte {
if resp.bodyStream != nil {
bb.Reset()
_, err := copyZeroAlloc(bb, resp.bodyStream)
resp.closeBodyStream() //nolint:errcheck
resp.closeBodyStream(err) //nolint:errcheck
if err != nil {
bb.Reset()
bb.SetString(err.Error())
Expand Down Expand Up @@ -1725,10 +1730,13 @@ func (resp *Response) brotliBody(level int) {
wf: zw,
bw: sw,
}
copyZeroAlloc(fw, bs) //nolint:errcheck
_, wErr := copyZeroAlloc(fw, bs)
releaseStacklessBrotliWriter(zw, level)
if bsc, ok := bs.(io.Closer); ok {
bsc.Close()
switch v := bs.(type) {
case io.Closer:
v.Close()
case ReadCloserWithError:
v.CloseWithError(wErr) //nolint:errcheck
}
})
} else {
Expand Down Expand Up @@ -1780,10 +1788,13 @@ func (resp *Response) gzipBody(level int) {
wf: zw,
bw: sw,
}
copyZeroAlloc(fw, bs) //nolint:errcheck
_, wErr := copyZeroAlloc(fw, bs)
releaseStacklessGzipWriter(zw, level)
if bsc, ok := bs.(io.Closer); ok {
bsc.Close()
switch v := bs.(type) {
case io.Closer:
v.Close()
case ReadCloserWithError:
v.CloseWithError(wErr) //nolint:errcheck
}
})
} else {
Expand Down Expand Up @@ -1835,10 +1846,13 @@ func (resp *Response) deflateBody(level int) {
wf: zw,
bw: sw,
}
copyZeroAlloc(fw, bs) //nolint:errcheck
_, wErr := copyZeroAlloc(fw, bs)
releaseStacklessDeflateWriter(zw, level)
if bsc, ok := bs.(io.Closer); ok {
bsc.Close()
switch v := bs.(type) {
case io.Closer:
v.Close()
case ReadCloserWithError:
v.CloseWithError(wErr) //nolint:errcheck
}
})
} else {
Expand Down Expand Up @@ -1887,10 +1901,13 @@ func (resp *Response) zstdBody(level int) {
wf: zw,
bw: sw,
}
copyZeroAlloc(fw, bs) //nolint:errcheck
_, wErr := copyZeroAlloc(fw, bs)
releaseStacklessZstdWriter(zw, level)
if bsc, ok := bs.(io.Closer); ok {
bsc.Close()
switch v := bs.(type) {
case io.Closer:
v.Close()
case ReadCloserWithError:
v.CloseWithError(wErr) //nolint:errcheck
}
})
} else {
Expand Down Expand Up @@ -2053,7 +2070,7 @@ func (resp *Response) writeBodyStream(w *bufio.Writer, sendBody bool) (err error
}
}
}
errc := resp.closeBodyStream()
errc := resp.closeBodyStream(err)
if err == nil {
err = errc
}
Expand All @@ -2075,14 +2092,17 @@ func (req *Request) closeBodyStream() error {
return err
}

func (resp *Response) closeBodyStream() error {
func (resp *Response) closeBodyStream(wErr error) error {
if resp.bodyStream == nil {
return nil
}
var err error
if bsc, ok := resp.bodyStream.(io.Closer); ok {
err = bsc.Close()
}
if bsc, ok := resp.bodyStream.(ReadCloserWithError); ok {
err = bsc.CloseWithError(wErr)
}
if bsr, ok := resp.bodyStream.(*requestStream); ok {
releaseRequestStream(bsr)
}
Expand Down
2 changes: 1 addition & 1 deletion http_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2943,7 +2943,7 @@ func TestResponseBodyStream(t *testing.T) {
t.Fatalf("parse response find err: %v", err)
}
defer func() {
if err := response.closeBodyStream(); err != nil {
if err := response.closeBodyStream(nil); err != nil {
t.Fatalf("close body stream err: %v", err)
}
}()
Expand Down
Loading