Eliminate redundant receive-side copy in gRPC protobuf decode path

Every protobuf message received over gRPC was copied twice: once by
the codec (Materialize, necessary) and again by the encoding layer
(bufferpool.ReadFrom, redundant). The second copy exists because the
transport wraps the already-materialized []byte in bytes.NewReader,
and the encoding layer has no way to extract the raw bytes.

Add a bytesBody type to the transport layer that implements
io.ReadCloser and exposes the underlying []byte via a Bytes() method.
The encoding layer's unmarshal function type-asserts for this method
and skips the bufferpool copy when available. Non-gRPC transports
(HTTP, TChannel) fall through to the existing ReadFrom path.

Changes:
- transport/grpc: Add bytesBody with Bytes() method for zero-copy access
- transport/grpc: Update all four receive paths (server stream, client
  stream, unary inbound, unary outbound) to use bytesBody
- encoding/protobuf: Add Bytes() fast-path in unmarshal (gogo and v2)

Safety:
- bytesBody is unexported and implements io.ReadCloser, the interface
  expected by transport.Request.Body and transport.StreamMessage.Body.
  It is a drop-in replacement for bytes.NewReader / ioutil.NopCloser.
- The Bytes() fast-path is opt-in via interface type assertion. Any
  reader that does not implement Bytes() (HTTP, TChannel) falls through
  to the existing bufferpool.ReadFrom slow path with zero behavior change.
- The fast-path calls unmarshalBytes, the same function the slow path
  calls after draining the reader — deserialization logic is identical.
- No public API changes; bytesBody is internal to transport/grpc.

Benchmark added in #2515

Benchmark (BenchmarkUnmarshalBytesReader, n=10, AMD EPYC 9B45, Go 1.26.1):

  Library  CPU Δ     Heap Δ   GC Cycles Δ   GC Pause Δ
  gogo     -16.1%    +1.1%    -3.6%          -9.3%
  v2       -15.1%    +0.9%    -2.8%          -7.6%
  average  -15.6%    +1.0%    -3.2%          -8.5%

  Heap increase at small payloads (<1KB) is due to bytesBody struct
  being marginally larger than bytes.Reader; at >=10KB the eliminated
  bufferpool copy dominates and heap decreases.

RELEASE NOTES: N/A (internal optimization, no API changes)

Made-with: Cursor
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: kshitijsuri90

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"
 
@@ -193,7 +192,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,7 +93,7 @@ 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
 }
@@ -106,6 +106,21 @@ func (r readCloser) Close() error {
 	return nil
 }
 
+// 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 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 +188,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