Skip to content

Commit

Permalink
fix: buffer reuse and release #917
Browse files Browse the repository at this point in the history
  • Loading branch information
jeevatkm committed Nov 22, 2024
1 parent d598157 commit 9a9930f
Show file tree
Hide file tree
Showing 4 changed files with 8 additions and 31 deletions.
1 change: 0 additions & 1 deletion client.go
Original file line number Diff line number Diff line change
Expand Up @@ -1236,7 +1236,6 @@ func (c *Client) executeBefore(req *Request) error {
return wrapNoRetryErr(err)
}

req.RawRequest.Body = newRequestBodyReleaser(req.RawRequest.Body, req.bodyBuf)
return nil
}

Expand Down
2 changes: 2 additions & 0 deletions request.go
Original file line number Diff line number Diff line change
Expand Up @@ -1030,6 +1030,7 @@ func (r *Request) Execute(method, url string) (*Response, error) {
r.Attempt = 1
resp, err = r.client.execute(r)
r.client.onErrorHooks(r, resp, unwrapNoRetryErr(err))
backToBufPool(r.bodyBuf)
return resp, unwrapNoRetryErr(err)
}

Expand Down Expand Up @@ -1059,6 +1060,7 @@ func (r *Request) Execute(method, url string) (*Response, error) {
}

r.client.onErrorHooks(r, resp, unwrapNoRetryErr(err))
backToBufPool(r.bodyBuf)
return resp, unwrapNoRetryErr(err)
}

Expand Down
6 changes: 2 additions & 4 deletions request_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2203,10 +2203,8 @@ func TestRequestGH917(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
b, err := io.ReadAll(r.Body)
assertError(t, err)
if len(b) > 0 {
// sometimes, the body is "testtest" instead of "test"
assertEqual(t, "test", string(b))
}
assertEqual(t, "test", string(b))

w.WriteHeader(http.StatusInternalServerError)
}))

Expand Down
30 changes: 4 additions & 26 deletions util.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import (
"runtime"
"sort"
"strings"
"sync"
)

//‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾‾
Expand Down Expand Up @@ -297,38 +296,17 @@ func acquireBuffer() *bytes.Buffer {

func releaseBuffer(buf *bytes.Buffer) {
if buf != nil {
buf.Reset()
bufPool.Put(buf)
}
}

// requestBodyReleaser wraps requests's body and implements custom Close for it.
// The Close method closes original body and releases request body back to sync.Pool.
type requestBodyReleaser struct {
releaseOnce sync.Once
reqBuf *bytes.Buffer
io.ReadCloser
}

func newRequestBodyReleaser(respBody io.ReadCloser, reqBuf *bytes.Buffer) io.ReadCloser {
if reqBuf == nil {
return respBody
}

return &requestBodyReleaser{
reqBuf: reqBuf,
ReadCloser: respBody,
func backToBufPool(buf *bytes.Buffer) {
if buf != nil {
bufPool.Put(buf)
}
}

func (rr *requestBodyReleaser) Close() error {
err := rr.ReadCloser.Close()
rr.releaseOnce.Do(func() {
releaseBuffer(rr.reqBuf)
})

return err
}

func closeq(v interface{}) {
if c, ok := v.(io.Closer); ok {
silently(c.Close())
Expand Down

0 comments on commit 9a9930f

Please sign in to comment.