grpc: support per-message compression and enforce embedding in ServerTransportStream implementations#8972
grpc: support per-message compression and enforce embedding in ServerTransportStream implementations#8972Dostonlv wants to merge 9 commits intogrpc:masterfrom
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #8972 +/- ##
==========================================
- Coverage 83.23% 81.77% -1.47%
==========================================
Files 410 413 +3
Lines 32572 33386 +814
==========================================
+ Hits 27111 27300 +189
- Misses 4066 4277 +211
- Partials 1395 1809 +414
🚀 New features to boost your workflow:
|
easwars
added
Type: API Change
| // 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) |
There was a problem hiding this comment.
There are a few things here that I'm not convinced about:
- We are adding the
compressKeyfor every stream context whether or not condition compression enabling/disabling will be used or not. I'm not sure how much of a performance overhead this would be. @arjan-bal : Do you have any thoughts? - We are modifying a value stored in a context here, which breaks the guarantee that contexts are immutable in Go. I agree that we do have this pattern in a few places already in our codebase, but if we can avoid it, that would be great. One option is to avoid it would be not store a pointer to
compressOptions, but instead store it by value, and change this function to return a new context that containscompressKeywith the given value. This would mean that the caller on the client and the server would have to use the newly returned context going forward. @dfawley : Do you have any objections to change this function to return a context instead?
There was a problem hiding this comment.
Thanks for the suggestion! I thought about this approach too, but returning a new context doesn't work for streams.
The reason is that serverStream.SendMsg reads from ss.ctx — the context that was captured when the stream was created. Even if SetMessageCompression returns a new context, ss.ctx inside the stream object doesn't change, so SendMsg would never pick it up. The same issue applies to clientStream and addrConnStream.
Since there's no way to update the stream's internal context from the outside, mutating the pointer was the approach I went with. That said, I'm open to other ideas if there's a cleaner way to handle this — happy to adjust if you have something in mind.
There was a problem hiding this comment.
As you mentioned, I wanted to share my thinking on the compressKey overhead. Since SetMessageCompression is already a no-op without a compressor, I considered adding it only when a compressor is actually set. But the tricky part is when SetSendCompressor is called first — the compressor isn't known at stream creation time. So I kept it as is for now. Happy to explore this further if you think it's worth optimizing.
There was a problem hiding this comment.
One more thought — to address the mutability concern, I'm planning to
switch from plain bool to atomic.Bool in compressOptions. Would that work?
There was a problem hiding this comment.
Presently we have ~24 allocs/op for streaming RPCs and ~140 allocs/op for unary RPCs for client and server side combined. Even a single extra heap allocation does show up as a <1% QPS drop in our benchmarks when sending very small messages. Here is how to run the benchmarks for unary RPCs:
$ git checkout master
$ RUN_NAME=unary-before
$ go run benchmark/benchmain/main.go -benchtime=60s -workloads=unary \
-compression=off -maxConcurrentCalls=200 -trace=off \
-reqSizeBytes=100 -respSizeBytes=100 -networkMode=Local -resultFile="${RUN_NAME}"
$ git checkout feature-branch
$ RUN_NAME=unary-after
$ go run benchmark/benchmain/main.go -benchtime=60s -workloads=unary \
-compression=off -maxConcurrentCalls=200 -trace=off \
-reqSizeBytes=100 -respSizeBytes=100 -networkMode=Local -resultFile="${RUN_NAME}"
# Compare the results
$ go run benchmark/benchresult/main.go unary-before unary-after
We should ensure that users not utilizing the new SetMessageCompression API avoid the extra allocation. Could we treat the absence of compressKey{} as "compression enabled" to match current behavior? This should should avoid heap allocations for existing use-cases.
There was a problem hiding this comment.
I ran the benchmarks and the results confirm your concern:
- Allocs/op: 143 → 147 (+2.79%)
- 99th percentile latency: 2.636ms → 3.051ms (+15.74%)
I'll fix this by treating the absence of compressKey as
"compression enabled", so users not using SetMessageCompression
avoid the extra allocation entirely.
There was a problem hiding this comment.
Updated the implementation to address the allocation concern.
Server side: moved doNotCompress into transport.ServerStream directly, no context allocation needed.
Client side: compressKey is added to context only when a compressor is set, so users not using SetMessageCompression get zero extra allocations.
Benchmark results (unary, compression=off):
Title Before After Percentage
TotalOps 8318091 8070904 -2.97%
Bytes/op 9927.29 9883.28 -0.44%
Allocs/op 143.22 143.03 0.00%
50th-Lat 1.210583ms 1.255084ms 3.68%
90th-Lat 2.083459ms 2.034416ms -2.35%
99th-Lat 4.868417ms 4.718292ms -3.08%
Avg-Lat 1.440822ms 1.483993ms 3.00%
Benchmark results (unary, compression=gzip):
Title Before After Percentage
TotalOps 2591595 2578328 -0.51%
Bytes/op 14332.39 14312.50 -0.14%
Allocs/op 183.36 185.21 +1.09%
50th-Lat 4.095041ms 4.1095ms 0.35%
90th-Lat 6.164916ms 6.204125ms 0.64%
99th-Lat 15.452875ms 15.49425ms 0.27%
Avg-Lat 4.624957ms 4.648262ms 0.50%
The +2 allocs in gzip case is expected since compressKey is only added when a compressor is set.
@arjan-bal @easwars
This is one of my first contributions to the project, so any guidance or feedback would be greatly appreciated.
There was a problem hiding this comment.
Thanks for optimizing this. The results seem good.
| // 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) |
There was a problem hiding this comment.
Thanks for optimizing this. The results seem good.
|
@arjan-bal Assalamu alaykum. Updated the PR based on the feedback. Would appreciate another review |
arjan-bal
left a comment
arjan-bal
left a comment
There was a problem hiding this comment.
Mostly looks good! Just left a few minor housekeeping comments. Great job on the tests, I especially like how you used stats handlers to verify compression.
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces experimental APIs to enable or disable per-message compression mid-stream for both gRPC clients and servers. The changes include updates to the internal transport and stream handling logic to respect these toggles, along with comprehensive tests. Feedback suggests ensuring the client-side compression flag is always initialized to align with documentation and making sure internal stream types like addrConnStream correctly support the new toggle mechanism.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces the ability to enable or disable per-message compression mid-stream for both client and server sides. It adds SetServerStreamMessageCompression and SetClientStreamMessageCompression functions, updates the ServerTransportStream interface, and ensures that compression settings are respected during message preparation. The review feedback highlights a discrepancy between the implementation and documentation of SetClientStreamMessageCompression, suggesting that the compression flag should be consistently present in the client stream context to allow the function to act as a no-op when no compressor is configured, rather than returning an error.
| if !ok || flag == nil { | ||
| return fmt.Errorf("grpc: SetClientStreamMessageCompression called on a non-client-stream context or a compressor is not configured on the stream") | ||
| } |
There was a problem hiding this comment.
The implementation of SetClientStreamMessageCompression returns an error if no compressor is configured on the stream, which contradicts the documentation on lines 79-80 stating that it should be a no-op in this case.
To align with the documentation and provide a consistent API (similar to the server-side implementation), the compressKey should always be present in the client stream context. This allows the function to distinguish between an invalid context and a valid stream context where compression is simply not configured. Additionally, the flag == nil check is redundant as the type assertion .(*bool) would only succeed with a non-nil value given how gRPC initializes this flag.
| if !ok || flag == nil { | |
| return fmt.Errorf("grpc: SetClientStreamMessageCompression called on a non-client-stream context or a compressor is not configured on the stream") | |
| } | |
| if !ok { | |
| return fmt.Errorf("grpc: SetClientStreamMessageCompression called on a non-client-stream context") | |
| } |
There was a problem hiding this comment.
We need to update this line in the godoc to mention that an error is returned if no compressor is configured:
Compression is enabled by default and is a no-op if no compressor is configured on the stream
|
@arjan-bal
PTAL |
arjan-bal
changed the title
arjan-bal
left a comment
arjan-bal
left a comment
There was a problem hiding this comment.
LGTM, just one minor comment to address. I'm also adding a second reviewer.
| if !ok || flag == nil { | ||
| return fmt.Errorf("grpc: SetClientStreamMessageCompression called on a non-client-stream context or a compressor is not configured on the stream") | ||
| } |
There was a problem hiding this comment.
We need to update this line in the godoc to mention that an error is returned if no compressor is configured:
Compression is enabled by default and is a no-op if no compressor is configured on the stream
| // | ||
| // Notice: This API is EXPERIMENTAL and may be changed or removed in a | ||
| // later release. | ||
| func SetClientStreamMessageCompression(ctx context.Context, enable bool) error { |
There was a problem hiding this comment.
How will this be used on the client-side? The reason I'm asking is because I want to see if we can avoid mutating the context. This is an anti-pattern that we already have enough occurrences of and it would be nice if we can avoid adding more.
What I want to see if whether having a function like:
func NewClientStreamWithCompression(cs ClientStream, enable bool) ClientSteam {
// Create a new context with the compression enabled flag set to the value of `enable`
// Wrap the passed in stream with the new one, with all methods except Context being a passthrough
// The Context() method will return the newly created context that has the correct value of `enable`
}There was a problem hiding this comment.
I'm coming back to how this will be used. I'm wondering if an approach similar to the one on the server side can be used, where this function will accept a ClientStream instead of a context. The implementation will then propagate the value of the enable flag down to the client stream in the transport whenever appropriate (the clientStream struct in the grpc package has a reference to a csAttempt which has a ref to the underlying ClientStream in the transport).
|
|
||
| // ServerStream implements streaming functionality for a gRPC server. | ||
| type ServerStream struct { | ||
| internal.EnforceServerTransportStreamEmbedding |
There was a problem hiding this comment.
Do we really need this? This is an internal package, so we control all implementations, and can handle breaking changes appropriately.
|
@Dostonlv : Apologies for the delay in getting back. |
|
This PR is labeled as requiring an update from the reporter, and no update has been received after 6 days. If no update is provided in the next 7 days, this issue will be automatically closed. |
4 participants
Fixes #8662
RELEASE NOTES:
SetServerStreamMessageCompressionandSetClientStreamMessageCompressionfunctions.ServerTransportStreaminterface to force implementors to embed a delegate implementation. This helps gRPC add new methods to the interface without breaking builds for users of this interface.