net,os: arrange zero-copy of os.File and net.TCPConn to net.UnixConn

Fixes #58808

goos: linux
goarch: amd64
pkg: net
cpu: DO-Premium-Intel
                             │      old      │                 new                  │
                             │    sec/op     │    sec/op     vs base                │
Splice/tcp-to-unix/1024-4       3.783µ ± 10%   3.201µ ±  7%  -15.40% (p=0.001 n=10)
Splice/tcp-to-unix/2048-4       3.967µ ± 13%   3.818µ ± 16%        ~ (p=0.971 n=10)
Splice/tcp-to-unix/4096-4       4.988µ ± 16%   4.590µ ± 11%        ~ (p=0.089 n=10)
Splice/tcp-to-unix/8192-4       6.981µ ± 13%   5.236µ ±  9%  -25.00% (p=0.000 n=10)
Splice/tcp-to-unix/16384-4     10.192µ ±  9%   7.350µ ±  7%  -27.89% (p=0.000 n=10)
Splice/tcp-to-unix/32768-4      19.65µ ± 13%   10.28µ ± 16%  -47.69% (p=0.000 n=10)
Splice/tcp-to-unix/65536-4      41.89µ ± 18%   15.70µ ± 13%  -62.52% (p=0.000 n=10)
Splice/tcp-to-unix/131072-4     90.05µ ± 11%   29.55µ ± 10%  -67.18% (p=0.000 n=10)
Splice/tcp-to-unix/262144-4    170.24µ ± 15%   52.66µ ±  4%  -69.06% (p=0.000 n=10)
Splice/tcp-to-unix/524288-4     326.4µ ± 13%   109.3µ ± 11%  -66.52% (p=0.000 n=10)
Splice/tcp-to-unix/1048576-4    651.4µ ±  9%   228.3µ ± 14%  -64.95% (p=0.000 n=10)
geomean                         29.42µ         15.62µ        -46.90%

                             │      old      │                  new                   │
                             │      B/s      │      B/s       vs base                 │
Splice/tcp-to-unix/1024-4      258.2Mi ± 11%   305.2Mi ±  8%   +18.21% (p=0.001 n=10)
Splice/tcp-to-unix/2048-4      492.5Mi ± 15%   511.7Mi ± 13%         ~ (p=0.971 n=10)
Splice/tcp-to-unix/4096-4      783.5Mi ± 14%   851.2Mi ± 12%         ~ (p=0.089 n=10)
Splice/tcp-to-unix/8192-4      1.093Gi ± 11%   1.458Gi ±  8%   +33.36% (p=0.000 n=10)
Splice/tcp-to-unix/16384-4     1.497Gi ±  9%   2.076Gi ±  7%   +38.67% (p=0.000 n=10)
Splice/tcp-to-unix/32768-4     1.553Gi ± 11%   2.969Gi ± 14%   +91.17% (p=0.000 n=10)
Splice/tcp-to-unix/65536-4     1.458Gi ± 23%   3.888Gi ± 11%  +166.69% (p=0.000 n=10)
Splice/tcp-to-unix/131072-4    1.356Gi ± 10%   4.131Gi ±  9%  +204.72% (p=0.000 n=10)
Splice/tcp-to-unix/262144-4    1.434Gi ± 13%   4.637Gi ±  4%  +223.32% (p=0.000 n=10)
Splice/tcp-to-unix/524288-4    1.497Gi ± 15%   4.468Gi ± 10%  +198.47% (p=0.000 n=10)
Splice/tcp-to-unix/1048576-4   1.501Gi ± 10%   4.277Gi ± 16%  +184.88% (p=0.000 n=10)
geomean                        1.038Gi         1.954Gi         +88.28%

                             │      old      │                   new                   │
                             │     B/op      │    B/op     vs base                     │
Splice/tcp-to-unix/1024-4      0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/2048-4      0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/4096-4      0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/8192-4      0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/16384-4     0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/32768-4     0.000 ±  0%     0.000 ± 0%         ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/65536-4     1.000 ±   ?     0.000 ± 0%  -100.00% (p=0.001 n=10)
Splice/tcp-to-unix/131072-4    2.000 ±  0%     0.000 ± 0%  -100.00% (p=0.000 n=10)
Splice/tcp-to-unix/262144-4    4.000 ± 25%     0.000 ± 0%  -100.00% (p=0.000 n=10)
Splice/tcp-to-unix/524288-4    7.500 ± 33%     0.000 ± 0%  -100.00% (p=0.000 n=10)
Splice/tcp-to-unix/1048576-4   17.00 ± 12%      0.00 ± 0%  -100.00% (p=0.000 n=10)
geomean                                    ²               ?                       ² ³
¹ all samples are equal
² summaries must be >0 to compute geomean
³ ratios must be >0 to compute geomean

                             │     old      │                 new                 │
                             │  allocs/op   │ allocs/op   vs base                 │
Splice/tcp-to-unix/1024-4      0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/2048-4      0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/4096-4      0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/8192-4      0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/16384-4     0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/32768-4     0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/65536-4     0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/131072-4    0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/262144-4    0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/524288-4    0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
Splice/tcp-to-unix/1048576-4   0.000 ± 0%     0.000 ± 0%       ~ (p=1.000 n=10) ¹
geomean                                   ²               +0.00%                ²
¹ all samples are equal
² summaries must be >0 to compute geomean

Change-Id: I829061b009a0929a8ef1a15c183793c0b9104dde
Reviewed-on: https://go-review.googlesource.com/c/go/+/472475
Reviewed-by: Damien Neil <[email protected]>
Reviewed-by: Bryan Mills <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>

Author: panjf2000

api/next/58808.txt

@@ -0,0 +1,2 @@
+pkg net, method (*TCPConn) WriteTo(io.Writer) (int64, error) #58808
+pkg os, method (*File) WriteTo(io.Writer) (int64, error) #58808

src/internal/poll/fd.go

@@ -81,3 +81,14 @@ func consume(v *[][]byte, n int64) {
 
 // TestHookDidWritev is a hook for testing writev.
 var TestHookDidWritev = func(wrote int) {}
+
+// String is an internal string definition for methods/functions
+// that is not intended for use outside the standard libraries.
+//
+// Other packages in std that import internal/poll and have some
+// exported APIs (now we've got some in net.rawConn) which are only used
+// internally and are not intended to be used outside the standard libraries,
+// Therefore, we make those APIs use internal types like poll.FD or poll.String
+// in their function signatures to disable the usability of these APIs from
+// external codebase.
+type String string

src/net/http/transfer_test.go

@@ -264,6 +264,12 @@ func TestTransferWriterWriteBodyReaderTypes(t *testing.T) {
 					actualReader = reflect.TypeOf(lr.R)
 				} else {
 					actualReader = reflect.TypeOf(mw.CalledReader)
+					// We have to handle this special case for genericWriteTo in os,
+					// this struct is introduced to support a zero-copy optimization,
+					// check out https://go.dev/issue/58808 for details.
+					if actualReader.Kind() == reflect.Struct && actualReader.PkgPath() == "os" && actualReader.Name() == "fileWithoutWriteTo" {
+						actualReader = actualReader.Field(1).Type
+					}
 				}
 
 				if tc.expectedReader != actualReader {

src/net/net.go

@@ -664,15 +664,53 @@ var errClosed = poll.ErrNetClosing
 // errors.Is(err, net.ErrClosed).
 var ErrClosed error = errClosed
 
-type writerOnly struct {
-	io.Writer
+// noReadFrom can be embedded alongside another type to
+// hide the ReadFrom method of that other type.
+type noReadFrom struct{}
+
+// ReadFrom hides another ReadFrom method.
+// It should never be called.
+func (noReadFrom) ReadFrom(io.Reader) (int64, error) {
+	panic("can't happen")
+}
+
+// tcpConnWithoutReadFrom implements all the methods of *TCPConn other
+// than ReadFrom. This is used to permit ReadFrom to call io.Copy
+// without leading to a recursive call to ReadFrom.
+type tcpConnWithoutReadFrom struct {
+	noReadFrom
+	*TCPConn
 }
 
 // Fallback implementation of io.ReaderFrom's ReadFrom, when sendfile isn't
 // applicable.
-func genericReadFrom(w io.Writer, r io.Reader) (n int64, err error) {
+func genericReadFrom(c *TCPConn, r io.Reader) (n int64, err error) {
 	// Use wrapper to hide existing r.ReadFrom from io.Copy.
-	return io.Copy(writerOnly{w}, r)
+	return io.Copy(tcpConnWithoutReadFrom{TCPConn: c}, r)
+}
+
+// noWriteTo can be embedded alongside another type to
+// hide the WriteTo method of that other type.
+type noWriteTo struct{}
+
+// WriteTo hides another WriteTo method.
+// It should never be called.
+func (noWriteTo) WriteTo(io.Writer) (int64, error) {
+	panic("can't happen")
+}
+
+// tcpConnWithoutWriteTo implements all the methods of *TCPConn other
+// than WriteTo. This is used to permit WriteTo to call io.Copy
+// without leading to a recursive call to WriteTo.
+type tcpConnWithoutWriteTo struct {
+	noWriteTo
+	*TCPConn
+}
+
+// Fallback implementation of io.WriterTo's WriteTo, when zero-copy isn't applicable.
+func genericWriteTo(c *TCPConn, w io.Writer) (n int64, err error) {
+	// Use wrapper to hide existing w.WriteTo from io.Copy.
+	return io.Copy(w, tcpConnWithoutWriteTo{TCPConn: c})
 }
 
 // Limit the number of concurrent cgo-using goroutines, because

src/net/rawconn.go

@@ -79,6 +79,17 @@ func newRawConn(fd *netFD) *rawConn {
 	return &rawConn{fd: fd}
 }
 
+// Network returns the network type of the underlying connection.
+//
+// Other packages in std that import internal/poll and are unable to
+// import net (such as os) can use a type assertion to access this
+// extension method so that they can distinguish different socket types.
+//
+// Network is not intended for use outside the standard library.
+func (c *rawConn) Network() poll.String {
+	return poll.String(c.fd.net)
+}
+
 type rawListener struct {
 	rawConn
 }

src/net/sendfile_linux_test.go

@@ -14,29 +14,36 @@ import (
 )
 
 func BenchmarkSendFile(b *testing.B) {
+	b.Run("file-to-tcp", func(b *testing.B) { benchmarkSendFile(b, "tcp") })
+	b.Run("file-to-unix", func(b *testing.B) { benchmarkSendFile(b, "unix") })
+}
+
+func benchmarkSendFile(b *testing.B, proto string) {
 	for i := 0; i <= 10; i++ {
 		size := 1 << (i + 10)
-		bench := sendFileBench{chunkSize: size}
+		bench := sendFileBench{
+			proto:     proto,
+			chunkSize: size,
+		}
 		b.Run(strconv.Itoa(size), bench.benchSendFile)
 	}
 }
 
 type sendFileBench struct {
+	proto     string
 	chunkSize int
 }
 
 func (bench sendFileBench) benchSendFile(b *testing.B) {
 	fileSize := b.N * bench.chunkSize
 	f := createTempFile(b, fileSize)
-	fileName := f.Name()
-	defer os.Remove(fileName)
-	defer f.Close()
 
-	client, server := spliceTestSocketPair(b, "tcp")
+	client, server := spliceTestSocketPair(b, bench.proto)
 	defer server.Close()
 
 	cleanUp, err := startSpliceClient(client, "r", bench.chunkSize, fileSize)
 	if err != nil {
+		client.Close()
 		b.Fatal(err)
 	}
 	defer cleanUp()
@@ -51,15 +58,18 @@ func (bench sendFileBench) benchSendFile(b *testing.B) {
 		b.Fatalf("failed to copy data with sendfile, error: %v", err)
 	}
 	if sent != int64(fileSize) {
-		b.Fatalf("bytes sent mismatch\n\texpect: %d\n\tgot: %d", fileSize, sent)
+		b.Fatalf("bytes sent mismatch, got: %d, want: %d", sent, fileSize)
 	}
 }
 
 func createTempFile(b *testing.B, size int) *os.File {
-	f, err := os.CreateTemp("", "linux-sendfile-test")
+	f, err := os.CreateTemp(b.TempDir(), "linux-sendfile-bench")
 	if err != nil {
 		b.Fatalf("failed to create temporary file: %v", err)
 	}
+	b.Cleanup(func() {
+		f.Close()
+	})
 
 	data := make([]byte, size)
 	if _, err := f.Write(data); err != nil {

src/net/splice_linux.go

@@ -9,12 +9,12 @@ import (
 	"io"
 )
 
-// splice transfers data from r to c using the splice system call to minimize
-// copies from and to userspace. c must be a TCP connection. Currently, splice
-// is only enabled if r is a TCP or a stream-oriented Unix connection.
+// spliceFrom transfers data from r to c using the splice system call to minimize
+// copies from and to userspace. c must be a TCP connection.
+// Currently, spliceFrom is only enabled if r is a TCP or a stream-oriented Unix connection.
 //
-// If splice returns handled == false, it has performed no work.
-func splice(c *netFD, r io.Reader) (written int64, err error, handled bool) {
+// If spliceFrom returns handled == false, it has performed no work.
+func spliceFrom(c *netFD, r io.Reader) (written int64, err error, handled bool) {
 	var remain int64 = 1<<63 - 1 // by default, copy until EOF
 	lr, ok := r.(*io.LimitedReader)
 	if ok {
@@ -25,14 +25,17 @@ func splice(c *netFD, r io.Reader) (written int64, err error, handled bool) {
 	}
 
 	var s *netFD
-	if tc, ok := r.(*TCPConn); ok {
-		s = tc.fd
-	} else if uc, ok := r.(*UnixConn); ok {
-		if uc.fd.net != "unix" {
+	switch v := r.(type) {
+	case *TCPConn:
+		s = v.fd
+	case tcpConnWithoutWriteTo:
+		s = v.fd
+	case *UnixConn:
+		if v.fd.net != "unix" {
 			return 0, nil, false
 		}
-		s = uc.fd
-	} else {
+		s = v.fd
+	default:
 		return 0, nil, false
 	}
 
@@ -42,3 +45,18 @@ func splice(c *netFD, r io.Reader) (written int64, err error, handled bool) {
 	}
 	return written, wrapSyscallError(sc, err), handled
 }
+
+// spliceTo transfers data from c to w using the splice system call to minimize
+// copies from and to userspace. c must be a TCP connection.
+// Currently, spliceTo is only enabled if w is a stream-oriented Unix connection.
+//
+// If spliceTo returns handled == false, it has performed no work.
+func spliceTo(w io.Writer, c *netFD) (written int64, err error, handled bool) {
+	uc, ok := w.(*UnixConn)
+	if !ok || uc.fd.net != "unix" {
+		return
+	}
+
+	written, handled, sc, err := poll.Splice(&uc.fd.pfd, &c.pfd, 1<<63-1)
+	return written, wrapSyscallError(sc, err), handled
+}

src/net/splice_stub.go

@@ -8,6 +8,10 @@ package net
 
 import "io"
 
-func splice(c *netFD, r io.Reader) (int64, error, bool) {
+func spliceFrom(_ *netFD, _ io.Reader) (int64, error, bool) {
+	return 0, nil, false
+}
+
+func spliceTo(_ io.Writer, _ *netFD) (int64, error, bool) {
 	return 0, nil, false
 }

src/net/splice_test.go

@@ -23,6 +23,7 @@ func TestSplice(t *testing.T) {
 		t.Skip("skipping unix-to-tcp tests")
 	}
 	t.Run("unix-to-tcp", func(t *testing.T) { testSplice(t, "unix", "tcp") })
+	t.Run("tcp-to-unix", func(t *testing.T) { testSplice(t, "tcp", "unix") })
 	t.Run("tcp-to-file", func(t *testing.T) { testSpliceToFile(t, "tcp", "file") })
 	t.Run("unix-to-file", func(t *testing.T) { testSpliceToFile(t, "unix", "file") })
 	t.Run("no-unixpacket", testSpliceNoUnixpacket)
@@ -159,23 +160,30 @@ func (tc spliceTestCase) testFile(t *testing.T) {
 }
 
 func testSpliceReaderAtEOF(t *testing.T, upNet, downNet string) {
+	// UnixConn doesn't implement io.ReaderFrom, which will fail
+	// the following test in asserting a UnixConn to be an io.ReaderFrom,
+	// so skip this test.
+	if upNet == "unix" || downNet == "unix" {
+		t.Skip("skipping test on unix socket")
+	}
+
 	clientUp, serverUp := spliceTestSocketPair(t, upNet)
 	defer clientUp.Close()
 	clientDown, serverDown := spliceTestSocketPair(t, downNet)
 	defer clientDown.Close()
 
 	serverUp.Close()
 
-	// We'd like to call net.splice here and check the handled return
+	// We'd like to call net.spliceFrom here and check the handled return
 	// value, but we disable splice on old Linux kernels.
 	//
-	// In that case, poll.Splice and net.splice return a non-nil error
+	// In that case, poll.Splice and net.spliceFrom return a non-nil error
 	// and handled == false. We'd ideally like to see handled == true
 	// because the source reader is at EOF, but if we're running on an old
-	// kernel, and splice is disabled, we won't see EOF from net.splice,
+	// kernel, and splice is disabled, we won't see EOF from net.spliceFrom,
 	// because we won't touch the reader at all.
 	//
-	// Trying to untangle the errors from net.splice and match them
+	// Trying to untangle the errors from net.spliceFrom and match them
 	// against the errors created by the poll package would be brittle,
 	// so this is a higher level test.
 	//
@@ -268,7 +276,7 @@ func testSpliceNoUnixpacket(t *testing.T) {
 	//
 	// What we want is err == nil and handled == false, i.e. we never
 	// called poll.Splice, because we know the unix socket's network.
-	_, err, handled := splice(serverDown.(*TCPConn).fd, serverUp)
+	_, err, handled := spliceFrom(serverDown.(*TCPConn).fd, serverUp)
 	if err != nil || handled != false {
 		t.Fatalf("got err = %v, handled = %t, want nil error, handled == false", err, handled)
 	}
@@ -289,7 +297,7 @@ func testSpliceNoUnixgram(t *testing.T) {
 	defer clientDown.Close()
 	defer serverDown.Close()
 	// Analogous to testSpliceNoUnixpacket.
-	_, err, handled := splice(serverDown.(*TCPConn).fd, up)
+	_, err, handled := spliceFrom(serverDown.(*TCPConn).fd, up)
 	if err != nil || handled != false {
 		t.Fatalf("got err = %v, handled = %t, want nil error, handled == false", err, handled)
 	}
@@ -300,6 +308,7 @@ func BenchmarkSplice(b *testing.B) {
 
 	b.Run("tcp-to-tcp", func(b *testing.B) { benchSplice(b, "tcp", "tcp") })
 	b.Run("unix-to-tcp", func(b *testing.B) { benchSplice(b, "unix", "tcp") })
+	b.Run("tcp-to-unix", func(b *testing.B) { benchSplice(b, "tcp", "unix") })
 }
 
 func benchSplice(b *testing.B, upNet, downNet string) {

src/net/tcpsock.go

@@ -134,6 +134,18 @@ func (c *TCPConn) ReadFrom(r io.Reader) (int64, error) {
 	return n, err
 }
 
+// WriteTo implements the io.WriterTo WriteTo method.
+func (c *TCPConn) WriteTo(w io.Writer) (int64, error) {
+	if !c.ok() {
+		return 0, syscall.EINVAL
+	}
+	n, err := c.writeTo(w)
+	if err != nil && err != io.EOF {
+		err = &OpError{Op: "writeto", Net: c.fd.net, Source: c.fd.laddr, Addr: c.fd.raddr, Err: err}
+	}
+	return n, err
+}
+
 // CloseRead shuts down the reading side of the TCP connection.
 // Most callers should just use Close.
 func (c *TCPConn) CloseRead() error {

src/net/tcpsock_plan9.go

@@ -14,6 +14,10 @@ func (c *TCPConn) readFrom(r io.Reader) (int64, error) {
 	return genericReadFrom(c, r)
 }
 
+func (c *TCPConn) writeTo(w io.Writer) (int64, error) {
+	return genericWriteTo(c, w)
+}
+
 func (sd *sysDialer) dialTCP(ctx context.Context, laddr, raddr *TCPAddr) (*TCPConn, error) {
 	if h := sd.testHookDialTCP; h != nil {
 		return h(ctx, sd.network, laddr, raddr)

src/net/tcpsock_posix.go

@@ -45,7 +45,7 @@ func (a *TCPAddr) toLocal(net string) sockaddr {
 }
 
 func (c *TCPConn) readFrom(r io.Reader) (int64, error) {
-	if n, err, handled := splice(c.fd, r); handled {
+	if n, err, handled := spliceFrom(c.fd, r); handled {
 		return n, err
 	}
 	if n, err, handled := sendFile(c.fd, r); handled {
@@ -54,6 +54,13 @@ func (c *TCPConn) readFrom(r io.Reader) (int64, error) {
 	return genericReadFrom(c, r)
 }
 
+func (c *TCPConn) writeTo(w io.Writer) (int64, error) {
+	if n, err, handled := spliceTo(w, c.fd); handled {
+		return n, err
+	}
+	return genericWriteTo(c, w)
+}
+
 func (sd *sysDialer) dialTCP(ctx context.Context, laddr, raddr *TCPAddr) (*TCPConn, error) {
 	if h := sd.testHookDialTCP; h != nil {
 		return h(ctx, sd.network, laddr, raddr)

src/os/export_linux_test.go

@@ -5,7 +5,8 @@
 package os
 
 var (
-	PollCopyFileRangeP = &pollCopyFileRange
-	PollSpliceFile     = &pollSplice
-	GetPollFDForTest   = getPollFD
+	PollCopyFileRangeP  = &pollCopyFileRange
+	PollSpliceFile      = &pollSplice
+	PollSendFile        = &pollSendFile
+	GetPollFDAndNetwork = getPollFDAndNetwork
 )