Allow connection close for custom streams (#1603)

* feat: allow connection close for custom streams

* fix: avoid req access since might already be released

* fix: fix aloc test fails

* fix: race condition when clossing body stream

2 files changed, 6 insertions, 3 deletions

diff --git a/client.go b/client.go
index 8e7b2f9..14321d1 100644
--- a/client.go
+++ b/client.go
@@ -2894,8 +2894,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
@@ -2906,10 +2906,11 @@ 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)
 			}
-			if closeConn {
+			if closeConn || resp.ConnectionClose() {
 				hc.closeConn(cc)
 			} else {
 				hc.releaseConn(cc)
@@ -2917,6 +2918,8 @@ func (t *transport) RoundTrip(hc *HostClient, req *Request, resp *Response) (ret
 			return nil
 		})
 		return false, nil
+	} else {
+		hc.releaseReader(br)
 	}
 
 	if closeConn {
diff --git a/http.go b/http.go
index 6640be2..0a5b446 100644
--- a/http.go
+++ b/http.go
@@ -1108,8 +1108,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