From 5686b2c944c896441aa03622465fc8284c90996e Mon Sep 17 00:00:00 2001 From: Suzy Mueller Date: Fri, 18 Sep 2026 22:27:18 +0000 Subject: [PATCH] feat(pubsub): pipe publish timeout to concurrent batch actor Forward the configured publish timeout from BasePublisher through PublisherPartialBuilder and Dispatcher to ConcurrentBatchActor. This enables publisher request hedging to clamp hedged attempt timeouts to min(remaining_total_time, 10s). --- src/pubsub/src/publisher/actor.rs | 22 +++ src/pubsub/src/publisher/base_publisher.rs | 60 +++++++- src/pubsub/src/publisher/builder.rs | 8 ++ src/pubsub/src/publisher/hedging.rs | 151 ++++++++++++++++++++- 4 files changed, 237 insertions(+), 4 deletions(-) diff --git a/src/pubsub/src/publisher/actor.rs b/src/pubsub/src/publisher/actor.rs index 96a2122714..ebad0c98ee 100644 --- a/src/pubsub/src/publisher/actor.rs +++ b/src/pubsub/src/publisher/actor.rs @@ -64,6 +64,7 @@ pub(crate) struct Dispatcher { client: GapicPublisher, batching_options: BatchingOptions, hedging_options: Option, + total_timeout: Option, rx: mpsc::UnboundedReceiver, } @@ -73,6 +74,7 @@ impl Dispatcher { client: GapicPublisher, batching_options: BatchingOptions, hedging_options: Option, + total_timeout: Option, rx: mpsc::UnboundedReceiver, ) -> Self { Self { @@ -81,6 +83,7 @@ impl Dispatcher { rx, batching_options, hedging_options, + total_timeout, } } @@ -94,6 +97,7 @@ impl Dispatcher { self.client.clone(), self.batching_options.clone(), self.hedging_options.clone(), + self.total_timeout, rx, ) .run(), @@ -232,6 +236,7 @@ impl BatchActorContext { struct ConcurrentBatchActor { context: BatchActorContext, hedging: Option, + total_timeout: Option, } impl ConcurrentBatchActor { @@ -240,12 +245,14 @@ impl ConcurrentBatchActor { client: GapicPublisher, batching_options: BatchingOptions, hedging_options: Option, + total_timeout: Option, rx: mpsc::UnboundedReceiver, ) -> Self { let hedging = hedging_options.map(HedgingScheduler::spawn); ConcurrentBatchActor { context: BatchActorContext::new(topic, client, batching_options, rx), hedging, + total_timeout, } } @@ -341,6 +348,7 @@ impl ConcurrentBatchActor { txs, self.context.client.clone(), self.context.topic.clone(), + self.total_timeout, inflight, ); } else { @@ -823,6 +831,7 @@ mod tests { client.clone(), batching_options.clone(), None, + None, rx, ); @@ -852,6 +861,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(2_u32), None, + None, actor_rx, ) .run(), @@ -930,6 +940,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(2_u32), None, + None, actor_rx, ) .run(), @@ -1012,6 +1023,7 @@ mod tests { GapicPublisher::from_stub(MockGapicPublisher::new()), BatchingOptions::default(), None, + None, actor_rx, ) .run(), @@ -1102,6 +1114,7 @@ mod tests { .set_byte_threshold(MAX_BYTES) .set_delay_threshold(std::time::Duration::MAX), None, + None, actor_rx, ) .run(), @@ -1162,6 +1175,7 @@ mod tests { .set_message_count_threshold(MAX_MESSAGES) .set_byte_threshold(25_u32), // The current test generates 24 byte single message batches. None, + None, actor_rx, ) .run(), @@ -1275,6 +1289,7 @@ mod tests { .set_message_count_threshold(MAX_MESSAGES) .set_byte_threshold(1_u32), // The current test generates 24 byte single message batches. None, + None, actor_rx, ) .run(), @@ -1328,6 +1343,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1377,6 +1393,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1434,6 +1451,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1494,6 +1512,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1553,6 +1572,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1607,6 +1627,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), @@ -1668,6 +1689,7 @@ mod tests { GapicPublisher::from_stub(mock), BatchingOptions::default().set_message_count_threshold(1_u32), Some(hedging_options), + None, actor_rx, ) .run(), diff --git a/src/pubsub/src/publisher/base_publisher.rs b/src/pubsub/src/publisher/base_publisher.rs index a5a56b4c3a..a0afb13c29 100644 --- a/src/pubsub/src/publisher/base_publisher.rs +++ b/src/pubsub/src/publisher/base_publisher.rs @@ -13,6 +13,7 @@ // limitations under the License. use crate::publisher::builder::PublisherPartialBuilder; +use std::time::Duration; /// Creates [`Publisher`](crate::client::Publisher) instances. /// @@ -43,6 +44,7 @@ use crate::publisher::builder::PublisherPartialBuilder; #[derive(Clone, Debug)] pub struct BasePublisher { pub(crate) inner: crate::generated::gapic_dataplane::client::Publisher, + pub(crate) total_timeout: Option, } pub use super::client_builder::BasePublisherBuilder; @@ -62,9 +64,16 @@ impl BasePublisher { /// Creates a new Pub/Sub publisher client with the given configuration. pub(crate) async fn new(builder: BasePublisherBuilder) -> crate::ClientBuilderResult { + let total_timeout = + builder.config.retry_policy.as_ref().and_then(|p| { + p.remaining_time(&google_cloud_gax::retry_state::RetryState::new(false)) + }); let inner = crate::generated::gapic_dataplane::client::Publisher::new(builder.config).await?; - std::result::Result::Ok(Self { inner }) + std::result::Result::Ok(Self { + inner, + total_timeout, + }) } /// Creates a new `Publisher` for a given topic. @@ -85,6 +94,7 @@ impl BasePublisher { T: Into, { PublisherPartialBuilder::new(self.inner.clone(), topic.into()) + .with_total_timeout(self.total_timeout) } } @@ -92,6 +102,8 @@ impl BasePublisher { mod tests { use super::BasePublisher; use google_cloud_auth::credentials::anonymous::Builder as Anonymous; + use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt}; + use std::time::Duration; #[tokio::test] async fn builder() -> anyhow::Result<()> { @@ -102,4 +114,50 @@ mod tests { let _ = client.publisher("projects/my-project/topics/my-topic".to_string()); Ok(()) } + + #[tokio::test] + async fn default_total_timeout() -> anyhow::Result<()> { + let client = BasePublisher::builder() + .with_credentials(Anonymous::new().build()) + .build() + .await?; + let timeout = client + .total_timeout + .expect("default total_timeout should be present"); + assert!(timeout <= Duration::from_secs(600) && timeout >= Duration::from_secs(590)); + + let partial_builder = client.publisher("projects/my-project/topics/my-topic"); + assert_eq!(partial_builder.total_timeout, client.total_timeout); + Ok(()) + } + + #[tokio::test] + async fn custom_total_timeout() -> anyhow::Result<()> { + let client = BasePublisher::builder() + .with_credentials(Anonymous::new().build()) + .with_retry_policy(AlwaysRetry.with_time_limit(Duration::from_secs(45))) + .build() + .await?; + let timeout = client + .total_timeout + .expect("custom total_timeout should be present"); + assert!(timeout <= Duration::from_secs(45) && timeout >= Duration::from_secs(40)); + + let partial_builder = client.publisher("projects/my-project/topics/my-topic"); + assert_eq!(partial_builder.total_timeout, client.total_timeout); + Ok(()) + } + + #[tokio::test] + async fn attempt_limit_only_has_no_total_timeout() -> anyhow::Result<()> { + let client = BasePublisher::builder() + .with_credentials(Anonymous::new().build()) + .with_retry_policy(AlwaysRetry.with_attempt_limit(3)) + .build() + .await?; + assert_eq!(client.total_timeout, None); + let partial_builder = client.publisher("projects/my-project/topics/my-topic"); + assert_eq!(partial_builder.total_timeout, None); + Ok(()) + } } diff --git a/src/pubsub/src/publisher/builder.rs b/src/pubsub/src/publisher/builder.rs index 65edce269b..5387251785 100644 --- a/src/pubsub/src/publisher/builder.rs +++ b/src/pubsub/src/publisher/builder.rs @@ -357,6 +357,7 @@ pub struct PublisherPartialBuilder { topic: String, batching_options: BatchingOptions, hedging_options: Option, + pub(crate) total_timeout: Option, } impl PublisherPartialBuilder { @@ -367,9 +368,15 @@ impl PublisherPartialBuilder { topic, batching_options: BatchingOptions::default(), hedging_options: None, + total_timeout: None, } } + pub(crate) fn with_total_timeout(mut self, total_timeout: Option) -> Self { + self.total_timeout = total_timeout; + self + } + /// Sets the message count threshold for batching. /// /// The publisher will send a batch of messages when the number of messages @@ -522,6 +529,7 @@ impl PublisherPartialBuilder { self.inner, batching_options.clone(), hedging_options.clone(), + self.total_timeout, rx, ); let handle = tokio::spawn(dispatcher.run()); diff --git a/src/pubsub/src/publisher/hedging.rs b/src/pubsub/src/publisher/hedging.rs index c2dc89cba4..c1040e7c0c 100644 --- a/src/pubsub/src/publisher/hedging.rs +++ b/src/pubsub/src/publisher/hedging.rs @@ -51,6 +51,8 @@ pub(crate) struct BatchState { pub topic: String, pub token_bucket: Arc, pub start_time: Option, + pub start_instant: tokio::time::Instant, + pub total_timeout: Option, /// Cancellation token shared across all attempts (initial and hedged) for this batch. /// Triggered as soon as any attempt succeeds (or the initial attempt fails permanently). pub cancel_token: CancellationToken, @@ -63,6 +65,7 @@ impl BatchState { client: GapicPublisher, topic: String, token_bucket: Arc, + total_timeout: Option, done_tx: tokio::sync::oneshot::Sender>, ) -> Self { Self { @@ -72,6 +75,8 @@ impl BatchState { topic, token_bucket, start_time: wkt::Timestamp::try_from(std::time::SystemTime::now()).ok(), + start_instant: tokio::time::Instant::now(), + total_timeout, cancel_token: CancellationToken::new(), } } @@ -105,9 +110,21 @@ impl BatchState { if self.cancel_token.is_cancelled() { return; } - - // TODO(#6776): clamp the timeout to the remaining time of the initial request. - let timeout = Duration::from_secs(10); + // The server allows a maximum timeout of 10s for publish RPC attempts. + // We cap individual hedged attempts to this limit, while also clamping + // to any remaining total retry timeout. + const MAX_ATTEMPT_TIMEOUT: Duration = Duration::from_secs(10); + let timeout = self + .total_timeout + .map(|t| { + let elapsed = self.start_instant.elapsed(); + t.saturating_sub(elapsed) + }) + .unwrap_or(MAX_ATTEMPT_TIMEOUT) + .min(MAX_ATTEMPT_TIMEOUT); + if timeout.is_zero() { + return; + } let request = self .client @@ -170,6 +187,7 @@ impl HedgingSchedulerHandle { txs: Vec>>, client: GapicPublisher, topic: String, + total_timeout: Option, inflight: &mut JoinSet>, ) { let (done_tx, done_rx) = oneshot::channel(); @@ -179,6 +197,7 @@ impl HedgingSchedulerHandle { client, topic, self.token_bucket.clone(), + total_timeout, done_tx, )); inflight.spawn(async move { @@ -334,6 +353,19 @@ mod tests { Arc, oneshot::Receiver>, oneshot::Receiver>, + ) { + test_batch_state_with_timeout(client, token_bucket, None) + } + + #[allow(clippy::type_complexity)] + fn test_batch_state_with_timeout( + client: GapicPublisher, + token_bucket: Arc, + total_timeout: Option, + ) -> ( + Arc, + oneshot::Receiver>, + oneshot::Receiver>, ) { let (tx, rx) = oneshot::channel(); let (done_tx, done_rx) = oneshot::channel(); @@ -343,6 +375,7 @@ mod tests { client, "topic".to_string(), token_bucket, + total_timeout, done_tx, )); (state, rx, done_rx) @@ -973,4 +1006,116 @@ mod tests { Ok(()) } + + #[tokio_test_no_panics(start_paused = true)] + async fn test_hedged_rpc_timeout_clamped_to_remaining_time() -> anyhow::Result<()> { + let mut mock = MockGapicPublisherWithFuture::new(); + mock.expect_publish().times(1).returning(|_, options| { + // Verify that the retry policy time limit on the hedged request was clamped + let policy = options.retry_policy(); + let policy = policy.as_ref().expect("retry policy must be set"); + // The policy should have a remaining time limit of exactly 3 seconds (5s total - 2s elapsed) + let state = google_cloud_gax::retry_state::RetryState::new(false) + .set_start(tokio::time::Instant::now().into_std()); + let remaining = policy.remaining_time(&state); + assert_eq!(remaining, Some(Duration::from_secs(3))); + Box::pin(async { mock_publish_response("msg-hedged") }) + }); + + let client = GapicPublisher::from_stub(mock); + let token_bucket = Arc::new(TokenBucket::new(10, 0.1)); + let (state, rx, done_rx) = + test_batch_state_with_timeout(client, token_bucket, Some(Duration::from_secs(5))); + + // Advance time by 2 seconds before hedged RPC is sent + tokio::time::advance(Duration::from_secs(2)).await; + + state.send_hedged_rpc(1).await; + + let msg_id = rx.await??; + assert_eq!(msg_id, "msg-hedged"); + assert!(state.cancel_token.is_cancelled()); + done_rx.await??; + + Ok(()) + } + + #[tokio_test_no_panics(start_paused = true)] + async fn test_hedged_rpc_timeout_capped_at_max_attempt_timeout() -> anyhow::Result<()> { + let mut mock = MockGapicPublisherWithFuture::new(); + mock.expect_publish().times(1).returning(|_, options| { + let policy = options.retry_policy(); + let policy = policy.as_ref().expect("retry policy must be set"); + // Even though 600s - 1s = 599s remains, the hedge is capped at MAX_ATTEMPT_TIMEOUT (10s) + let state = google_cloud_gax::retry_state::RetryState::new(false) + .set_start(tokio::time::Instant::now().into_std()); + let remaining = policy.remaining_time(&state); + assert_eq!(remaining, Some(Duration::from_secs(10))); + Box::pin(async { mock_publish_response("msg-hedged") }) + }); + + let client = GapicPublisher::from_stub(mock); + let token_bucket = Arc::new(TokenBucket::new(10, 0.1)); + let (state, rx, done_rx) = + test_batch_state_with_timeout(client, token_bucket, Some(Duration::from_secs(600))); + + tokio::time::advance(Duration::from_secs(1)).await; + + state.send_hedged_rpc(1).await; + + let msg_id = rx.await??; + assert_eq!(msg_id, "msg-hedged"); + assert!(state.cancel_token.is_cancelled()); + done_rx.await??; + + Ok(()) + } + + #[tokio_test_no_panics(start_paused = true)] + async fn test_hedged_rpc_skipped_when_total_timeout_expired() -> anyhow::Result<()> { + let mock = MockGapicPublisher::new(); + // Gapic client should not be called because remaining time is zero + let client = GapicPublisher::from_stub(mock); + let token_bucket = Arc::new(TokenBucket::new(10, 0.1)); + let (state, _rx, _done_rx) = + test_batch_state_with_timeout(client, token_bucket, Some(Duration::from_millis(500))); + + // Advance time past total timeout + tokio::time::advance(Duration::from_secs(1)).await; + + state.send_hedged_rpc(1).await; + + // cancel_token is not cancelled, batch is not completed + assert!(!state.cancel_token.is_cancelled()); + + Ok(()) + } + + #[tokio_test_no_panics(start_paused = true)] + async fn test_hedged_rpc_without_total_timeout() -> anyhow::Result<()> { + let mut mock = MockGapicPublisherWithFuture::new(); + mock.expect_publish().times(1).returning(|_, options| { + let policy = options.retry_policy(); + let policy = policy.as_ref().expect("retry policy must be set"); + // With None total_timeout, remaining time defaults to MAX_ATTEMPT_TIMEOUT (10s) + let state = google_cloud_gax::retry_state::RetryState::new(false) + .set_start(tokio::time::Instant::now().into_std()); + let remaining = policy.remaining_time(&state); + assert_eq!(remaining, Some(Duration::from_secs(10))); + Box::pin(async { mock_publish_response("msg-hedged") }) + }); + + let client = GapicPublisher::from_stub(mock); + let token_bucket = Arc::new(TokenBucket::new(10, 0.1)); + let (state, rx, done_rx) = test_batch_state_with_timeout(client, token_bucket, None); + + state.send_hedged_rpc(1).await; + + let msg_id = rx.await??; + assert_eq!(msg_id, "msg-hedged"); + assert!(state.cancel_token.is_cancelled()); + done_rx.await??; + + Ok(()) + } }