net/http: make Transport retry GetBody requests if nothing written

This is another attempt at the change attempted in
https://golang.org/cl/27117 and rolled back in https://golang.org/cl/34134

The difference between this and the previous attempt is that this version only
retries if the new field GetBody is set on the Request.

Additionally, this allows retries of requests with idempotent methods even if
they have bodies, as long as GetBody is defined.

This also fixes an existing bug where readLoop could make a redundant call to
setReqCanceler for DELETE/POST/PUT/etc requests with no body with zero bytes
written.

This clarifies the existing TestRetryIdempotentRequestsOnError test (and changes
it into a test with 4 subtests).  When that test was written, it was in fact
testing "retry idempotent requests" logic, but the logic had changed since then,
and it was actually testing "retry requests with no body when no bytes have been
written". (You can confirm this by changing the existing test from a GET to a
DELETE; it passes without the changes in this CL.) We now test for the no-Body
and GetBody cases for both idempotent and nothing-written-non-idempotent
requests.

Fixes #18241
Fixes #17844

Change-Id: I69a48691796f6dc08c31f7aa7887b7dfd67e278a
Reviewed-on: https://go-review.googlesource.com/42142
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>

6 files changed, 186 insertions, 82 deletions

diff --git a/src/net/http/client_test.go b/src/net/http/client_test.go
index 73f22212f6..b9a1c31e43 100644
--- a/src/net/http/client_test.go
+++ b/src/net/http/client_test.go
@@ -1727,8 +1727,8 @@ func (b issue18239Body) Close() error {
 	return nil
 }
 
-// Issue 18239: make sure the Transport doesn't retry requests with bodies.
-// (Especially if Request.GetBody is not defined.)
+// Issue 18239: make sure the Transport doesn't retry requests with bodies
+// if Request.GetBody is not defined.
 func TestTransportBodyReadError(t *testing.T) {
 	setParallel(t)
 	defer afterTest(t)
diff --git a/src/net/http/export_test.go b/src/net/http/export_test.go
index 98fb0834dd..2ef145e534 100644
--- a/src/net/http/export_test.go
+++ b/src/net/http/export_test.go
@@ -25,6 +25,7 @@ var (
 	ExportCloseWriteAndWait           = (*conn).closeWriteAndWait
 	ExportErrRequestCanceled          = errRequestCanceled
 	ExportErrRequestCanceledConn      = errRequestCanceledConn
+	ExportErrServerClosedIdle         = errServerClosedIdle
 	ExportServeFile                   = serveFile
 	ExportScanETag                    = scanETag
 	ExportHttp2ConfigureServer        = http2ConfigureServer
diff --git a/src/net/http/request.go b/src/net/http/request.go
index 82466d9b36..7f473dd15d 100644
--- a/src/net/http/request.go
+++ b/src/net/http/request.go
@@ -1317,7 +1317,7 @@ func (r *Request) closeBody() {
 }
 
 func (r *Request) isReplayable() bool {
-	if r.Body == nil {
+	if r.Body == nil || r.Body == NoBody || r.GetBody != nil {
 		switch valueOrDefault(r.Method, "GET") {
 		case "GET", "HEAD", "OPTIONS", "TRACE":
 			return true
diff --git a/src/net/http/transport.go b/src/net/http/transport.go
index abb22d4f8d..9dedc22272 100644
--- a/src/net/http/transport.go
+++ b/src/net/http/transport.go
@@ -419,6 +419,18 @@ func (t *Transport) RoundTrip(req *Request) (*Response, error) {
 			return nil, err
 		}
 		testHookRoundTripRetried()
+
+		// Rewind the body if we're able to.  (HTTP/2 does this itself so we only
+		// need to do it for HTTP/1.1 connections.)
+		if req.GetBody != nil && pconn.alt == nil {
+			newReq := *req
+			var err error
+			newReq.Body, err = req.GetBody()
+			if err != nil {
+				return nil, err
+			}
+			req = &newReq
+		}
 	}
 }
 
@@ -450,8 +462,9 @@ func (pc *persistConn) shouldRetryRequest(req *Request, err error) bool {
 		return false
 	}
 	if _, ok := err.(nothingWrittenError); ok {
-		// We never wrote anything, so it's safe to retry.
-		return true
+		// We never wrote anything, so it's safe to retry, if there's no body or we
+		// can "rewind" the body with GetBody.
+		return req.outgoingLength() == 0 || req.GetBody != nil
 	}
 	if !req.isReplayable() {
 		// Don't retry non-idempotent requests.
@@ -1475,7 +1488,7 @@ func (pc *persistConn) mapRoundTripError(req *transportRequest, startBytesWritte
 	}
 	if pc.isBroken() {
 		<-pc.writeLoopDone
-		if pc.nwrite == startBytesWritten && req.outgoingLength() == 0 {
+		if pc.nwrite == startBytesWritten {
 			return nothingWrittenError{err}
 		}
 		return fmt.Errorf("net/http: HTTP/1.x transport connection broken: %v", err)
@@ -1544,16 +1557,6 @@ func (pc *persistConn) readLoop() {
 				err = fmt.Errorf("net/http: server response headers exceeded %d bytes; aborted", pc.maxHeaderResponseSize())
 			}
 
-			// If we won't be able to retry this request later (from the
-			// roundTrip goroutine), mark it as done now.
-			// BEFORE the send on rc.ch, as the client might re-use the
-			// same *Request pointer, and we don't want to set call
-			// t.setReqCanceler from this persistConn while the Transport
-			// potentially spins up a different persistConn for the
-			// caller's subsequent request.
-			if !pc.shouldRetryRequest(rc.req, err) {
-				pc.t.setReqCanceler(rc.req, nil)
-			}
 			select {
 			case rc.ch <- responseAndError{err: err}:
 			case <-rc.callerGone:
@@ -1768,7 +1771,7 @@ func (pc *persistConn) writeLoop() {
 			}
 			if err != nil {
 				wr.req.Request.closeBody()
-				if pc.nwrite == startBytesWritten && wr.req.outgoingLength() == 0 {
+				if pc.nwrite == startBytesWritten {
 					err = nothingWrittenError{err}
 				}
 			}
diff --git a/src/net/http/transport_internal_test.go b/src/net/http/transport_internal_test.go
index 262d8b4ac5..594bf6e2c8 100644
--- a/src/net/http/transport_internal_test.go
+++ b/src/net/http/transport_internal_test.go
@@ -9,6 +9,7 @@ package http
 import (
 	"errors"
 	"net"
+	"strings"
 	"testing"
 )
 
@@ -81,6 +82,19 @@ func dummyRequest(method string) *Request {
 	}
 	return req
 }
+func dummyRequestWithBody(method string) *Request {
+	req, err := NewRequest(method, "http://fake.tld/", strings.NewReader("foo"))
+	if err != nil {
+		panic(err)
+	}
+	return req
+}
+
+func dummyRequestWithBodyNoGetBody(method string) *Request {
+	req := dummyRequestWithBody(method)
+	req.GetBody = nil
+	return req
+}
 
 func TestTransportShouldRetryRequest(t *testing.T) {
 	tests := []struct {
@@ -132,6 +146,18 @@ func TestTransportShouldRetryRequest(t *testing.T) {
 			err:  errServerClosedIdle,
 			want: true,
 		},
+		7: {
+			pc:   &persistConn{reused: true},
+			req:  dummyRequestWithBody("POST"),
+			err:  nothingWrittenError{},
+			want: true,
+		},
+		8: {
+			pc:   &persistConn{reused: true},
+			req:  dummyRequestWithBodyNoGetBody("POST"),
+			err:  nothingWrittenError{},
+			want: false,
+		},
 	}
 	for i, tt := range tests {
 		got := tt.pc.shouldRetryRequest(tt.req, tt.err)
diff --git a/src/net/http/transport_test.go b/src/net/http/transport_test.go
index c516380990..27b55dca2f 100644
--- a/src/net/http/transport_test.go
+++ b/src/net/http/transport_test.go
@@ -2601,86 +2601,160 @@ type writerFuncConn struct {
 
 func (c writerFuncConn) Write(p []byte) (n int, err error) { return c.write(p) }
 
-// Issue 4677. If we try to reuse a connection that the server is in the
-// process of closing, we may end up successfully writing out our request (or a
-// portion of our request) only to find a connection error when we try to read
-// from (or finish writing to) the socket.
+// Issues 4677, 18241, and 17844. If we try to reuse a connection that the
+// server is in the process of closing, we may end up successfully writing out
+// our request (or a portion of our request) only to find a connection error
+// when we try to read from (or finish writing to) the socket.
 //
-// NOTE: we resend a request only if the request is idempotent, we reused a
-// keep-alive connection, and we haven't yet received any header data. This
-// automatically prevents an infinite resend loop because we'll run out of the
-// cached keep-alive connections eventually.
-func TestRetryIdempotentRequestsOnError(t *testing.T) {
-	defer afterTest(t)
+// NOTE: we resend a request only if:
+//   - we reused a keep-alive connection
+//   - we haven't yet received any header data
+//   - either we wrote no bytes to the server, or the request is idempotent
+// This automatically prevents an infinite resend loop because we'll run out of
+// the cached keep-alive connections eventually.
+func TestRetryRequestsOnError(t *testing.T) {
+	newRequest := func(method, urlStr string, body io.Reader) *Request {
+		req, err := NewRequest(method, urlStr, body)
+		if err != nil {
+			t.Fatal(err)
+		}
+		return req
+	}
 
-	var (
-		mu     sync.Mutex
-		logbuf bytes.Buffer
-	)
-	logf := func(format string, args ...interface{}) {
-		mu.Lock()
-		defer mu.Unlock()
-		fmt.Fprintf(&logbuf, format, args...)
-		logbuf.WriteByte('\n')
+	testCases := []struct {
+		name       string
+		failureN   int
+		failureErr error
+		// Note that we can't just re-use the Request object across calls to c.Do
+		// because we need to rewind Body between calls.  (GetBody is only used to
+		// rewind Body on failure and redirects, not just because it's done.)
+		req       func() *Request
+		reqString string
+	}{
+		{
+			name: "IdempotentNoBodySomeWritten",
+			// Believe that we've written some bytes to the server, so we know we're
+			// not just in the "retry when no bytes sent" case".
+			failureN: 1,
+			// Use the specific error that shouldRetryRequest looks for with idempotent requests.
+			failureErr: ExportErrServerClosedIdle,
+			req: func() *Request {
+				return newRequest("GET", "http://fake.golang", nil)
+			},
+			reqString: `GET / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nAccept-Encoding: gzip\r\n\r\n`,
+		},
+		{
+			name: "IdempotentGetBodySomeWritten",
+			// Believe that we've written some bytes to the server, so we know we're
+			// not just in the "retry when no bytes sent" case".
+			failureN: 1,
+			// Use the specific error that shouldRetryRequest looks for with idempotent requests.
+			failureErr: ExportErrServerClosedIdle,
+			req: func() *Request {
+				return newRequest("GET", "http://fake.golang", strings.NewReader("foo\n"))
+			},
+			reqString: `GET / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nContent-Length: 4\r\nAccept-Encoding: gzip\r\n\r\nfoo\n`,
+		},
+		{
+			name: "NothingWrittenNoBody",
+			// It's key that we return 0 here -- that's what enables Transport to know
+			// that nothing was written, even though this is a non-idempotent request.
+			failureN:   0,
+			failureErr: errors.New("second write fails"),
+			req: func() *Request {
+				return newRequest("DELETE", "http://fake.golang", nil)
+			},
+			reqString: `DELETE / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nAccept-Encoding: gzip\r\n\r\n`,
+		},
+		{
+			name: "NothingWrittenGetBody",
+			// It's key that we return 0 here -- that's what enables Transport to know
+			// that nothing was written, even though this is a non-idempotent request.
+			failureN:   0,
+			failureErr: errors.New("second write fails"),
+			// Note that NewRequest will set up GetBody for strings.Reader, which is
+			// required for the retry to occur
+			req: func() *Request {
+				return newRequest("POST", "http://fake.golang", strings.NewReader("foo\n"))
+			},
+			reqString: `POST / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nContent-Length: 4\r\nAccept-Encoding: gzip\r\n\r\nfoo\n`,
+		},
 	}
 
-	ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
-		logf("Handler")
-		w.Header().Set("X-Status", "ok")
-	}))
-	defer ts.Close()
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			defer afterTest(t)
 
-	var writeNumAtomic int32
-	c := ts.Client()
-	c.Transport.(*Transport).Dial = func(network, addr string) (net.Conn, error) {
-		logf("Dial")
-		c, err := net.Dial(network, ts.Listener.Addr().String())
-		if err != nil {
-			logf("Dial error: %v", err)
-			return nil, err
-		}
-		return &writerFuncConn{
-			Conn: c,
-			write: func(p []byte) (n int, err error) {
-				if atomic.AddInt32(&writeNumAtomic, 1) == 2 {
-					logf("intentional write failure")
-					return 0, errors.New("second write fails")
+			var (
+				mu     sync.Mutex
+				logbuf bytes.Buffer
+			)
+			logf := func(format string, args ...interface{}) {
+				mu.Lock()
+				defer mu.Unlock()
+				fmt.Fprintf(&logbuf, format, args...)
+				logbuf.WriteByte('\n')
+			}
+
+			ts := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
+				logf("Handler")
+				w.Header().Set("X-Status", "ok")
+			}))
+			defer ts.Close()
+
+			var writeNumAtomic int32
+			c := ts.Client()
+			c.Transport.(*Transport).Dial = func(network, addr string) (net.Conn, error) {
+				logf("Dial")
+				c, err := net.Dial(network, ts.Listener.Addr().String())
+				if err != nil {
+					logf("Dial error: %v", err)
+					return nil, err
 				}
-				logf("Write(%q)", p)
-				return c.Write(p)
-			},
-		}, nil
-	}
+				return &writerFuncConn{
+					Conn: c,
+					write: func(p []byte) (n int, err error) {
+						if atomic.AddInt32(&writeNumAtomic, 1) == 2 {
+							logf("intentional write failure")
+							return tc.failureN, tc.failureErr
+						}
+						logf("Write(%q)", p)
+						return c.Write(p)
+					},
+				}, nil
+			}
 
-	SetRoundTripRetried(func() {
-		logf("Retried.")
-	})
-	defer SetRoundTripRetried(nil)
+			SetRoundTripRetried(func() {
+				logf("Retried.")
+			})
+			defer SetRoundTripRetried(nil)
 
-	for i := 0; i < 3; i++ {
-		res, err := c.Get("http://fake.golang/")
-		if err != nil {
-			t.Fatalf("i=%d: Get = %v", i, err)
-		}
-		res.Body.Close()
-	}
+			for i := 0; i < 3; i++ {
+				res, err := c.Do(tc.req())
+				if err != nil {
+					t.Fatalf("i=%d: Do = %v", i, err)
+				}
+				res.Body.Close()
+			}
 
-	mu.Lock()
-	got := logbuf.String()
-	mu.Unlock()
-	const want = `Dial
-Write("GET / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nAccept-Encoding: gzip\r\n\r\n")
+			mu.Lock()
+			got := logbuf.String()
+			mu.Unlock()
+			want := fmt.Sprintf(`Dial
+Write("%s")
 Handler
 intentional write failure
 Retried.
 Dial
-Write("GET / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nAccept-Encoding: gzip\r\n\r\n")
+Write("%s")
 Handler
-Write("GET / HTTP/1.1\r\nHost: fake.golang\r\nUser-Agent: Go-http-client/1.1\r\nAccept-Encoding: gzip\r\n\r\n")
+Write("%s")
 Handler
-`
-	if got != want {
-		t.Errorf("Log of events differs. Got:\n%s\nWant:\n%s", got, want)
+`, tc.reqString, tc.reqString, tc.reqString)
+			if got != want {
+				t.Errorf("Log of events differs. Got:\n%s\nWant:\n%s", got, want)
+			}
+		})
 	}
 }