leyline/core/config/
host_limits.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex, PoisonError};
4use std::time::Duration;
5
6use tokio::sync::{OwnedSemaphorePermit, Semaphore};
7use tokio::time::Instant;
8
9#[derive(Debug, Clone, PartialEq, Eq, Hash)]
10pub(crate) struct Origin {
11    scheme: String,
12    host: String,
13    port: u16,
14}
15
16impl std::fmt::Display for Origin {
17    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
18        write!(f, "{}://{}:{}", self.scheme, self.host, self.port)
19    }
20}
21
22impl Origin {
23    pub(crate) fn of(url: &url::Url) -> Option<Self> {
24        Some(Self {
25            scheme: url.scheme().to_ascii_lowercase(),
26            host: url.host_str()?.trim_end_matches('.').to_ascii_lowercase(),
27            port: url.port_or_known_default()?,
28        })
29    }
30}
31
32#[derive(Debug, Clone, Copy, Default)]
33struct Rule {
34    max_in_flight: Option<usize>,
35    interval: Option<Duration>,
36}
37
38impl Rule {
39    fn is_unlimited(self) -> bool {
40        self.max_in_flight.is_none() && self.interval.is_none()
41    }
42}
43
44#[derive(Debug, Clone, PartialEq, Eq)]
45enum HostPattern {
46    Exact(String),
47    Subdomains(String),
48}
49
50impl HostPattern {
51    fn parse(host: &str) -> Self {
52        let host = host.trim().trim_end_matches('.').to_ascii_lowercase();
53        match host.strip_prefix("*.") {
54            Some(domain) => Self::Subdomains(domain.to_owned()),
55            None => Self::Exact(host),
56        }
57    }
58
59    fn rank(&self, host: &str) -> Option<usize> {
60        match self {
61            Self::Exact(exact) => (exact == host).then_some(usize::MAX),
62            Self::Subdomains(domain) => host
63                .strip_suffix(domain.as_str())
64                .is_some_and(|rest| rest.ends_with('.') && rest.len() > 1)
65                .then_some(domain.len()),
66        }
67    }
68}
69
70#[derive(Debug, Clone)]
71#[non_exhaustive]
72pub struct HostLimits {
73    rule: Rule,
74    overrides: Vec<(HostPattern, Rule)>,
75    max_total: Option<usize>,
76    pause: Pause,
77    gates: Arc<Gates>,
78}
79
80#[derive(Debug, Clone)]
81struct Pause {
82    statuses: Vec<http::StatusCode>,
83    invalid: Option<u16>,
84    default: Duration,
85    max: Duration,
86}
87
88impl Pause {
89    fn length(&self, status: http::StatusCode, wait: Option<Duration>) -> Option<Duration> {
90        self.statuses
91            .contains(&status)
92            .then(|| wait.unwrap_or(self.default).min(self.max))
93    }
94}
95
96impl Default for HostLimits {
97    fn default() -> Self {
98        Self {
99            rule: Rule::default(),
100            overrides: Vec::new(),
101            max_total: None,
102            pause: Pause {
103                statuses: Vec::new(),
104                invalid: None,
105                default: Duration::from_secs(60),
106                max: Duration::from_secs(24 * 60 * 60),
107            },
108            gates: Arc::default(),
109        }
110    }
111}
112
113impl HostLimits {
114    pub fn new() -> Self {
115        Self::default()
116    }
117
118    pub fn max_in_flight(mut self, n: usize) -> Self {
119        self.rule.max_in_flight = Some(n.max(1));
120        self
121    }
122
123    pub fn per_second(mut self, rate: f64) -> Self {
124        self.rule.interval = if rate > 0.0 {
125            Duration::try_from_secs_f64(rate.recip()).ok()
126        } else {
127            None
128        };
129        self
130    }
131
132    pub fn host(mut self, host: &str, limits: HostLimits) -> Self {
133        let pattern = HostPattern::parse(host);
134        self.overrides.retain(|(held, _)| *held != pattern);
135        self.overrides.push((pattern, limits.rule));
136        self
137    }
138
139    pub fn max_total_in_flight(mut self, n: usize) -> Self {
140        self.max_total = Some(n.max(1));
141        self
142    }
143
144    pub fn pause_on(mut self, statuses: impl IntoIterator<Item = u16>) -> Self {
145        self.pause.statuses.clear();
146        self.pause.invalid = None;
147        for code in statuses {
148            match http::StatusCode::from_u16(code) {
149                Ok(status) => self.pause.statuses.push(status),
150                Err(_) => {
151                    self.pause.invalid.get_or_insert(code);
152                }
153            }
154        }
155        self
156    }
157
158    pub fn pause_for(mut self, default: Duration) -> Self {
159        self.pause.default = default;
160        self
161    }
162
163    pub fn max_pause(mut self, max: Duration) -> Self {
164        self.pause.max = max;
165        self
166    }
167
168    pub(crate) fn config_error(&self) -> Option<String> {
169        self.pause
170            .invalid
171            .map(|code| format!("HostLimits::pause_on: {code} is not a valid HTTP status code"))
172    }
173
174    pub(crate) fn armed(mut self) -> Self {
175        self.gates = Arc::new(Gates::new(self.max_total));
176        self
177    }
178
179    pub(crate) fn observe(&self, url: &url::Url, status: http::StatusCode, wait: Option<Duration>) {
180        let Some(length) = self.pause.length(status, wait) else {
181            return;
182        };
183        let Some(origin) = Origin::of(url) else {
184            return;
185        };
186        let rule = self.rule_for(&origin.host);
187        let now = Instant::now();
188        let end = now.checked_add(length).unwrap_or(now);
189        self.gates.gate(origin, rule.max_in_flight).pause_until(end);
190    }
191
192    pub(crate) fn is_unlimited(&self) -> bool {
193        self.rule.is_unlimited()
194            && self.overrides.is_empty()
195            && self.max_total.is_none()
196            && self.pause.statuses.is_empty()
197    }
198
199    fn rule_for(&self, host: &str) -> Rule {
200        self.overrides
201            .iter()
202            .filter_map(|(pattern, rule)| Some((pattern.rank(host)?, *rule)))
203            .max_by_key(|(rank, _)| *rank)
204            .map_or(self.rule, |(_, rule)| rule)
205    }
206
207    pub(crate) async fn admit(&self, url: &url::Url) -> Option<HostPass> {
208        if self.is_unlimited() {
209            return None;
210        }
211        let origin = Origin::of(url)?;
212        let rule = self.rule_for(&origin.host);
213        let gate = self.gates.gate(origin, rule.max_in_flight);
214        let waiting = Count::enter(&gate.waiting);
215        let permit = acquire(gate.slots.as_ref()).await;
216        let total = loop {
217            gate.wait_unpaused().await;
218            gate.wait_spacing().await;
219            let total = acquire(self.gates.total.as_ref()).await;
220            if gate.try_admit(rule.interval) {
221                break total;
222            }
223        };
224        drop(waiting);
225        Some(HostPass {
226            _in_flight: Count::enter(&gate.in_flight),
227            _permit: permit,
228            _total: total,
229            _gate: gate,
230        })
231    }
232
233    pub(crate) fn stats(&self) -> Vec<HostStats> {
234        self.gates
235            .by_origin
236            .lock()
237            .unwrap_or_else(PoisonError::into_inner)
238            .iter()
239            .map(|(origin, gate)| HostStats {
240                origin: origin.to_string(),
241                in_flight: gate.in_flight.load(Ordering::Relaxed),
242                waiting: gate.waiting.load(Ordering::Relaxed),
243            })
244            .collect()
245    }
246}
247
248#[derive(Debug, Clone, PartialEq, Eq)]
249#[non_exhaustive]
250pub struct HostStats {
251    origin: String,
252    in_flight: usize,
253    waiting: usize,
254}
255
256impl HostStats {
257    pub fn origin(&self) -> &str {
258        &self.origin
259    }
260
261    pub fn in_flight(&self) -> usize {
262        self.in_flight
263    }
264
265    pub fn waiting(&self) -> usize {
266        self.waiting
267    }
268}
269
270struct Count(Arc<AtomicUsize>);
271
272impl Count {
273    fn enter(counter: &Arc<AtomicUsize>) -> Self {
274        counter.fetch_add(1, Ordering::Relaxed);
275        Self(Arc::clone(counter))
276    }
277}
278
279impl Drop for Count {
280    fn drop(&mut self) {
281        self.0.fetch_sub(1, Ordering::Relaxed);
282    }
283}
284
285async fn acquire(slots: Option<&Arc<Semaphore>>) -> Option<OwnedSemaphorePermit> {
286    Arc::clone(slots?).acquire_owned().await.ok()
287}
288
289pub(crate) struct HostPass {
290    _in_flight: Count,
291    _permit: Option<OwnedSemaphorePermit>,
292    _total: Option<OwnedSemaphorePermit>,
293    _gate: Arc<HostGate>,
294}
295
296#[derive(Debug, Default)]
297struct Gates {
298    by_origin: Mutex<HashMap<Origin, Arc<HostGate>>>,
299    total: Option<Arc<Semaphore>>,
300}
301
302impl Gates {
303    fn new(max_total: Option<usize>) -> Self {
304        Self {
305            by_origin: Mutex::default(),
306            total: max_total.map(|n| Arc::new(Semaphore::new(n))),
307        }
308    }
309
310    fn gate(&self, origin: Origin, max_in_flight: Option<usize>) -> Arc<HostGate> {
311        let mut map = self
312            .by_origin
313            .lock()
314            .unwrap_or_else(PoisonError::into_inner);
315        if let Some(gate) = map.get(&origin) {
316            return Arc::clone(gate);
317        }
318        let now = Instant::now();
319        map.retain(|_, gate| Arc::strong_count(gate) > 1 || gate.busy_after(now));
320        let gate = Arc::new(HostGate::new(max_in_flight));
321        map.insert(origin, Arc::clone(&gate));
322        gate
323    }
324}
325
326#[derive(Debug)]
327struct HostGate {
328    slots: Option<Arc<Semaphore>>,
329    next: Mutex<Instant>,
330    paused: Mutex<Instant>,
331    in_flight: Arc<AtomicUsize>,
332    waiting: Arc<AtomicUsize>,
333}
334
335impl HostGate {
336    fn new(max_in_flight: Option<usize>) -> Self {
337        Self {
338            slots: max_in_flight.map(|n| Arc::new(Semaphore::new(n))),
339            next: Mutex::new(Instant::now()),
340            paused: Mutex::new(Instant::now()),
341            in_flight: Arc::default(),
342            waiting: Arc::default(),
343        }
344    }
345
346    fn try_admit(&self, interval: Option<Duration>) -> bool {
347        let mut next = self.next.lock().unwrap_or_else(PoisonError::into_inner);
348        let now = Instant::now();
349        if *next > now || self.paused_until() > now {
350            return false;
351        }
352        if let Some(interval) = interval {
353            *next = now.checked_add(interval).unwrap_or(now);
354        }
355        true
356    }
357
358    async fn wait_spacing(&self) {
359        let next = *self.next.lock().unwrap_or_else(PoisonError::into_inner);
360        tokio::time::sleep_until(next).await;
361    }
362
363    fn busy_after(&self, now: Instant) -> bool {
364        *self.next.lock().unwrap_or_else(PoisonError::into_inner) > now || self.paused_until() > now
365    }
366
367    fn paused_until(&self) -> Instant {
368        *self.paused.lock().unwrap_or_else(PoisonError::into_inner)
369    }
370
371    fn pause_until(&self, end: Instant) {
372        let mut paused = self.paused.lock().unwrap_or_else(PoisonError::into_inner);
373        *paused = (*paused).max(end);
374    }
375
376    async fn wait_unpaused(&self) {
377        loop {
378            let end = self.paused_until();
379            if end <= Instant::now() {
380                return;
381            }
382            tokio::time::sleep_until(end).await;
383        }
384    }
385}
leyline/core/config/host_limits.rs

leyline/core/config/
host_limits.rs
