Conversation
Implement request hedging for batch publishing in the Pub/Sub publisher. When request hedging is enabled, if an initial publish RPC does not complete within the configured delay threshold and tokens are available in the token bucket, a hedged publish RPC is dispatched with `NeverRetry`. The first attempt to succeed resolves the pending publish futures, refills the token bucket, and cancels any remaining in-flight attempts via a shared CancellationToken. Rate-limited hedges (where no token is available at deadline expiry) are discarded to shed load and avoid sending excessive requests during sustained backend latency spikes. Hedged request timeouts are capped at 10 seconds pending retry policy timeout integration.
![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 request hedging to the Pub/Sub publisher by adding a dedicated HedgingScheduler and BatchState management. It refactors the batching logic to support concurrent and sequential actors with hedging options, and adds comprehensive test coverage for various hedging scenarios. Feedback on the changes suggests removing the unused base64 dependency and RequestOptionsBuilder import, as well as releasing the Mutex lock early in BatchState::complete to minimize lock contention during non-atomic operations.
| async-trait.workspace = true | ||
| base64.workspace = true | ||
| google-cloud-auth = { workspace = true } |
There was a problem hiding this comment.
The base64 dependency is added to Cargo.toml but is not used anywhere in the changes or the crate. We should remove it to avoid unnecessary dependency bloat and keep the build clean.
| async-trait.workspace = true | |
| base64.workspace = true | |
| google-cloud-auth = { workspace = true } | |
| async-trait.workspace = true | |
| google-cloud-auth = { workspace = true } |
References
- Expose Hidden Costs: Question new dependencies or patterns that add significant boilerplate for minimal gain. (link)
| use google_cloud_gax::options::RequestOptionsBuilder; | ||
| use google_cloud_gax::retry_policy::{NeverRetry, RetryPolicyExt}; |
There was a problem hiding this comment.
The RequestOptionsBuilder import is unused in this file and can be safely removed.
| use google_cloud_gax::options::RequestOptionsBuilder; | |
| use google_cloud_gax::retry_policy::{NeverRetry, RetryPolicyExt}; | |
| use google_cloud_gax::retry_policy::{NeverRetry, RetryPolicyExt}; |
References
- Formatting: Code must be formatted with cargo fmt. (link)
| fn complete(&self, resp: crate::Result<PublishResponse>) { | ||
| let mut lock = self.txs.lock().unwrap(); | ||
| if let Some((txs, done_tx)) = lock.take() { | ||
| self.cancel_token.cancel(); | ||
| if resp.is_ok() { | ||
| self.token_bucket.refill(); | ||
| } | ||
| let _ = done_tx.send(batch_resolve_publish_futures(resp, txs)); | ||
| } | ||
| } |
There was a problem hiding this comment.
Holding a Mutex lock while performing potentially blocking or complex operations (such as cancelling a token, refilling a token bucket, or sending messages over channels) can lead to unnecessary lock contention and potential deadlocks. Since we only need to atomically take the senders from self.txs, we can lock the mutex, take() the option, and release the lock immediately before executing the rest of the completion logic.
| fn complete(&self, resp: crate::Result<PublishResponse>) { | |
| let mut lock = self.txs.lock().unwrap(); | |
| if let Some((txs, done_tx)) = lock.take() { | |
| self.cancel_token.cancel(); | |
| if resp.is_ok() { | |
| self.token_bucket.refill(); | |
| } | |
| let _ = done_tx.send(batch_resolve_publish_futures(resp, txs)); | |
| } | |
| } | |
| fn complete(&self, resp: crate::Result<PublishResponse>) { | |
| let txs_and_done_tx = self.txs.lock().unwrap().take(); | |
| if let Some((txs, done_tx)) = txs_and_done_tx { | |
| self.cancel_token.cancel(); | |
| if resp.is_ok() { | |
| self.token_bucket.refill(); | |
| } | |
| let _ = done_tx.send(batch_resolve_publish_futures(resp, txs)); | |
| } | |
| } |
References
- Unnecessary locks: Scrutinize any use of Mutex<>. Do multiple threads really need to access this data? Could we avoid locks with a different model? (link)
|
I think this PR is too big for me to provide good feedback. Can we break it down, e.g. start with some of the types from |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #6809 +/- ##
========================================
Coverage 97.02% 97.03%
========================================
Files 326 327 +1
Lines 107511 108204 +693
========================================
+ Hits 104310 104992 +682
- Misses 3201 3212 +11 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
2 participants
Implement request hedging for batch publishing in the Pub/Sub publisher. When request hedging is enabled, if an initial publish RPC does not complete within the configured delay threshold and tokens are available in the token bucket, a hedged publish RPC is dispatched with
NeverRetry. The first attempt to succeed resolves the pending publish futures, refills the token bucket, and cancels any remaining in-flight attempts via a shared CancellationToken. Only errors from the initial request are returned to the user, all other errors are discarded.Rate-limited hedges (where no token is available at deadline expiry) are discarded to shed load and avoid sending excessive requests during sustained backend latency spikes.
Hedged request timeouts are capped at 10 seconds pending retry policy timeout integration (it should be the minimum of 10 seconds and total_timeout of the original request).
For #6776