millipede_storage_memory/
policy.rs

1use millipede_core::request::Request;
2use std::collections::{HashMap, VecDeque};
3
4/// Ordering policy used by an in-memory request queue.
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
6#[non_exhaustive]
7pub enum MemoryQueuePolicy {
8    /// Process requests in first-in, first-out order.
9    #[default]
10    Fifo,
11    /// Alternate between URL hosts while preserving per-host order.
12    DomainRoundRobin,
13}
14
15#[derive(Debug)]
16pub(crate) enum Frontier {
17    Fifo(VecDeque<Request>),
18    DomainRoundRobin(DomainRoundRobin),
19}
20
21impl Frontier {
22    pub(crate) fn new(policy: MemoryQueuePolicy) -> Self {
23        match policy {
24            MemoryQueuePolicy::Fifo => Self::Fifo(VecDeque::new()),
25            MemoryQueuePolicy::DomainRoundRobin => {
26                Self::DomainRoundRobin(DomainRoundRobin::default())
27            }
28        }
29    }
30
31    pub(crate) fn push_back(&mut self, request: Request) {
32        match self {
33            Self::Fifo(requests) => requests.push_back(request),
34            Self::DomainRoundRobin(requests) => requests.push_back(request),
35        }
36    }
37
38    pub(crate) fn push_front(&mut self, request: Request) {
39        match self {
40            Self::Fifo(requests) => requests.push_front(request),
41            Self::DomainRoundRobin(requests) => requests.push_front(request),
42        }
43    }
44
45    pub(crate) fn pop_front(&mut self) -> Option<Request> {
46        match self {
47            Self::Fifo(requests) => requests.pop_front(),
48            Self::DomainRoundRobin(requests) => requests.pop_front(),
49        }
50    }
51
52    pub(crate) fn len(&self) -> usize {
53        match self {
54            Self::Fifo(requests) => requests.len(),
55            Self::DomainRoundRobin(requests) => requests.len(),
56        }
57    }
58
59    pub(crate) fn is_empty(&self) -> bool {
60        match self {
61            Self::Fifo(requests) => requests.is_empty(),
62            Self::DomainRoundRobin(requests) => requests.is_empty(),
63        }
64    }
65}
66
67/// A pure frontier that rotates fairly between URL hosts.
68#[derive(Debug, Default)]
69pub struct DomainRoundRobin {
70    by_host: HashMap<String, VecDeque<Request>>,
71    rotation: VecDeque<String>,
72}
73
74impl DomainRoundRobin {
75    fn host_key(request: &Request) -> String {
76        request.url.host_str().unwrap_or("").to_ascii_lowercase()
77    }
78
79    fn push_back(&mut self, request: Request) {
80        let host = Self::host_key(&request);
81        if let Some(requests) = self.by_host.get_mut(&host) {
82            requests.push_back(request);
83        } else {
84            self.by_host.insert(host.clone(), VecDeque::from([request]));
85            self.rotation.push_back(host);
86        }
87    }
88
89    fn push_front(&mut self, request: Request) {
90        let host = Self::host_key(&request);
91        if let Some(requests) = self.by_host.get_mut(&host) {
92            requests.push_front(request);
93        } else {
94            self.by_host.insert(host.clone(), VecDeque::from([request]));
95            self.rotation.push_front(host);
96        }
97    }
98
99    fn pop_front(&mut self) -> Option<Request> {
100        let host = self.rotation.pop_front()?;
101        let requests = self
102            .by_host
103            .get_mut(&host)
104            .expect("rotation only contains active hosts");
105        let request = requests.pop_front();
106        if requests.is_empty() {
107            self.by_host.remove(&host);
108        } else {
109            self.rotation.push_back(host);
110        }
111        request
112    }
113
114    fn len(&self) -> usize {
115        self.by_host.values().map(VecDeque::len).sum()
116    }
117
118    fn is_empty(&self) -> bool {
119        self.by_host.is_empty()
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::DomainRoundRobin;
126    use millipede_core::request::Request;
127
128    fn request(url: &str) -> Request {
129        Request::get(url).build().unwrap()
130    }
131
132    fn host(frontier: &mut DomainRoundRobin) -> String {
133        frontier
134            .pop_front()
135            .unwrap()
136            .url
137            .host_str()
138            .unwrap()
139            .to_owned()
140    }
141
142    #[test]
143    fn strictly_rotates_between_hosts() {
144        let mut frontier = DomainRoundRobin::default();
145        for url in [
146            "https://a.com/1",
147            "https://a.com/2",
148            "https://a.com/3",
149            "https://b.com/1",
150            "https://b.com/2",
151            "https://c.com/1",
152        ] {
153            frontier.push_back(request(url));
154        }
155
156        let hosts = (0..6).map(|_| host(&mut frontier)).collect::<Vec<_>>();
157        assert_eq!(
158            hosts,
159            ["a.com", "b.com", "c.com", "a.com", "b.com", "a.com"]
160        );
161    }
162
163    #[test]
164    fn push_front_heads_new_host_queue_and_rotation() {
165        let mut frontier = DomainRoundRobin::default();
166        frontier.push_back(request("https://a.com/old"));
167        frontier.push_front(request("https://b.com/front"));
168
169        assert_eq!(frontier.pop_front().unwrap().url.path(), "/front");
170        assert_eq!(frontier.pop_front().unwrap().url.path(), "/old");
171    }
172
173    #[test]
174    fn push_front_precedes_items_for_an_existing_host() {
175        let mut frontier = DomainRoundRobin::default();
176        frontier.push_back(request("https://a.com/old-1"));
177        frontier.push_back(request("https://a.com/old-2"));
178        frontier.push_front(request("https://a.com/front"));
179
180        assert_eq!(frontier.pop_front().unwrap().url.path(), "/front");
181        assert_eq!(frontier.pop_front().unwrap().url.path(), "/old-1");
182        assert_eq!(frontier.pop_front().unwrap().url.path(), "/old-2");
183    }
184
185    #[test]
186    fn exhausted_host_leaves_rotation() {
187        let mut frontier = DomainRoundRobin::default();
188        for url in [
189            "https://a.com/1",
190            "https://b.com/1",
191            "https://c.com/1",
192            "https://a.com/2",
193            "https://b.com/2",
194            "https://a.com/3",
195            "https://b.com/3",
196        ] {
197            frontier.push_back(request(url));
198        }
199
200        assert_eq!(host(&mut frontier), "a.com");
201        assert_eq!(host(&mut frontier), "b.com");
202        assert_eq!(host(&mut frontier), "c.com");
203        let remaining = (0..4).map(|_| host(&mut frontier)).collect::<Vec<_>>();
204        assert_eq!(remaining, ["a.com", "b.com", "a.com", "b.com"]);
205    }
206
207    #[test]
208    fn tracks_length_and_emptiness() {
209        let mut frontier = DomainRoundRobin::default();
210        assert!(frontier.is_empty());
211        assert_eq!(frontier.len(), 0);
212        frontier.push_back(request("https://a.com/1"));
213        frontier.push_back(request("https://b.com/1"));
214        assert!(!frontier.is_empty());
215        assert_eq!(frontier.len(), 2);
216        frontier.pop_front();
217        assert_eq!(frontier.len(), 1);
218        frontier.pop_front();
219        assert!(frontier.is_empty());
220    }
221}
millipede_storage_memory/policy.rs

millipede_storage_memory/
policy.rs
