dope-session 0.2.3

use dope_transport::Transport;
use std::time::{Duration, Instant};

pub enum ConnectAction<T: Transport> {
    Connect { addr: T::Addr, tag: u32 },
    Backoff { until: Instant },
    Idle,
}

pub trait ConnectSource<T: Transport>: 'static {
    fn poll_connect(&mut self, now: Instant) -> ConnectAction<T>;

    fn on_connect_outcome(&mut self, tag: u32, success: bool, now: Instant);

    fn on_disconnect(&mut self, tag: u32, now: Instant) {
        let _ = (tag, now);
    }
}

#[derive(Default, Clone, Copy)]
struct BackoffSlot {
    failed_at: Option<Instant>,
    attempt: u8,
}

const MAX_BACKOFF_SHIFT: u8 = 10;

pub struct StaticUpstreams<T: Transport> {
    upstreams: Vec<T::Addr>,
    backoff: Vec<BackoffSlot>,
    base_window: Duration,
    next: u32,
    rng_state: u64,
}

impl<T: Transport> StaticUpstreams<T>
where
    T::Addr: Clone,
{
    pub fn new(upstreams: Vec<T::Addr>, base_window: Duration) -> Self {
        let n = upstreams.len();
        let seed = (n as u64).wrapping_mul(0x9E3779B97F4A7C15) ^ base_window.as_nanos() as u64;
        Self {
            upstreams,
            backoff: vec![BackoffSlot::default(); n],
            base_window,
            next: 0,
            rng_state: seed.max(1),
        }
    }

    #[inline(always)]
    fn next_rand(&mut self) -> u64 {
        let mut x = self.rng_state;
        x ^= x >> 12;
        x ^= x << 25;
        x ^= x >> 27;
        self.rng_state = x;
        x.wrapping_mul(0x2545F4914F6CDD1D)
    }

    fn retry_at(&mut self, slot: BackoffSlot) -> Option<Instant> {
        let failed_at = slot.failed_at?;
        let shift = slot.attempt.min(MAX_BACKOFF_SHIFT) as u32;
        let scaled = self.base_window.saturating_mul(1u32 << shift);
        let r = self.next_rand();
        let jitter_quarters = (r % 4) as u32;
        let jitter = scaled / 4 * jitter_quarters;
        Some(failed_at + scaled / 2 + jitter)
    }
}

impl<T: Transport> ConnectSource<T> for StaticUpstreams<T>
where
    T::Addr: Clone,
{
    fn poll_connect(&mut self, now: Instant) -> ConnectAction<T> {
        if self.upstreams.is_empty() {
            return ConnectAction::Idle;
        }
        let n = self.upstreams.len() as u32;
        let mut earliest_retry: Option<Instant> = None;
        for offset in 0..n {
            let idx = ((self.next + offset) % n) as usize;
            let slot = self.backoff[idx];
            if slot.failed_at.is_none() {
                let addr = self.upstreams[idx].clone();
                self.next = (idx as u32 + 1) % n;
                return ConnectAction::Connect {
                    addr,
                    tag: idx as u32,
                };
            }
            let retry = self.retry_at(slot).expect("failed_at present");
            if retry <= now {
                let addr = self.upstreams[idx].clone();
                self.backoff[idx].failed_at = None;
                self.next = (idx as u32 + 1) % n;
                return ConnectAction::Connect {
                    addr,
                    tag: idx as u32,
                };
            }
            earliest_retry = Some(earliest_retry.map_or(retry, |e| e.min(retry)));
        }
        match earliest_retry {
            Some(until) => ConnectAction::Backoff { until },
            None => ConnectAction::Idle,
        }
    }

    fn on_connect_outcome(&mut self, tag: u32, success: bool, now: Instant) {
        let idx = tag as usize;
        if idx >= self.backoff.len() {
            return;
        }
        if success {
            self.backoff[idx] = BackoffSlot::default();
        } else {
            let attempt = self.backoff[idx]
                .attempt
                .saturating_add(1)
                .min(MAX_BACKOFF_SHIFT);
            self.backoff[idx] = BackoffSlot {
                failed_at: Some(now),
                attempt,
            };
        }
    }

    fn on_disconnect(&mut self, tag: u32, now: Instant) {
        let idx = tag as usize;
        if idx >= self.backoff.len() {
            return;
        }
        self.backoff[idx] = BackoffSlot {
            failed_at: Some(now),
            attempt: 1,
        };
    }
}