perf: improve copyZeroAlloc for os.File and net.TCPConn

http.go

@@ -2219,20 +2219,109 @@ func writeBodyFixedSize(w *bufio.Writer, r io.Reader, size int64) error {
 	return err
 }
 
+// copyZeroAlloc optimizes io.Copy by calling ReadFrom or WriteTo only when
+// copying between os.File and net.TCPConn. If the reader has a WriteTo
+// method, it uses WriteTo for copying; if the writer has a ReadFrom method,
+// it uses ReadFrom for copying. If neither method is available, it gets a
+// buffer from sync.Pool to perform the copy.
+//
+// io.CopyBuffer always uses the WriterTo or ReadFrom interface if it's
+// available. however, os.File and net.TCPConn unfortunately have a
+// fallback in their WriterTo that calls io.Copy if sendfile isn't possible.
+//
+// See issue: https://github.com/valyala/fasthttp/issues/1889
+//
+// sendfile can only be triggered when copying between os.File and net.TCPConn.
+// Since the function confirming zero-copy is a private function, we use
+// ReadFrom only in this specific scenario. For all other cases, we prioritize
+// using our own copyBuffer method.
+//
+// o: our copyBuffer
+// r: readFrom
+// w: writeTo
+//
+// write\read *File  *TCPConn  writeTo  other
+// *File        o       r         w       o
+// *TCPConn    w,r      o         w       o
+// readFrom     r       r         w       r
+// other        o       o         w       o
+//
+//nolint:dupword
 func copyZeroAlloc(w io.Writer, r io.Reader) (int64, error) {
-	if wt, ok := r.(io.WriterTo); ok {
-		return wt.WriteTo(w)
-	}
-	if rt, ok := w.(io.ReaderFrom); ok {
-		return rt.ReadFrom(r)
+	var readerIsFile, readerIsConn bool
+
+	switch r := r.(type) {
+	case *os.File:
+		readerIsFile = true
+	case *net.TCPConn:
+		readerIsConn = true
+	case io.WriterTo:
+		return r.WriteTo(w)
+	}
+
+	switch w := w.(type) {
+	case *os.File:
+		if readerIsConn {
+			return w.ReadFrom(r)
+		}
+	case *net.TCPConn:
+		if readerIsFile {
+			// net.WriteTo requires go1.22 or later
+			// Benchmark tests show that on Windows, WriteTo performs
+			// significantly better than ReadFrom. On Linux, however,
+			// ReadFrom slightly outperforms WriteTo. When possible,
+			// copyZeroAlloc aims to perform  better than or as well
+			// as io.Copy, so we use WriteTo whenever possible for
+			// optimal performance.
+			if rt, ok := r.(io.WriterTo); ok {
+				return rt.WriteTo(w)
+			}
+			return w.ReadFrom(r)
+		}
+	case io.ReaderFrom:
+		return w.ReadFrom(r)
 	}
+
 	vbuf := copyBufPool.Get()
 	buf := vbuf.([]byte)
-	n, err := io.CopyBuffer(w, r, buf)
+	n, err := copyBuffer(w, r, buf)
 	copyBufPool.Put(vbuf)
 	return n, err
 }
 
+// copyBuffer is rewritten from io.copyBuffer. We do not check if src has a
+// WriteTo method, if dst has a ReadFrom method, or if buf is empty.
+func copyBuffer(dst io.Writer, src io.Reader, buf []byte) (written int64, err error) {
+	for {
+		nr, er := src.Read(buf)
+		if nr > 0 {
+			nw, ew := dst.Write(buf[0:nr])
+			if nw < 0 || nr < nw {
+				nw = 0
+				if ew == nil {
+					ew = errors.New("invalid write result")
+				}
+			}
+			written += int64(nw)
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
+		}
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
+		}
+	}
+	return written, err
+}
+
 var copyBufPool = sync.Pool{
 	New: func() any {
 		return make([]byte, 4096)

http_timing_test.go

@@ -0,0 +1,283 @@
+package fasthttp
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"os"
+	"strings"
+	"testing"
+)
+
+func BenchmarkCopyZeroAllocOSFileToBytesBuffer(b *testing.B) {
+	r, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer r.Close()
+
+	buf := &bytes.Buffer{}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		buf.Reset()
+		_, err = copyZeroAlloc(buf, r)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocBytesBufferToOSFile(b *testing.B) {
+	f, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer f.Close()
+
+	buf := &bytes.Buffer{}
+	_, err = io.Copy(buf, f)
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	tmp, err := os.CreateTemp(os.TempDir(), "test_*")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer os.Remove(tmp.Name())
+
+	w, err := os.OpenFile(tmp.Name(), os.O_WRONLY, 0o444)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer w.Close()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := w.Seek(0, 0)
+		if err != nil {
+			b.Fatal(err)
+		}
+		_, err = copyZeroAlloc(w, buf)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocOSFileToStringsBuilder(b *testing.B) {
+	r, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatalf("Failed to open testing file: %v", err)
+	}
+	defer r.Close()
+
+	w := &strings.Builder{}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		w.Reset()
+		_, err = copyZeroAlloc(w, r)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocIOLimitedReaderToOSFile(b *testing.B) {
+	f, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer f.Close()
+
+	r := io.LimitReader(f, 1024)
+
+	tmp, err := os.CreateTemp(os.TempDir(), "test_*")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer os.Remove(tmp.Name())
+
+	w, err := os.OpenFile(tmp.Name(), os.O_WRONLY, 0o444)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer w.Close()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := w.Seek(0, 0)
+		if err != nil {
+			b.Fatal(err)
+		}
+		_, err = copyZeroAlloc(w, r)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocOSFileToOSFile(b *testing.B) {
+	r, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer r.Close()
+
+	f, err := os.CreateTemp(os.TempDir(), "test_*")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer os.Remove(f.Name())
+
+	w, err := os.OpenFile(f.Name(), os.O_WRONLY, 0o444)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer w.Close()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := w.Seek(0, 0)
+		if err != nil {
+			b.Fatal(err)
+		}
+		_, err = copyZeroAlloc(w, r)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocOSFileToNetConn(b *testing.B) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	addr := ln.Addr().String()
+	defer ln.Close()
+
+	done := make(chan struct{})
+	defer close(done)
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			b.Error(err)
+			return
+		}
+		defer conn.Close()
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				_, err := io.Copy(io.Discard, conn)
+				if err != nil {
+					b.Error(err)
+					return
+				}
+			}
+		}
+	}()
+
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer conn.Close()
+
+	file, err := os.Open("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer file.Close()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		if _, err := copyZeroAlloc(conn, file); err != nil {
+			b.Fatal(err)
+		}
+	}
+}
+
+func BenchmarkCopyZeroAllocNetConnToOSFile(b *testing.B) {
+	data, err := os.ReadFile("./README.md")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		b.Fatal(err)
+	}
+
+	addr := ln.Addr().String()
+	defer ln.Close()
+
+	done := make(chan struct{})
+	defer close(done)
+
+	writeDone := make(chan struct{})
+	go func() {
+		for {
+			select {
+			case <-done:
+				return
+			default:
+				conn, err := ln.Accept()
+				if err != nil {
+					b.Error(err)
+					return
+				}
+				_, err = conn.Write(data)
+				if err != nil {
+					b.Error(err)
+				}
+				conn.Close()
+				writeDone <- struct{}{}
+			}
+		}
+	}()
+
+	tmp, err := os.CreateTemp(os.TempDir(), "test_*")
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer os.Remove(tmp.Name())
+
+	file, err := os.OpenFile(tmp.Name(), os.O_WRONLY, 0o444)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer file.Close()
+
+	conn, err := net.Dial("tcp", addr)
+	if err != nil {
+		b.Fatal(err)
+	}
+	defer conn.Close()
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		b.StopTimer()
+		<-writeDone
+		_, err = file.Seek(0, 0)
+		if err != nil {
+			b.Fatal(err)
+		}
+		b.StartTimer()
+		_, err = copyZeroAlloc(file, conn)
+		if err != nil {
+			b.Fatal(err)
+		}
+		b.StopTimer()
+		conn, err = net.Dial("tcp", addr)
+		if err != nil {
+			b.Fatal(err)
+		}
+	}
+}