net/http: avoid leaking writer goroutines in tests

In TestTransportPrefersResponseOverWriteError and TestMaxBytesHandler,
the server may respond to an incoming request without ever reading the
request body. The client's Do method will return as soon as the
server's response headers are read, but the Transport will remain
active until it notices that the server has closed the connection,
which may be arbitrarily later.

When the server has closed the connection, it will call the Close
method on the request body (if it has such a method). So we can use
that method to find out when the Transport is close enough to done for
the test to complete without interfering too much with other tests.

For #57612.
For #59526.

Change-Id: Iddc7a3b7b09429113ad76ccc1c090ebc9e1835a1
Reviewed-on: https://go-review.googlesource.com/c/go/+/483895
Run-TryBot: Bryan Mills <[email protected]>
Commit-Queue: Bryan Mills <[email protected]>
Auto-Submit: Bryan Mills <[email protected]>
TryBot-Result: Gopher Robot <[email protected]>
Reviewed-by: Michael Pratt <[email protected]>

Author: Bryan C. Mills

src/net/http/serve_test.go

@@ -6546,9 +6546,30 @@ func testMaxBytesHandler(t *testing.T, mode testMode, maxSize, requestSize int64
 	defer ts.Close()
 
 	c := ts.Client()
+
+	body := strings.Repeat("a", int(requestSize))
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	getBody := func() (io.ReadCloser, error) {
+		wg.Add(1)
+		body := &wgReadCloser{
+			Reader: strings.NewReader(body),
+			wg:     &wg,
+		}
+		return body, nil
+	}
+	reqBody, _ := getBody()
+	req, err := NewRequest("POST", ts.URL, reqBody)
+	if err != nil {
+		reqBody.Close()
+		t.Fatal(err)
+	}
+	req.ContentLength = int64(len(body))
+	req.GetBody = getBody
+	req.Header.Set("Content-Type", "text/plain")
+
 	var buf strings.Builder
-	body := strings.NewReader(strings.Repeat("a", int(requestSize)))
-	res, err := c.Post(ts.URL, "text/plain", body)
+	res, err := c.Do(req)
 	if err != nil {
 		t.Errorf("unexpected connection error: %v", err)
 	} else {

src/net/http/transport_test.go

@@ -4250,6 +4250,21 @@ func testTransportFlushesRequestHeader(t *testing.T, mode testMode) {
 	<-gotRes
 }
 
+type wgReadCloser struct {
+	io.Reader
+	wg     *sync.WaitGroup
+	closed bool
+}
+
+func (c *wgReadCloser) Close() error {
+	if c.closed {
+		return net.ErrClosed
+	}
+	c.closed = true
+	c.wg.Done()
+	return nil
+}
+
 // Issue 11745.
 func TestTransportPrefersResponseOverWriteError(t *testing.T) {
 	run(t, testTransportPrefersResponseOverWriteError)
@@ -4271,12 +4286,29 @@ func testTransportPrefersResponseOverWriteError(t *testing.T, mode testMode) {
 
 	fail := 0
 	count := 100
+
 	bigBody := strings.Repeat("a", contentLengthLimit*2)
+	var wg sync.WaitGroup
+	defer wg.Wait()
+	getBody := func() (io.ReadCloser, error) {
+		wg.Add(1)
+		body := &wgReadCloser{
+			Reader: strings.NewReader(bigBody),
+			wg:     &wg,
+		}
+		return body, nil
+	}
+
 	for i := 0; i < count; i++ {
-		req, err := NewRequest("PUT", ts.URL, strings.NewReader(bigBody))
+		reqBody, _ := getBody()
+		req, err := NewRequest("PUT", ts.URL, reqBody)
 		if err != nil {
+			reqBody.Close()
 			t.Fatal(err)
 		}
+		req.ContentLength = int64(len(bigBody))
+		req.GetBody = getBody
+
 		resp, err := c.Do(req)
 		if err != nil {
 			fail++