stochastic-routing-extended 1.0.2

//! Non-linear retry strategy.
//!
//! Retries are sent through different transports than the original attempt,
//! preventing signature matching by DPI systems.

use crate::transport::TransportKind;

/// A non-linear retry strategy that varies the transport on each attempt.
pub struct RetryStrategy {
    max_attempts: u32,
    current_attempt: u32,
}

impl RetryStrategy {
    pub fn new(max_attempts: u32) -> Self {
        Self {
            max_attempts,
            current_attempt: 0,
        }
    }

    /// Get the next transport to try, or [`None`] if attempts exhausted.
    ///
    /// Prefers any transport **other than** `exclude`; if that would empty the pool,
    /// falls back to the full `available` list (e.g. single-transport deployments).
    pub fn next_transport(
        &mut self,
        available: &[TransportKind],
        exclude: TransportKind,
    ) -> Option<TransportKind> {
        if self.current_attempt >= self.max_attempts {
            return None;
        }
        if available.is_empty() {
            return None;
        }

        let preferred: Vec<TransportKind> = available
            .iter()
            .copied()
            .filter(|&t| t != exclude)
            .collect();

        let pool = if preferred.is_empty() {
            available.to_vec()
        } else {
            preferred
        };

        let idx = self.current_attempt as usize % pool.len();
        let choice = pool[idx];
        self.current_attempt += 1;
        Some(choice)
    }

    /// Reset the retry counter.
    pub fn reset(&mut self) {
        self.current_attempt = 0;
    }

    /// Remaining attempts after the last [`next_transport`] call.
    #[must_use]
    pub fn attempts_used(&self) -> u32 {
        self.current_attempt
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn skips_excluded_until_fallback() {
        let kinds = [TransportKind::Tcp, TransportKind::Udp];
        let mut r = RetryStrategy::new(5);
        let first = r.next_transport(&kinds, TransportKind::Tcp).unwrap();
        assert_eq!(first, TransportKind::Udp);
    }

    #[test]
    fn cycles_when_multiple_allowed() {
        let kinds = [TransportKind::Tcp, TransportKind::Udp, TransportKind::Quic];
        let mut r = RetryStrategy::new(10);
        let _ = r.next_transport(&kinds, TransportKind::Grpc); // exclude not in list → use pool
        let a = r.next_transport(&kinds, TransportKind::Grpc).unwrap();
        let b = r.next_transport(&kinds, TransportKind::Grpc).unwrap();
        assert_eq!(a, kinds[1]);
        assert_eq!(b, kinds[2]);
    }

    #[test]
    fn exhausts_max() {
        let kinds = [TransportKind::Tcp];
        let mut r = RetryStrategy::new(2);
        assert!(r.next_transport(&kinds, TransportKind::Udp).is_some());
        assert!(r.next_transport(&kinds, TransportKind::Udp).is_some());
        assert!(r.next_transport(&kinds, TransportKind::Udp).is_none());
    }
}