From 4a40b302509d7293435acc924d8be939907eb1cf Mon Sep 17 00:00:00 2001 From: Armin Date: Fri, 11 Aug 2023 11:47:06 +0200 Subject: [PATCH 1/4] feat: allow connection close for custom streams --- client.go | 1 + 1 file changed, 1 insertion(+) diff --git a/client.go b/client.go index 52230316b6..358ea9ef7a 100644 --- a/client.go +++ b/client.go @@ -2905,6 +2905,7 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret if customStreamBody && resp.bodyStream != nil { rbs := resp.bodyStream resp.bodyStream = newCloseReader(rbs, func() error { + closeConn = closeConn || req.ConnectionClose() || resp.ConnectionClose() if r, ok := rbs.(*requestStream); ok { releaseRequestStream(r) } From caa857bb83b0ee09178098597220c4f97dd28c37 Mon Sep 17 00:00:00 2001 From: Armin Date: Fri, 11 Aug 2023 14:58:19 +0200 Subject: [PATCH 2/4] fix: avoid req access since might already be released --- client.go | 2 +- http.go | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/client.go b/client.go index 358ea9ef7a..1db54e49d3 100644 --- a/client.go +++ b/client.go @@ -2905,7 +2905,7 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret if customStreamBody && resp.bodyStream != nil { rbs := resp.bodyStream resp.bodyStream = newCloseReader(rbs, func() error { - closeConn = closeConn || req.ConnectionClose() || resp.ConnectionClose() + closeConn = closeConn || resp.ConnectionClose() if r, ok := rbs.(*requestStream); ok { releaseRequestStream(r) } diff --git a/http.go b/http.go index 595eafa9c2..788f67d06d 100644 --- a/http.go +++ b/http.go @@ -1106,8 +1106,8 @@ func (resp *Response) Reset() { if responseBodyPoolSizeLimit >= 0 && resp.body != nil { resp.ReleaseBody(responseBodyPoolSizeLimit) } - resp.Header.Reset() resp.resetSkipHeader() + resp.Header.Reset() resp.SkipBody = false resp.raddr = nil resp.laddr = nil From 7ccc02ab3524ced46c101ef5c8562f1b5b9c4ec7 Mon Sep 17 00:00:00 2001 From: Armin Date: Mon, 28 Aug 2023 09:53:32 +0200 Subject: [PATCH 3/4] fix: fix aloc test fails --- client.go | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/client.go b/client.go index 1db54e49d3..fa6be5e1ae 100644 --- a/client.go +++ b/client.go @@ -2905,11 +2905,10 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret if customStreamBody && resp.bodyStream != nil { rbs := resp.bodyStream resp.bodyStream = newCloseReader(rbs, func() error { - closeConn = closeConn || resp.ConnectionClose() if r, ok := rbs.(*requestStream); ok { releaseRequestStream(r) } - if closeConn { + if closeConn || resp.ConnectionClose() { hc.closeConn(cc) } else { hc.releaseConn(cc) From 56c46eac498053e31c621b38dbc6e2eff1ff2d4d Mon Sep 17 00:00:00 2001 From: Armin Date: Mon, 28 Aug 2023 14:29:59 +0200 Subject: [PATCH 4/4] fix: race condition when clossing body stream --- client.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/client.go b/client.go index fa6be5e1ae..003d8aaf24 100644 --- a/client.go +++ b/client.go @@ -2893,8 +2893,8 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret br := hc.acquireReader(conn) err = resp.ReadLimitBody(br, hc.MaxResponseBodySize) - hc.releaseReader(br) if err != nil { + hc.releaseReader(br) hc.closeConn(cc) // Don't retry in case of ErrBodyTooLarge since we will just get the same again. needRetry := err != ErrBodyTooLarge @@ -2905,6 +2905,7 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret if customStreamBody && resp.bodyStream != nil { rbs := resp.bodyStream resp.bodyStream = newCloseReader(rbs, func() error { + hc.releaseReader(br) if r, ok := rbs.(*requestStream); ok { releaseRequestStream(r) } @@ -2916,6 +2917,8 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret return nil }) return false, nil + } else { + hc.releaseReader(br) } if closeConn {