Skip to main content

google_cloud_pubsub/publisher/
options.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
15/// Configure publisher batching behavior.
16#[derive(Clone, Debug)]
17#[non_exhaustive]
18pub struct BatchingOptions {
19    pub message_count_threshold: u32,
20    pub byte_threshold: u32,
21    pub delay_threshold: std::time::Duration,
22}
23
24impl BatchingOptions {
25    /// Create a new instance.
26    pub fn new() -> Self {
27        Self::default()
28    }
29
30    /// Set the [BatchingOptions][Self::message_count_threshold] field.
31    pub fn set_message_count_threshold<V: Into<u32>>(mut self, v: V) -> Self {
32        self.message_count_threshold = v.into();
33        self
34    }
35
36    /// Set the [BatchingOptions][Self::byte_threshold] field.
37    pub(crate) fn set_byte_threshold<V: Into<u32>>(mut self, v: V) -> Self {
38        self.byte_threshold = v.into();
39        self
40    }
41
42    /// Set the [BatchingOptions][Self::delay_threshold] field.
43    pub fn set_delay_threshold<V: Into<std::time::Duration>>(mut self, v: V) -> Self {
44        self.delay_threshold = v.into();
45        self
46    }
47}
48
49impl std::default::Default for BatchingOptions {
50    fn default() -> Self {
51        Self {
52            message_count_threshold: 100_u32,
53            byte_threshold: 1_000_000_u32, // 1 MB
54            delay_threshold: std::time::Duration::from_millis(10),
55        }
56    }
57}
58
59use super::constants::*;
60
61/// Configure publisher request hedging behavior.
62///
63/// Request hedging sends a duplicate publish request when an in-flight batch publish
64/// RPC exceeds a configured delay threshold, mitigating tail latency caused by slow backend
65/// tasks or transient network stalls.
66///
67/// Hedging uses a token bucket to rate-limit hedged RPCs. Successful publish RPCs refill
68/// fractional tokens, and sending a hedged RPC decrements 1 token.
69///
70/// Request hedging is only active for messages published without an ordering key. For
71/// ordered publishing (messages with an ordering key), hedging is disabled to preserve
72/// strict ordering guarantees.
73#[derive(Clone, Debug, PartialEq)]
74#[non_exhaustive]
75pub struct HedgingOptions {
76    /// The delay before sending a hedged request for an outstanding batch.
77    ///
78    /// Clamped between 100ms and 10s. Defaults to 1s.
79    pub(crate) delay: std::time::Duration,
80    /// The maximum number of tokens in the token bucket.
81    ///
82    /// Represents the maximum burst capacity of hedged requests. Clamped between 1 and 250.
83    /// Defaults to 50.
84    pub(crate) max_tokens: u32,
85    /// The fraction of a token added to the bucket for each successful publish RPC.
86    ///
87    /// Clamped between 0.001 and 0.2. Defaults to 0.1 (1 full token per 10 successful RPCs).
88    pub(crate) refill_ratio: f32,
89}
90
91impl HedgingOptions {
92    /// Create a new instance.
93    pub fn new() -> Self {
94        Self::default()
95    }
96
97    /// Set the delay before sending a hedged request.
98    ///
99    /// Clamped between 100ms and 10s.
100    pub fn set_delay<V: Into<std::time::Duration>>(mut self, v: V) -> Self {
101        self.delay = v.into().clamp(MIN_HEDGING_DELAY, MAX_HEDGING_DELAY);
102        self
103    }
104
105    /// Set the maximum number of tokens in the token bucket.
106    ///
107    /// Clamped between 1 and 250.
108    pub fn set_max_tokens<V: Into<u32>>(mut self, v: V) -> Self {
109        self.max_tokens = v
110            .into()
111            .clamp(MIN_HEDGING_MAX_TOKENS, MAX_HEDGING_MAX_TOKENS);
112        self
113    }
114
115    /// Set the fraction of a token refilled per successful publish RPC.
116    ///
117    /// Clamped between 0.001 and 0.2.
118    pub fn set_refill_ratio<V: Into<f32>>(mut self, v: V) -> Self {
119        let val = v.into();
120        self.refill_ratio = if val.is_nan() {
121            DEFAULT_HEDGING_REFILL_RATIO
122        } else {
123            val.clamp(MIN_HEDGING_REFILL_RATIO, MAX_HEDGING_REFILL_RATIO)
124        };
125        self
126    }
127}
128
129impl std::default::Default for HedgingOptions {
130    fn default() -> Self {
131        Self {
132            delay: DEFAULT_HEDGING_DELAY,
133            max_tokens: DEFAULT_HEDGING_MAX_TOKENS,
134            refill_ratio: DEFAULT_HEDGING_REFILL_RATIO,
135        }
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::*;
142    use std::time::Duration;
143
144    #[tokio::test]
145    async fn batching_options() -> anyhow::Result<()> {
146        let options = BatchingOptions::new()
147            .set_byte_threshold(1_234_u32)
148            .set_message_count_threshold(123_u32)
149            .set_delay_threshold(std::time::Duration::from_millis(12));
150        assert_eq!(options.byte_threshold, 1_234_u32);
151        assert_eq!(options.message_count_threshold, 123_u32);
152        assert_eq!(
153            options.delay_threshold,
154            std::time::Duration::from_millis(12)
155        );
156        Ok(())
157    }
158
159    #[test]
160    fn hedging_options_defaults_and_builder() {
161        let default_opts = HedgingOptions::default();
162        assert_eq!(default_opts.delay, Duration::from_secs(1));
163        assert_eq!(default_opts.max_tokens, 50);
164        assert_eq!(default_opts.refill_ratio, 0.1);
165        assert_eq!(HedgingOptions::new(), default_opts);
166
167        let custom_opts = HedgingOptions::new()
168            .set_delay(Duration::from_millis(500))
169            .set_max_tokens(100_u32)
170            .set_refill_ratio(0.05_f32);
171        assert_eq!(custom_opts.delay, Duration::from_millis(500));
172        assert_eq!(custom_opts.max_tokens, 100);
173        assert_eq!(custom_opts.refill_ratio, 0.05);
174    }
175
176    #[test]
177    fn hedging_options_clamps_values() {
178        let under_opts = HedgingOptions::default()
179            .set_delay(Duration::from_millis(10))
180            .set_max_tokens(0_u32)
181            .set_refill_ratio(0.0001_f32);
182        assert_eq!(under_opts.delay, MIN_HEDGING_DELAY);
183        assert_eq!(under_opts.max_tokens, MIN_HEDGING_MAX_TOKENS);
184        assert_eq!(under_opts.refill_ratio, MIN_HEDGING_REFILL_RATIO);
185
186        let over_opts = HedgingOptions::default()
187            .set_delay(Duration::from_secs(60))
188            .set_max_tokens(500_u32)
189            .set_refill_ratio(0.5_f32);
190        assert_eq!(over_opts.delay, MAX_HEDGING_DELAY);
191        assert_eq!(over_opts.max_tokens, MAX_HEDGING_MAX_TOKENS);
192        assert_eq!(over_opts.refill_ratio, MAX_HEDGING_REFILL_RATIO);
193    }
194
195    #[test_case::test_case(f32::MAX, MAX_HEDGING_REFILL_RATIO)]
196    #[test_case::test_case(f32::MIN, MIN_HEDGING_REFILL_RATIO)]
197    #[test_case::test_case(f32::NAN, DEFAULT_HEDGING_REFILL_RATIO)]
198    fn refill_ratio_clamps_values(val: f32, want: f32) {
199        let opts = HedgingOptions::default().set_refill_ratio(val);
200        assert_eq!(opts.refill_ratio, want);
201    }
202}