encoding: add support for compressor options in Compress methods

Author: Dostonlv

encoding/compressor_test.go

@@ -45,9 +45,9 @@ type wrapCompressor struct {
 	compressInvokes int32
 }
 
-func (wc *wrapCompressor) Compress(w io.Writer) (io.WriteCloser, error) {
+func (wc *wrapCompressor) Compress(w io.Writer, opts ...any) (io.WriteCloser, error) {
 	atomic.AddInt32(&wc.compressInvokes, 1)
-	return wc.Compressor.Compress(w)
+	return wc.Compressor.Compress(w, opts...)
 }
 
 func setupGzipWrapCompressor(t *testing.T) *wrapCompressor {
@@ -186,7 +186,7 @@ type fakeCompressor struct {
 	decompressedMessageSize int
 }
 
-func (f *fakeCompressor) Compress(w io.Writer) (io.WriteCloser, error) {
+func (f *fakeCompressor) Compress(w io.Writer, _ ...any) (io.WriteCloser, error) {
 	return nopWriteCloser{w}, nil
 }
 

encoding/encoding.go

@@ -61,8 +61,10 @@ func init() {
 type Compressor interface {
 	// Compress writes the data written to wc to w after compressing it.  If an
 	// error occurs while initializing the compressor, that error is returned
-	// instead.
-	Compress(w io.Writer) (io.WriteCloser, error)
+	// instead.  opts passes caller-provided context to the compressor (e.g.
+	// dictionary IDs for trained compression formats).  Unknown options must
+	// be silently ignored.
+	Compress(w io.Writer, opts ...any) (io.WriteCloser, error)
 	// Decompress reads data from r, decompresses it, and provides the
 	// uncompressed data via the returned io.Reader.  If an error occurs while
 	// initializing the decompressor, that error is returned instead.

encoding/gzip/gzip.go

@@ -70,7 +70,7 @@ func SetLevel(level int) error {
 	return nil
 }
 
-func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
+func (c *compressor) Compress(w io.Writer, _ ...any) (io.WriteCloser, error) {
 	z := c.poolCompressor.Get().(*writer)
 	z.Writer.Reset(w)
 	return z, nil

internal/transport/server_stream.go

@@ -50,6 +50,8 @@ type ServerStream struct {
 	headerSent atomic.Bool // atomically set when the headers are sent out.
 
 	headerWireLength int
+
+	sendCompressOptions []any
 }
 
 // Read reads an n byte message from the input stream.
@@ -108,6 +110,17 @@ func (s *ServerStream) ContentSubtype() string {
 	return s.contentSubtype
 }
 
+// SetSendCompressOptions sets the compressor options to be passed to the
+// compressor's Compress method.
+func (s *ServerStream) SetSendCompressOptions(opts []any) {
+	s.sendCompressOptions = opts
+}
+
+// SendCompressOptions returns the compressor options set for the stream.
+func (s *ServerStream) SendCompressOptions() []any {
+	return s.sendCompressOptions
+}
+
 // SetSendCompress sets the compression algorithm to the stream.
 func (s *ServerStream) SetSendCompress(name string) error {
 	if s.isHeaderSent() || s.getState() == streamDone {

rpc_util.go

@@ -168,6 +168,7 @@ type callInfo struct {
 	onFinish                    []func(err error)
 	authority                   string
 	acceptedResponseCompressors []string
+	compressorOptions           []any
 }
 
 func acceptedCompressorAllows(allowed []string, name string) bool {
@@ -490,14 +491,16 @@ func (o PerRPCCredsCallOption) after(*callInfo, *csAttempt) {}
 
 // UseCompressor returns a CallOption which sets the compressor used when
 // sending the request.  If WithCompressor is also set, UseCompressor has
-// higher priority.
+// higher priority.  The optional compressorOptions are forwarded to the
+// compressor's Compress method, allowing callers to pass additional context
+// such as dictionary IDs for trained compression formats.
 //
 // # Experimental
 //
 // Notice: This API is EXPERIMENTAL and may be changed or removed in a
 // later release.
-func UseCompressor(name string) CallOption {
-	return CompressorCallOption{CompressorType: name}
+func UseCompressor(name string, compressorOptions ...any) CallOption {
+	return CompressorCallOption{CompressorType: name, CompressorOptions: compressorOptions}
 }
 
 // CompressorCallOption is a CallOption that indicates the compressor to use.
@@ -507,11 +510,13 @@ func UseCompressor(name string) CallOption {
 // Notice: This type is EXPERIMENTAL and may be changed or removed in a
 // later release.
 type CompressorCallOption struct {
-	CompressorType string
+	CompressorType    string
+	CompressorOptions []any
 }
 
 func (o CompressorCallOption) before(c *callInfo) error {
 	c.compressorName = o.CompressorType
+	c.compressorOptions = o.CompressorOptions
 	return nil
 }
 func (o CompressorCallOption) after(*callInfo, *csAttempt) {}
@@ -817,7 +822,7 @@ func encode(c baseCodec, msg any) (mem.BufferSlice, error) {
 // indicating no compression was done.
 //
 // TODO(dfawley): eliminate cp parameter by wrapping Compressor in an encoding.Compressor.
-func compress(in mem.BufferSlice, cp Compressor, compressor encoding.Compressor, pool mem.BufferPool) (mem.BufferSlice, payloadFormat, error) {
+func compress(in mem.BufferSlice, cp Compressor, compressor encoding.Compressor, pool mem.BufferPool, compressorOptions ...any) (mem.BufferSlice, payloadFormat, error) {
 	if (compressor == nil && cp == nil) || in.Len() == 0 {
 		return nil, compressionNone, nil
 	}
@@ -828,7 +833,7 @@ func compress(in mem.BufferSlice, cp Compressor, compressor encoding.Compressor,
 		return status.Errorf(codes.Internal, "grpc: error while compressing: %v", err.Error())
 	}
 	if compressor != nil {
-		z, err := compressor.Compress(w)
+		z, err := compressor.Compress(w, compressorOptions...)
 		if err != nil {
 			return nil, 0, wrapErr(err)
 		}

rpc_util_test.go

@@ -52,7 +52,7 @@ type testCompressorForRegistry struct {
 	name string
 }
 
-func (c *testCompressorForRegistry) Compress(w io.Writer) (io.WriteCloser, error) {
+func (c *testCompressorForRegistry) Compress(w io.Writer, _ ...any) (io.WriteCloser, error) {
 	return &testWriteCloser{w}, nil
 }
 
@@ -541,7 +541,7 @@ type mockCompressor struct {
 	ch chan<- struct{}
 }
 
-func (m *mockCompressor) Compress(io.Writer) (io.WriteCloser, error) {
+func (m *mockCompressor) Compress(io.Writer, ...any) (io.WriteCloser, error) {
 	panic("unimplemented")
 }
 

server.go

@@ -2146,11 +2146,15 @@ func SendHeader(ctx context.Context, md metadata.MD) error {
 // It is not safe to call SetSendCompressor concurrently with SendHeader and
 // SendMsg.
 //
+// The optional compressorOptions are forwarded to the compressor's Compress
+// method on each SendMsg call, allowing callers to pass additional context
+// such as dictionary IDs for trained compression formats.
+//
 // # Experimental
 //
 // Notice: This function is EXPERIMENTAL and may be changed or removed in a
 // later release.
-func SetSendCompressor(ctx context.Context, name string) error {
+func SetSendCompressor(ctx context.Context, name string, compressorOptions ...any) error {
 	stream, ok := ServerTransportStreamFromContext(ctx).(*transport.ServerStream)
 	if !ok || stream == nil {
 		return fmt.Errorf("failed to fetch the stream from the given context")
@@ -2160,6 +2164,7 @@ func SetSendCompressor(ctx context.Context, name string) error {
 		return fmt.Errorf("unable to set send compressor: %w", err)
 	}
 
+	stream.SetSendCompressOptions(compressorOptions)
 	return stream.SetSendCompress(name)
 }
 

stream.go

@@ -633,6 +633,8 @@ type clientStream struct {
 	// nameResolutionDelay indicates if there was a delay in the name resolution.
 	// This field is only valid on client side, it's always false on server side.
 	nameResolutionDelay bool
+
+	compressorOptions []any
 }
 
 type replayOp struct {
@@ -964,7 +966,7 @@ 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)
+	hdr, data, payload, pf, err := prepareMsg(m, cs.codec, cs.compressorV0, cs.compressorV1, cs.cc.dopts.copts.BufferPool, cs.compressorOptions...)
 	if err != nil {
 		return err
 	}
@@ -1471,7 +1473,7 @@ 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)
+	hdr, data, payload, pf, err := prepareMsg(m, as.codec, as.sendCompressorV0, as.sendCompressorV1, as.ac.dopts.copts.BufferPool, as.callInfo.compressorOptions...)
 	if err != nil {
 		return err
 	}
@@ -1669,7 +1671,8 @@ type serverStream struct {
 	// synchronized.
 	serverHeaderBinlogged bool
 
-	mu sync.Mutex // protects trInfo.tr after the service handler runs.
+	mu                    sync.Mutex // protects trInfo.tr after the service handler runs.
+	sendCompressorOptions []any
 }
 
 func (ss *serverStream) Context() context.Context {
@@ -1748,10 +1751,11 @@ func (ss *serverStream) SendMsg(m any) (err error) {
 	if sendCompressorsName := ss.s.SendCompress(); sendCompressorsName != ss.sendCompressorName {
 		ss.compressorV1 = encoding.GetCompressor(sendCompressorsName)
 		ss.sendCompressorName = sendCompressorsName
+		ss.sendCompressorOptions = ss.s.SendCompressOptions()
 	}
 
 	// load hdr, payload, data
-	hdr, data, payload, pf, err := prepareMsg(m, ss.codec, ss.compressorV0, ss.compressorV1, ss.p.bufferPool)
+	hdr, data, payload, pf, err := prepareMsg(m, ss.codec, ss.compressorV0, ss.compressorV1, ss.p.bufferPool, ss.sendCompressorOptions...)
 	if err != nil {
 		return err
 	}
@@ -1893,7 +1897,7 @@ func MethodFromServerStream(stream ServerStream) (string, bool) {
 // compression was made and therefore whether the payload needs to be freed in
 // addition to the returned data. Freeing the payload if the returned boolean is
 // false can lead to undefined behavior.
-func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor, pool mem.BufferPool) (hdr []byte, data, payload mem.BufferSlice, pf payloadFormat, err error) {
+func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor, pool mem.BufferPool, compressorOptions ...any) (hdr []byte, data, payload mem.BufferSlice, pf payloadFormat, err error) {
 	if preparedMsg, ok := m.(*PreparedMsg); ok {
 		return preparedMsg.hdr, preparedMsg.encodedData, preparedMsg.payload, preparedMsg.pf, nil
 	}
@@ -1903,7 +1907,7 @@ func prepareMsg(m any, codec baseCodec, cp Compressor, comp encoding.Compressor,
 	if err != nil {
 		return nil, nil, nil, 0, err
 	}
-	compData, pf, err := compress(data, cp, comp, pool)
+	compData, pf, err := compress(data, cp, comp, pool, compressorOptions...)
 	if err != nil {
 		data.Free()
 		return nil, nil, nil, 0, err