google_cloud_pubsub/publisher/
options.rs1#[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 pub fn new() -> Self {
27 Self::default()
28 }
29
30 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 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 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, delay_threshold: std::time::Duration::from_millis(10),
55 }
56 }
57}
58
59use super::constants::*;
60
61#[derive(Clone, Debug, PartialEq)]
74#[non_exhaustive]
75pub struct HedgingOptions {
76 pub(crate) delay: std::time::Duration,
80 pub(crate) max_tokens: u32,
85 pub(crate) refill_ratio: f32,
89}
90
91impl HedgingOptions {
92 pub fn new() -> Self {
94 Self::default()
95 }
96
97 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 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 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}