use std::sync::atomic::{AtomicUsize, Ordering};
use serde::{Deserialize, Serialize};
use crate::check::UncertainReason;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum TransportTier {
Http,
Impersonate,
Browser,
}
impl TransportTier {
#[must_use]
pub const fn as_str(self) -> &'static str {
match self {
Self::Http => "http",
Self::Impersonate => "impersonate",
Self::Browser => "browser",
}
}
}
impl core::fmt::Display for TransportTier {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug)]
pub struct EscalationBudget {
used: AtomicUsize,
cap: usize,
}
impl EscalationBudget {
#[must_use]
pub const fn new(cap: usize) -> Self {
Self {
used: AtomicUsize::new(0),
cap,
}
}
#[must_use]
pub const fn unlimited() -> Self {
Self::new(usize::MAX)
}
pub fn try_consume(&self) -> bool {
let mut cur = self.used.load(Ordering::Acquire);
loop {
if cur >= self.cap {
return false;
}
match self
.used
.compare_exchange_weak(cur, cur + 1, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => return true,
Err(actual) => cur = actual,
}
}
}
#[must_use]
pub fn used(&self) -> usize {
self.used.load(Ordering::Acquire)
}
#[must_use]
pub const fn cap(&self) -> usize {
self.cap
}
}
pub(crate) const fn should_escalate(reason: &UncertainReason) -> bool {
matches!(
reason,
UncertainReason::CloudflareChallenge | UncertainReason::RateLimited
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn escalates_on_cloudflare_and_rate_limited_only() {
assert!(should_escalate(&UncertainReason::CloudflareChallenge));
assert!(should_escalate(&UncertainReason::RateLimited));
assert!(!should_escalate(&UncertainReason::Captcha));
assert!(!should_escalate(&UncertainReason::RobotsDisallowed));
assert!(!should_escalate(&UncertainReason::Deadline));
assert!(!should_escalate(&UncertainReason::SchedulerClosed));
assert!(!should_escalate(&UncertainReason::Network(
"refused".into()
)));
assert!(!should_escalate(&UncertainReason::BodyRead("eof".into())));
assert!(!should_escalate(&UncertainReason::BrowserBudget));
assert!(!should_escalate(&UncertainReason::UsernameNotAllowed));
assert!(!should_escalate(&UncertainReason::BrowserFailed(
"timeout".into()
)));
assert!(!should_escalate(&UncertainReason::GeoUnavailable));
assert!(!should_escalate(&UncertainReason::SessionRequired));
assert!(!should_escalate(&UncertainReason::Other("?".into())));
}
#[test]
fn budget_consumes_up_to_cap() {
let b = EscalationBudget::new(2);
assert!(b.try_consume());
assert!(b.try_consume());
assert!(!b.try_consume());
assert_eq!(b.used(), 2);
assert_eq!(b.cap(), 2);
}
#[test]
fn budget_zero_denies_all() {
let b = EscalationBudget::new(0);
assert!(!b.try_consume());
}
#[test]
fn budget_unlimited_never_denies() {
let b = EscalationBudget::unlimited();
for _ in 0..1024 {
assert!(b.try_consume());
}
}
#[test]
fn transport_tier_as_str_matches_serde() {
assert_eq!(TransportTier::Http.as_str(), "http");
assert_eq!(TransportTier::Impersonate.as_str(), "impersonate");
assert_eq!(TransportTier::Browser.as_str(), "browser");
let json = serde_json::to_string(&TransportTier::Impersonate).unwrap();
assert_eq!(json, r#""impersonate""#);
let back: TransportTier = serde_json::from_str(&json).unwrap();
assert_eq!(back, TransportTier::Impersonate);
}
}