From 6e738868a7a943d7d4fd6bb1963e7f6d78111726 Mon Sep 17 00:00:00 2001
From: Damien Neil <dneil@google.com>
Date: Tue, 3 Aug 2021 19:38:37 -0700
Subject: [PATCH] net/http: speed up and deflake
 TestCancelRequestWhenSharingConnection

This test made many requests over the same connection for 10
seconds, trusting that this will exercise the request cancelation
race from #41600.

Change the test to exhibit the specific race in a targeted fashion
with only two requests.

Updates #41600.
Updates #47016.

Change-Id: If99c9b9331ff645f6bb67fe9fb79b8aab8784710
Reviewed-on: https://go-review.googlesource.com/c/go/+/339594
Trust: Damien Neil <dneil@google.com>
Run-TryBot: Damien Neil <dneil@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
---
 src/net/http/transport_test.go | 81 ++++++++++++++++++++++------------
 1 file changed, 53 insertions(+), 28 deletions(-)

diff --git a/src/net/http/transport_test.go b/src/net/http/transport_test.go
index 690e0c299d..eeaa492644 100644
--- a/src/net/http/transport_test.go
+++ b/src/net/http/transport_test.go
@@ -6441,10 +6441,11 @@ func TestErrorWriteLoopRace(t *testing.T) {
 // Test that a new request which uses the connection of an active request
 // cannot cause it to be canceled as well.
 func TestCancelRequestWhenSharingConnection(t *testing.T) {
-	if testing.Short() {
-		t.Skip("skipping in short mode")
-	}
+	reqc := make(chan chan struct{}, 2)
 	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, req *Request) {
+		ch := make(chan struct{}, 1)
+		reqc <- ch
+		<-ch
 		w.Header().Add("Content-Length", "0")
 	}))
 	defer ts.Close()
@@ -6456,34 +6457,58 @@ func TestCancelRequestWhenSharingConnection(t *testing.T) {
 
 	var wg sync.WaitGroup
 
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-
-	for i := 0; i < 10; i++ {
-		wg.Add(1)
-		go func() {
-			defer wg.Done()
-			for ctx.Err() == nil {
-				reqctx, reqcancel := context.WithCancel(ctx)
-				go reqcancel()
-				req, _ := NewRequestWithContext(reqctx, "GET", ts.URL, nil)
-				res, err := client.Do(req)
-				if err == nil {
-					res.Body.Close()
-				}
-			}
-		}()
-	}
-
-	for ctx.Err() == nil {
-		req, _ := NewRequest("GET", ts.URL, nil)
-		if res, err := client.Do(req); err != nil {
-			t.Errorf("unexpected: %p %v", req, err)
-			break
-		} else {
+	wg.Add(1)
+	putidlec := make(chan chan struct{})
+	go func() {
+		defer wg.Done()
+		ctx := httptrace.WithClientTrace(context.Background(), &httptrace.ClientTrace{
+			PutIdleConn: func(error) {
+				// Signal that the idle conn has been returned to the pool,
+				// and wait for the order to proceed.
+				ch := make(chan struct{})
+				putidlec <- ch
+				<-ch
+			},
+		})
+		req, _ := NewRequestWithContext(ctx, "GET", ts.URL, nil)
+		res, err := client.Do(req)
+		if err == nil {
 			res.Body.Close()
 		}
-	}
+		if err != nil {
+			t.Errorf("request 1: got err %v, want nil", err)
+		}
+	}()
 
+	// Wait for the first request to receive a response and return the
+	// connection to the idle pool.
+	r1c := <-reqc
+	close(r1c)
+	idlec := <-putidlec
+
+	wg.Add(1)
+	cancelctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		defer wg.Done()
+		req, _ := NewRequestWithContext(cancelctx, "GET", ts.URL, nil)
+		res, err := client.Do(req)
+		if err == nil {
+			res.Body.Close()
+		}
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("request 2: got err %v, want Canceled", err)
+		}
+	}()
+
+	// Wait for the second request to arrive at the server, and then cancel
+	// the request context.
+	r2c := <-reqc
 	cancel()
+
+	// Give the cancelation a moment to take effect, and then unblock the first request.
+	time.Sleep(1 * time.Millisecond)
+	close(idlec)
+
+	close(r2c)
 	wg.Wait()
 }