#
google_cloud_pubsub/publisher/
worker.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::batch::Batch;
    18use std::collections::{HashMap, VecDeque};
    19use tokio::sync::{mpsc, oneshot};
    20use tokio::task::JoinSet;
    21
    22/// A command sent from the `Publisher` to the background `Worker`.
    23pub(crate) enum ToWorker {
    24    /// A request to publish a single message.
    25    Publish(BundledMessage),
    26    /// A request to flush all outstanding messages.
    27    Flush(oneshot::Sender<()>),
    28    // TODO(#4015): Add a resume function to allow resume Publishing on a ordering key after a
    29    // failure.
    30}
    31
    32/// A command sent from the `Worker` to the background `batch` worker.
    33pub(crate) enum ToBatchWorker {
    34    /// A request to publish a single message.
    35    Publish(BundledMessage),
    36    /// A request to flush all outstanding messages.
    37    Flush(oneshot::Sender<()>),
    38    // TODO(#4015): Add a resume function to allow resume Publishing on a ordering key after a
    39    // failure.
    40}
    41
    42/// Object that is passed to the worker task over the
    43/// main channel. This represents a single message and the sender
    44/// half of the channel to resolve the [PublishHandle].
    45#[derive(Debug)]
    46pub(crate) struct BundledMessage {
    47    pub msg: crate::model::PubsubMessage,
    48    pub tx: oneshot::Sender<crate::Result<String>>,
    49}
    50
    51/// The worker is spawned in a background task and handles
    52/// batching and publishing all messages that are sent to the publisher.
    53#[derive(Debug)]
    54pub(crate) struct Worker {
    55    topic_name: String,
    56    client: GapicPublisher,
    57    #[allow(dead_code)]
    58    batching_options: BatchingOptions,
    59    rx: mpsc::UnboundedReceiver<ToWorker>,
    60}
    61
    62impl Worker {
    63    pub(crate) fn new(
    64        topic_name: String,
    65        client: GapicPublisher,
    66        batching_options: BatchingOptions,
    67        rx: mpsc::UnboundedReceiver<ToWorker>,
    68    ) -> Self {
    69        Self {
    70            topic_name,
    71            client,
    72            rx,
    73            batching_options,
    74        }
    75    }
    76
    77    /// The main loop of the background worker.
    78    ///
    79    /// This method concurrently handles four main events:
    80    ///
    81    /// 1. Publish command from the `Publisher` is demultiplexed to the batch worker
    82    ///    for its ordering key.
    83    /// 2. A Flush command from the `Publisher` causes all batch workers to flush
    84    ///    all pending messages.
    85    /// 2. A timer fire causes the Worker to flush all pending batches.
    86    ///
    87    /// The loop terminates when the `rx` channel is closed, which happens when all
    88    /// `Publisher` clones have been dropped.
    89    pub(crate) async fn run(mut self) {
    90        // A dictionary of ordering key to outstanding publish operations.
    91        // We batch publish operations on the same ordering key together.
    92        // Publish without ordering keys are treated as having the key "".
    93        // TODO(#4012): Remove batch workers when there are no outstanding operations on the ordering key.
    94        let mut batch_workers: HashMap<String, mpsc::UnboundedSender<ToBatchWorker>> =
    95            HashMap::new();
    96        let delay = self.batching_options.delay_threshold;
    97        let batch_worker_error_msg = "Batch worker should not close the channel";
    98
    99        let timer = tokio::time::sleep(delay);
    100        // Pin the timer to the stack.
    101        tokio::pin!(timer);
    102        loop {
    103            tokio::select! {
    104                // Currently, the batch worker periodically flush on a shared timer. If needed,
    105                // this can be moved into the batch worker such that each are running on a
    106                // separate timer.
    107                _ = &mut timer => {
    108                    for (_, batch_worker) in batch_workers.iter_mut() {
    109                        let (tx, _) = oneshot::channel();
    110                        batch_worker
    111                            .send(ToBatchWorker::Flush(tx))
    112                            .expect(batch_worker_error_msg);
    113                    }
    114                    timer.as_mut().reset(tokio::time::Instant::now() + delay);
    115                }
    116                // Handle receiving a message from the channel.
    117                msg = self.rx.recv() => {
    118                    match msg {
    119                        Some(ToWorker::Publish(msg)) => {
    120                            let ordering_key = msg.msg.ordering_key.clone();
    121                            let batch_worker =
    122                                batch_workers
    123                                    .entry(ordering_key.clone())
    124                                    .or_insert({
    125                                        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
    126                                        let batch_worker = BatchWorker::new(
    127                                                        self.topic_name.clone(),
    128                                                        ordering_key,
    129                                                        self.client.clone(),
    130                                                        self.batching_options.clone(),
    131                                                        rx,
    132                                                );
    133                                        tokio::spawn(batch_worker.run());
    134                                        tx
    135                                });
    136                            batch_worker
    137                                .send(ToBatchWorker::Publish(msg))
    138                                .expect(batch_worker_error_msg);
    139                        },
    140                        Some(ToWorker::Flush(tx)) => {
    141                            let mut flush_set = JoinSet::new();
    142                            for (_, batch_worker) in batch_workers.iter_mut() {
    143                                let (tx, rx) = oneshot::channel();
    144                                batch_worker
    145                                    .send(ToBatchWorker::Flush(tx))
    146                                    .expect(batch_worker_error_msg);
    147                                flush_set.spawn(rx);
    148                            }
    149                            // Wait on all the tasks that exist right now.
    150                            // We could instead tokio::spawn this as well so the publisher
    151                            // can keep working on additional messages. The worker would
    152                            // also need to keep track of any pending flushes, and make sure
    153                            // all of those resolve as well.
    154                            flush_set.join_all().await;
    155                            let _ = tx.send(());
    156                        },
    157                        None => {
    158                            // The sender has been dropped send batch and stop running.
    159                            // This isn't guaranteed to execute if a user does not .await on the
    160                            // corresponding PublishHandles for the batch and the program ends.
    161                            let mut flush_set = JoinSet::new();
    162                            for (_, batch_worker) in batch_workers.iter_mut() {
    163                                let (tx, rx) = oneshot::channel();
    164                                batch_worker
    165                                    .send(ToBatchWorker::Flush(tx))
    166                                    .expect(batch_worker_error_msg);
    167                                flush_set.spawn(rx);
    168                            }
    169                            flush_set.join_all().await;
    170                            break;
    171                        }
    172                    }
    173                }
    174
    175            }
    176        }
    177    }
    178}
    179
    180/// A background worker that continuously handles Publisher commands for a specific ordering key.
    181#[derive(Debug)]
    182pub(crate) struct BatchWorker {
    183    topic: String,
    184    ordering_key: String,
    185    client: GapicPublisher,
    186    batching_options: BatchingOptions,
    187    rx: mpsc::UnboundedReceiver<ToBatchWorker>,
    188    pending_batch: Batch,
    189    pending_msgs: VecDeque<BundledMessage>,
    190}
    191
    192impl BatchWorker {
    193    pub(crate) fn new(
    194        topic: String,
    195        ordering_key: String,
    196        client: GapicPublisher,
    197        batching_options: BatchingOptions,
    198        rx: mpsc::UnboundedReceiver<ToBatchWorker>,
    199    ) -> Self {
    200        BatchWorker {
    201            pending_batch: Batch::new(topic.len() as u32),
    202            topic,
    203            ordering_key,
    204            client,
    205            batching_options,
    206            rx,
    207            pending_msgs: VecDeque::new(),
    208        }
    209    }
    210
    211    fn at_batch_threshold(&mut self) -> bool {
    212        // TODO(#4012): When the message increases the batch size beyond pubsub server acceptable
    213        // message size (10MB), we should flush existing batch first.
    214        self.pending_batch.len() as u32 >= self.batching_options.message_count_threshold
    215            || self.pending_batch.size() >= self.batching_options.byte_threshold
    216    }
    217
    218    // Move pending messages to the pending batch respecting batch thresholds.
    219    pub(crate) fn move_to_batch(&mut self) {
    220        while let Some(publish) = self.pending_msgs.pop_front() {
    221            self.pending_batch.push(publish);
    222            if self.at_batch_threshold() {
    223                break;
    224            }
    225        }
    226    }
    227
    228    /// The main loop of the batch worker.
    229    ///
    230    /// This method concurrently handles the following events:
    231    ///
    232    /// 1. A Publish command from the `Worker` causes the new message to be
    233    ///    added a pending message queue. If there is enough message to create
    234    ///    a full batch, then also flush the batch.
    235    /// 2. A `Flush` command from the `Publisher` causes all the pending messages
    236    ///    to be flushed respecting the configured batch size and message ordering.
    237    /// 4. A `inflight` batch completion causes the next batch to send if it satisfies
    238    ///    the configured batch threshold.
    239    ///
    240    /// The loop terminates when the `rx` channel is closed, which happens when the
    241    /// `Worker` drops the Sender.
    242    pub(crate) async fn run(mut self) {
    243        // While it is possible to use Some(JoinHandle) here as there is at max
    244        // a single inflight task at any given time, the use of JoinSet
    245        // simplify the managing the inflight JoinHandle.
    246        let mut inflight = JoinSet::new();
    247        // For empty ordering key, we do not need to limit the number of inflight batches.
    248        let unlimited_inflight = self.ordering_key.is_empty();
    249        loop {
    250            tokio::select! {
    251                _ = inflight.join_next(), if !inflight.is_empty() => {
    252                    if unlimited_inflight {
    253                        // When there are no limits to the number of inflight, we can skip sending the
    254                        // next batch as it is handled by the Publish and Flush operations.
    255                        continue;
    256                    }
    257                    self.move_to_batch();
    258                    if self.at_batch_threshold() {
    259                        self.pending_batch.flush(self.client.clone(), self.topic.clone(), &mut inflight);
    260                    }
    261                }
    262                msg = self.rx.recv() => {
    263                    match msg {
    264                        Some(ToBatchWorker::Publish(msg)) => {
    265                            self.pending_msgs.push_back(msg);
    266                            self.move_to_batch();
    267                            if self.at_batch_threshold() && (unlimited_inflight || inflight.is_empty()) {
    268                                self.pending_batch.flush(self.client.clone(), self.topic.clone(), &mut inflight);
    269                            }
    270                        },
    271                        Some(ToBatchWorker::Flush(tx)) => {
    272                            if unlimited_inflight {
    273                                // Send all the batches concurrently.
    274                                while !self.pending_batch.is_empty() || !self.pending_msgs.is_empty() {
    275                                    self.move_to_batch();
    276                                    self.pending_batch.flush(self.client.clone(), self.topic.clone(), &mut inflight);
    277                                }
    278                                inflight.join_all().await;
    279                                inflight = JoinSet::new();
    280                            } else {
    281                                // Send batches sequentially.
    282                                if !inflight.is_empty() {
    283                                    inflight.join_next().await;
    284                                }
    285                                while !self.pending_batch.is_empty() || !self.pending_msgs.is_empty() {
    286                                    self.move_to_batch();
    287                                    self.pending_batch.flush(self.client.clone(), self.topic.clone(), &mut inflight);
    288                                    inflight.join_next().await;
    289                                }
    290                            }
    291                            let _ = tx.send(());
    292                        },
    293                        None => {
    294                            break;
    295                        }
    296                    }
    297                }
    298            }
    299        }
    300    }
    301}