Merge add4546 into de2b514

internal/iopool/copy.go

@@ -0,0 +1,26 @@
+package iopool
+
+import (
+	"io"
+	"sync"
+)
+
+type buffer struct {
+	b []byte
+}
+
+var _copyBufSize = 1024 * 32
+var _pool = sync.Pool{
+	New: func() interface{} {
+		return &buffer{make([]byte, _copyBufSize)}
+	},
+}
+
+// Copy wraps the io library's CopyBuffer func with a preallocated buffer from a
+// sync.Pool we maintain.
+func Copy(dst io.Writer, src io.Reader) (int64, error) {
+	buf := _pool.Get().(*buffer)
+	written, err := io.CopyBuffer(dst, src, buf.b)
+	_pool.Put(buf)
+	return written, err
+}

internal/iopool/copy_test.go

@@ -0,0 +1,32 @@
+package iopool
+
+import (
+	"bytes"
+	"math/rand"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestBuffers(t *testing.T) {
+	var wg sync.WaitGroup
+	for g := 0; g < 10; g++ {
+		wg.Add(1)
+		go func() {
+			for i := 0; i < 100; i++ {
+				inputBytes := make([]byte, rand.Intn(5000)+20)
+				_, err := rand.Read(inputBytes)
+				assert.NoError(t, err, "Unexpected error from rand.Read")
+				reader := bytes.NewReader(inputBytes)
+
+				outputBytes := make([]byte, 0, len(inputBytes))
+				writer := bytes.NewBuffer(outputBytes)
+
+				Copy(writer, reader)
+			}
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+}

transport/http/handler.go

@@ -23,12 +23,12 @@ package http
 import (
 	"bytes"
 	"context"
-	"io"
 	"net/http"
 	"time"
 
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/internal/errors"
+	"go.uber.org/yarpc/internal/iopool"
 	"go.uber.org/yarpc/internal/request"
 
 	"github.com/opentracing/opentracing-go"
@@ -130,7 +130,7 @@ func handleOnewayRequest(
 	// we will lose access to the body unless we read all the bytes before
 	// returning from the request
 	var buff bytes.Buffer
-	if _, err := io.Copy(&buff, treq.Body); err != nil {
+	if _, err := iopool.Copy(&buff, treq.Body); err != nil {
 		return err
 	}
 	treq.Body = &buff

transport/http/inbound_example_test.go

@@ -2,12 +2,12 @@ package http_test
 
 import (
 	"fmt"
-	"io"
 	"log"
 	nethttp "net/http"
 	"os"
 
 	"go.uber.org/yarpc"
+	"go.uber.org/yarpc/internal/iopool"
 	"go.uber.org/yarpc/transport/http"
 )
 
@@ -58,7 +58,7 @@ func ExampleMux() {
 	}
 	defer res.Body.Close()
 
-	if _, err := io.Copy(os.Stdout, res.Body); err != nil {
+	if _, err := iopool.Copy(os.Stdout, res.Body); err != nil {
 		log.Fatal(err)
 	}
 	// Output: hello from /health

transport/tchannel/channel_outbound.go

@@ -28,9 +28,10 @@ import (
 	"go.uber.org/yarpc/internal/encoding"
 	"go.uber.org/yarpc/internal/errors"
 	"go.uber.org/yarpc/internal/introspection"
+	"go.uber.org/yarpc/internal/iopool"
+	"go.uber.org/yarpc/internal/sync"
 
 	"github.com/uber/tchannel-go"
-	"go.uber.org/yarpc/internal/sync"
 )
 
 var (
@@ -207,7 +208,7 @@ func writeBody(body io.Reader, call *tchannel.OutboundCall) error {
 		return err
 	}
 
-	if _, err := io.Copy(w, body); err != nil {
+	if _, err := iopool.Copy(w, body); err != nil {
 		return err
 	}