net/http: deflake TestTimeoutHandlerContextCanceled

Fixes #49448

Change-Id: Ie2acff7dedbca9bd1cc0b1b3dd0a01573c7befee
Reviewed-on: https://go-review.googlesource.com/c/go/+/361920
Trust: Damien Neil <dneil@google.com>
Run-TryBot: Damien Neil <dneil@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Bryan C. Mills <bcmills@google.com>

1 files changed, 15 insertions, 9 deletions

diff --git a/src/net/http/serve_test.go b/src/net/http/serve_test.go
index 27dff2bf45..30a6555d30 100644
--- a/src/net/http/serve_test.go
+++ b/src/net/http/serve_test.go
@@ -2520,22 +2520,28 @@ func TestTimeoutHandlerStartTimerWhenServing(t *testing.T) {
 func TestTimeoutHandlerContextCanceled(t *testing.T) {
 	setParallel(t)
 	defer afterTest(t)
-	sendHi := make(chan bool, 1)
 	writeErrors := make(chan error, 1)
 	sayHi := HandlerFunc(func(w ResponseWriter, r *Request) {
 		w.Header().Set("Content-Type", "text/plain")
-		<-sendHi
-		_, werr := w.Write([]byte("hi"))
-		writeErrors <- werr
+		var err error
+		// The request context has already been canceled, but
+		// retry the write for a while to give the timeout handler
+		// a chance to notice.
+		for i := 0; i < 100; i++ {
+			_, err = w.Write([]byte("a"))
+			if err != nil {
+				break
+			}
+			time.Sleep(1 * time.Millisecond)
+		}
+		writeErrors <- err
 	})
-	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Hour)
-	h := NewTestTimeoutHandler(sayHi, ctx)
+	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
+	h := NewTestTimeoutHandler(sayHi, ctx)
 	cst := newClientServerTest(t, h1Mode, h)
 	defer cst.close()
 
-	// Succeed without timing out:
-	sendHi <- true
 	res, err := cst.c.Get(cst.ts.URL)
 	if err != nil {
 		t.Error(err)
@@ -2548,7 +2554,7 @@ func TestTimeoutHandlerContextCanceled(t *testing.T) {
 		t.Errorf("got body %q; expected %q", g, e)
 	}
 	if g, e := <-writeErrors, context.Canceled; g != e {
-		t.Errorf("got unexpected Write error on first request: %v", g)
+		t.Errorf("got unexpected Write in handler: %v, want %g", g, e)
 	}
 }