stealthscraper_rs/state/
store.rs

1//! The [`StateStore`] output port and an in-memory adapter.
2
3use std::collections::HashMap;
4use std::sync::Mutex;
5
6use crate::Error;
7
8use super::model::DomainState;
9
10/// Output port for persisting [`DomainState`] keyed by host.
11///
12/// Implementations must be cheap to share across threads (`Send + Sync`); the
13/// scraper holds one behind an `Arc`. Methods are synchronous to match embedded
14/// stores like `redb`.
15pub trait StateStore: Send + Sync {
16    /// Fetch the stored state for `host`, if any.
17    fn get(&self, host: &str) -> Result<Option<DomainState>, Error>;
18    /// Insert or replace the state for `state.host`.
19    fn put(&self, state: &DomainState) -> Result<(), Error>;
20    /// Delete any stored state for `host` (no-op if absent).
21    fn remove(&self, host: &str) -> Result<(), Error>;
22
23    /// Atomically read-modify-write the state for `host`.
24    ///
25    /// `update` receives the current state (or a fresh [`DomainState`] for
26    /// `host` if none exists) and returns the new value to persist. The returned
27    /// state is stored and handed back.
28    ///
29    /// The default implementation is a non-atomic `get` + `put` and is therefore
30    /// subject to lost updates under concurrency; adapters that can do better
31    /// (a single lock or transaction) should override it. The built-in
32    /// [`InMemoryStateStore`] and `RedbStateStore` do.
33    fn update(
34        &self,
35        host: &str,
36        update: &mut dyn FnMut(DomainState) -> DomainState,
37    ) -> Result<DomainState, Error> {
38        let current = self.get(host)?.unwrap_or_else(|| DomainState::new(host));
39        let next = update(current);
40        self.put(&next)?;
41        Ok(next)
42    }
43}
44
45/// Ephemeral, in-process [`StateStore`]. The default when no persistent backend
46/// is configured; also handy in tests.
47#[derive(Debug, Default)]
48pub struct InMemoryStateStore {
49    inner: Mutex<HashMap<String, DomainState>>,
50}
51
52impl InMemoryStateStore {
53    /// Create an empty in-memory store.
54    pub fn new() -> Self {
55        Self::default()
56    }
57}
58
59impl StateStore for InMemoryStateStore {
60    fn get(&self, host: &str) -> Result<Option<DomainState>, Error> {
61        Ok(self
62            .inner
63            .lock()
64            .expect("state store lock poisoned")
65            .get(host)
66            .cloned())
67    }
68
69    fn put(&self, state: &DomainState) -> Result<(), Error> {
70        self.inner
71            .lock()
72            .expect("state store lock poisoned")
73            .insert(state.host.clone(), state.clone());
74        Ok(())
75    }
76
77    fn remove(&self, host: &str) -> Result<(), Error> {
78        self.inner
79            .lock()
80            .expect("state store lock poisoned")
81            .remove(host);
82        Ok(())
83    }
84
85    fn update(
86        &self,
87        host: &str,
88        update: &mut dyn FnMut(DomainState) -> DomainState,
89    ) -> Result<DomainState, Error> {
90        let mut guard = self.inner.lock().expect("state store lock poisoned");
91        let current = guard
92            .get(host)
93            .cloned()
94            .unwrap_or_else(|| DomainState::new(host));
95        let next = update(current);
96        guard.insert(next.host.clone(), next.clone());
97        Ok(next)
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use crate::state::Outcome;
105    use std::time::Duration;
106
107    #[test]
108    fn in_memory_round_trips_state() {
109        let store = InMemoryStateStore::new();
110        assert_eq!(store.get("example.com").unwrap(), None);
111
112        let state = DomainState::new("example.com").record(
113            Outcome::Success,
114            Some("http://p:1".into()),
115            10,
116            Duration::ZERO,
117        );
118        store.put(&state).unwrap();
119
120        let loaded = store.get("example.com").unwrap().unwrap();
121        assert_eq!(loaded, state);
122    }
123
124    #[test]
125    fn in_memory_update_creates_then_modifies() {
126        let store = InMemoryStateStore::new();
127
128        // First update on a missing host starts from a fresh state.
129        let s1 = store
130            .update("example.com", &mut |cur| {
131                cur.record(Outcome::Blocked, None, 1, Duration::ZERO)
132            })
133            .unwrap();
134        assert_eq!(s1.failures, 1);
135        assert_eq!(s1.host, "example.com");
136
137        // Second update sees the persisted value.
138        let s2 = store
139            .update("example.com", &mut |cur| {
140                cur.record(Outcome::Success, None, 2, Duration::ZERO)
141            })
142            .unwrap();
143        assert_eq!(s2.failures, 1);
144        assert_eq!(s2.successes, 1);
145        assert_eq!(store.get("example.com").unwrap().unwrap(), s2);
146    }
147
148    #[test]
149    fn in_memory_remove_deletes() {
150        let store = InMemoryStateStore::new();
151        store.put(&DomainState::new("h")).unwrap();
152        store.remove("h").unwrap();
153        assert_eq!(store.get("h").unwrap(), None);
154        // Removing a missing host is a no-op.
155        store.remove("h").unwrap();
156    }
157}
stealthscraper_rs/state/store.rs

stealthscraper_rs/state/
store.rs
