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(()) + } }