| // Copyright 2018 The Go Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style |
| // license that can be found in the LICENSE file. |
| |
| // +build linux |
| |
| package net |
| |
| import ( |
| "bytes" |
| "fmt" |
| "io" |
| "io/ioutil" |
| "sync" |
| "testing" |
| ) |
| |
| func TestSplice(t *testing.T) { |
| t.Run("simple", testSpliceSimple) |
| t.Run("multipleWrite", testSpliceMultipleWrite) |
| t.Run("big", testSpliceBig) |
| t.Run("honorsLimitedReader", testSpliceHonorsLimitedReader) |
| t.Run("readerAtEOF", testSpliceReaderAtEOF) |
| t.Run("issue25985", testSpliceIssue25985) |
| } |
| |
| func testSpliceSimple(t *testing.T) { |
| srv, err := newSpliceTestServer() |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer srv.Close() |
| copyDone := srv.Copy() |
| msg := []byte("splice test") |
| if _, err := srv.Write(msg); err != nil { |
| t.Fatal(err) |
| } |
| got := make([]byte, len(msg)) |
| if _, err := io.ReadFull(srv, got); err != nil { |
| t.Fatal(err) |
| } |
| if !bytes.Equal(got, msg) { |
| t.Errorf("got %q, wrote %q", got, msg) |
| } |
| srv.CloseWrite() |
| srv.CloseRead() |
| if err := <-copyDone; err != nil { |
| t.Errorf("splice: %v", err) |
| } |
| } |
| |
| func testSpliceMultipleWrite(t *testing.T) { |
| srv, err := newSpliceTestServer() |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer srv.Close() |
| copyDone := srv.Copy() |
| msg1 := []byte("splice test part 1 ") |
| msg2 := []byte(" splice test part 2") |
| if _, err := srv.Write(msg1); err != nil { |
| t.Fatalf("Write: %v", err) |
| } |
| if _, err := srv.Write(msg2); err != nil { |
| t.Fatal(err) |
| } |
| got := make([]byte, len(msg1)+len(msg2)) |
| if _, err := io.ReadFull(srv, got); err != nil { |
| t.Fatal(err) |
| } |
| want := append(msg1, msg2...) |
| if !bytes.Equal(got, want) { |
| t.Errorf("got %q, wrote %q", got, want) |
| } |
| srv.CloseWrite() |
| srv.CloseRead() |
| if err := <-copyDone; err != nil { |
| t.Errorf("splice: %v", err) |
| } |
| } |
| |
| func testSpliceBig(t *testing.T) { |
| size := 1<<31 - 1 |
| if testing.Short() { |
| size = 1 << 25 |
| } |
| srv, err := newSpliceTestServer() |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer srv.Close() |
| big := make([]byte, size) |
| copyDone := srv.Copy() |
| type readResult struct { |
| b []byte |
| err error |
| } |
| readDone := make(chan readResult) |
| go func() { |
| got := make([]byte, len(big)) |
| _, err := io.ReadFull(srv, got) |
| readDone <- readResult{got, err} |
| }() |
| if _, err := srv.Write(big); err != nil { |
| t.Fatal(err) |
| } |
| res := <-readDone |
| if res.err != nil { |
| t.Fatal(res.err) |
| } |
| got := res.b |
| if !bytes.Equal(got, big) { |
| t.Errorf("input and output differ") |
| } |
| srv.CloseWrite() |
| srv.CloseRead() |
| if err := <-copyDone; err != nil { |
| t.Errorf("splice: %v", err) |
| } |
| } |
| |
| func testSpliceHonorsLimitedReader(t *testing.T) { |
| t.Run("stopsAfterN", testSpliceStopsAfterN) |
| t.Run("updatesN", testSpliceUpdatesN) |
| } |
| |
| func testSpliceStopsAfterN(t *testing.T) { |
| clientUp, serverUp, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientUp.Close() |
| defer serverUp.Close() |
| clientDown, serverDown, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientDown.Close() |
| defer serverDown.Close() |
| count := 128 |
| copyDone := make(chan error) |
| lr := &io.LimitedReader{ |
| N: int64(count), |
| R: serverUp, |
| } |
| go func() { |
| _, err := io.Copy(serverDown, lr) |
| serverDown.Close() |
| copyDone <- err |
| }() |
| msg := make([]byte, 2*count) |
| if _, err := clientUp.Write(msg); err != nil { |
| t.Fatal(err) |
| } |
| clientUp.Close() |
| var buf bytes.Buffer |
| if _, err := io.Copy(&buf, clientDown); err != nil { |
| t.Fatal(err) |
| } |
| if buf.Len() != count { |
| t.Errorf("splice transferred %d bytes, want to stop after %d", buf.Len(), count) |
| } |
| clientDown.Close() |
| if err := <-copyDone; err != nil { |
| t.Errorf("splice: %v", err) |
| } |
| } |
| |
| func testSpliceUpdatesN(t *testing.T) { |
| clientUp, serverUp, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientUp.Close() |
| defer serverUp.Close() |
| clientDown, serverDown, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientDown.Close() |
| defer serverDown.Close() |
| count := 128 |
| copyDone := make(chan error) |
| lr := &io.LimitedReader{ |
| N: int64(100 + count), |
| R: serverUp, |
| } |
| go func() { |
| _, err := io.Copy(serverDown, lr) |
| copyDone <- err |
| }() |
| msg := make([]byte, count) |
| if _, err := clientUp.Write(msg); err != nil { |
| t.Fatal(err) |
| } |
| clientUp.Close() |
| got := make([]byte, count) |
| if _, err := io.ReadFull(clientDown, got); err != nil { |
| t.Fatal(err) |
| } |
| clientDown.Close() |
| if err := <-copyDone; err != nil { |
| t.Errorf("splice: %v", err) |
| } |
| wantN := int64(100) |
| if lr.N != wantN { |
| t.Errorf("lr.N = %d, want %d", lr.N, wantN) |
| } |
| } |
| |
| func testSpliceReaderAtEOF(t *testing.T) { |
| clientUp, serverUp, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientUp.Close() |
| defer serverUp.Close() |
| clientDown, serverDown, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer clientDown.Close() |
| defer serverDown.Close() |
| |
| serverUp.Close() |
| _, err, handled := splice(serverDown.(*TCPConn).fd, serverUp) |
| if !handled { |
| t.Errorf("closed connection: got err = %v, handled = %t, want handled = true", err, handled) |
| } |
| lr := &io.LimitedReader{ |
| N: 0, |
| R: serverUp, |
| } |
| _, err, handled = splice(serverDown.(*TCPConn).fd, lr) |
| if !handled { |
| t.Errorf("exhausted LimitedReader: got err = %v, handled = %t, want handled = true", err, handled) |
| } |
| } |
| |
| func testSpliceIssue25985(t *testing.T) { |
| front, err := newLocalListener("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer front.Close() |
| back, err := newLocalListener("tcp") |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer back.Close() |
| |
| var wg sync.WaitGroup |
| wg.Add(2) |
| |
| proxy := func() { |
| src, err := front.Accept() |
| if err != nil { |
| return |
| } |
| dst, err := Dial("tcp", back.Addr().String()) |
| if err != nil { |
| return |
| } |
| defer dst.Close() |
| defer src.Close() |
| go func() { |
| io.Copy(src, dst) |
| wg.Done() |
| }() |
| go func() { |
| io.Copy(dst, src) |
| wg.Done() |
| }() |
| } |
| |
| go proxy() |
| |
| toFront, err := Dial("tcp", front.Addr().String()) |
| if err != nil { |
| t.Fatal(err) |
| } |
| |
| io.WriteString(toFront, "foo") |
| toFront.Close() |
| |
| fromProxy, err := back.Accept() |
| if err != nil { |
| t.Fatal(err) |
| } |
| defer fromProxy.Close() |
| |
| _, err = ioutil.ReadAll(fromProxy) |
| if err != nil { |
| t.Fatal(err) |
| } |
| |
| wg.Wait() |
| } |
| |
| func BenchmarkTCPReadFrom(b *testing.B) { |
| testHookUninstaller.Do(uninstallTestHooks) |
| |
| var chunkSizes []int |
| for i := uint(10); i <= 20; i++ { |
| chunkSizes = append(chunkSizes, 1<<i) |
| } |
| // To benchmark the genericReadFrom code path, set this to false. |
| useSplice := true |
| for _, chunkSize := range chunkSizes { |
| b.Run(fmt.Sprint(chunkSize), func(b *testing.B) { |
| benchmarkSplice(b, chunkSize, useSplice) |
| }) |
| } |
| } |
| |
| func benchmarkSplice(b *testing.B, chunkSize int, useSplice bool) { |
| srv, err := newSpliceTestServer() |
| if err != nil { |
| b.Fatal(err) |
| } |
| defer srv.Close() |
| var copyDone <-chan error |
| if useSplice { |
| copyDone = srv.Copy() |
| } else { |
| copyDone = srv.CopyNoSplice() |
| } |
| chunk := make([]byte, chunkSize) |
| discardDone := make(chan struct{}) |
| go func() { |
| for { |
| buf := make([]byte, chunkSize) |
| _, err := srv.Read(buf) |
| if err != nil { |
| break |
| } |
| } |
| discardDone <- struct{}{} |
| }() |
| b.SetBytes(int64(chunkSize)) |
| b.ResetTimer() |
| for i := 0; i < b.N; i++ { |
| srv.Write(chunk) |
| } |
| srv.CloseWrite() |
| <-copyDone |
| srv.CloseRead() |
| <-discardDone |
| } |
| |
| type spliceTestServer struct { |
| clientUp io.WriteCloser |
| clientDown io.ReadCloser |
| serverUp io.ReadCloser |
| serverDown io.WriteCloser |
| } |
| |
| func newSpliceTestServer() (*spliceTestServer, error) { |
| // For now, both networks are hard-coded to TCP. |
| // If splice is enabled for non-tcp upstream connections, |
| // newSpliceTestServer will need to take a network parameter. |
| clientUp, serverUp, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| return nil, err |
| } |
| clientDown, serverDown, err := spliceTestSocketPair("tcp") |
| if err != nil { |
| clientUp.Close() |
| serverUp.Close() |
| return nil, err |
| } |
| return &spliceTestServer{clientUp, clientDown, serverUp, serverDown}, nil |
| } |
| |
| // Read reads from the downstream connection. |
| func (srv *spliceTestServer) Read(b []byte) (int, error) { |
| return srv.clientDown.Read(b) |
| } |
| |
| // Write writes to the upstream connection. |
| func (srv *spliceTestServer) Write(b []byte) (int, error) { |
| return srv.clientUp.Write(b) |
| } |
| |
| // Close closes the server. |
| func (srv *spliceTestServer) Close() error { |
| err := srv.closeUp() |
| err1 := srv.closeDown() |
| if err == nil { |
| return err1 |
| } |
| return err |
| } |
| |
| // CloseWrite closes the client side of the upstream connection. |
| func (srv *spliceTestServer) CloseWrite() error { |
| return srv.clientUp.Close() |
| } |
| |
| // CloseRead closes the client side of the downstream connection. |
| func (srv *spliceTestServer) CloseRead() error { |
| return srv.clientDown.Close() |
| } |
| |
| // Copy copies from the server side of the upstream connection |
| // to the server side of the downstream connection, in a separate |
| // goroutine. Copy is done when the first send on the returned |
| // channel succeeds. |
| func (srv *spliceTestServer) Copy() <-chan error { |
| ch := make(chan error) |
| go func() { |
| _, err := io.Copy(srv.serverDown, srv.serverUp) |
| ch <- err |
| close(ch) |
| }() |
| return ch |
| } |
| |
| // CopyNoSplice is like Copy, but ensures that the splice code path |
| // is not reached. |
| func (srv *spliceTestServer) CopyNoSplice() <-chan error { |
| type onlyReader struct { |
| io.Reader |
| } |
| ch := make(chan error) |
| go func() { |
| _, err := io.Copy(srv.serverDown, onlyReader{srv.serverUp}) |
| ch <- err |
| close(ch) |
| }() |
| return ch |
| } |
| |
| func (srv *spliceTestServer) closeUp() error { |
| var err, err1 error |
| if srv.serverUp != nil { |
| err = srv.serverUp.Close() |
| } |
| if srv.clientUp != nil { |
| err1 = srv.clientUp.Close() |
| } |
| if err == nil { |
| return err1 |
| } |
| return err |
| } |
| |
| func (srv *spliceTestServer) closeDown() error { |
| var err, err1 error |
| if srv.serverDown != nil { |
| err = srv.serverDown.Close() |
| } |
| if srv.clientDown != nil { |
| err1 = srv.clientDown.Close() |
| } |
| if err == nil { |
| return err1 |
| } |
| return err |
| } |
| |
| func spliceTestSocketPair(net string) (client, server Conn, err error) { |
| ln, err := newLocalListener(net) |
| if err != nil { |
| return nil, nil, err |
| } |
| defer ln.Close() |
| var cerr, serr error |
| acceptDone := make(chan struct{}) |
| go func() { |
| server, serr = ln.Accept() |
| acceptDone <- struct{}{} |
| }() |
| client, cerr = Dial(ln.Addr().Network(), ln.Addr().String()) |
| <-acceptDone |
| if cerr != nil { |
| if server != nil { |
| server.Close() |
| } |
| return nil, nil, cerr |
| } |
| if serr != nil { |
| if client != nil { |
| client.Close() |
| } |
| return nil, nil, serr |
| } |
| return client, server, nil |
| } |