net/http: optimize Transport request body writes

net.TCPConn has the ability to send data out using system calls such as
sendfile when the source data comes from an *os.File. However, the way
that I/O has been laid out in the transport means that the File is
actually wrapped behind two outer io.Readers, and as such the TCP stack
cannot properly type-assert the reader, ensuring that it falls back to
genericReadFrom.

This commit does the following:

* Removes transferBodyReader and moves its functionality to a new
doBodyCopy helper. This is not an io.Reader implementation, but no
functionality is lost this way, and it allows us to unwrap one layer
from the body.

* The second layer of the body is unwrapped if the original writer
was wrapped with ioutil.NopCloser, which is what NewRequest wraps the
body in if it's not a ReadCloser on its own. The unwrap operation
passes through the existing body if there's no nopCloser.

Note that this depends on change https://golang.org/cl/163599 to
properly function, as the lack of ReaderFrom implementation otherwise
means that this functionality is essentially walled off.

Updates #30377.

Author: vancluever

src/net/http/transfer.go

@@ -53,19 +53,6 @@ func (br *byteReader) Read(p []byte) (n int, err error) {
 	return 1, io.EOF
 }
 
-// transferBodyReader is an io.Reader that reads from tw.Body
-// and records any non-EOF error in tw.bodyReadError.
-// It is exactly 1 pointer wide to avoid allocations into interfaces.
-type transferBodyReader struct{ tw *transferWriter }
-
-func (br transferBodyReader) Read(p []byte) (n int, err error) {
-	n, err = br.tw.Body.Read(p)
-	if err != nil && err != io.EOF {
-		br.tw.bodyReadError = err
-	}
-	return
-}
-
 // transferWriter inspects the fields of a user-supplied Request or Response,
 // sanitizes them without changing the user object and provides methods for
 // writing the respective header, body and trailer in wire format.
@@ -347,15 +334,18 @@ func (t *transferWriter) writeBody(w io.Writer) error {
 	var err error
 	var ncopy int64
 
-	// Write body
+	// Write body. We "unwrap" the body first if it was wrapped in a
+	// nopCloser. This is to ensure that we can take advantage of
+	// OS-level optimizations in the event that the body is an
+	// *os.File.
 	if t.Body != nil {
-		var body = transferBodyReader{t}
+		var body = t.unwrapBody()
 		if chunked(t.TransferEncoding) {
 			if bw, ok := w.(*bufio.Writer); ok && !t.IsResponse {
 				w = &internal.FlushAfterChunkWriter{Writer: bw}
 			}
 			cw := internal.NewChunkedWriter(w)
-			_, err = io.Copy(cw, body)
+			_, err = t.doBodyCopy(cw, body)
 			if err == nil {
 				err = cw.Close()
 			}
@@ -364,14 +354,14 @@ func (t *transferWriter) writeBody(w io.Writer) error {
 			if t.Method == "CONNECT" {
 				dst = bufioFlushWriter{dst}
 			}
-			ncopy, err = io.Copy(dst, body)
+			ncopy, err = t.doBodyCopy(dst, body)
 		} else {
-			ncopy, err = io.Copy(w, io.LimitReader(body, t.ContentLength))
+			ncopy, err = t.doBodyCopy(w, io.LimitReader(body, t.ContentLength))
 			if err != nil {
 				return err
 			}
 			var nextra int64
-			nextra, err = io.Copy(ioutil.Discard, body)
+			nextra, err = t.doBodyCopy(ioutil.Discard, body)
 			ncopy += nextra
 		}
 		if err != nil {
@@ -402,6 +392,31 @@ func (t *transferWriter) writeBody(w io.Writer) error {
 	return err
 }
 
+// doBodyCopy wraps a copy operation, with any resulting error also
+// being saved in bodyReadError.
+//
+// This function is only intended for use in writeBody.
+func (t *transferWriter) doBodyCopy(dst io.Writer, src io.Reader) (n int64, err error) {
+	n, err = io.Copy(dst, src)
+	if err != nil && err != io.EOF {
+		t.bodyReadError = err
+	}
+	return
+}
+
+// unwrapBodyReader unwraps the body's inner reader if it's a
+// nopCloser. This is to ensure that body writes sourced from local
+// files (*os.File types) are properly optimized.
+//
+// This function is only intended for use in writeBody.
+func (t *transferWriter) unwrapBody() io.Reader {
+	if reflect.TypeOf(t.Body) == nopCloserType {
+		return reflect.ValueOf(t.Body).Field(0).Interface().(io.Reader)
+	}
+
+	return t.Body
+}
+
 type transferReader struct {
 	// Input
 	Header        Header

src/net/http/transfer_test.go

@@ -9,6 +9,8 @@ import (
 	"bytes"
 	"io"
 	"io/ioutil"
+	"os"
+	"reflect"
 	"strings"
 	"testing"
 )
@@ -90,3 +92,187 @@ func TestDetectInMemoryReaders(t *testing.T) {
 		}
 	}
 }
+
+type mockTransferWriterBodyWriter struct {
+	CalledReader io.Reader
+	WriteCalled  bool
+}
+
+func (w *mockTransferWriterBodyWriter) ReadFrom(r io.Reader) (int64, error) {
+	w.CalledReader = r
+	return io.Copy(ioutil.Discard, r)
+}
+
+func (w *mockTransferWriterBodyWriter) Write(p []byte) (int, error) {
+	w.WriteCalled = true
+	return ioutil.Discard.Write(p)
+}
+
+func TestTransferWriterWriteBodyReaderTypes(t *testing.T) {
+	fileTyp := reflect.TypeOf(&os.File{})
+	bufferTyp := reflect.TypeOf(&bytes.Buffer{})
+
+	newFileFunc := func() (io.Reader, func(), error) {
+		f, err := ioutil.TempFile("", "net-http-testtransferwriterwritebodyreadertypes")
+		if err != nil {
+			return nil, nil, err
+		}
+
+		// 1K zeros just to get a file that we can read
+		_, err = f.Write(make([]byte, 1024))
+		f.Close()
+		if err != nil {
+			return nil, nil, err
+		}
+
+		f, err = os.Open(f.Name())
+		if err != nil {
+			return nil, nil, err
+		}
+
+		return f, func() {
+			f.Close()
+			os.Remove(f.Name())
+		}, nil
+	}
+
+	newBufferFunc := func() (io.Reader, func(), error) {
+		return bytes.NewBuffer(make([]byte, 1024)), func() {}, nil
+	}
+
+	cases := []struct {
+		Name             string
+		BodyFunc         func() (io.Reader, func(), error)
+		Method           string
+		ContentLength    int64
+		TransferEncoding []string
+		LimitedReader    bool
+		ExpectedReader   reflect.Type
+		ExpectedWrite    bool
+	}{
+		{
+			Name:           "file, non-chunked, size set",
+			BodyFunc:       newFileFunc,
+			Method:         "PUT",
+			ContentLength:  1024,
+			LimitedReader:  true,
+			ExpectedReader: fileTyp,
+		},
+		{
+			Name:   "file, non-chunked, size set, nopCloser wrapped",
+			Method: "PUT",
+			BodyFunc: func() (io.Reader, func(), error) {
+				r, cleanup, err := newFileFunc()
+				return ioutil.NopCloser(r), cleanup, err
+			},
+			ContentLength:  1024,
+			LimitedReader:  true,
+			ExpectedReader: fileTyp,
+		},
+		{
+			Name:           "file, non-chunked, negative size",
+			Method:         "PUT",
+			BodyFunc:       newFileFunc,
+			ContentLength:  -1,
+			ExpectedReader: fileTyp,
+		},
+		{
+			Name:           "file, non-chunked, CONNECT, negative size",
+			Method:         "CONNECT",
+			BodyFunc:       newFileFunc,
+			ContentLength:  -1,
+			ExpectedReader: fileTyp,
+		},
+		{
+			Name:             "file, chunked",
+			Method:           "PUT",
+			BodyFunc:         newFileFunc,
+			TransferEncoding: []string{"chunked"},
+			ExpectedWrite:    true,
+		},
+		{
+			Name:           "buffer, non-chunked, size set",
+			BodyFunc:       newBufferFunc,
+			Method:         "PUT",
+			ContentLength:  1024,
+			LimitedReader:  true,
+			ExpectedReader: bufferTyp,
+		},
+		{
+			Name:   "buffer, non-chunked, size set, nopCloser wrapped",
+			Method: "PUT",
+			BodyFunc: func() (io.Reader, func(), error) {
+				r, cleanup, err := newBufferFunc()
+				return ioutil.NopCloser(r), cleanup, err
+			},
+			ContentLength:  1024,
+			LimitedReader:  true,
+			ExpectedReader: bufferTyp,
+		},
+		{
+			Name:          "buffer, non-chunked, negative size",
+			Method:        "PUT",
+			BodyFunc:      newBufferFunc,
+			ContentLength: -1,
+			ExpectedWrite: true,
+		},
+		{
+			Name:          "buffer, non-chunked, CONNECT, negative size",
+			Method:        "CONNECT",
+			BodyFunc:      newBufferFunc,
+			ContentLength: -1,
+			ExpectedWrite: true,
+		},
+		{
+			Name:             "buffer, chunked",
+			Method:           "PUT",
+			BodyFunc:         newBufferFunc,
+			TransferEncoding: []string{"chunked"},
+			ExpectedWrite:    true,
+		},
+	}
+
+	for _, tc := range cases {
+		t.Run(tc.Name, func(t *testing.T) {
+			body, cleanup, err := tc.BodyFunc()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			defer cleanup()
+
+			mw := &mockTransferWriterBodyWriter{}
+			tw := &transferWriter{
+				Body:             body,
+				ContentLength:    tc.ContentLength,
+				TransferEncoding: tc.TransferEncoding,
+			}
+
+			if err := tw.writeBody(mw); err != nil {
+				t.Fatal(err)
+			}
+
+			if tc.ExpectedReader != nil {
+				if mw.CalledReader == nil {
+					t.Fatal("expected ReadFrom to be called, but it wasn't")
+				}
+
+				var actualReader reflect.Type
+				lr, ok := mw.CalledReader.(*io.LimitedReader)
+				if ok && tc.LimitedReader {
+					actualReader = reflect.TypeOf(lr.R)
+				} else {
+					actualReader = reflect.TypeOf(mw.CalledReader)
+				}
+
+				if tc.ExpectedReader != actualReader {
+					t.Fatalf("expected reader to be %s, got %s", tc.ExpectedReader, actualReader)
+				}
+			}
+
+			if tc.ExpectedWrite && !mw.WriteCalled {
+				t.Fatal("expected Read to be called, but it wasn't")
+			}
+		})
+	}
+}