#
google_cloud_pubsub/publisher/
batch.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 tokio::task::JoinSet;
    16
    17use crate::generated::gapic_dataplane::client::Publisher as GapicPublisher;
    18use crate::publisher::worker::BundledMessage;
    19use std::sync::Arc;
    20
    21#[derive(Debug, Default)]
    22pub(crate) struct Batch {
    23    messages: Vec<BundledMessage>,
    24    initial_size: u32,
    25    messages_byte_size: u32,
    26}
    27
    28impl Batch {
    29    pub(crate) fn new(initial_size: u32) -> Self {
    30        Batch {
    31            initial_size,
    32            messages_byte_size: initial_size,
    33            ..Batch::default()
    34        }
    35    }
    36
    37    pub(crate) fn is_empty(&self) -> bool {
    38        self.messages.is_empty()
    39    }
    40
    41    pub(crate) fn len(&self) -> usize {
    42        self.messages.len()
    43    }
    44
    45    pub(crate) fn size(&self) -> u32 {
    46        self.messages_byte_size
    47    }
    48
    49    pub(crate) fn push(&mut self, msg: BundledMessage) {
    50        self.messages_byte_size += Self::message_size(&msg.msg) as u32;
    51        self.messages.push(msg);
    52    }
    53
    54    fn message_size(msg: &crate::model::PubsubMessage) -> usize {
    55        // This is only an estimate and not the wire length.
    56        // TODO(#3963): If we move on to use protobuf crate, then it may be
    57        // possible to use compute_size to find the wire length.
    58        msg.attributes
    59            .iter()
    60            .fold(msg.data.len() + msg.ordering_key.len(), |acc, (k, v)| {
    61                acc + k.len() + v.len()
    62            })
    63    }
    64
    65    /// Drains the batch and spawns a task to send the messages.
    66    ///
    67    /// This method mutably drains the messages from the current batch, leaving it
    68    /// empty, and returns a `JoinHandle` for the spawned send operation. This allows
    69    /// the `Worker` to immediately begin creating a new batch while the old one is
    70    /// being sent in the background.
    71    pub(crate) fn flush(
    72        &mut self,
    73        client: GapicPublisher,
    74        topic: String,
    75        inflight: &mut JoinSet<()>,
    76    ) {
    77        let batch_to_send = Self {
    78            initial_size: self.initial_size,
    79            messages: self.messages.drain(..).collect(),
    80            messages_byte_size: self.messages_byte_size,
    81        };
    82        self.messages_byte_size = self.initial_size;
    83        inflight.spawn(batch_to_send.send(client, topic));
    84    }
    85
    86    /// Send the batch to the service and process the results.
    87    async fn send(self, client: GapicPublisher, topic: String) {
    88        let (msgs, txs): (Vec<_>, Vec<_>) = self
    89            .messages
    90            .into_iter()
    91            .map(|msg| (msg.msg, msg.tx))
    92            .unzip();
    93        let request = client.publish().set_topic(topic).set_messages(msgs);
    94
    95        // Handle the response by extracting the message ID on success.
    96        match request.send().await {
    97            Err(e) => {
    98                // TODO(#4013): To support message ordering retry, we need to correctly handle
    99                // the send error here with either retry or propagate to the user.
    100                let e = Arc::new(e);
    101                for tx in txs {
    102                    // The user may have dropped the handle, so it is ok if this fails.
    103                    // TODO(#3689): The error type for this is incorrect, will need to handle
    104                    // this error propagation more fully.
    105                    let _ = tx.send(Err(gax::error::Error::io(e.clone())));
    106                }
    107            }
    108            Ok(result) => {
    109                txs.into_iter()
    110                    .zip(result.message_ids.into_iter())
    111                    .for_each(|(tx, result)| {
    112                        // The user may have dropped the handle, so it is ok if this fails.
    113                        let _ = tx.send(Ok(result));
    114                    });
    115            }
    116        }
    117    }
    118}
    119
    120#[cfg(test)]
    121mod tests {
    122    use crate::{
    123        generated::gapic_dataplane::client::Publisher as GapicPublisher,
    124        model::{PublishResponse, PubsubMessage},
    125        publisher::batch::Batch,
    126        publisher::worker::BundledMessage,
    127    };
    128    use tokio::task::JoinSet;
    129
    130    mockall::mock! {
    131        #[derive(Debug)]
    132        GapicPublisher {}
    133        impl crate::generated::gapic_dataplane::stub::Publisher for GapicPublisher {
    134            async fn publish(&self, req: crate::model::PublishRequest, _options: gax::options::RequestOptions) -> gax::Result<gax::response::Response<crate::model::PublishResponse>>;
    135        }
    136    }
    137
    138    #[tokio::test]
    139    async fn test_push_and_flush_batch() {
    140        let mut batch = Batch::new("topic".len() as u32);
    141        assert!(batch.is_empty());
    142
    143        let (message_a, _rx_a) = create_bundled_message_from_bytes("hello");
    144        batch.push(message_a);
    145        assert_eq!(batch.len(), 1);
    146
    147        let (message_b, _rx_b) = create_bundled_message_from_bytes(", ");
    148        batch.push(message_b);
    149        assert_eq!(batch.len(), 2);
    150
    151        let (message_c, _rx_c) = create_bundled_message_from_bytes("world");
    152        batch.push(message_c);
    153        assert_eq!(batch.len(), 3);
    154
    155        let mut mock = MockGapicPublisher::new();
    156        mock.expect_publish().return_once({
    157            |r, _| {
    158                assert_eq!(r.topic, "topic");
    159                assert_eq!(r.messages.len(), 3);
    160                Ok(gax::response::Response::from(PublishResponse::new()))
    161            }
    162        });
    163        let client = GapicPublisher::from_stub(mock);
    164        let mut inflight = JoinSet::new();
    165        batch.flush(client, "topic".to_string(), &mut inflight);
    166        assert_eq!(batch.len(), 0);
    167    }
    168
    169    #[tokio::test]
    170    async fn test_size() {
    171        use std::collections::HashMap;
    172
    173        let topic = "topic";
    174        let mut batch: Batch = Batch::new(topic.len() as u32);
    175        let mut expected_encoded_len = topic.len();
    176        assert_eq!(batch.size(), expected_encoded_len as u32);
    177
    178        let (message_data_only, _rx) = create_bundled_message_from_bytes("message_data_only");
    179        expected_encoded_len += message_data_only.msg.data.len();
    180        batch.push(message_data_only);
    181        assert_eq!(batch.size(), expected_encoded_len as u32);
    182
    183        let (message_with_ordering, _rx) = create_bundled_message_from_pubsub_message(
    184            PubsubMessage::new().set_ordering_key("ordering_key"),
    185        );
    186        expected_encoded_len += message_with_ordering.msg.ordering_key.len();
    187        batch.push(message_with_ordering);
    188        assert_eq!(batch.size(), expected_encoded_len as u32);
    189
    190        let attributes = HashMap::from([("k1", "v1"), ("key2", "value2")]);
    191        let (message_with_attributes, _rx) = create_bundled_message_from_pubsub_message(
    192            PubsubMessage::new().set_attributes(attributes),
    193        );
    194        expected_encoded_len += 14;
    195        batch.push(message_with_attributes);
    196        assert_eq!(batch.size(), expected_encoded_len as u32);
    197
    198        let mut mock = MockGapicPublisher::new();
    199        mock.expect_publish().return_once({
    200            |r, _| {
    201                assert_eq!(r.topic, "topic");
    202                assert_eq!(r.messages.len(), 3);
    203                Ok(gax::response::Response::from(PublishResponse::new()))
    204            }
    205        });
    206        let client = GapicPublisher::from_stub(mock);
    207        let mut inflight = JoinSet::new();
    208        batch.flush(client, "topic".to_string(), &mut inflight);
    209        assert_eq!(batch.size(), "topic".len() as u32);
    210    }
    211
    212    fn create_bundled_message_from_bytes<T: Into<::bytes::Bytes>>(
    213        data: T,
    214    ) -> (
    215        BundledMessage,
    216        tokio::sync::oneshot::Receiver<crate::Result<String>>,
    217    ) {
    218        create_bundled_message_from_pubsub_message(PubsubMessage::new().set_data(data.into()))
    219    }
    220
    221    fn create_bundled_message_from_pubsub_message(
    222        msg: PubsubMessage,
    223    ) -> (
    224        BundledMessage,
    225        tokio::sync::oneshot::Receiver<crate::Result<String>>,
    226    ) {
    227        let (tx, rx) = tokio::sync::oneshot::channel();
    228        (BundledMessage { tx, msg }, rx)
    229    }
    230}