offline_protocol_reliability/
ack_optimization.rs

1//! ACK optimization to prevent feedback loop amplification in large networks.
2//!
3//! This module implements several techniques to reduce ACK-related network load:
4//! - Probabilistic ACK relaying: not all nodes relay every ACK
5//! - ACK aggregation: batch multiple ACKs into single messages
6//! - Piggyback ACKs: attach ACKs to data messages going the same direction
7
8use crate::constants::{
9    ESTIMATED_ACK_SIZE_BYTES, RELAY_DENSITY_FACTOR_NUMERATOR, SMALL_NETWORK_RELAY_THRESHOLD,
10};
11use chrono::{DateTime, Utc};
12use offline_protocol_core::MessageId;
13use std::collections::hash_map::DefaultHasher;
14use std::collections::HashMap;
15use std::hash::{Hash, Hasher};
16
17/// Configuration for ACK optimization.
18#[derive(Debug, Clone)]
19pub struct AckOptimizationConfig {
20    /// Enable probabilistic ACK relaying.
21    pub probabilistic_relay_enabled: bool,
22    /// Base probability for relaying an ACK (0.0-1.0).
23    pub base_relay_probability: f32,
24    /// Minimum relay probability (ensures some ACKs always propagate).
25    pub min_relay_probability: f32,
26    /// Enable ACK aggregation (batching).
27    pub aggregation_enabled: bool,
28    /// Maximum number of ACKs to aggregate before flushing.
29    pub max_aggregation_batch: usize,
30    /// Maximum time to wait for more ACKs before flushing (milliseconds).
31    pub aggregation_timeout_ms: u64,
32    /// Enable piggybacking ACKs on data messages.
33    pub piggyback_enabled: bool,
34    /// Maximum ACK data size for piggybacking (bytes).
35    pub max_piggyback_size: usize,
36}
37
38impl Default for AckOptimizationConfig {
39    fn default() -> Self {
40        Self {
41            probabilistic_relay_enabled: true,
42            base_relay_probability: 0.5, // Only 50% of nodes relay each ACK
43            min_relay_probability: 0.2,  // At least 20% always relay
44            aggregation_enabled: true,
45            max_aggregation_batch: 10,
46            aggregation_timeout_ms: 500,
47            piggyback_enabled: true,
48            max_piggyback_size: 64, // Keep piggyback data small
49        }
50    }
51}
52
53/// Aggregated ACK entry ready for transmission.
54#[derive(Debug, Clone)]
55pub struct AggregatedAck {
56    /// Original message ID this ACK is for.
57    pub original_message_id: MessageId,
58    /// Hop count when ACK was generated.
59    pub hop_count: u8,
60    /// Timestamp when ACK was received.
61    pub received_at: DateTime<Utc>,
62}
63
64/// Pending piggyback ACK data.
65#[derive(Debug, Clone)]
66pub struct PiggybackAckData {
67    /// Destination (sender of original message).
68    pub destination: String,
69    /// ACKs to piggyback on messages to this destination.
70    pub acks: Vec<AggregatedAck>,
71    /// When the first ACK was added.
72    pub first_added_at: DateTime<Utc>,
73}
74
75/// ACK optimizer for reducing network load from acknowledgments.
76pub struct AckOptimizer {
77    config: AckOptimizationConfig,
78    /// Local device ID for deterministic randomness.
79    local_device_id: String,
80    /// Pending ACKs waiting for aggregation, keyed by destination.
81    pending_acks: HashMap<String, Vec<AggregatedAck>>,
82    /// Timestamp of first pending ACK per destination.
83    pending_since: HashMap<String, DateTime<Utc>>,
84}
85
86impl AckOptimizer {
87    /// Creates a new ACK optimizer with default configuration.
88    pub fn new(local_device_id: impl Into<String>) -> Self {
89        Self::with_config(local_device_id, AckOptimizationConfig::default())
90    }
91
92    /// Creates a new ACK optimizer with custom configuration.
93    pub fn with_config(local_device_id: impl Into<String>, config: AckOptimizationConfig) -> Self {
94        Self {
95            config,
96            local_device_id: local_device_id.into(),
97            pending_acks: HashMap::new(),
98            pending_since: HashMap::new(),
99        }
100    }
101
102    /// Determines if this node should relay a received ACK.
103    ///
104    /// Uses deterministic pseudo-random selection based on ACK ID and local device ID
105    /// to ensure consistent relay decisions across the network.
106    ///
107    /// # Arguments
108    ///
109    /// * `ack_message_id` - ID of the ACK message being relayed
110    /// * `visible_peer_count` - Number of visible peers in the network
111    ///
112    /// # Returns
113    ///
114    /// `true` if this node should relay the ACK, `false` otherwise.
115    pub fn should_relay_ack(&self, ack_message_id: &str, visible_peer_count: usize) -> bool {
116        if !self.config.probabilistic_relay_enabled {
117            return true; // Always relay if disabled
118        }
119
120        // Higher peer count = lower relay probability to prevent storms
121        let probability = self.compute_relay_probability(visible_peer_count);
122
123        // Deterministic selection based on ACK ID and device ID
124        let mut hasher = DefaultHasher::new();
125        ack_message_id.hash(&mut hasher);
126        self.local_device_id.hash(&mut hasher);
127        let hash = hasher.finish();
128
129        let hash_probability = (hash as f64 / u64::MAX as f64) as f32;
130        hash_probability < probability
131    }
132
133    /// Computes the relay probability based on network density.
134    fn compute_relay_probability(&self, visible_peer_count: usize) -> f32 {
135        if visible_peer_count <= SMALL_NETWORK_RELAY_THRESHOLD {
136            return 1.0; // Small network: always relay
137        }
138
139        // As network grows, reduce probability
140        // At 50 peers: ~25% probability
141        // At 100 peers: ~15% probability
142        let density_factor = RELAY_DENSITY_FACTOR_NUMERATOR / (visible_peer_count as f32);
143        let raw_probability = self.config.base_relay_probability * density_factor;
144
145        raw_probability
146            .max(self.config.min_relay_probability)
147            .min(1.0)
148    }
149
150    /// Adds an ACK to the pending aggregation buffer.
151    ///
152    /// # Arguments
153    ///
154    /// * `destination` - User ID of the original message sender (ACK destination)
155    /// * `original_message_id` - ID of the message being acknowledged
156    /// * `hop_count` - Number of hops the original message traveled
157    ///
158    /// # Returns
159    ///
160    /// `Some(Vec<AggregatedAck>)` if batch is ready for transmission, `None` otherwise.
161    pub fn add_ack_for_aggregation(
162        &mut self,
163        destination: &str,
164        original_message_id: MessageId,
165        hop_count: u8,
166    ) -> Option<Vec<AggregatedAck>> {
167        if !self.config.aggregation_enabled {
168            // No aggregation: return immediately
169            return Some(vec![AggregatedAck {
170                original_message_id,
171                hop_count,
172                received_at: Utc::now(),
173            }]);
174        }
175
176        let now = Utc::now();
177        let ack = AggregatedAck {
178            original_message_id,
179            hop_count,
180            received_at: now,
181        };
182
183        let pending = self
184            .pending_acks
185            .entry(destination.to_string())
186            .or_default();
187        pending.push(ack);
188
189        // Record when we started collecting for this destination
190        self.pending_since
191            .entry(destination.to_string())
192            .or_insert(now);
193
194        // Check if batch is full
195        if pending.len() >= self.config.max_aggregation_batch {
196            return self.flush_pending_acks(destination);
197        }
198
199        None
200    }
201
202    /// Flushes pending ACKs for a destination.
203    pub fn flush_pending_acks(&mut self, destination: &str) -> Option<Vec<AggregatedAck>> {
204        self.pending_since.remove(destination);
205        self.pending_acks.remove(destination)
206    }
207
208    /// Gets destinations with ACKs that have timed out and should be flushed.
209    ///
210    /// # Returns
211    ///
212    /// A vector of destination user IDs whose ACK buffers should be flushed.
213    pub fn get_timed_out_destinations(&self) -> Vec<String> {
214        let now = Utc::now();
215        let timeout = chrono::Duration::milliseconds(self.config.aggregation_timeout_ms as i64);
216
217        self.pending_since
218            .iter()
219            .filter_map(|(dest, since)| {
220                // Inclusive: a batch that has waited *at least* the configured
221                // timeout is due. The strict form made a zero timeout mean
222                // "never, until the clock happens to tick", so a batch
223                // configured to flush immediately only flushed when two
224                // `Utc::now()` calls landed in different microseconds.
225                if now.signed_duration_since(*since) >= timeout {
226                    Some(dest.clone())
227                } else {
228                    None
229                }
230            })
231            .collect()
232    }
233
234    /// Drains all timed-out ACK batches.
235    ///
236    /// # Returns
237    ///
238    /// A vector of (destination, acks) tuples for destinations that timed out.
239    pub fn drain_timed_out(&mut self) -> Vec<(String, Vec<AggregatedAck>)> {
240        let destinations = self.get_timed_out_destinations();
241        destinations
242            .into_iter()
243            .filter_map(|dest| {
244                let acks = self.flush_pending_acks(&dest)?;
245                Some((dest, acks))
246            })
247            .collect()
248    }
249
250    /// Checks if there's a pending ACK to piggyback on a data message.
251    ///
252    /// # Arguments
253    ///
254    /// * `message_destination` - Destination of the outgoing data message
255    ///
256    /// # Returns
257    ///
258    /// `Some(PiggybackAckData)` if there are ACKs to piggyback, `None` otherwise.
259    pub fn get_piggyback_acks(&mut self, message_destination: &str) -> Option<PiggybackAckData> {
260        if !self.config.piggyback_enabled {
261            return None;
262        }
263
264        // Get pending ACKs for this destination
265        let mut acks = self.flush_pending_acks(message_destination)?;
266        if acks.is_empty() {
267            return None;
268        }
269
270        // Limit size for piggybacking; return overflow back to the queue
271        let limit = self.estimate_piggyback_limit();
272        if acks.len() > limit {
273            let overflow = acks.split_off(limit);
274            self.pending_acks
275                .insert(message_destination.to_string(), overflow);
276            // Preserve pending_since so the overflow batch retains its age
277            // (it was already removed by flush_pending_acks; re-set to now
278            // so the overflow doesn't immediately time-out on the next tick).
279            self.pending_since
280                .insert(message_destination.to_string(), Utc::now());
281        }
282
283        if acks.is_empty() {
284            return None;
285        }
286
287        Some(PiggybackAckData {
288            destination: message_destination.to_string(),
289            acks,
290            first_added_at: Utc::now(),
291        })
292    }
293
294    /// Estimates how many ACKs can be piggybacked within size limit.
295    fn estimate_piggyback_limit(&self) -> usize {
296        (self.config.max_piggyback_size / ESTIMATED_ACK_SIZE_BYTES).max(1)
297    }
298
299    /// Returns the number of pending ACK destinations.
300    pub fn pending_destination_count(&self) -> usize {
301        self.pending_acks.len()
302    }
303
304    /// Returns the total number of pending ACKs.
305    pub fn pending_ack_count(&self) -> usize {
306        self.pending_acks.values().map(|v| v.len()).sum()
307    }
308
309    /// Clears all pending ACKs.
310    pub fn clear(&mut self) {
311        self.pending_acks.clear();
312        self.pending_since.clear();
313    }
314
315    /// Gets the configuration.
316    pub fn config(&self) -> &AckOptimizationConfig {
317        &self.config
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324
325    #[test]
326    fn test_probabilistic_relay_small_network() {
327        let optimizer = AckOptimizer::new("device1");
328
329        // In small networks, should always relay
330        let should_relay = optimizer.should_relay_ack("ack123", 3);
331        assert!(should_relay);
332    }
333
334    #[test]
335    fn test_probabilistic_relay_deterministic() {
336        let optimizer = AckOptimizer::new("device1");
337
338        // Same inputs should always produce same result
339        let result1 = optimizer.should_relay_ack("ack123", 50);
340        let result2 = optimizer.should_relay_ack("ack123", 50);
341        assert_eq!(result1, result2);
342    }
343
344    #[test]
345    fn test_probabilistic_relay_varies_by_device() {
346        let optimizer1 = AckOptimizer::new("device1");
347        let optimizer2 = AckOptimizer::new("device2");
348
349        // Different devices may have different relay decisions
350        let results: Vec<bool> = (0..10)
351            .map(|i| optimizer1.should_relay_ack(&format!("ack{}", i), 100))
352            .collect();
353        let results2: Vec<bool> = (0..10)
354            .map(|i| optimizer2.should_relay_ack(&format!("ack{}", i), 100))
355            .collect();
356
357        // Results should differ for at least some ACKs (probabilistically)
358        // This test might rarely fail due to randomness, but it's a good sanity check
359        assert_ne!(results, results2);
360    }
361
362    #[test]
363    fn test_aggregation_single_ack() {
364        let mut optimizer = AckOptimizer::new("device1");
365
366        let result = optimizer.add_ack_for_aggregation("alice", MessageId::new(), 3);
367
368        // Single ACK shouldn't flush (waiting for more)
369        assert!(result.is_none());
370        assert_eq!(optimizer.pending_ack_count(), 1);
371    }
372
373    #[test]
374    fn test_aggregation_batch_full() {
375        let config = AckOptimizationConfig {
376            max_aggregation_batch: 3,
377            ..Default::default()
378        };
379        let mut optimizer = AckOptimizer::with_config("device1", config);
380
381        // Add ACKs until batch is full
382        optimizer.add_ack_for_aggregation("alice", MessageId::new(), 1);
383        optimizer.add_ack_for_aggregation("alice", MessageId::new(), 2);
384        let result = optimizer.add_ack_for_aggregation("alice", MessageId::new(), 3);
385
386        // Should flush when batch is full
387        assert!(result.is_some());
388        assert_eq!(result.unwrap().len(), 3);
389        assert_eq!(optimizer.pending_ack_count(), 0);
390    }
391
392    #[test]
393    fn test_aggregation_disabled() {
394        let config = AckOptimizationConfig {
395            aggregation_enabled: false,
396            ..Default::default()
397        };
398        let mut optimizer = AckOptimizer::with_config("device1", config);
399
400        let result = optimizer.add_ack_for_aggregation("alice", MessageId::new(), 3);
401
402        // Should return immediately when aggregation disabled
403        assert!(result.is_some());
404        assert_eq!(result.unwrap().len(), 1);
405    }
406
407    #[test]
408    fn test_piggyback_acks() {
409        let mut optimizer = AckOptimizer::new("device1");
410
411        // Add an ACK for alice
412        optimizer.add_ack_for_aggregation("alice", MessageId::new(), 3);
413
414        // Get piggybacked ACKs when sending to alice
415        let piggyback = optimizer.get_piggyback_acks("alice");
416        assert!(piggyback.is_some());
417        assert_eq!(piggyback.unwrap().acks.len(), 1);
418
419        // Should be cleared now
420        assert_eq!(optimizer.pending_ack_count(), 0);
421    }
422
423    #[test]
424    fn test_piggyback_no_acks() {
425        let mut optimizer = AckOptimizer::new("device1");
426
427        // No pending ACKs for bob
428        let piggyback = optimizer.get_piggyback_acks("bob");
429        assert!(piggyback.is_none());
430    }
431
432    #[test]
433    fn test_drain_timed_out() {
434        let config = AckOptimizationConfig {
435            aggregation_timeout_ms: 0, // Immediate timeout for testing
436            ..Default::default()
437        };
438        let mut optimizer = AckOptimizer::with_config("device1", config);
439
440        optimizer.add_ack_for_aggregation("alice", MessageId::new(), 1);
441        optimizer.add_ack_for_aggregation("bob", MessageId::new(), 2);
442
443        // All should be timed out immediately
444        let timed_out = optimizer.drain_timed_out();
445        assert_eq!(timed_out.len(), 2);
446    }
447}
offline_protocol_reliability/ack_optimization.rs

offline_protocol_reliability/
ack_optimization.rs
