stealthscraper_rs/state/
model.rs1use std::time::Duration;
12
13use serde::{Deserialize, Serialize};
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
17pub enum Outcome {
18 Success,
20 Challenged,
22 Blocked,
24 RateLimited,
26}
27
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
30#[serde(deny_unknown_fields)]
31pub struct DomainState {
32 pub host: String,
34 pub last_outcome: Option<Outcome>,
36 pub last_proxy: Option<String>,
38 pub successes: u32,
40 pub failures: u32,
42 pub cooldown_until: Option<u64>,
44 pub updated_at: u64,
46}
47
48impl DomainState {
49 pub fn new(host: impl Into<String>) -> Self {
51 Self {
52 host: host.into(),
53 last_outcome: None,
54 last_proxy: None,
55 successes: 0,
56 failures: 0,
57 cooldown_until: None,
58 updated_at: 0,
59 }
60 }
61
62 pub fn record(
75 &self,
76 outcome: Outcome,
77 proxy: Option<String>,
78 now: u64,
79 rate_limit_cooldown: Duration,
80 ) -> Self {
81 let mut next = self.clone();
82 next.last_outcome = Some(outcome);
83 if proxy.is_some() {
84 next.last_proxy = proxy;
85 }
86 next.updated_at = now;
87 match outcome {
88 Outcome::Success | Outcome::Challenged => {
89 next.successes = next.successes.saturating_add(1);
90 next.cooldown_until = None;
91 }
92 Outcome::RateLimited => {
93 next.failures = next.failures.saturating_add(1);
94 next.cooldown_until = Some(now.saturating_add(rate_limit_cooldown.as_secs()));
95 }
96 Outcome::Blocked => {
97 next.failures = next.failures.saturating_add(1);
98 next.cooldown_until = None;
99 }
100 }
101 next
102 }
103
104 pub fn in_cooldown(&self, now: u64) -> bool {
106 self.cooldown_until.is_some_and(|until| now < until)
107 }
108
109 pub fn cooldown_remaining(&self, now: u64) -> Option<Duration> {
111 self.cooldown_until
112 .and_then(|until| until.checked_sub(now))
113 .filter(|&secs| secs > 0)
114 .map(Duration::from_secs)
115 }
116}
117
118#[cfg(test)]
119mod tests {
120 use super::*;
121
122 #[test]
123 fn record_success_increments_and_clears_cooldown() {
124 let s = DomainState::new("example.com").record(
125 Outcome::RateLimited,
126 None,
127 100,
128 Duration::from_secs(60),
129 );
130 assert!(s.in_cooldown(120));
131
132 let s = s.record(
133 Outcome::Success,
134 Some("http://p:1".into()),
135 200,
136 Duration::ZERO,
137 );
138 assert_eq!(s.successes, 1);
139 assert_eq!(s.last_outcome, Some(Outcome::Success));
140 assert_eq!(s.last_proxy.as_deref(), Some("http://p:1"));
141 assert!(!s.in_cooldown(200));
142 assert_eq!(s.cooldown_until, None);
143 }
144
145 #[test]
146 fn record_rate_limited_sets_cooldown() {
147 let s = DomainState::new("h").record(
148 Outcome::RateLimited,
149 None,
150 1_000,
151 Duration::from_secs(30),
152 );
153 assert_eq!(s.failures, 1);
154 assert!(s.in_cooldown(1_029));
155 assert!(!s.in_cooldown(1_030));
156 assert_eq!(s.cooldown_remaining(1_010), Some(Duration::from_secs(20)));
157 assert_eq!(s.cooldown_remaining(1_030), None);
158 }
159
160 #[test]
161 fn record_keeps_proxy_when_none_passed() {
162 let s = DomainState::new("h")
163 .record(Outcome::Success, Some("http://a".into()), 1, Duration::ZERO)
164 .record(Outcome::Blocked, None, 2, Duration::ZERO);
165 assert_eq!(s.last_proxy.as_deref(), Some("http://a"));
166 assert_eq!(s.failures, 1);
167 assert_eq!(s.last_outcome, Some(Outcome::Blocked));
168 }
169
170 #[test]
171 fn no_cooldown_by_default() {
172 let s = DomainState::new("h");
173 assert!(!s.in_cooldown(0));
174 assert_eq!(s.cooldown_remaining(0), None);
175 }
176
177 #[test]
178 fn challenged_counts_as_success_and_clears_cooldown() {
179 let s =
181 DomainState::new("h").record(Outcome::RateLimited, None, 100, Duration::from_secs(60));
182 assert!(s.in_cooldown(120));
183
184 let s = s.record(Outcome::Challenged, None, 200, Duration::ZERO);
185 assert_eq!(s.successes, 1);
186 assert_eq!(s.failures, 1); assert_eq!(s.last_outcome, Some(Outcome::Challenged));
188 assert!(!s.in_cooldown(200));
189 }
190
191 #[test]
192 fn blocked_counts_as_failure_and_clears_stale_cooldown() {
193 let s =
194 DomainState::new("h").record(Outcome::RateLimited, None, 100, Duration::from_secs(60));
195 assert!(s.in_cooldown(120));
196
197 let s = s.record(Outcome::Blocked, None, 130, Duration::ZERO);
200 assert_eq!(s.failures, 2);
201 assert_eq!(s.successes, 0);
202 assert!(!s.in_cooldown(130));
203 assert_eq!(s.cooldown_until, None);
204 }
205}