t.Skip("skipping unix-to-tcp tests")
}
t.Run("unix-to-tcp", func(t *testing.T) { testSplice(t, "unix", "tcp") })
+ 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)
t.Run("no-unixgram", testSpliceNoUnixgram)
}
+func testSpliceToFile(t *testing.T, upNet, downNet string) {
+ t.Run("simple", spliceTestCase{upNet, downNet, 128, 128, 0}.testFile)
+ t.Run("multipleWrite", spliceTestCase{upNet, downNet, 4096, 1 << 20, 0}.testFile)
+ t.Run("big", spliceTestCase{upNet, downNet, 5 << 20, 1 << 30, 0}.testFile)
+ t.Run("honorsLimitedReader", spliceTestCase{upNet, downNet, 4096, 1 << 20, 1 << 10}.testFile)
+ t.Run("updatesLimitedReaderN", spliceTestCase{upNet, downNet, 1024, 4096, 4096 + 100}.testFile)
+ t.Run("limitedReaderAtLimit", spliceTestCase{upNet, downNet, 32, 128, 128}.testFile)
+}
+
func testSplice(t *testing.T, upNet, downNet string) {
t.Run("simple", spliceTestCase{upNet, downNet, 128, 128, 0}.test)
t.Run("multipleWrite", spliceTestCase{upNet, downNet, 4096, 1 << 20, 0}.test)
}
}
+func (tc spliceTestCase) testFile(t *testing.T) {
+ f, err := os.CreateTemp(t.TempDir(), "linux-splice-to-file")
+ if err != nil {
+ t.Fatal(err)
+ }
+ defer f.Close()
+
+ client, server := spliceTestSocketPair(t, tc.upNet)
+ defer server.Close()
+
+ cleanup, err := startSpliceClient(client, "w", tc.chunkSize, tc.totalSize)
+ if err != nil {
+ client.Close()
+ t.Fatal("failed to start splice client:", err)
+ }
+ defer cleanup()
+
+ var (
+ r io.Reader = server
+ actualSize = tc.totalSize
+ )
+ if tc.limitReadSize > 0 {
+ if tc.limitReadSize < actualSize {
+ actualSize = tc.limitReadSize
+ }
+
+ r = &io.LimitedReader{
+ N: int64(tc.limitReadSize),
+ R: r,
+ }
+ }
+
+ got, err := io.Copy(f, r)
+ if err != nil {
+ t.Fatalf("failed to ReadFrom with error: %v", err)
+ }
+ if want := int64(actualSize); got != want {
+ t.Errorf("got %d bytes, want %d", got, want)
+ }
+ if tc.limitReadSize > 0 {
+ wantN := 0
+ if tc.limitReadSize > actualSize {
+ wantN = tc.limitReadSize - actualSize
+ }
+
+ if gotN := r.(*io.LimitedReader).N; gotN != int64(wantN) {
+ t.Errorf("r.N = %d, want %d", gotN, wantN)
+ }
+ }
+}
+
func testSpliceReaderAtEOF(t *testing.T, upNet, downNet string) {
clientUp, serverUp := spliceTestSocketPair(t, upNet)
defer clientUp.Close()
}
}
}
+
+func BenchmarkSpliceFile(b *testing.B) {
+ b.Run("tcp-to-file", func(b *testing.B) { benchmarkSpliceFile(b, "tcp") })
+ b.Run("unix-to-file", func(b *testing.B) { benchmarkSpliceFile(b, "unix") })
+}
+
+func benchmarkSpliceFile(b *testing.B, proto string) {
+ for i := 0; i <= 10; i++ {
+ size := 1 << (i + 10)
+ bench := spliceFileBench{
+ proto: proto,
+ chunkSize: size,
+ }
+ b.Run(strconv.Itoa(size), bench.benchSpliceFile)
+ }
+}
+
+type spliceFileBench struct {
+ proto string
+ chunkSize int
+}
+
+func (bench spliceFileBench) benchSpliceFile(b *testing.B) {
+ f, err := os.OpenFile(os.DevNull, os.O_WRONLY, 0)
+ if err != nil {
+ b.Fatal(err)
+ }
+ defer f.Close()
+
+ totalSize := b.N * bench.chunkSize
+
+ client, server := spliceTestSocketPair(b, bench.proto)
+ defer server.Close()
+
+ cleanup, err := startSpliceClient(client, "w", bench.chunkSize, totalSize)
+ if err != nil {
+ client.Close()
+ b.Fatalf("failed to start splice client: %v", err)
+ }
+ defer cleanup()
+
+ b.ReportAllocs()
+ b.SetBytes(int64(bench.chunkSize))
+ b.ResetTimer()
+
+ got, err := io.Copy(f, server)
+ if err != nil {
+ b.Fatalf("failed to ReadFrom with error: %v", err)
+ }
+ if want := int64(totalSize); got != want {
+ b.Errorf("bytes sent mismatch, got: %d, want: %d", got, want)
+ }
+}
import (
"internal/poll"
"io"
+ "syscall"
)
-var pollCopyFileRange = poll.CopyFileRange
+var (
+ pollCopyFileRange = poll.CopyFileRange
+ pollSplice = poll.Splice
+)
func (f *File) readFrom(r io.Reader) (written int64, handled bool, err error) {
+ written, handled, err = f.copyFileRange(r)
+ if handled {
+ return
+ }
+ return f.spliceToFile(r)
+}
+
+func (f *File) spliceToFile(r io.Reader) (written int64, handled bool, err error) {
+ var (
+ remain int64
+ lr *io.LimitedReader
+ )
+ if lr, r, remain = tryLimitedReader(r); remain <= 0 {
+ return 0, true, nil
+ }
+
+ pfd := getPollFD(r)
+ // TODO(panjf2000): run some tests to see if we should unlock the non-streams for splice.
+ // Streams benefit the most from the splice(2), non-streams are not even supported in old kernels
+ // where splice(2) will just return EINVAL; newer kernels support non-streams like UDP, but I really
+ // doubt that splice(2) could help non-streams, cuz they usually send small frames respectively
+ // and one splice call would result in one frame.
+ // splice(2) is suitable for large data but the generation of fragments defeats its edge here.
+ // Therefore, don't bother to try splice if the r is not a streaming descriptor.
+ if pfd == nil || !pfd.IsStream {
+ return
+ }
+
+ var syscallName string
+ written, handled, syscallName, err = pollSplice(&f.pfd, pfd, remain)
+
+ if lr != nil {
+ lr.N = remain - written
+ }
+
+ return written, handled, NewSyscallError(syscallName, err)
+}
+
+// getPollFD tries to get the poll.FD from the given io.Reader by expecting
+// the underlying type of r to be the implementation of syscall.Conn that contains
+// a *net.rawConn.
+func getPollFD(r io.Reader) *poll.FD {
+ sc, ok := r.(syscall.Conn)
+ if !ok {
+ return nil
+ }
+ rc, err := sc.SyscallConn()
+ if err != nil {
+ return nil
+ }
+ ipfd, ok := rc.(interface{ PollFD() *poll.FD })
+ if !ok {
+ return nil
+ }
+ return ipfd.PollFD()
+}
+
+func (f *File) copyFileRange(r io.Reader) (written int64, handled bool, err error) {
// copy_file_range(2) does not support destinations opened with
// O_APPEND, so don't even try.
if f.appendMode {
return 0, false, nil
}
- remain := int64(1 << 62)
-
- lr, ok := r.(*io.LimitedReader)
- if ok {
- remain, r = lr.N, lr.R
- if remain <= 0 {
- return 0, true, nil
- }
+ var (
+ remain int64
+ lr *io.LimitedReader
+ )
+ if lr, r, remain = tryLimitedReader(r); remain <= 0 {
+ return 0, true, nil
}
src, ok := r.(*File)
}
return written, handled, NewSyscallError("copy_file_range", err)
}
+
+// tryLimitedReader tries to assert the io.Reader to io.LimitedReader, it returns the io.LimitedReader,
+// the underlying io.Reader and the remaining amount of bytes if the assertion succeeds,
+// otherwise it just returns the original io.Reader and the theoretical unlimited remaining amount of bytes.
+func tryLimitedReader(r io.Reader) (*io.LimitedReader, io.Reader, int64) {
+ remain := int64(1 << 62)
+
+ lr, ok := r.(*io.LimitedReader)
+ if !ok {
+ return nil, r, remain
+ }
+
+ remain = lr.N
+ return lr, lr.R, remain
+}
"internal/poll"
"io"
"math/rand"
+ "net"
"os"
. "os"
"path/filepath"
"syscall"
"testing"
"time"
+
+ "golang.org/x/net/nettest"
)
func TestCopyFileRange(t *testing.T) {
})
}
+func TestSpliceFile(t *testing.T) {
+ sizes := []int{
+ 1,
+ 42,
+ 1025,
+ syscall.Getpagesize() + 1,
+ 32769,
+ }
+ t.Run("Basic-TCP", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "tcp", int64(size), -1)
+ })
+ }
+ })
+ t.Run("Basic-Unix", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "unix", int64(size), -1)
+ })
+ }
+ })
+ t.Run("Limited", func(t *testing.T) {
+ t.Run("OneLess-TCP", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "tcp", int64(size), int64(size)-1)
+ })
+ }
+ })
+ t.Run("OneLess-Unix", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "unix", int64(size), int64(size)-1)
+ })
+ }
+ })
+ t.Run("Half-TCP", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "tcp", int64(size), int64(size)/2)
+ })
+ }
+ })
+ t.Run("Half-Unix", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "unix", int64(size), int64(size)/2)
+ })
+ }
+ })
+ t.Run("More-TCP", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "tcp", int64(size), int64(size)+1)
+ })
+ }
+ })
+ t.Run("More-Unix", func(t *testing.T) {
+ for _, size := range sizes {
+ t.Run(strconv.Itoa(size), func(t *testing.T) {
+ testSpliceFile(t, "unix", int64(size), int64(size)+1)
+ })
+ }
+ })
+ })
+}
+
+func testSpliceFile(t *testing.T, proto string, size, limit int64) {
+ dst, src, data, hook, cleanup := newSpliceFileTest(t, proto, size)
+ defer cleanup()
+
+ // If we have a limit, wrap the reader.
+ var (
+ r io.Reader
+ lr *io.LimitedReader
+ )
+ if limit >= 0 {
+ lr = &io.LimitedReader{N: limit, R: src}
+ r = lr
+ if limit < int64(len(data)) {
+ data = data[:limit]
+ }
+ } else {
+ r = src
+ }
+ // Now call ReadFrom (through io.Copy), which will hopefully call poll.Splice
+ n, err := io.Copy(dst, r)
+ if err != nil {
+ t.Fatal(err)
+ }
+
+ // We should have called poll.Splice with the right file descriptor arguments.
+ if n > 0 && !hook.called {
+ t.Fatal("expected to called poll.Splice")
+ }
+ if hook.called && hook.dstfd != int(dst.Fd()) {
+ t.Fatalf("wrong destination file descriptor: got %d, want %d", hook.dstfd, dst.Fd())
+ }
+ sc, ok := src.(syscall.Conn)
+ if !ok {
+ t.Fatalf("server Conn is not a syscall.Conn")
+ }
+ rc, err := sc.SyscallConn()
+ if err != nil {
+ t.Fatalf("server Conn SyscallConn error: %v", err)
+ }
+ if err = rc.Control(func(fd uintptr) {
+ if hook.called && hook.srcfd != int(fd) {
+ t.Fatalf("wrong source file descriptor: got %d, want %d", hook.srcfd, int(fd))
+ }
+ }); err != nil {
+ t.Fatalf("server Conn Control error: %v", err)
+ }
+
+ // Check that the offsets after the transfer make sense, that the size
+ // of the transfer was reported correctly, and that the destination
+ // file contains exactly the bytes we expect it to contain.
+ dstoff, err := dst.Seek(0, io.SeekCurrent)
+ if err != nil {
+ t.Fatal(err)
+ }
+ if dstoff != int64(len(data)) {
+ t.Errorf("dstoff = %d, want %d", dstoff, len(data))
+ }
+ if n != int64(len(data)) {
+ t.Errorf("short ReadFrom: wrote %d bytes, want %d", n, len(data))
+ }
+ mustSeekStart(t, dst)
+ mustContainData(t, dst, data)
+
+ // If we had a limit, check that it was updated.
+ if lr != nil {
+ if want := limit - n; lr.N != want {
+ t.Fatalf("didn't update limit correctly: got %d, want %d", lr.N, want)
+ }
+ }
+}
+
func testCopyFileRange(t *testing.T, size int64, limit int64) {
dst, src, data, hook := newCopyFileRangeTest(t, size)
return dst, src, data, hook
}
+// newSpliceFileTest initializes a new test for splice.
+//
+// It creates source sockets and destination file, and populates the source sockets
+// with random data of the specified size. It also hooks package os' call
+// to poll.Splice and returns the hook so it can be inspected.
+func newSpliceFileTest(t *testing.T, proto string, size int64) (*File, net.Conn, []byte, *spliceFileHook, func()) {
+ t.Helper()
+
+ hook := hookSpliceFile(t)
+
+ client, server := createSocketPair(t, proto)
+
+ dst, err := CreateTemp(t.TempDir(), "dst-splice-file-test")
+ if err != nil {
+ t.Fatal(err)
+ }
+ t.Cleanup(func() { dst.Close() })
+
+ randSeed := time.Now().Unix()
+ t.Logf("random data seed: %d\n", randSeed)
+ prng := rand.New(rand.NewSource(randSeed))
+ data := make([]byte, size)
+ prng.Read(data)
+
+ done := make(chan struct{})
+ go func() {
+ client.Write(data)
+ client.Close()
+ close(done)
+ }()
+
+ return dst, server, data, hook, func() { <-done }
+}
+
// mustContainData ensures that the specified file contains exactly the
// specified data.
func mustContainData(t *testing.T, f *File, data []byte) {
*PollCopyFileRangeP = h.original
}
+func hookSpliceFile(t *testing.T) *spliceFileHook {
+ h := new(spliceFileHook)
+ h.install()
+ t.Cleanup(h.uninstall)
+ return h
+}
+
+type spliceFileHook struct {
+ called bool
+ dstfd int
+ srcfd int
+ remain int64
+
+ written int64
+ handled bool
+ sc string
+ err error
+
+ original func(dst, src *poll.FD, remain int64) (int64, bool, string, error)
+}
+
+func (h *spliceFileHook) install() {
+ h.original = *PollSpliceFile
+ *PollSpliceFile = func(dst, src *poll.FD, remain int64) (int64, bool, string, error) {
+ h.called = true
+ h.dstfd = dst.Sysfd
+ h.srcfd = src.Sysfd
+ h.remain = remain
+ h.written, h.handled, h.sc, h.err = h.original(dst, src, remain)
+ return h.written, h.handled, h.sc, h.err
+ }
+}
+
+func (h *spliceFileHook) uninstall() {
+ *PollSpliceFile = h.original
+}
+
// On some kernels copy_file_range fails on files in /proc.
func TestProcCopy(t *testing.T) {
t.Parallel()
t.Errorf("copy of %q got %q want %q\n", cmdlineFile, copy, cmdline)
}
}
+
+func TestGetPollFDFromReader(t *testing.T) {
+ t.Run("tcp", func(t *testing.T) { testGetPollFromReader(t, "tcp") })
+ t.Run("unix", func(t *testing.T) { testGetPollFromReader(t, "unix") })
+}
+
+func testGetPollFromReader(t *testing.T, proto string) {
+ _, server := createSocketPair(t, proto)
+ sc, ok := server.(syscall.Conn)
+ if !ok {
+ t.Fatalf("server Conn is not a syscall.Conn")
+ }
+ rc, err := sc.SyscallConn()
+ if err != nil {
+ t.Fatalf("server SyscallConn error: %v", err)
+ }
+ if err = rc.Control(func(fd uintptr) {
+ pfd := os.GetPollFDForTest(server)
+ if pfd == nil {
+ t.Fatalf("GetPollFDForTest didn't return poll.FD")
+ }
+ if pfd.Sysfd != int(fd) {
+ t.Fatalf("GetPollFDForTest returned wrong poll.FD, got: %d, want: %d", pfd.Sysfd, int(fd))
+ }
+ if !pfd.IsStream {
+ t.Fatalf("expected IsStream to be true")
+ }
+ if err = pfd.Init(proto, true); err == nil {
+ t.Fatalf("Init should have failed with the initialized poll.FD and return EEXIST error")
+ }
+ }); err != nil {
+ t.Fatalf("server Control error: %v", err)
+ }
+}
+
+func createSocketPair(t *testing.T, proto string) (client, server net.Conn) {
+ t.Helper()
+
+ ln, err := nettest.NewLocalListener(proto)
+ if err != nil {
+ t.Fatalf("NewLocalListener error: %v", err)
+ }
+ t.Cleanup(func() {
+ if ln != nil {
+ ln.Close()
+ }
+ if client != nil {
+ client.Close()
+ }
+ if server != nil {
+ server.Close()
+ }
+ })
+ ch := make(chan struct{})
+ go func() {
+ var err error
+ server, err = ln.Accept()
+ if err != nil {
+ t.Errorf("Accept new connection error: %v", err)
+ }
+ ch <- struct{}{}
+ }()
+ client, err = net.Dial(proto, ln.Addr().String())
+ <-ch
+ if err != nil {
+ t.Fatalf("Dial new connection error: %v", err)
+ }
+ return client, server
+}