Eliminate redundant receive-side copy in gRPC protobuf decode path

encoding/protobuf/marshal.go

@@ -54,6 +54,19 @@ func newCodec(anyResolver jsonpb.AnyResolver) *codec {
 }
 
 func unmarshal(encoding transport.Encoding, reader io.Reader, message proto.Message, c *codec) error {
+	// Fast path: if the reader directly exposes raw bytes, skip the redundant
+	// copy through bufferpool. This is the case for gRPC transports where the
+	// bytes have already been materialized from the wire.
+	if br, ok := reader.(interface{ Bytes() []byte }); ok {
+		body := br.Bytes()
+		if len(body) == 0 {
+			return nil
+		}
+		return unmarshalBytes(encoding, body, message, c)
+	}
+
+	// Slow path: drain the reader into a pooled buffer. Used by non-gRPC
+	// transports (HTTP, TChannel) that provide a plain io.Reader.
 	buf := bufferpool.Get()
 	defer bufferpool.Put(buf)
 	if _, err := buf.ReadFrom(reader); err != nil {

encoding/protobuf/marshal_test.go

@@ -21,10 +21,14 @@
 package protobuf
 
 import (
+	"bytes"
 	"strings"
 	"testing"
 
+	"github.com/gogo/protobuf/proto"
+	"github.com/gogo/protobuf/types"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/yarpcerrors"
 )
@@ -38,3 +42,71 @@ func TestUnhandledEncoding(t *testing.T) {
 	assert.NotNil(t, cleanup, "cleanup function should never be nil")
 	assert.NotPanics(t, func() { cleanup() }, "cleanup should be safe to call even on error")
 }
+
+// testBytesReader simulates a body from gRPC transport that exposes raw bytes
+// via a Bytes() method, allowing the unmarshal fast path to skip bufferpool copies.
+type testBytesReader struct {
+	data        []byte
+	reader      *bytes.Reader
+	bytesCalled bool
+	readCalled  bool
+}
+
+func newTestBytesReader(data []byte) *testBytesReader {
+	return &testBytesReader{data: data, reader: bytes.NewReader(data)}
+}
+
+func (r *testBytesReader) Read(p []byte) (int, error) {
+	r.readCalled = true
+	return r.reader.Read(p)
+}
+
+func (r *testBytesReader) Bytes() []byte {
+	r.bytesCalled = true
+	return r.data
+}
+
+// TestUnmarshalFastPath verifies that when the reader exposes a Bytes() method,
+// unmarshal uses the zero-copy path and never calls Read().
+func TestUnmarshalFastPath(t *testing.T) {
+	c := newCodec(nil)
+
+	t.Run("Bytes called and Read not called", func(t *testing.T) {
+		original := &types.StringValue{Value: "hello"}
+		data, err := proto.Marshal(original)
+		require.NoError(t, err)
+
+		reader := newTestBytesReader(data)
+		got := &types.StringValue{}
+
+		err = unmarshal(Encoding, reader, got, c)
+		assert.NoError(t, err)
+		assert.Equal(t, original.Value, got.Value, "Message should be deserialized correctly via fast path")
+		assert.True(t, reader.bytesCalled, "Bytes() should be called on fast path")
+		assert.False(t, reader.readCalled, "Read() should not be called on fast path")
+	})
+
+	t.Run("empty body returns nil", func(t *testing.T) {
+		reader := newTestBytesReader([]byte{})
+
+		err := unmarshal(Encoding, reader, nil, c)
+		assert.NoError(t, err, "Empty body on fast path should return nil")
+		assert.True(t, reader.bytesCalled, "Bytes() should still be called for empty body")
+		assert.False(t, reader.readCalled, "Read() should not be called for empty body")
+	})
+
+	t.Run("invalid encoding returns error", func(t *testing.T) {
+		reader := newTestBytesReader([]byte("data"))
+
+		err := unmarshal(transport.Encoding("unknown"), reader, nil, c)
+		assert.Equal(t, yarpcerrors.CodeInternal, yarpcerrors.FromError(err).Code(),
+			"Fast path should still return encoding error for unrecognized encoding")
+	})
+
+	t.Run("malformed protobuf returns unmarshal error", func(t *testing.T) {
+		reader := newTestBytesReader([]byte{0xff, 0xff, 0xff})
+
+		err := unmarshal(Encoding, reader, nil, c)
+		assert.Error(t, err, "Malformed protobuf should return an error on fast path")
+	})
+}

encoding/protobuf/v2/marshal.go

@@ -70,6 +70,19 @@ func newCodec(anyResolver AnyResolver) *codec {
 }
 
 func unmarshal(encoding transport.Encoding, reader io.Reader, message proto.Message, codec *codec) error {
+	// Fast path: if the reader directly exposes raw bytes, skip the redundant
+	// copy through bufferpool. This is the case for gRPC transports where the
+	// bytes have already been materialized from the wire.
+	if br, ok := reader.(interface{ Bytes() []byte }); ok {
+		body := br.Bytes()
+		if len(body) == 0 {
+			return nil
+		}
+		return unmarshalBytes(encoding, body, message, codec)
+	}
+
+	// Slow path: drain the reader into a pooled buffer. Used by non-gRPC
+	// transports (HTTP, TChannel) that provide a plain io.Reader.
 	buf := bufferpool.Get()
 	defer bufferpool.Put(buf)
 	if _, err := buf.ReadFrom(reader); err != nil {

encoding/protobuf/v2/marshal_test.go

@@ -21,12 +21,16 @@
 package v2
 
 import (
+	"bytes"
 	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/yarpc/api/transport"
 	"go.uber.org/yarpc/yarpcerrors"
+	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
 func TestUnhandledEncoding(t *testing.T) {
@@ -35,3 +39,71 @@ func TestUnhandledEncoding(t *testing.T) {
 	_, _, err := marshal(transport.Encoding("foo"), nil, newCodec(nil))
 	assert.Equal(t, yarpcerrors.CodeInternal, yarpcerrors.FromError(err).Code())
 }
+
+// testBytesReader simulates a body from gRPC transport that exposes raw bytes
+// via a Bytes() method, allowing the unmarshal fast path to skip bufferpool copies.
+type testBytesReader struct {
+	data        []byte
+	reader      *bytes.Reader
+	bytesCalled bool
+	readCalled  bool
+}
+
+func newTestBytesReader(data []byte) *testBytesReader {
+	return &testBytesReader{data: data, reader: bytes.NewReader(data)}
+}
+
+func (r *testBytesReader) Read(p []byte) (int, error) {
+	r.readCalled = true
+	return r.reader.Read(p)
+}
+
+func (r *testBytesReader) Bytes() []byte {
+	r.bytesCalled = true
+	return r.data
+}
+
+// TestUnmarshalFastPath verifies that when the reader exposes a Bytes() method,
+// unmarshal uses the zero-copy path and never calls Read().
+func TestUnmarshalFastPath(t *testing.T) {
+	c := newCodec(nil)
+
+	t.Run("Bytes called and Read not called", func(t *testing.T) {
+		original := &wrapperspb.StringValue{Value: "hello"}
+		data, err := proto.Marshal(original)
+		require.NoError(t, err)
+
+		reader := newTestBytesReader(data)
+		got := &wrapperspb.StringValue{}
+
+		err = unmarshal(Encoding, reader, got, c)
+		assert.NoError(t, err)
+		assert.Equal(t, original.Value, got.Value, "Message should be deserialized correctly via fast path")
+		assert.True(t, reader.bytesCalled, "Bytes() should be called on fast path")
+		assert.False(t, reader.readCalled, "Read() should not be called on fast path")
+	})
+
+	t.Run("empty body returns nil", func(t *testing.T) {
+		reader := newTestBytesReader([]byte{})
+
+		err := unmarshal(Encoding, reader, nil, c)
+		assert.NoError(t, err, "Empty body on fast path should return nil")
+		assert.True(t, reader.bytesCalled, "Bytes() should still be called for empty body")
+		assert.False(t, reader.readCalled, "Read() should not be called for empty body")
+	})
+
+	t.Run("invalid encoding returns error", func(t *testing.T) {
+		reader := newTestBytesReader([]byte("data"))
+
+		err := unmarshal(transport.Encoding("unknown"), reader, nil, c)
+		assert.Equal(t, yarpcerrors.CodeInternal, yarpcerrors.FromError(err).Code(),
+			"Fast path should still return encoding error for unrecognized encoding")
+	})
+
+	t.Run("malformed protobuf returns unmarshal error", func(t *testing.T) {
+		reader := newTestBytesReader([]byte{0xff, 0xff, 0xff})
+
+		err := unmarshal(Encoding, reader, &wrapperspb.StringValue{}, c)
+		assert.Error(t, err, "Malformed protobuf should return an error on fast path")
+	})
+}

transport/grpc/handler.go

@@ -21,7 +21,6 @@
 package grpc
 
 import (
-	"bytes"
 	"strings"
 	"time"
 
@@ -194,7 +193,7 @@ func (h *handler) handleUnary(
 		return err
 	}
 
-	transportRequest.Body = bytes.NewReader(requestData)
+	transportRequest.Body = newBytesBody(requestData)
 	transportRequest.BodySize = len(requestData)
 
 	responseWriter := newResponseWriter()

transport/grpc/outbound.go

@@ -21,7 +21,6 @@
 package grpc
 
 import (
-	"bytes"
 	"context"
 	"io/ioutil"
 	"strings"
@@ -139,7 +138,7 @@ func (o *Outbound) DirectCall(ctx context.Context, request *transport.Request) (
 		return nil, err
 	}
 	return &transport.Response{
-		Body:                 ioutil.NopCloser(bytes.NewReader(responseBody)),
+		Body:                 newBytesBody(responseBody),
 		BodySize:             len(responseBody),
 		Headers:              responseHeaders,
 		ApplicationError:     metadataToIsApplicationError(responseMD),

transport/grpc/stream.go

@@ -93,19 +93,26 @@ func (ss *serverStream) ReceiveMessage(_ context.Context) (*transport.StreamMess
 		return nil, toYARPCStreamError(err)
 	}
 	return &transport.StreamMessage{
-		Body:     readCloser{bytes.NewReader(msg)},
+		Body:     newBytesBody(msg),
 		BodySize: len(msg),
 	}, nil
 }
 
-type readCloser struct {
-	*bytes.Reader
+// bytesBody wraps a []byte and exposes direct access to
+// the underlying bytes via the Bytes method.
+type bytesBody struct {
+	data   []byte
+	reader *bytes.Reader
 }
 
-func (r readCloser) Close() error {
-	return nil
+func newBytesBody(data []byte) *bytesBody {
+	return &bytesBody{data: data, reader: bytes.NewReader(data)}
 }
 
+func (b *bytesBody) Read(p []byte) (int, error) { return b.reader.Read(p) }
+func (b *bytesBody) Bytes() []byte              { return b.data }
+func (b *bytesBody) Close() error               { return nil }
+
 func (ss *serverStream) SendHeaders(headers transport.Headers) error {
 	md := make(metadata.MD, headers.Len())
 	for k, v := range headers.Items() {
@@ -173,7 +180,7 @@ func (cs *clientStream) ReceiveMessage(context.Context) (*transport.StreamMessag
 	if err := cs.stream.RecvMsg(&msg); err != nil {
 		return nil, toYARPCStreamError(cs.closeWithErr(err))
 	}
-	return &transport.StreamMessage{Body: ioutil.NopCloser(bytes.NewReader(msg))}, nil
+	return &transport.StreamMessage{Body: newBytesBody(msg)}, nil
 }
 
 func (cs *clientStream) Close(context.Context) error {

transport/grpc/stream_test.go

@@ -144,6 +144,62 @@ func TestServerStreamSendMessage(t *testing.T) {
 	})
 }
 
+// TestStreamReceiveMessage verifies bytesBody wiring: both server and client
+// receive paths wrap payloads in bytesBody (exposes Bytes() for zero-copy).
+// The actual fast-path dispatch is tested in TestUnmarshalFastPath.
+func TestStreamReceiveMessage(t *testing.T) {
+	t.Run("server stream wraps body in bytesBody", func(t *testing.T) {
+		payload := []byte("server receive test")
+
+		mockStream := &mockGRPCServerStream{
+			recvMsgFunc: func(m any) error {
+				ptr := m.(*[]byte)
+				*ptr = payload
+				return nil
+			},
+		}
+
+		ss := &serverStream{
+			ctx:    context.Background(),
+			stream: mockStream,
+			req:    &transport.StreamRequest{Meta: &transport.RequestMeta{}},
+		}
+
+		msg, err := ss.ReceiveMessage(context.Background())
+		assert.NoError(t, err)
+
+		br, ok := msg.Body.(interface{ Bytes() []byte })
+		assert.True(t, ok, "Body should implement Bytes() []byte")
+		assert.Equal(t, payload, br.Bytes())
+		assert.Equal(t, len(payload), msg.BodySize)
+	})
+
+	t.Run("client stream wraps body in bytesBody", func(t *testing.T) {
+		payload := []byte("client receive test")
+
+		mockStream := &mockGRPCClientStream{
+			recvMsgFunc: func(m any) error {
+				ptr := m.(*[]byte)
+				*ptr = payload
+				return nil
+			},
+		}
+
+		cs := &clientStream{
+			ctx:    context.Background(),
+			stream: mockStream,
+			req:    &transport.StreamRequest{Meta: &transport.RequestMeta{}},
+		}
+
+		msg, err := cs.ReceiveMessage(context.Background())
+		assert.NoError(t, err)
+
+		br, ok := msg.Body.(interface{ Bytes() []byte })
+		assert.True(t, ok, "Body should implement Bytes() []byte")
+		assert.Equal(t, payload, br.Bytes())
+	})
+}
+
 func TestClientStreamSendMessage(t *testing.T) {
 	t.Run("passes mem.Buffer to gRPC without conversion", func(t *testing.T) {
 		data := []byte("client test")
@@ -212,6 +268,7 @@ func TestClientStreamSendMessage(t *testing.T) {
 type mockGRPCServerStream struct {
 	grpc.ServerStream
 	sendMsgFunc func(m any) error
+	recvMsgFunc func(m any) error
 }
 
 func (m *mockGRPCServerStream) SendMsg(msg any) error {
@@ -221,6 +278,13 @@ func (m *mockGRPCServerStream) SendMsg(msg any) error {
 	return nil
 }
 
+func (m *mockGRPCServerStream) RecvMsg(msg any) error {
+	if m.recvMsgFunc != nil {
+		return m.recvMsgFunc(msg)
+	}
+	return nil
+}
+
 type mockGRPCClientStream struct {
 	grpc.ClientStream
 	sendMsgFunc func(m any) error