1use crate::constants::*;
7
8#[derive(Debug, Clone)]
10pub struct AdaptiveTtlConfig {
11 pub base_ttl: u8,
13 pub ttl_per_100_devices: u8,
15 pub max_ttl: u8,
17 pub min_ttl: u8,
19 pub congestion_boost: u8,
21 pub small_network_threshold: usize,
23 pub enabled: bool,
25}
26
27impl Default for AdaptiveTtlConfig {
28 fn default() -> Self {
29 Self {
30 base_ttl: ADAPTIVE_TTL_BASE,
31 ttl_per_100_devices: ADAPTIVE_TTL_PER_100_DEVICES,
32 max_ttl: ADAPTIVE_TTL_MAX,
33 min_ttl: ADAPTIVE_TTL_MIN,
34 congestion_boost: ADAPTIVE_TTL_CONGESTION_BOOST,
35 small_network_threshold: ADAPTIVE_TTL_SMALL_NETWORK_THRESHOLD,
36 enabled: true,
37 }
38 }
39}
40
41#[derive(Debug, Clone, Default)]
43pub struct NetworkSizeEstimate {
44 pub connected_peers: usize,
46 pub visible_peers: usize,
48 pub estimated_total: Option<usize>,
50 pub congestion_level: f32,
52 pub has_delivery_failures: bool,
54}
55
56impl NetworkSizeEstimate {
57 pub fn new(connected_peers: usize, visible_peers: usize) -> Self {
59 Self {
60 connected_peers,
61 visible_peers,
62 estimated_total: None,
63 congestion_level: 0.0,
64 has_delivery_failures: false,
65 }
66 }
67
68 pub fn with_estimated_total(mut self, total: usize) -> Self {
70 self.estimated_total = Some(total);
71 self
72 }
73
74 pub fn with_congestion(mut self, level: f32) -> Self {
76 self.congestion_level = level.clamp(0.0, 1.0);
77 self
78 }
79
80 pub fn with_delivery_failures(mut self) -> Self {
82 self.has_delivery_failures = true;
83 self
84 }
85
86 fn estimated_size(&self) -> usize {
93 if let Some(total) = self.estimated_total {
94 return total;
95 }
96
97 let min_estimate = self.visible_peers.max(self.connected_peers);
101 let max_estimate = self.visible_peers.saturating_mul(4);
102
103 ((min_estimate as f64 * max_estimate as f64).sqrt() as usize).max(min_estimate)
105 }
106}
107
108pub struct AdaptiveTtlCalculator {
110 config: AdaptiveTtlConfig,
111}
112
113impl AdaptiveTtlCalculator {
114 pub fn new() -> Self {
116 Self::with_config(AdaptiveTtlConfig::default())
117 }
118
119 pub fn with_config(config: AdaptiveTtlConfig) -> Self {
121 Self { config }
122 }
123
124 pub fn compute_ttl(
136 &self,
137 network_estimate: &NetworkSizeEstimate,
138 is_priority_message: bool,
139 is_queued: bool,
140 ) -> u8 {
141 if !self.config.enabled {
142 return self.config.base_ttl;
143 }
144
145 let estimated_size = network_estimate.estimated_size();
146
147 let mut ttl = self.config.base_ttl;
149
150 if estimated_size > self.config.small_network_threshold {
155 let extra_hundreds = (estimated_size - self.config.small_network_threshold) / 100;
156 let additional_ttl = u8::try_from(extra_hundreds)
157 .unwrap_or(u8::MAX)
158 .saturating_mul(self.config.ttl_per_100_devices);
159 ttl = ttl.saturating_add(additional_ttl);
160 }
161
162 if is_queued {
164 ttl = ttl.saturating_add(self.config.congestion_boost);
165 }
166
167 if network_estimate.has_delivery_failures {
169 ttl = ttl.saturating_add(1);
170 }
171
172 if is_priority_message {
174 ttl = ttl.saturating_add(2);
175 }
176
177 if network_estimate.congestion_level > 0.7 && !is_priority_message {
179 ttl = ttl.saturating_sub(1);
180 }
181
182 ttl.clamp(self.config.min_ttl, self.config.max_ttl)
184 }
185
186 pub fn compute_reply_ttl(
191 &self,
192 incoming_hop_count: u8,
193 network_estimate: &NetworkSizeEstimate,
194 ) -> u8 {
195 let min_reply_ttl = incoming_hop_count.saturating_add(2); let normal_ttl = self.compute_ttl(network_estimate, false, false);
200
201 min_reply_ttl.max(normal_ttl).min(self.config.max_ttl)
203 }
204
205 pub fn recommend_ttl_boost(
209 &self,
210 current_ttl: u8,
211 network_estimate: &NetworkSizeEstimate,
212 ) -> u8 {
213 let optimal_ttl = self.compute_ttl(network_estimate, false, false);
214
215 if current_ttl >= optimal_ttl {
216 0
217 } else {
218 optimal_ttl.saturating_sub(current_ttl)
219 }
220 }
221
222 pub fn config(&self) -> &AdaptiveTtlConfig {
224 &self.config
225 }
226}
227
228impl Default for AdaptiveTtlCalculator {
229 fn default() -> Self {
230 Self::new()
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use super::*;
237
238 #[test]
239 fn test_base_ttl_for_small_network() {
240 let calc = AdaptiveTtlCalculator::new();
241 let estimate = NetworkSizeEstimate::new(4, 10);
242
243 let ttl = calc.compute_ttl(&estimate, false, false);
244 assert_eq!(ttl, 8); }
246
247 #[test]
248 fn test_increased_ttl_for_large_network() {
249 let calc = AdaptiveTtlCalculator::new();
250 let estimate = NetworkSizeEstimate::new(4, 100).with_estimated_total(500);
251
252 let ttl = calc.compute_ttl(&estimate, false, false);
253 assert!(ttl > 8);
256 assert!(ttl <= 24);
257 }
258
259 #[test]
260 fn test_ttl_boost_for_queued_message() {
261 let calc = AdaptiveTtlCalculator::new();
262 let estimate = NetworkSizeEstimate::new(4, 20);
263
264 let normal_ttl = calc.compute_ttl(&estimate, false, false);
265 let queued_ttl = calc.compute_ttl(&estimate, false, true);
266
267 assert!(queued_ttl > normal_ttl);
268 }
269
270 #[test]
271 fn test_ttl_boost_for_priority_message() {
272 let calc = AdaptiveTtlCalculator::new();
273 let estimate = NetworkSizeEstimate::new(4, 20);
274
275 let normal_ttl = calc.compute_ttl(&estimate, false, false);
276 let priority_ttl = calc.compute_ttl(&estimate, true, false);
277
278 assert!(priority_ttl > normal_ttl);
279 }
280
281 #[test]
282 fn test_ttl_reduction_in_congestion() {
283 let calc = AdaptiveTtlCalculator::new();
284 let normal_estimate = NetworkSizeEstimate::new(4, 20);
285 let congested_estimate = NetworkSizeEstimate::new(4, 20).with_congestion(0.8);
286
287 let normal_ttl = calc.compute_ttl(&normal_estimate, false, false);
288 let congested_ttl = calc.compute_ttl(&congested_estimate, false, false);
289
290 assert!(congested_ttl <= normal_ttl);
291 }
292
293 #[test]
294 fn test_ttl_clamped_to_max() {
295 let calc = AdaptiveTtlCalculator::new();
296 let huge_estimate = NetworkSizeEstimate::new(4, 1000).with_estimated_total(10000);
297
298 let ttl = calc.compute_ttl(&huge_estimate, true, true);
299 assert!(ttl <= 24); }
301
302 #[test]
303 fn test_ttl_no_wrap_at_u8_boundary() {
304 let calc = AdaptiveTtlCalculator::new();
305
306 let boundary = NetworkSizeEstimate::new(4, 100).with_estimated_total(25_650);
310 assert_eq!(calc.compute_ttl(&boundary, false, false), 24);
311
312 let huge = NetworkSizeEstimate::new(4, 100).with_estimated_total(10_000_000);
314 assert_eq!(calc.compute_ttl(&huge, false, false), 24);
315 }
316
317 #[test]
318 fn test_ttl_clamped_to_min() {
319 let config = AdaptiveTtlConfig {
320 base_ttl: 2, ..Default::default()
322 };
323 let calc = AdaptiveTtlCalculator::with_config(config);
324 let estimate = NetworkSizeEstimate::new(2, 5).with_congestion(0.9);
325
326 let ttl = calc.compute_ttl(&estimate, false, false);
327 assert!(ttl >= 4); }
329
330 #[test]
331 fn test_reply_ttl_matches_hop_count() {
332 let calc = AdaptiveTtlCalculator::new();
333 let estimate = NetworkSizeEstimate::new(4, 20);
334
335 let reply_ttl = calc.compute_reply_ttl(6, &estimate);
337 assert!(reply_ttl >= 8);
338 }
339
340 #[test]
341 fn test_disabled_adaptive_ttl() {
342 let config = AdaptiveTtlConfig {
343 enabled: false,
344 base_ttl: 10,
345 ..Default::default()
346 };
347 let calc = AdaptiveTtlCalculator::with_config(config);
348
349 let estimate = NetworkSizeEstimate::new(4, 1000).with_estimated_total(5000);
350 let ttl = calc.compute_ttl(&estimate, true, true);
351
352 assert_eq!(ttl, 10); }
354
355 #[test]
356 fn test_network_size_estimation_heuristic() {
357 let estimate = NetworkSizeEstimate::new(4, 50);
358 let size = estimate.estimated_size();
359
360 assert!(size >= 50);
362 assert!(size <= 200);
363 }
364}