google_cloud_pubsub/publisher/
publisher.rs

1// Copyright 2025 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use super::options::BatchingOptions;
16use crate::generated::gapic_dataplane::client::Publisher as GapicPublisher;
17use crate::publisher::worker::BundledMessage;
18use crate::publisher::worker::ToWorker;
19use crate::publisher::worker::Worker;
20use std::time::Duration;
21use tokio::sync::mpsc::UnboundedSender;
22use tokio::sync::oneshot;
23
24const MAX_DELAY: Duration = Duration::from_secs(60 * 60 * 24); // 1 day
25// These limits come from https://cloud.google.com/pubsub/docs/batch-messaging#quotas_and_limits_on_batch_messaging.
26// Client libraries are expected to enforce these limits on batch siziing.
27const MAX_MESSAGES: u32 = 1000;
28const MAX_BYTES: u32 = 1e7 as u32; // 10MB
29
30/// A `Publisher` sends messages to a specific topic. It manages message batching
31/// and sending in a background task.
32///
33/// Publishers are created via a [`Client`](crate::client::Client).
34///
35/// ```
36/// # async fn sample() -> anyhow::Result<()> {
37/// # use google_cloud_pubsub::*;
38/// # use client::Client;
39/// # use model::PubsubMessage;
40/// let client = Client::builder().build().await?;
41/// let publisher = client.publisher("projects/my-project/topics/my-topic").build();
42/// let message_id = publisher.publish(PubsubMessage::new().set_data("Hello, World"));
43/// # Ok(()) }
44/// ```
45#[derive(Debug, Clone)]
46pub struct Publisher {
47    #[allow(dead_code)]
48    pub(crate) batching_options: BatchingOptions,
49    tx: UnboundedSender<ToWorker>,
50}
51
52impl Publisher {
53    /// Publishes a message to the topic.
54    ///
55    /// ```
56    /// # use google_cloud_pubsub::client::Publisher;
57    /// # async fn sample(publisher: Publisher) -> anyhow::Result<()> {
58    /// # use google_cloud_pubsub::model::PubsubMessage;
59    /// let message_id = publisher.publish(PubsubMessage::new().set_data("Hello, World")).await?;
60    /// # Ok(()) }
61    /// ```
62    pub fn publish(&self, msg: crate::model::PubsubMessage) -> crate::model_ext::PublishHandle {
63        let (tx, rx) = tokio::sync::oneshot::channel();
64
65        // If this fails, the worker is gone, which indicates something bad has happened.
66        // The PublishHandle will automatically receive an error when `tx` is dropped.
67        if self
68            .tx
69            .send(ToWorker::Publish(BundledMessage { msg, tx }))
70            .is_err()
71        {
72            // `tx` is dropped here if the send errors.
73        }
74        crate::model_ext::PublishHandle { rx }
75    }
76
77    /// Flushes all outstanding messages.
78    ///
79    /// This method sends any messages that have been published but not yet sent,
80    /// regardless of the configured batching options (`delay_threshold`, etc.).
81    ///
82    /// This method is `async` and returns only after all publish attempts for the
83    /// messages in the snapshot have completed. A "completed" attempt means the
84    /// message has either been successfully sent, or has failed permanently after
85    /// exhausting any applicable retry policies.
86    ///
87    /// After flush()` returns, the final result of each individual publish
88    /// operation (i.e., a success with a message ID or a terminal error) will
89    /// be available on its corresponding [PublishHandle](crate::model_ext::PublishHandle).
90    ///
91    /// Messages published after `flush()` is called will be buffered for a
92    /// subsequent batch and are not included in this flush operation.
93    ///
94    /// # Example
95    ///
96    /// ```
97    /// # use google_cloud_pubsub::model::PubsubMessage;
98    /// # async fn sample(publisher: google_cloud_pubsub::client::Publisher) -> anyhow::Result<()> {
99    /// // Publish some messages. They will be buffered according to batching options.
100    /// let handle1 = publisher.publish(PubsubMessage::new().set_data("foo".to_string()));
101    /// let handle2 = publisher.publish(PubsubMessage::new().set_data("bar".to_string()));
102    ///
103    /// // Flush ensures that these messages are sent immediately and waits for
104    /// // the send to complete.
105    /// publisher.flush().await;
106    ///
107    /// // The results for handle1 and handle2 are available.
108    /// let id1 = handle1.await?;
109    /// let id2 = handle2.await?;
110    /// # Ok(())
111    /// # }
112    /// ```
113    pub async fn flush(&self) {
114        let (tx, rx) = oneshot::channel();
115        if self.tx.send(ToWorker::Flush(tx)).is_err() {
116            // `tx` is dropped here if the send errors.
117        }
118        rx.await
119            .expect("the client library should not release the sender");
120    }
121}
122
123/// Creates `Publisher`s.
124///
125/// Publishers are created via a [`Client`][crate::client::Client].
126///
127/// # Example
128///
129/// ```
130/// # async fn sample() -> anyhow::Result<()> {
131/// # use google_cloud_pubsub::*;
132/// # use builder::publisher::ClientBuilder;
133/// # use client::Client;
134/// let client = Client::builder().build().await?;
135/// let publisher = client.publisher("projects/my-project/topics/topic").build();
136/// # Ok(()) }
137/// ```
138#[derive(Clone, Debug)]
139pub struct PublisherBuilder {
140    pub(crate) inner: GapicPublisher,
141    topic: String,
142    batching_options: BatchingOptions,
143}
144
145impl PublisherBuilder {
146    /// Creates a new Pub/Sub publisher builder for topic.
147    pub(crate) fn new(client: GapicPublisher, topic: String) -> Self {
148        Self {
149            inner: client,
150            topic,
151            batching_options: BatchingOptions::default(),
152        }
153    }
154
155    /// Sets the maximum number of messages to be batched together for a single `Publish` call.
156    /// When this number is reached, the batch is sent.
157    ///
158    /// Setting this to `1` disables batching by message count.
159    ///
160    /// # Example
161    /// ```
162    /// # use google_cloud_pubsub::client::Client;
163    /// # async fn sample() -> anyhow::Result<()> {
164    /// # let client = Client::builder().build().await?;
165    /// let publisher = client.publisher("projects/my-project/topics/my-topic")
166    ///     .set_message_count_threshold(100)
167    ///     .build();
168    /// # Ok(()) }
169    /// ```
170    pub fn set_message_count_threshold(mut self, threshold: u32) -> PublisherBuilder {
171        self.batching_options = self.batching_options.set_message_count_threshold(threshold);
172        self
173    }
174
175    /// Sets the byte threshold for batching in a single `Publish` call.
176    /// When this many bytes are accumulated, the batch is sent.
177    ///
178    /// # Example
179    /// ```
180    /// # use google_cloud_pubsub::client::Client;
181    /// # async fn sample() -> anyhow::Result<()> {
182    /// # let client = Client::builder().build().await?;
183    /// let publisher = client.publisher("projects/my-project/topics/my-topic")
184    ///     .set_byte_threshold(100)
185    ///     .build();
186    /// # Ok(()) }
187    /// ```
188    pub fn set_byte_threshold(mut self, threshold: u32) -> PublisherBuilder {
189        self.batching_options = self.batching_options.set_byte_threshold(threshold);
190        self
191    }
192
193    /// Sets the maximum amount of time the publisher will wait before sending a
194    /// batch. When this delay is reached, the current batch is sent, regardless
195    /// of the number of messages or total byte size.
196    ///
197    /// # Example
198    /// ```
199    /// # use google_cloud_pubsub::client::Client;
200    /// # use std::time::Duration;
201    /// # async fn sample() -> anyhow::Result<()> {
202    /// # let client = Client::builder().build().await?;
203    /// let publisher = client.publisher("projects/my-project/topics/my-topic")
204    ///     .set_delay_threshold(Duration::from_millis(50))
205    ///     .build();
206    /// # Ok(()) }
207    /// ```
208    pub fn set_delay_threshold(mut self, threshold: Duration) -> PublisherBuilder {
209        self.batching_options = self.batching_options.set_delay_threshold(threshold);
210        self
211    }
212
213    /// Creates a new [`Publisher`] from the builder's configuration.
214    // This method starts a background task to manage the batching
215    // and sending of messages. The returned `Publisher` is a
216    // lightweight handle for sending messages to that background task
217    // over a channel.
218    pub fn build(self) -> Publisher {
219        // Enforce limits by clamping the user-provided options.
220        let batching_options = BatchingOptions::new()
221            .set_delay_threshold(
222                self.batching_options
223                    .delay_threshold
224                    .clamp(Duration::ZERO, MAX_DELAY),
225            )
226            .set_message_count_threshold(
227                self.batching_options
228                    .message_count_threshold
229                    .clamp(0, MAX_MESSAGES),
230            )
231            .set_byte_threshold(self.batching_options.byte_threshold.clamp(0, MAX_BYTES));
232
233        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
234        // Create the batching worker that will run in the background.
235        // We don't need to keep track of a handle to the worker.
236        // Dropping the Publisher will drop the only sender to the channel.
237        // This will cause worker.run() to read None from the channel and close.
238        let worker = Worker::new(self.topic, self.inner, batching_options.clone(), rx);
239        tokio::spawn(worker.run());
240
241        Publisher {
242            batching_options,
243            tx,
244        }
245    }
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251    use crate::{client::Client, publisher::options::BatchingOptions};
252    use crate::{
253        generated::gapic_dataplane::client::Publisher as GapicPublisher,
254        model::{PublishResponse, PubsubMessage},
255    };
256    use mockall::Sequence;
257
258    mockall::mock! {
259        #[derive(Debug)]
260        GapicPublisher {}
261        impl crate::generated::gapic_dataplane::stub::Publisher for GapicPublisher {
262            async fn publish(&self, req: crate::model::PublishRequest, _options: gax::options::RequestOptions) -> gax::Result<gax::response::Response<crate::model::PublishResponse>>;
263        }
264    }
265
266    // Similar to GapicPublisher but returns impl Future instead.
267    // This is useful for mocking a response with delays/timeouts.
268    // See https://github.com/asomers/mockall/issues/189 for more
269    // detail on why this is needed.
270    // While this can used inplace of GapicPublisher, it makes the
271    // normal usage without async closure much more cumbersome.
272    mockall::mock! {
273        #[derive(Debug)]
274        GapicPublisherWithFuture {}
275        impl crate::generated::gapic_dataplane::stub::Publisher for GapicPublisherWithFuture {
276            fn publish(&self, req: crate::model::PublishRequest, _options: gax::options::RequestOptions) -> impl Future<Output=gax::Result<gax::response::Response<crate::model::PublishResponse>>> + Send;
277        }
278    }
279
280    #[tokio::test]
281    async fn test_worker_success() {
282        let mut mock = MockGapicPublisher::new();
283        mock.expect_publish()
284            .returning({
285                |r, _| {
286                    assert_eq!(r.topic, "my-topic");
287                    assert_eq!(r.messages.len(), 1);
288                    let id = String::from_utf8(r.messages[0].data.to_vec()).unwrap();
289                    Ok(gax::response::Response::from(
290                        PublishResponse::new().set_message_ids(vec![id]),
291                    ))
292                }
293            })
294            .times(2);
295
296        let client = GapicPublisher::from_stub(mock);
297        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
298            .set_message_count_threshold(1_u32)
299            .build();
300
301        let messages = vec![
302            PubsubMessage::new().set_data("hello".to_string()),
303            PubsubMessage::new().set_data("world".to_string()),
304        ];
305        let mut handles = Vec::new();
306        for msg in messages {
307            let handle = publisher.publish(msg.clone());
308            handles.push((msg, handle));
309        }
310
311        for (id, rx) in handles.into_iter() {
312            let got = rx.await.expect("expected message id");
313            let id = String::from_utf8(id.data.to_vec()).unwrap();
314            assert_eq!(got, id);
315        }
316    }
317
318    #[tokio::test(start_paused = true)]
319    async fn test_drop_publisher() {
320        // If we hold on to the handles returned from the publisher, it should
321        // be safe to drop the publisher and .await on the handles.
322        let mut mock = MockGapicPublisher::new();
323        mock.expect_publish().return_once({
324            |r, _| {
325                assert_eq!(r.topic, "my-topic");
326                let ids = r
327                    .messages
328                    .iter()
329                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
330                Ok(gax::response::Response::from(
331                    PublishResponse::new().set_message_ids(ids),
332                ))
333            }
334        });
335        let client = GapicPublisher::from_stub(mock);
336        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
337            .set_message_count_threshold(1000_u32)
338            .set_delay_threshold(Duration::from_secs(60))
339            .build();
340
341        let start = tokio::time::Instant::now();
342        let messages = vec![
343            PubsubMessage::new().set_data("hello".to_string()),
344            PubsubMessage::new().set_data("world".to_string()),
345        ];
346        let mut handles = Vec::new();
347        for msg in messages {
348            let handle = publisher.publish(msg.clone());
349            handles.push((msg, handle));
350        }
351        drop(publisher); // This should trigger the batch to send, no delay.
352
353        for (id, rx) in handles.into_iter() {
354            let got = rx.await.expect("expected message id");
355            let id = String::from_utf8(id.data.to_vec()).unwrap();
356            assert_eq!(got, id);
357            assert_eq!(start.elapsed(), Duration::ZERO);
358        }
359    }
360
361    #[tokio::test]
362    async fn test_worker_error() {
363        let mut mock = MockGapicPublisher::new();
364        mock.expect_publish()
365            .returning({
366                |r, _| {
367                    assert_eq!(r.topic, "my-topic");
368                    assert_eq!(r.messages.len(), 1);
369                    Err(gax::error::Error::io("io error"))
370                }
371            })
372            .times(2);
373
374        let client = GapicPublisher::from_stub(mock);
375        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
376            .set_message_count_threshold(1_u32)
377            .build();
378
379        let messages = vec![
380            PubsubMessage::new().set_data("hello".to_string()),
381            PubsubMessage::new().set_data("world".to_string()),
382        ];
383
384        let mut handles = Vec::new();
385        for msg in messages {
386            let handle = publisher.publish(msg.clone());
387            handles.push(handle);
388        }
389
390        for rx in handles.into_iter() {
391            let got = rx.await;
392            assert!(got.is_err());
393        }
394    }
395
396    #[tokio::test(start_paused = true)]
397    async fn test_worker_flush() {
398        let mut mock = MockGapicPublisher::new();
399        mock.expect_publish().returning({
400            |r, _| {
401                assert_eq!(r.topic, "my-topic");
402                let ids = r
403                    .messages
404                    .iter()
405                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
406                Ok(gax::response::Response::from(
407                    PublishResponse::new().set_message_ids(ids),
408                ))
409            }
410        });
411
412        let client = GapicPublisher::from_stub(mock);
413        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
414            // Set a long delay.
415            .set_message_count_threshold(1000_u32)
416            .set_delay_threshold(Duration::from_secs(60))
417            .build();
418
419        let start = tokio::time::Instant::now();
420        let messages = vec![
421            PubsubMessage::new().set_data("hello".to_string()),
422            PubsubMessage::new().set_data("world".to_string()),
423        ];
424        let mut handles = Vec::new();
425        for msg in messages {
426            let handle = publisher.publish(msg.clone());
427            handles.push((msg, handle));
428        }
429
430        publisher.flush().await;
431        assert_eq!(start.elapsed(), Duration::ZERO);
432
433        let post = publisher.publish(PubsubMessage::new().set_data("after".to_string()));
434        for (id, rx) in handles.into_iter() {
435            let got = rx.await.expect("expected message id");
436            let id = String::from_utf8(id.data.to_vec()).unwrap();
437            assert_eq!(got, id);
438            assert_eq!(start.elapsed(), Duration::ZERO);
439        }
440
441        // The last message is only sent after the next timeout
442        // (worker does not continue to flush).
443        let got = post.await.expect("expected message id");
444        assert_eq!(got, "after");
445        assert_eq!(start.elapsed(), Duration::from_secs(60));
446    }
447
448    // User's should be able to drop handles and the messages will still send.
449    #[tokio::test(start_paused = true)]
450    async fn test_worker_drop_handles() {
451        let mut mock = MockGapicPublisher::new();
452        mock.expect_publish().return_once({
453            move |r, _| {
454                assert_eq!(r.topic, "my-topic");
455                let ids = r
456                    .messages
457                    .iter()
458                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
459                assert_eq!(ids.len(), 2);
460                let ids = ids.collect::<Vec<_>>();
461                assert_eq!(ids.clone(), vec!["hello", "world"]);
462                Ok(gax::response::Response::from(
463                    PublishResponse::new().set_message_ids(ids),
464                ))
465            }
466        });
467
468        let client = GapicPublisher::from_stub(mock);
469        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
470            // Set a long delay.
471            .set_message_count_threshold(1000_u32)
472            .set_delay_threshold(Duration::from_secs(60))
473            .build();
474
475        let start = tokio::time::Instant::now();
476        let messages = vec![
477            PubsubMessage::new().set_data("hello".to_string()),
478            PubsubMessage::new().set_data("world".to_string()),
479        ];
480        for msg in messages {
481            publisher.publish(msg.clone());
482        }
483
484        publisher.flush().await;
485        assert_eq!(start.elapsed(), Duration::ZERO);
486    }
487
488    #[tokio::test(start_paused = true)]
489    async fn test_empty_flush() {
490        let mock = MockGapicPublisher::new();
491
492        let client = GapicPublisher::from_stub(mock);
493        let publisher = PublisherBuilder::new(client, "my-topic".to_string()).build();
494
495        let start = tokio::time::Instant::now();
496        publisher.flush().await;
497        assert_eq!(start.elapsed(), Duration::ZERO);
498    }
499
500    #[tokio::test]
501    async fn test_batching_send_on_message_count_threshold_success() {
502        // Make sure all messages in a batch receive the correct message ID.
503        let mut mock = MockGapicPublisher::new();
504        mock.expect_publish().return_once({
505            |r, _| {
506                assert_eq!(r.topic, "my-topic");
507                assert_eq!(r.messages.len(), 2);
508                let ids = r
509                    .messages
510                    .iter()
511                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
512                Ok(gax::response::Response::from(
513                    PublishResponse::new().set_message_ids(ids),
514                ))
515            }
516        });
517
518        let client = GapicPublisher::from_stub(mock);
519        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
520            .set_message_count_threshold(2_u32)
521            .set_byte_threshold(MAX_BYTES)
522            .set_delay_threshold(std::time::Duration::MAX)
523            .build();
524
525        let messages = vec![
526            PubsubMessage::new().set_data("hello".to_string()),
527            PubsubMessage::new().set_data("world".to_string()),
528        ];
529        let mut handles = Vec::new();
530        for msg in messages {
531            let handle = publisher.publish(msg.clone());
532            handles.push((msg, handle));
533        }
534
535        for (id, rx) in handles.into_iter() {
536            let got = rx.await.expect("expected message id");
537            let id = String::from_utf8(id.data.to_vec()).unwrap();
538            assert_eq!(got, id);
539        }
540    }
541
542    #[tokio::test]
543    async fn test_batching_send_on_message_count_threshold_error() {
544        // Make sure all messages in a batch receive an error.
545        let mut mock = MockGapicPublisher::new();
546        mock.expect_publish().return_once({
547            |r, _| {
548                assert_eq!(r.topic, "my-topic");
549                assert_eq!(r.messages.len(), 2);
550                Err(gax::error::Error::io("io error"))
551            }
552        });
553
554        let client = GapicPublisher::from_stub(mock);
555        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
556            .set_message_count_threshold(2_u32)
557            .set_byte_threshold(MAX_BYTES)
558            .set_delay_threshold(std::time::Duration::MAX)
559            .build();
560
561        let messages = vec![
562            PubsubMessage::new().set_data("hello".to_string()),
563            PubsubMessage::new().set_data("world".to_string()),
564        ];
565        let mut handles = Vec::new();
566        for msg in messages {
567            let handle = publisher.publish(msg.clone());
568            handles.push(handle);
569        }
570
571        for rx in handles.into_iter() {
572            let got = rx.await;
573            assert!(got.is_err());
574        }
575    }
576
577    #[tokio::test]
578    async fn test_batching_send_on_byte_threshold() {
579        // Make sure all messages in a batch receive the correct message ID.
580        let mut mock = MockGapicPublisher::new();
581        mock.expect_publish().return_once({
582            |r, _| {
583                assert_eq!(r.topic, "my-topic");
584                assert_eq!(r.messages.len(), 2);
585                let ids = r
586                    .messages
587                    .iter()
588                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
589                Ok(gax::response::Response::from(
590                    PublishResponse::new().set_message_ids(ids),
591                ))
592            }
593        });
594
595        let client = GapicPublisher::from_stub(mock);
596        // Ensure that the first message does not pass the threshold.
597        let byte_threshold = "my-topic".len() + "hello".len() + 1;
598        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
599            .set_message_count_threshold(MAX_MESSAGES)
600            .set_byte_threshold(byte_threshold as u32)
601            .set_delay_threshold(std::time::Duration::MAX)
602            .build();
603
604        let messages = vec![
605            PubsubMessage::new().set_data("hello".to_string()),
606            PubsubMessage::new().set_data("world".to_string()),
607        ];
608        let mut handles = Vec::new();
609        for msg in messages {
610            let handle = publisher.publish(msg.clone());
611            handles.push((msg, handle));
612        }
613
614        for (id, rx) in handles.into_iter() {
615            let got = rx.await.expect("expected message id");
616            let id = String::from_utf8(id.data.to_vec()).unwrap();
617            assert_eq!(got, id);
618        }
619    }
620
621    #[tokio::test(start_paused = true)]
622    async fn test_batching_send_on_delay_threshold() {
623        let mut mock = MockGapicPublisher::new();
624        mock.expect_publish().returning({
625            |r, _| {
626                assert_eq!(r.topic, "my-topic");
627                let ids = r
628                    .messages
629                    .iter()
630                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
631                Ok(gax::response::Response::from(
632                    PublishResponse::new().set_message_ids(ids),
633                ))
634            }
635        });
636
637        let client = GapicPublisher::from_stub(mock);
638        let delay = std::time::Duration::from_millis(10);
639        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
640            .set_message_count_threshold(u32::MAX)
641            .set_byte_threshold(MAX_BYTES)
642            .set_delay_threshold(delay)
643            .build();
644
645        // Test that messages send after delay.
646        for _ in 0..3 {
647            let start = tokio::time::Instant::now();
648            let messages = vec![
649                PubsubMessage::new().set_data("hello 0".to_string()),
650                PubsubMessage::new().set_data("hello 1".to_string()),
651                PubsubMessage::new()
652                    .set_data("hello 2".to_string())
653                    .set_ordering_key("ordering key 1"),
654                PubsubMessage::new()
655                    .set_data("hello 3".to_string())
656                    .set_ordering_key("ordering key 2"),
657            ];
658            let mut handles = Vec::new();
659            for msg in messages {
660                let handle = publisher.publish(msg.clone());
661                handles.push((msg, handle));
662            }
663
664            for (id, rx) in handles.into_iter() {
665                let got = rx.await.expect("expected message id");
666                let id = String::from_utf8(id.data.to_vec()).unwrap();
667                assert_eq!(got, id);
668                assert_eq!(
669                    start.elapsed(),
670                    delay,
671                    "batch of messages should have sent after {:?}",
672                    delay
673                )
674            }
675        }
676    }
677
678    #[tokio::test(start_paused = true)]
679    #[allow(clippy::get_first)]
680    async fn test_batching_on_ordering_key() {
681        // Publish messages with different ordering key and validate that they are in different batches.
682        let mut mock = MockGapicPublisher::new();
683        mock.expect_publish().returning({
684            |r, _| {
685                assert_eq!(r.topic, "my-topic");
686                assert_eq!(r.messages.len(), 2);
687                assert_eq!(
688                    r.messages.get(0).unwrap().ordering_key,
689                    r.messages.get(1).unwrap().ordering_key
690                );
691                let ids = r
692                    .messages
693                    .iter()
694                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
695                Ok(gax::response::Response::from(
696                    PublishResponse::new().set_message_ids(ids),
697                ))
698            }
699        });
700
701        let client = GapicPublisher::from_stub(mock);
702        // Use a low message count to trigger batch sends.
703        let message_count_threshold = 2_u32;
704        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
705            .set_message_count_threshold(message_count_threshold)
706            .set_byte_threshold(MAX_BYTES)
707            .set_delay_threshold(std::time::Duration::MAX)
708            .build();
709
710        let num_ordering_keys = 3;
711        let mut messages = Vec::new();
712        // We want the number of messages to be a multiple of num_ordering_keys
713        // and message_count_threshold. Otherwise, the final batch of each
714        // ordering key may fail the message len assertion.
715        for i in 0..(2 * message_count_threshold * num_ordering_keys) {
716            messages.push(
717                PubsubMessage::new()
718                    .set_data(format!("test message {}", i))
719                    .set_ordering_key(format!("ordering key: {}", i % num_ordering_keys)),
720            );
721        }
722        let mut handles = Vec::new();
723        for msg in messages {
724            let handle = publisher.publish(msg.clone());
725            handles.push((msg, handle));
726        }
727
728        for (id, rx) in handles.into_iter() {
729            let got = rx.await.expect("expected message id");
730            let id = String::from_utf8(id.data.to_vec()).unwrap();
731            assert_eq!(got, id);
732        }
733    }
734
735    #[tokio::test(start_paused = true)]
736    #[allow(clippy::get_first)]
737    async fn test_batching_empty_ordering_key() {
738        // Publish messages with different ordering key and validate that they are in different batches.
739        let mut mock = MockGapicPublisher::new();
740        mock.expect_publish().returning({
741            |r, _| {
742                assert_eq!(r.topic, "my-topic");
743                assert_eq!(r.messages.len(), 2);
744                assert_eq!(
745                    r.messages.get(0).unwrap().ordering_key,
746                    r.messages.get(1).unwrap().ordering_key
747                );
748                let ids = r
749                    .messages
750                    .iter()
751                    .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
752                Ok(gax::response::Response::from(
753                    PublishResponse::new().set_message_ids(ids),
754                ))
755            }
756        });
757
758        let client = GapicPublisher::from_stub(mock);
759        // Use a low message count to trigger batch sends.
760        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
761            .set_message_count_threshold(2_u32)
762            .set_byte_threshold(MAX_BYTES)
763            .set_delay_threshold(std::time::Duration::MAX)
764            .build();
765
766        let messages = vec![
767            PubsubMessage::new().set_data("hello 1".to_string()),
768            PubsubMessage::new()
769                .set_data("hello 2".to_string())
770                .set_ordering_key(""),
771            PubsubMessage::new()
772                .set_data("hello 3".to_string())
773                .set_ordering_key("ordering key :1"),
774            PubsubMessage::new()
775                .set_data("hello 4".to_string())
776                .set_ordering_key("ordering key :1"),
777        ];
778
779        let mut handles = Vec::new();
780        for msg in messages {
781            let handle = publisher.publish(msg.clone());
782            handles.push((msg, handle));
783        }
784
785        for (id, rx) in handles.into_iter() {
786            let got = rx.await.expect("expected message id");
787            let id = String::from_utf8(id.data.to_vec()).unwrap();
788            assert_eq!(got, id);
789        }
790    }
791
792    #[tokio::test(start_paused = true)]
793    #[allow(clippy::get_first)]
794    async fn test_ordering_key_only_one_outstanding_batch() {
795        // Verify that Publisher must only have 1 outstanding batch inflight at a time.
796        // This is done by validating that the 2 expected publish calls are done in sequence
797        // with a sleep delay in the first Publish reply.
798        let mut seq = Sequence::new();
799        let mut mock = MockGapicPublisherWithFuture::new();
800        mock.expect_publish()
801            .times(1)
802            .in_sequence(&mut seq)
803            .returning({
804                |r, _| {
805                    Box::pin(async move {
806                        tokio::time::sleep(Duration::from_millis(10)).await;
807                        assert_eq!(r.topic, "my-topic");
808                        assert_eq!(r.messages.len(), 1);
809                        let ids = r
810                            .messages
811                            .iter()
812                            .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
813                        Ok(gax::response::Response::from(
814                            PublishResponse::new().set_message_ids(ids),
815                        ))
816                    })
817                }
818            });
819
820        mock.expect_publish()
821            .times(1)
822            .in_sequence(&mut seq)
823            .returning({
824                |r, _| {
825                    Box::pin(async move {
826                        assert_eq!(r.topic, "my-topic");
827                        assert_eq!(r.messages.len(), 1);
828                        let ids = r
829                            .messages
830                            .iter()
831                            .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
832                        Ok(gax::response::Response::from(
833                            PublishResponse::new().set_message_ids(ids),
834                        ))
835                    })
836                }
837            });
838
839        let client = GapicPublisher::from_stub(mock);
840        // Use a low message count to trigger batch sends.
841        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
842            .set_message_count_threshold(1_u32)
843            .set_byte_threshold(MAX_BYTES)
844            .set_delay_threshold(std::time::Duration::MAX)
845            .build();
846
847        let messages = [
848            PubsubMessage::new()
849                .set_data("hello 1".to_string())
850                .set_ordering_key("ordering key"),
851            PubsubMessage::new()
852                .set_data("hello 2".to_string())
853                .set_ordering_key("ordering key"),
854        ];
855
856        let start = tokio::time::Instant::now();
857        let msg1_handle = publisher.publish(messages.get(0).unwrap().clone());
858        let msg2_handle = publisher.publish(messages.get(1).unwrap().clone());
859        assert_eq!(msg2_handle.await.expect("expected message id"), "hello 2");
860        assert_eq!(
861            start.elapsed(),
862            Duration::from_millis(10),
863            "the second batch of messages should have sent after the first which is has been delayed by {:?}",
864            Duration::from_millis(10)
865        );
866        // Also validate the content of the first publish.
867        assert_eq!(msg1_handle.await.expect("expected message id"), "hello 1");
868    }
869
870    #[tokio::test(start_paused = true)]
871    #[allow(clippy::get_first)]
872    async fn test_empty_ordering_key_concurrent_batches() {
873        // Verify that for empty ordering key, the Publisher will send multiple batches without
874        // awaiting for the results.
875        // This is done by adding a delay in the first Publish reply and validating that
876        // the second batch does not await for the first batch.
877        let mut seq = Sequence::new();
878        let mut mock = MockGapicPublisherWithFuture::new();
879        mock.expect_publish()
880            .times(1)
881            .in_sequence(&mut seq)
882            .returning({
883                |r, _| {
884                    Box::pin(async move {
885                        tokio::time::sleep(Duration::from_millis(10)).await;
886                        assert_eq!(r.topic, "my-topic");
887                        assert_eq!(r.messages.len(), 1);
888                        let ids = r
889                            .messages
890                            .iter()
891                            .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
892                        Ok(gax::response::Response::from(
893                            PublishResponse::new().set_message_ids(ids),
894                        ))
895                    })
896                }
897            });
898
899        mock.expect_publish()
900            .times(1)
901            .in_sequence(&mut seq)
902            .returning({
903                |r, _| {
904                    Box::pin(async move {
905                        assert_eq!(r.topic, "my-topic");
906                        assert_eq!(r.messages.len(), 1);
907                        let ids = r
908                            .messages
909                            .iter()
910                            .map(|m| String::from_utf8(m.data.to_vec()).unwrap());
911                        Ok(gax::response::Response::from(
912                            PublishResponse::new().set_message_ids(ids),
913                        ))
914                    })
915                }
916            });
917
918        let client = GapicPublisher::from_stub(mock);
919        // Use a low message count to trigger batch sends.
920        let publisher = PublisherBuilder::new(client, "my-topic".to_string())
921            .set_message_count_threshold(1_u32)
922            .set_byte_threshold(MAX_BYTES)
923            .set_delay_threshold(std::time::Duration::MAX)
924            .build();
925
926        let messages = [
927            PubsubMessage::new()
928                .set_data("hello 1".to_string())
929                .set_ordering_key(""),
930            PubsubMessage::new()
931                .set_data("hello 2".to_string())
932                .set_ordering_key(""),
933        ];
934
935        let start = tokio::time::Instant::now();
936        let msg1_handle = publisher.publish(messages.get(0).unwrap().clone());
937        let msg2_handle = publisher.publish(messages.get(1).unwrap().clone());
938        assert_eq!(msg2_handle.await.expect("expected message id"), "hello 2");
939        assert_eq!(
940            start.elapsed(),
941            Duration::from_millis(0),
942            "the second batch of messages should have sent without any delay"
943        );
944        // Also validate the content of the first publish.
945        assert_eq!(msg1_handle.await.expect("expected message id"), "hello 1");
946    }
947
948    #[tokio::test]
949    async fn builder() -> anyhow::Result<()> {
950        let client = Client::builder().build().await?;
951        let builder = client.publisher("projects/my-project/topics/my-topic".to_string());
952        let publisher = builder.set_message_count_threshold(1_u32).build();
953        assert_eq!(publisher.batching_options.message_count_threshold, 1_u32);
954        Ok(())
955    }
956
957    #[tokio::test]
958    async fn default_batching() -> anyhow::Result<()> {
959        let client = Client::builder().build().await?;
960        let publisher = client
961            .publisher("projects/my-project/topics/my-topic".to_string())
962            .build();
963
964        assert_eq!(
965            publisher.batching_options.message_count_threshold,
966            BatchingOptions::default().message_count_threshold
967        );
968        assert_eq!(
969            publisher.batching_options.byte_threshold,
970            BatchingOptions::default().byte_threshold
971        );
972        assert_eq!(
973            publisher.batching_options.delay_threshold,
974            BatchingOptions::default().delay_threshold
975        );
976        Ok(())
977    }
978
979    #[tokio::test]
980    async fn test_builder_clamping() -> anyhow::Result<()> {
981        // Test values that are too high and should be clamped.
982        let oversized_options = BatchingOptions::new()
983            .set_delay_threshold(MAX_DELAY + Duration::from_secs(1))
984            .set_message_count_threshold(MAX_MESSAGES + 1)
985            .set_byte_threshold(MAX_BYTES + 1);
986
987        let client = Client::builder().build().await?;
988        let publisher = client
989            .publisher("projects/my-project/topics/my-topic".to_string())
990            .set_delay_threshold(oversized_options.delay_threshold)
991            .set_message_count_threshold(oversized_options.message_count_threshold)
992            .set_byte_threshold(oversized_options.byte_threshold)
993            .build();
994        let got = publisher.batching_options;
995
996        assert_eq!(got.delay_threshold, MAX_DELAY);
997        assert_eq!(got.message_count_threshold, MAX_MESSAGES);
998        assert_eq!(got.byte_threshold, MAX_BYTES);
999
1000        // Test values that are within limits and should not be changed.
1001        let normal_options = BatchingOptions::new()
1002            .set_delay_threshold(Duration::from_secs(10))
1003            .set_message_count_threshold(10_u32)
1004            .set_byte_threshold(100_u32);
1005
1006        let publisher = client
1007            .publisher("projects/my-project/topics/my-topic".to_string())
1008            .set_delay_threshold(normal_options.delay_threshold)
1009            .set_message_count_threshold(normal_options.message_count_threshold)
1010            .set_byte_threshold(normal_options.byte_threshold)
1011            .build();
1012        let got = publisher.batching_options;
1013
1014        assert_eq!(got.delay_threshold, normal_options.delay_threshold);
1015        assert_eq!(
1016            got.message_count_threshold,
1017            normal_options.message_count_threshold
1018        );
1019        assert_eq!(got.byte_threshold, normal_options.byte_threshold);
1020
1021        Ok(())
1022    }
1023}
google_cloud_pubsub/publisher/publisher.rs

google_cloud_pubsub/publisher/
publisher.rs
