impl(pubsub): pipe publish timeout to concurrent batch actor

src/pubsub/src/publisher/actor.rs

@@ -19,6 +19,7 @@ use crate::model::{Message, PublishResponse};
 use crate::publisher::batch::Batch;
 use crate::publisher::hedging::{HedgingScheduler, HedgingSchedulerHandle};
 use crate::publisher::options::HedgingOptions;
+use google_cloud_gax::retry_policy::RetryPolicy;
 use std::collections::{HashMap, VecDeque};
 use std::pin::Pin;
 use std::sync::Arc;
@@ -64,6 +65,7 @@ pub(crate) struct Dispatcher {
     client: GapicPublisher,
     batching_options: BatchingOptions,
     hedging_options: Option<HedgingOptions>,
+    retry_policy: Option<Arc<dyn RetryPolicy>>,
     rx: mpsc::UnboundedReceiver<ToDispatcher>,
 }
 
@@ -73,6 +75,7 @@ impl Dispatcher {
         client: GapicPublisher,
         batching_options: BatchingOptions,
         hedging_options: Option<HedgingOptions>,
+        retry_policy: Option<Arc<dyn RetryPolicy>>,
         rx: mpsc::UnboundedReceiver<ToDispatcher>,
     ) -> Self {
         Self {
@@ -81,6 +84,7 @@ impl Dispatcher {
             rx,
             batching_options,
             hedging_options,
+            retry_policy,
         }
     }
 
@@ -94,6 +98,7 @@ impl Dispatcher {
                     self.client.clone(),
                     self.batching_options.clone(),
                     self.hedging_options.clone(),
+                    self.retry_policy.clone(),
                     rx,
                 )
                 .run(),
@@ -232,6 +237,8 @@ impl BatchActorContext {
 struct ConcurrentBatchActor {
     context: BatchActorContext,
     hedging: Option<HedgingSchedulerHandle>,
+    #[allow(dead_code)]
+    retry_policy: Option<Arc<dyn RetryPolicy>>,
 }
 
 impl ConcurrentBatchActor {
@@ -240,12 +247,14 @@ impl ConcurrentBatchActor {
         client: GapicPublisher,
         batching_options: BatchingOptions,
         hedging_options: Option<HedgingOptions>,
+        retry_policy: Option<Arc<dyn RetryPolicy>>,
         rx: mpsc::UnboundedReceiver<ToBatchActor>,
     ) -> Self {
         let hedging = hedging_options.map(HedgingScheduler::spawn);
         ConcurrentBatchActor {
             context: BatchActorContext::new(topic, client, batching_options, rx),
             hedging,
+            retry_policy,
         }
     }
 
@@ -819,6 +828,7 @@ mod tests {
             client.clone(),
             batching_options.clone(),
             None,
+            None,
             rx,
         );
 
@@ -848,6 +858,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(2_u32),
                 None,
+                None,
                 actor_rx,
             )
             .run(),
@@ -926,6 +937,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(2_u32),
                 None,
+                None,
                 actor_rx,
             )
             .run(),
@@ -1008,6 +1020,7 @@ mod tests {
                 GapicPublisher::from_stub(MockGapicPublisher::new()),
                 BatchingOptions::default(),
                 None,
+                None,
                 actor_rx,
             )
             .run(),
@@ -1098,6 +1111,7 @@ mod tests {
                     .set_byte_threshold(MAX_BYTES)
                     .set_delay_threshold(std::time::Duration::MAX),
                 None,
+                None,
                 actor_rx,
             )
             .run(),
@@ -1158,6 +1172,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(),
@@ -1271,6 +1286,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(),
@@ -1324,6 +1340,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1373,6 +1390,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1430,6 +1448,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1490,6 +1509,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1549,6 +1569,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1603,6 +1624,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),
@@ -1664,6 +1686,7 @@ mod tests {
                 GapicPublisher::from_stub(mock),
                 BatchingOptions::default().set_message_count_threshold(1_u32),
                 Some(hedging_options),
+                None,
                 actor_rx,
             )
             .run(),

src/pubsub/src/publisher/base_publisher.rs

@@ -13,12 +13,14 @@
 // limitations under the License.
 
 use crate::publisher::builder::PublisherPartialBuilder;
+use google_cloud_gax::retry_policy::RetryPolicy;
+use std::sync::Arc;
 
-/// Creates [`Publisher`](crate::client::Publisher) instances.
+/// A client for the Cloud Pub/Sub publish API.
 ///
-/// A single `BasePublisher` can be used to create multiple `Publisher` clients
-/// for different topics. It manages the underlying gRPC connection and
-/// authentication.
+/// This is the low-level client for the Cloud Pub/Sub publish API. Applications
+/// that need more control over the batching of messages can use this client to
+/// create a [`Publisher`] with custom batching options.
 ///
 /// # Example
 ///
@@ -28,21 +30,20 @@ use crate::publisher::builder::PublisherPartialBuilder;
 /// # use google_cloud_pubsub::model::Message;
 ///
 /// // Create a client.
-/// let client: BasePublisher = BasePublisher::builder().build().await?;
+/// let client = BasePublisher::builder().build().await?;
 ///
 /// // Create a publisher for a specific topic.
 /// let publisher = client.publisher("projects/my-project/topics/my-topic").build();
 ///
 /// // Publish a message.
-/// let handle = publisher.publish(Message::new().set_data("hello world"));
-/// let message_id = handle.await?;
-/// println!("Message sent with ID: {}", message_id);
+/// let message_id = publisher.publish(Message::new().set_data("Hello, World")).await?;
 /// # Ok(())
 /// # }
 /// ```
 #[derive(Clone, Debug)]
 pub struct BasePublisher {
     pub(crate) inner: crate::generated::gapic_dataplane::client::Publisher,
+    pub(crate) retry_policy: Option<Arc<dyn RetryPolicy>>,
 }
 
 pub use super::client_builder::BasePublisherBuilder;
@@ -62,9 +63,13 @@ impl BasePublisher {
 
     /// Creates a new Pub/Sub publisher client with the given configuration.
     pub(crate) async fn new(builder: BasePublisherBuilder) -> crate::ClientBuilderResult<Self> {
+        let retry_policy = builder.config.retry_policy.clone();
         let inner =
             crate::generated::gapic_dataplane::client::Publisher::new(builder.config).await?;
-        std::result::Result::Ok(Self { inner })
+        std::result::Result::Ok(Self {
+            inner,
+            retry_policy,
+        })
     }
 
     /// Creates a new `Publisher` for a given topic.
@@ -85,13 +90,18 @@ impl BasePublisher {
         T: Into<String>,
     {
         PublisherPartialBuilder::new(self.inner.clone(), topic.into())
+            .with_retry_policy(self.retry_policy.clone())
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::BasePublisher;
     use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
+    use google_cloud_gax::retry_policy::{AlwaysRetry, RetryPolicyExt};
+    use google_cloud_gax::retry_state::RetryState;
+    use std::sync::Arc;
+    use std::time::Duration;
 
     #[tokio::test]
     async fn builder() -> anyhow::Result<()> {
@@ -102,4 +112,57 @@ mod tests {
         let _ = client.publisher("projects/my-project/topics/my-topic".to_string());
         Ok(())
     }
+
+    #[tokio::test]
+    async fn default_retry_policy() -> anyhow::Result<()> {
+        let client = BasePublisher::builder()
+            .with_credentials(Anonymous::new().build())
+            .build()
+            .await?;
+        let policy = client
+            .retry_policy
+            .as_ref()
+            .expect("default retry_policy should be present");
+        let timeout = policy
+            .remaining_time(&RetryState::new(false))
+            .expect("default policy should have time limit");
+        assert!(timeout <= Duration::from_secs(600) && timeout >= Duration::from_secs(590));
+
+        let partial_builder = client.publisher("projects/my-project/topics/my-topic");
+        assert!(Arc::ptr_eq(
+            policy,
+            partial_builder
+                .retry_policy
+                .as_ref()
+                .expect("partial builder should have retry policy")
+        ));
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn custom_retry_policy() -> 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 policy = client
+            .retry_policy
+            .as_ref()
+            .expect("custom retry_policy should be present");
+        let timeout = policy
+            .remaining_time(&RetryState::new(false))
+            .expect("custom policy should have time limit");
+        assert!(timeout <= Duration::from_secs(45) && timeout >= Duration::from_secs(40));
+
+        let partial_builder = client.publisher("projects/my-project/topics/my-topic");
+        assert!(Arc::ptr_eq(
+            policy,
+            partial_builder
+                .retry_policy
+                .as_ref()
+                .expect("partial builder should have retry policy")
+        ));
+        Ok(())
+    }
 }

src/pubsub/src/publisher/builder.rs

@@ -18,9 +18,11 @@ use crate::generated::gapic_dataplane::client::Publisher as GapicPublisher;
 use crate::publisher::actor::Dispatcher;
 use crate::publisher::base_publisher::BasePublisher;
 use google_cloud_gax::{
-    backoff_policy::BackoffPolicyArg, retry_policy::RetryPolicyArg,
+    backoff_policy::BackoffPolicyArg,
+    retry_policy::{RetryPolicy, RetryPolicyArg},
     retry_throttler::RetryThrottlerArg,
 };
+use std::sync::Arc;
 use std::time::Duration;
 
 pub use super::base_publisher::BasePublisherBuilder;
@@ -360,6 +362,7 @@ pub struct PublisherPartialBuilder {
     topic: String,
     batching_options: BatchingOptions,
     hedging_options: Option<HedgingOptions>,
+    pub(crate) retry_policy: Option<Arc<dyn RetryPolicy>>,
 }
 
 impl PublisherPartialBuilder {
@@ -370,9 +373,15 @@ impl PublisherPartialBuilder {
             topic,
             batching_options: BatchingOptions::default(),
             hedging_options: None,
+            retry_policy: None,
         }
     }
 
+    pub(crate) fn with_retry_policy(mut self, retry_policy: Option<Arc<dyn RetryPolicy>>) -> Self {
+        self.retry_policy = retry_policy;
+        self
+    }
+
     /// Sets the message count threshold for batching.
     ///
     /// The publisher will send a batch of messages when the number of messages
@@ -528,6 +537,7 @@ impl PublisherPartialBuilder {
             self.inner,
             batching_options.clone(),
             hedging_options.clone(),
+            self.retry_policy,
             rx,
         );
         let handle = tokio::spawn(dispatcher.run());