Skip to main content

dope_session/pool/client/
connect_source.rs

1use dope_transport::Transport;
2use std::time::{Duration, Instant};
3
4pub enum ConnectAction<T: Transport> {
5    Connect { addr: T::Addr, tag: u32 },
6    Backoff { until: Instant },
7    Idle,
8}
9
10pub trait ConnectSource<T: Transport>: 'static {
11    fn poll_connect(&mut self, now: Instant) -> ConnectAction<T>;
12
13    fn on_connect_outcome(&mut self, tag: u32, success: bool, now: Instant);
14
15    fn on_disconnect(&mut self, tag: u32, now: Instant) {
16        let _ = (tag, now);
17    }
18}
19
20#[derive(Default, Clone, Copy)]
21struct BackoffSlot {
22    failed_at: Option<Instant>,
23    attempt: u8,
24}
25
26const MAX_BACKOFF_SHIFT: u8 = 10;
27
28pub struct StaticUpstreams<T: Transport> {
29    upstreams: Vec<T::Addr>,
30    backoff: Vec<BackoffSlot>,
31    base_window: Duration,
32    next: u32,
33    rng_state: u64,
34}
35
36impl<T: Transport> StaticUpstreams<T>
37where
38    T::Addr: Clone,
39{
40    pub fn new(upstreams: Vec<T::Addr>, base_window: Duration) -> Self {
41        let n = upstreams.len();
42        let seed = (n as u64).wrapping_mul(0x9E3779B97F4A7C15) ^ base_window.as_nanos() as u64;
43        Self {
44            upstreams,
45            backoff: vec![BackoffSlot::default(); n],
46            base_window,
47            next: 0,
48            rng_state: seed.max(1),
49        }
50    }
51
52    #[inline(always)]
53    fn next_rand(&mut self) -> u64 {
54        let mut x = self.rng_state;
55        x ^= x >> 12;
56        x ^= x << 25;
57        x ^= x >> 27;
58        self.rng_state = x;
59        x.wrapping_mul(0x2545F4914F6CDD1D)
60    }
61
62    fn retry_at(&mut self, slot: BackoffSlot) -> Option<Instant> {
63        let failed_at = slot.failed_at?;
64        let shift = slot.attempt.min(MAX_BACKOFF_SHIFT) as u32;
65        let scaled = self.base_window.saturating_mul(1u32 << shift);
66        let r = self.next_rand();
67        let jitter_quarters = (r % 4) as u32;
68        let jitter = scaled / 4 * jitter_quarters;
69        Some(failed_at + scaled / 2 + jitter)
70    }
71}
72
73impl<T: Transport> ConnectSource<T> for StaticUpstreams<T>
74where
75    T::Addr: Clone,
76{
77    fn poll_connect(&mut self, now: Instant) -> ConnectAction<T> {
78        if self.upstreams.is_empty() {
79            return ConnectAction::Idle;
80        }
81        let n = self.upstreams.len() as u32;
82        let mut earliest_retry: Option<Instant> = None;
83        for offset in 0..n {
84            let idx = ((self.next + offset) % n) as usize;
85            let slot = self.backoff[idx];
86            if slot.failed_at.is_none() {
87                let addr = self.upstreams[idx].clone();
88                self.next = (idx as u32 + 1) % n;
89                return ConnectAction::Connect {
90                    addr,
91                    tag: idx as u32,
92                };
93            }
94            let retry = self.retry_at(slot).expect("failed_at present");
95            if retry <= now {
96                let addr = self.upstreams[idx].clone();
97                self.backoff[idx].failed_at = None;
98                self.next = (idx as u32 + 1) % n;
99                return ConnectAction::Connect {
100                    addr,
101                    tag: idx as u32,
102                };
103            }
104            earliest_retry = Some(earliest_retry.map_or(retry, |e| e.min(retry)));
105        }
106        match earliest_retry {
107            Some(until) => ConnectAction::Backoff { until },
108            None => ConnectAction::Idle,
109        }
110    }
111
112    fn on_connect_outcome(&mut self, tag: u32, success: bool, now: Instant) {
113        let idx = tag as usize;
114        if idx >= self.backoff.len() {
115            return;
116        }
117        if success {
118            self.backoff[idx] = BackoffSlot::default();
119        } else {
120            let attempt = self.backoff[idx]
121                .attempt
122                .saturating_add(1)
123                .min(MAX_BACKOFF_SHIFT);
124            self.backoff[idx] = BackoffSlot {
125                failed_at: Some(now),
126                attempt,
127            };
128        }
129    }
130
131    fn on_disconnect(&mut self, tag: u32, now: Instant) {
132        let idx = tag as usize;
133        if idx >= self.backoff.len() {
134            return;
135        }
136        self.backoff[idx] = BackoffSlot {
137            failed_at: Some(now),
138            attempt: 1,
139        };
140    }
141}