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