grpc: support per-message compression and enforce embedding in ServerTransportStream implementations

server.go

@@ -1423,6 +1423,7 @@ func (s *Server) processUnaryRPC(ctx context.Context, stream *transport.ServerSt
 		return nil
 	}
 	ctx = NewContextWithServerTransportStream(ctx, stream)
+	ctx = context.WithValue(ctx, compressKey{}, &compressOptions{})
 	reply, appErr := md.Handler(info.serviceImpl, ctx, df, s.opts.unaryInt)
 	if appErr != nil {
 		appStatus, ok := status.FromError(appErr)
@@ -1470,7 +1471,11 @@ func (s *Server) processUnaryRPC(ctx context.Context, stream *transport.ServerSt
 	if stream.SendCompress() != sendCompressorName {
 		comp = encoding.GetCompressor(stream.SendCompress())
 	}
-	if err := s.sendResponse(ctx, stream, reply, cp, opts, comp); err != nil {
+	compV0, compV1 := cp, comp
+	if opts, ok := ctx.Value(compressKey{}).(*compressOptions); ok && opts.DoNotCompress {
+		compV0, compV1 = nil, nil
+	}
+	if err := s.sendResponse(ctx, stream, reply, compV0, opts, compV1); err != nil {
 		if err == io.EOF {
 			// The entire stream is done (for unary RPC only).
 			return err
@@ -1585,6 +1590,7 @@ func (s *Server) processStreamingRPC(ctx context.Context, stream *transport.Serv
 		sh.HandleRPC(ctx, statsBegin)
 	}
 	ctx = NewContextWithServerTransportStream(ctx, stream)
+	ctx = context.WithValue(ctx, compressKey{}, &compressOptions{})
 	ss := &serverStream{
 		ctx:                   ctx,
 		s:                     stream,

stream.go

@@ -21,6 +21,7 @@ package grpc
 import (
 	"context"
 	"errors"
+	"fmt"
 	"io"
 	"math"
 	rand "math/rand/v2"
@@ -51,6 +52,33 @@ import (
 
 var metadataFromOutgoingContextRaw = internal.FromOutgoingContextRaw.(func(context.Context) (metadata.MD, [][]string, bool))
 
+// SetMessageCompression enables or disables per-message compression on a stream
+// if a compressor is specified for the stream (e.g. using UseCompressor or
+// SetSendCompressor) and if the encoding type is supported by the receiver.
+// By default, message compression is enabled, but is a no-op if compression
+// is not enabled on the stream.
+//
+// This method must not be called concurrently with SendMsg.
+//
+// # Experimental
+//
+// Notice: This API is EXPERIMENTAL and may be changed or removed in a
+// later release.
+func SetMessageCompression(ctx context.Context, enable bool) error {
+	opts, ok := ctx.Value(compressKey{}).(*compressOptions)
+	if !ok || opts == nil {
+		return fmt.Errorf("grpc: SetMessageCompression called on an uninitialized or non-stream context")
+	}
+	opts.DoNotCompress = !enable
+	return nil
+}
+
+type compressOptions struct {
+	DoNotCompress bool
+}
+
+type compressKey struct{}
+
 // StreamHandler defines the handler called by gRPC server to complete the
 // execution of a streaming RPC. srv is the service implementation on which the
 // RPC was invoked.
@@ -299,6 +327,7 @@ func newClientStreamWithParams(ctx context.Context, desc *StreamDesc, cc *Client
 	} else {
 		ctx, cancel = context.WithCancel(ctx)
 	}
+	ctx = context.WithValue(ctx, compressKey{}, &compressOptions{})
 	defer func() {
 		if err != nil {
 			cancel()
@@ -954,7 +983,11 @@ func (cs *clientStream) SendMsg(m any) (err error) {
 	}
 
 	// load hdr, payload, data
-	hdr, data, payload, pf, err := prepareMsg(m, cs.codec, cs.compressorV0, cs.compressorV1, cs.cc.dopts.copts.BufferPool)
+	compV0, compV1 := cs.compressorV0, cs.compressorV1
+	if opts, ok := cs.ctx.Value(compressKey{}).(*compressOptions); ok && opts.DoNotCompress {
+		compV0, compV1 = nil, nil
+	}
+	hdr, data, payload, pf, err := prepareMsg(m, cs.codec, compV0, compV1, cs.cc.dopts.copts.BufferPool)
 	if err != nil {
 		return err
 	}
@@ -1461,7 +1494,11 @@ func (as *addrConnStream) SendMsg(m any) (err error) {
 	}
 
 	// load hdr, payload, data
-	hdr, data, payload, pf, err := prepareMsg(m, as.codec, as.sendCompressorV0, as.sendCompressorV1, as.ac.dopts.copts.BufferPool)
+	compV0, compV1 := as.sendCompressorV0, as.sendCompressorV1
+	if opts, ok := as.ctx.Value(compressKey{}).(*compressOptions); ok && opts.DoNotCompress {
+		compV0, compV1 = nil, nil
+	}
+	hdr, data, payload, pf, err := prepareMsg(m, as.codec, compV0, compV1, as.ac.dopts.copts.BufferPool)
 	if err != nil {
 		return err
 	}
@@ -1741,7 +1778,11 @@ func (ss *serverStream) SendMsg(m any) (err error) {
 	}
 
 	// load hdr, payload, data
-	hdr, data, payload, pf, err := prepareMsg(m, ss.codec, ss.compressorV0, ss.compressorV1, ss.p.bufferPool)
+	compV0, compV1 := ss.compressorV0, ss.compressorV1
+	if opts, ok := ss.ctx.Value(compressKey{}).(*compressOptions); ok && opts.DoNotCompress {
+		compV0, compV1 = nil, nil
+	}
+	hdr, data, payload, pf, err := prepareMsg(m, ss.codec, compV0, compV1, ss.p.bufferPool)
 	if err != nil {
 		return err
 	}

test/compressor_test.go

@@ -25,6 +25,7 @@ import (
 	"io"
 	"reflect"
 	"strings"
+	"sync/atomic"
 	"testing"
 
 	"google.golang.org/grpc"
@@ -34,6 +35,7 @@ import (
 	"google.golang.org/grpc/internal/grpcutil"
 	"google.golang.org/grpc/internal/stubserver"
 	"google.golang.org/grpc/metadata"
+	"google.golang.org/grpc/stats"
 	"google.golang.org/grpc/status"
 
 	testgrpc "google.golang.org/grpc/interop/grpc_testing"
@@ -692,3 +694,137 @@ func (s) TestGzipBadChecksum(t *testing.T) {
 		t.Errorf("ss.Client.UnaryCall(_) = _, %v\n\twant: _, status(codes.Internal, contains %q)", err, gzip.ErrChecksum)
 	}
 }
+
+type statsHandler struct {
+	stats.Handler
+	compress   atomic.Int32
+	decompress atomic.Int32
+}
+
+func (h *statsHandler) TagRPC(ctx context.Context, _ *stats.RPCTagInfo) context.Context   { return ctx }
+func (h *statsHandler) TagConn(ctx context.Context, _ *stats.ConnTagInfo) context.Context { return ctx }
+func (h *statsHandler) HandleConn(context.Context, stats.ConnStats)                       {}
+func (h *statsHandler) HandleRPC(_ context.Context, s stats.RPCStats) {
+	switch st := s.(type) {
+	case *stats.OutPayload:
+		if st.CompressedLength < st.Length {
+			h.compress.Add(1)
+		}
+	case *stats.InPayload:
+		if st.CompressedLength < st.Length {
+			h.decompress.Add(1)
+		}
+	}
+}
+
+func (s) TestMessageCompression_Stream(t *testing.T) {
+
+	sh := &statsHandler{}
+	ss := &stubserver.StubServer{
+		FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
+			if _, err := stream.Recv(); err != nil {
+				return err
+			}
+			if err := stream.Send(&testpb.StreamingOutputCallResponse{
+				Payload: &testpb.Payload{Body: make([]byte, 1000)},
+			}); err != nil {
+				return err
+			}
+			if _, err := stream.Recv(); err != nil {
+				return err
+			}
+			if err := grpc.SetMessageCompression(stream.Context(), false); err != nil {
+				return err
+			}
+			if err := stream.Send(&testpb.StreamingOutputCallResponse{
+				Payload: &testpb.Payload{Body: make([]byte, 1000)},
+			}); err != nil {
+				return err
+			}
+			if _, err := stream.Recv(); err != nil {
+				return err
+			}
+			if err := grpc.SetMessageCompression(stream.Context(), true); err != nil {
+				return err
+			}
+			return stream.Send(&testpb.StreamingOutputCallResponse{
+				Payload: &testpb.Payload{Body: make([]byte, 1000)},
+			})
+		},
+	}
+
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting server: %v", err)
+	}
+	defer ss.Stop()
+	if err := ss.StartClient(grpc.WithStatsHandler(sh)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	stream, err := ss.Client.FullDuplexCall(ctx, grpc.UseCompressor("gzip"))
+	if err != nil {
+		t.Fatalf("FullDuplexCall failed: %v", err)
+	}
+
+	// 1. Send first compressed message
+	stream.Send(&testpb.StreamingOutputCallRequest{Payload: &testpb.Payload{Body: make([]byte, 1000)}})
+	stream.Recv()
+	if sh.compress.Load() != 1 || sh.decompress.Load() != 1 {
+		t.Fatalf("After Call 1, expected 1, 1. got %d, %d", sh.compress.Load(), sh.decompress.Load())
+	}
+
+	// 2. Disable message compression and send second message
+	grpc.SetMessageCompression(stream.Context(), false)
+	stream.Send(&testpb.StreamingOutputCallRequest{Payload: &testpb.Payload{Body: make([]byte, 1000)}})
+	stream.Recv()
+	if sh.compress.Load() != 1 || sh.decompress.Load() != 1 {
+		t.Fatalf("After Call 2, expected 1, 1. got %d, %d", sh.compress.Load(), sh.decompress.Load())
+	}
+
+	// 3. Enable message compression and send third message
+	grpc.SetMessageCompression(stream.Context(), true)
+	stream.Send(&testpb.StreamingOutputCallRequest{Payload: &testpb.Payload{Body: make([]byte, 1000)}})
+	stream.Recv()
+	if sh.compress.Load() != 2 || sh.decompress.Load() != 2 {
+		t.Fatalf("After Call 3, expected 2, 2. got %d, %d", sh.compress.Load(), sh.decompress.Load())
+	}
+}
+
+func (s) TestMessageCompression_Unary(t *testing.T) {
+	sh := &statsHandler{}
+	ss := &stubserver.StubServer{
+		UnaryCallF: func(ctx context.Context, in *testpb.SimpleRequest) (*testpb.SimpleResponse, error) {
+			grpc.SetSendCompressor(ctx, "gzip")
+			if in.ResponseSize == 0 {
+				grpc.SetMessageCompression(ctx, false)
+			}
+			return &testpb.SimpleResponse{Payload: &testpb.Payload{Body: make([]byte, 10000)}}, nil
+		},
+	}
+
+	if err := ss.Start(nil); err != nil {
+		t.Fatalf("Error starting server: %v", err)
+	}
+	defer ss.Stop()
+	if err := ss.StartClient(grpc.WithStatsHandler(sh)); err != nil {
+		t.Fatalf("Error starting client: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
+	defer cancel()
+
+	// Call 1: Compression ON
+	ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{ResponseSize: 1, Payload: &testpb.Payload{Body: make([]byte, 1000)}}, grpc.UseCompressor("gzip"))
+	if sh.compress.Load() != 1 || sh.decompress.Load() != 1 {
+		t.Fatalf("Expected 1/1, got %d/%d", sh.compress.Load(), sh.decompress.Load())
+	}
+
+	// Call 2: Compression OFF (for response, but request is still compressed by UseCompressor)
+	ss.Client.UnaryCall(ctx, &testpb.SimpleRequest{ResponseSize: 0, Payload: &testpb.Payload{Body: make([]byte, 1000)}}, grpc.UseCompressor("gzip"))
+	if sh.compress.Load() != 2 || sh.decompress.Load() != 1 {
+		t.Fatalf("Expected 2/1, got %d/%d", sh.compress.Load(), sh.decompress.Load())
+	}
+}