Optimize and fix regression on recvAndDecompress method when utilizing compression

Signed-off-by: alanprot <[email protected]>

Author: alanprot

rpc_util.go

@@ -876,31 +876,27 @@ func decompress(compressor encoding.Compressor, d mem.BufferSlice, maxReceiveMes
 	if err != nil {
 		return nil, 0, err
 	}
-
-	// TODO: Can/should this still be preserved with the new BufferSlice API? Are
-	//  there any actual benefits to allocating a single large buffer instead of
-	//  multiple smaller ones?
-	//if sizer, ok := compressor.(interface {
-	//	DecompressedSize(compressedBytes []byte) int
-	//}); ok {
-	//	if size := sizer.DecompressedSize(d); size >= 0 {
-	//		if size > maxReceiveMessageSize {
-	//			return nil, size, nil
-	//		}
-	//		// size is used as an estimate to size the buffer, but we
-	//		// will read more data if available.
-	//		// +MinRead so ReadFrom will not reallocate if size is correct.
-	//		//
-	//		// TODO: If we ensure that the buffer size is the same as the DecompressedSize,
-	//		// we can also utilize the recv buffer pool here.
-	//		buf := bytes.NewBuffer(make([]byte, 0, size+bytes.MinRead))
-	//		bytesRead, err := buf.ReadFrom(io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
-	//		return buf.Bytes(), int(bytesRead), err
-	//	}
-	//}
+	bufferSize := 32 * 1024
+
+	if sizer, ok := compressor.(interface {
+		DecompressedSize(compressedBytes []byte) int
+	}); ok {
+		if size := sizer.DecompressedSize(d.Materialize()); size >= 0 {
+			if size > maxReceiveMessageSize {
+				return nil, size, nil
+			}
+			// size is used as an estimate to size the buffer, but we
+			// will read more data if available
+			// io.CopyBuffer will not reallocate if size is correct.
+			bufferSize = size
+		}
+	}
 
 	var out mem.BufferSlice
-	_, err = io.Copy(mem.NewWriter(&out, pool), io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1))
+	buff := pool.Get(bufferSize)
+	b := mem.NewBuffer(buff, pool)
+	_, err = io.CopyBuffer(mem.NewWriter(&out, pool), io.LimitReader(dcReader, int64(maxReceiveMessageSize)+1), *buff)
+	b.Free()
 	if err != nil {
 		out.Free()
 		return nil, 0, err

server_ext_test.go

@@ -20,6 +20,9 @@ package grpc_test
 
 import (
 	"context"
+	"fmt"
+	"google.golang.org/grpc/encoding"
+	"google.golang.org/grpc/encoding/gzip"
 	"io"
 	"runtime"
 	"sync"
@@ -330,3 +333,70 @@ func (s) TestServer_GracefulStopWaits(t *testing.T) {
 		t.Fatalf("Timed out waiting for second RPC to start on server.")
 	}
 }
+
+type nopCloserWriter struct{ io.Writer }
+
+func (n nopCloserWriter) Close() error { return nil }
+
+type noopCompressor struct {
+}
+
+func (i *noopCompressor) Compress(w io.Writer) (io.WriteCloser, error) {
+	return &nopCloserWriter{
+		Writer: w,
+	}, nil
+}
+
+func (i *noopCompressor) Decompress(r io.Reader) (io.Reader, error) {
+	return r, nil
+}
+
+func (i *noopCompressor) Name() string {
+	return "noop"
+}
+
+func BenchmarkRPCCompressor(b *testing.B) {
+	encoding.RegisterCompressor(&noopCompressor{})
+
+	for _, comp := range []string{gzip.Name, "noop"} {
+		for _, payloadSize := range []int{1024, 10 * 1024, 500 * 1024} {
+			b.Run(fmt.Sprintf("comp=%v,payloadSize=%v", comp, payloadSize), func(b *testing.B) {
+				ss := stubserver.StubServer{
+					UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+						return &testpb.SimpleResponse{
+							Username: "test",
+						}, nil
+					},
+				}
+
+				// Start one RPC to the server.
+				if err := ss.Start(nil, grpc.WithDefaultCallOptions(grpc.UseCompressor(comp))); err != nil {
+					b.Fatal("Error starting server:", err)
+				}
+				defer ss.Stop()
+				paylaod := make([]byte, payloadSize)
+				for i := 0; i < payloadSize; i++ {
+					paylaod[i] = byte(i)
+				}
+
+				for i := 0; i < b.N; i++ {
+					b.ReportAllocs()
+					ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+					defer cancel()
+					_, err := ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{
+						ResponseType: testpb.PayloadType_COMPRESSABLE,
+						ResponseSize: 1024,
+						Payload: &testpb.Payload{
+							Type: testpb.PayloadType_COMPRESSABLE,
+							Body: paylaod,
+						},
+					})
+
+					if err != nil {
+						b.Fatal("Error staring call:", err)
+					}
+				}
+			})
+		}
+	}
+}