rpc_util: added a byte slice return to recvBufferPool

experimental/shared_buffer_pool_test.go

@@ -21,6 +21,7 @@ package experimental_test
 import (
 	"bytes"
 	"context"
+	"fmt"
 	"io"
 	"testing"
 	"time"
@@ -100,3 +101,86 @@ func (s) TestRecvBufferPool(t *testing.T) {
 		t.Errorf("Got replies %q; want %q", got, want)
 	}
 }
+
+func newPayload(t testpb.PayloadType, size int32) (*testpb.Payload, error) {
+	if size < 0 {
+		return nil, fmt.Errorf("requested a response with invalid length %d", size)
+	}
+	body := make([]byte, size)
+	switch t {
+	case testpb.PayloadType_COMPRESSABLE:
+	default:
+		return nil, fmt.Errorf("unsupported payload type: %d", t)
+	}
+	return &testpb.Payload{
+		Type: t,
+		Body: body,
+	}, nil
+}
+
+func (s) TestRecvBufferPoolUnary(t *testing.T) {
+	const largeSize = 1024
+	const bufSize = 1030
+
+	ss := &stubserver.StubServer{
+		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			payload, err := newPayload(testpb.PayloadType_COMPRESSABLE, largeSize)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			return &testpb.SimpleResponse{
+				Payload: payload,
+			}, nil
+		},
+	}
+
+	pool := &checkBufferPool{}
+
+	if err := ss.Start(
+		[]grpc.ServerOption{grpc.RecvBufferPool(pool)},
+		grpc.WithRecvBufferPool(pool),
+	); err != nil {
+		t.Fatalf("Error starting endpoint server: %v", err)
+	}
+	defer ss.Stop()
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	const reqCount = 10
+	for i := 0; i < reqCount; i++ {
+		payload, err := newPayload(testpb.PayloadType_COMPRESSABLE, largeSize)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{Payload: payload})
+		if err != nil {
+			t.Fatalf("ss.Client.UnaryCall failed: %f", err)
+		}
+	}
+
+	const bufferCount = reqCount * 2 // req + resp
+	if len(pool.puts) != bufferCount {
+		t.Fatalf("Expected 10 buffers to be returned to the pool, got %d", len(pool.puts))
+	}
+
+	for _, bs := range pool.puts {
+		if len(bs) != bufSize {
+			t.Fatalf("Expected buffer size %d, got %d", bufSize, len(bs))
+		}
+	}
+}
+
+type checkBufferPool struct {
+	puts [][]byte
+}
+
+func (p *checkBufferPool) Get(size int) []byte {
+	return make([]byte, size)
+}
+
+func (p *checkBufferPool) Put(bs *[]byte) {
+	p.puts = append(p.puts, *bs)
+}

server.go

@@ -1330,6 +1330,9 @@ func (s *Server) processUnaryRPC(ctx context.Context, t transport.ServerTranspor
 		t.IncrMsgRecv()
 	}
 	df := func(v any) error {
+		if len(shs) == 0 && len(binlogs) == 0 {
+			defer s.opts.recvBufferPool.Put(&d)
+		}
 		if err := s.getCodec(stream.ContentSubtype()).Unmarshal(d, v); err != nil {
 			return status.Errorf(codes.Internal, "grpc: error unmarshalling request: %v", err)
 		}