net/http: allow reuse connection if entire unread body is in buffer, v2

Author: canoriz

src/net/http/client_test.go

@@ -845,6 +845,7 @@ func testClientInsecureTransport(t *testing.T, mode testMode) {
 		if res != nil {
 			res.Body.Close()
 		}
+		c.CloseIdleConnections()
 	}
 
 	cst.close()

src/net/http/transport.go

@@ -2359,13 +2359,27 @@ func (pc *persistConn) readLoop() {
 		}
 
 		waitForBodyRead := make(chan bool, 2)
+
 		body := &bodyEOFSignal{
 			body: resp.Body,
-			earlyCloseFn: func() error {
-				waitForBodyRead <- false
+			earlyCloseFn: func(r io.ReadCloser) error {
+				isEOF := false
+				if b, ok := r.(*body); ok {
+					if lr, ok := b.src.(*io.LimitedReader); ok {
+						if br, ok := (lr.R).(*bufio.Reader); ok {
+							if lr.N == int64(br.Buffered()) {
+								// if bufio.Reader buffer have all bytes remaining in LimitReader,
+								// discard the buffer then reuse connection, set EOF flag.
+								b.sawEOF = true
+								br.Discard(br.Buffered())
+								isEOF = true
+							}
+						}
+					}
+				}
+				waitForBodyRead <- isEOF
 				<-eofc // will be closed by deferred call at the end of the function
 				return nil
-
 			},
 			fn: func(err error) error {
 				isEOF := err == io.EOF
@@ -2981,11 +2995,11 @@ func canonicalAddr(url *url.URL) string {
 // the return value from Close.
 type bodyEOFSignal struct {
 	body         io.ReadCloser
-	mu           sync.Mutex        // guards following 4 fields
-	closed       bool              // whether Close has been called
-	rerr         error             // sticky Read error
-	fn           func(error) error // err will be nil on Read io.EOF
-	earlyCloseFn func() error      // optional alt Close func used if io.EOF not seen
+	mu           sync.Mutex                // guards following 4 fields
+	closed       bool                      // whether Close has been called
+	rerr         error                     // sticky Read error
+	fn           func(error) error         // err will be nil on Read io.EOF
+	earlyCloseFn func(io.ReadCloser) error // optional alt Close func used if io.EOF not seen
 }
 
 var errReadOnClosedResBody = errors.New("http: read on closed response body")
@@ -3022,7 +3036,7 @@ func (es *bodyEOFSignal) Close() error {
 	}
 	es.closed = true
 	if es.earlyCloseFn != nil && es.rerr != io.EOF {
-		return es.earlyCloseFn()
+		return es.earlyCloseFn(es.body)
 	}
 	err := es.body.Close()
 	return es.condfn(err)

src/net/http/transport_test.go

@@ -474,6 +474,56 @@ func testTransportReadToEndReusesConn(t *testing.T, mode testMode) {
 	}
 }
 
+// Tests that the HTTP transport re-uses connections when a client
+// early closes a response Body but the content is fully read into the underlying
+// buffer. So we can discard the body buffer and reuse the connection.
+func TestTransportReusesEarlyCloseButAllReceivedConn(t *testing.T) {
+	run(t, testTransportReusesEarlyCloseButAllReceivedConn)
+}
+func testTransportReusesEarlyCloseButAllReceivedConn(t *testing.T, mode testMode) {
+	const msg = "foobar"
+
+	var addrSeen map[string]int
+	ts := newClientServerTest(t, mode, HandlerFunc(func(w ResponseWriter, r *Request) {
+		addrSeen[r.RemoteAddr]++
+		w.Header().Set("Content-Length", strconv.Itoa(len(msg)))
+		w.WriteHeader(200)
+		w.Write([]byte(msg))
+	})).ts
+
+	wantLen := len(msg)
+	addrSeen = make(map[string]int)
+	total := 5
+	for i := 0; i < total; i++ {
+		res, err := ts.Client().Get(ts.URL)
+		if err != nil {
+			t.Errorf("Get /: %v", err)
+			continue
+		}
+
+		if res.ContentLength != int64(wantLen) {
+			t.Errorf("res.ContentLength = %d; want %d", res.ContentLength, wantLen)
+		}
+
+		if i+1 < total {
+			// Close body directly. The body is small enough, so probably the underlying bufio.Reader
+			// has read entire body into buffer. Thus even if the body is not read, the buffer is discarded
+			// then connection is reused.
+			res.Body.Close()
+		} else {
+			// when reading body, everything should be same.
+			got, err := io.ReadAll(res.Body)
+			if string(got) != msg || err != nil {
+				t.Errorf("ReadAll(Body) = %q, %v; want %q, nil", string(got), err, msg)
+			}
+		}
+	}
+
+	if len(addrSeen) != 1 {
+		t.Errorf("server saw %d distinct client addresses; want 1", len(addrSeen))
+	}
+}
+
 func TestTransportMaxPerHostIdleConns(t *testing.T) {
 	run(t, testTransportMaxPerHostIdleConns, []testMode{http1Mode})
 }