spire_core/frontier/
memory.rs

1//! In-memory storage: [`InMemFrontier`], a FIFO, LIFO or priority queue.
2
3use std::cmp::{Ordering, Reverse};
4use std::collections::{BinaryHeap, HashMap, VecDeque};
5use std::fmt;
6use std::future::{Future, ready};
7use std::sync::Mutex;
8
9use tokio::time::Instant;
10
11use crate::Result;
12use crate::frontier::seed::{Fingerprint, Seed};
13use crate::frontier::{Ack, Admission, Frontier, NextSeed, Report, Schedule};
14use crate::sync::lock;
15
16/// An in-memory [`Frontier`] that crawls in insertion, reverse or priority
17/// order.
18pub struct InMemFrontier {
19    seeds: Mutex<Store>,
20    /// Created on the first schedule, so the constructors stay `const`.
21    scheduled: Mutex<Option<Scheduled>>,
22}
23
24/// Seeds held back until a time, one per fingerprint, see [`Schedule`].
25#[derive(Default)]
26struct Scheduled {
27    /// Every scheduled seed, with when it is due and the sequence of the
28    /// heap entry that is current for it.
29    seeds: HashMap<Fingerprint, (Instant, u64, Seed)>,
30    /// Due times, earliest first. Entries a later schedule replaced stay
31    /// until they come up, and are skipped then.
32    due: BinaryHeap<Reverse<(Instant, u64, Fingerprint)>>,
33    sequence: u64,
34}
35
36impl Scheduled {
37    fn insert(&mut self, seed: Seed, due: Instant) {
38        self.sequence += 1;
39        let fingerprint = seed.fingerprint();
40        self.due.push(Reverse((due, self.sequence, fingerprint)));
41        self.seeds.insert(fingerprint, (due, self.sequence, seed));
42    }
43
44    /// Drops heap entries that a later schedule replaced, so the top is
45    /// current.
46    fn settle(&mut self) {
47        while let Some(Reverse((_, sequence, fingerprint))) = self.due.peek() {
48            let current = self
49                .seeds
50                .get(fingerprint)
51                .is_some_and(|(_, current, _)| current == sequence);
52            if current {
53                break;
54            }
55            self.due.pop();
56        }
57    }
58
59    fn pop_due(&mut self, now: Instant) -> Option<Seed> {
60        self.settle();
61        let Reverse((due, _, fingerprint)) = *self.due.peek()?;
62        if due > now {
63            return None;
64        }
65        self.due.pop();
66        self.seeds.remove(&fingerprint).map(|(_, _, seed)| seed)
67    }
68
69    fn next_due(&mut self) -> Option<Instant> {
70        self.settle();
71        self.due.peek().map(|Reverse((due, _, _))| *due)
72    }
73}
74
75enum Store {
76    Fifo(VecDeque<Seed>),
77    Lifo(VecDeque<Seed>),
78    Priority(BinaryHeap<Prioritized>, u64),
79}
80
81/// A seed ordered by priority, then by insertion.
82struct Prioritized {
83    priority: i32,
84    sequence: u64,
85    seed: Seed,
86}
87
88impl Ord for Prioritized {
89    fn cmp(&self, other: &Self) -> Ordering {
90        self.priority
91            .cmp(&other.priority)
92            .then_with(|| other.sequence.cmp(&self.sequence))
93    }
94}
95
96impl PartialOrd for Prioritized {
97    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
98        Some(self.cmp(other))
99    }
100}
101
102impl PartialEq for Prioritized {
103    fn eq(&self, other: &Self) -> bool {
104        self.cmp(other).is_eq()
105    }
106}
107
108impl Eq for Prioritized {}
109
110impl Store {
111    fn len(&self) -> usize {
112        match self {
113            Self::Fifo(seeds) | Self::Lifo(seeds) => seeds.len(),
114            Self::Priority(seeds, _) => seeds.len(),
115        }
116    }
117
118    const fn order(&self) -> &'static str {
119        match self {
120            Self::Fifo(_) => "fifo",
121            Self::Lifo(_) => "lifo",
122            Self::Priority(..) => "priority",
123        }
124    }
125}
126
127impl InMemFrontier {
128    /// Crawls seeds in the order they were pushed: breadth-first.
129    #[must_use]
130    pub const fn fifo() -> Self {
131        Self::with_store(Store::Fifo(VecDeque::new()))
132    }
133
134    /// Crawls the most recently pushed seed first: depth-first.
135    #[must_use]
136    pub const fn lifo() -> Self {
137        Self::with_store(Store::Lifo(VecDeque::new()))
138    }
139
140    /// Crawls the seed with the highest [`priority`](Seed::with_priority)
141    /// first, and seeds of equal priority in the order they were pushed.
142    #[must_use]
143    pub const fn priority() -> Self {
144        Self::with_store(Store::Priority(BinaryHeap::new(), 0))
145    }
146
147    const fn with_store(store: Store) -> Self {
148        Self {
149            seeds: Mutex::new(store),
150            scheduled: Mutex::new(None),
151        }
152    }
153
154    fn insert(&self, seed: Seed) {
155        match &mut *lock(&self.seeds) {
156            Store::Fifo(seeds) | Store::Lifo(seeds) => seeds.push_back(seed),
157            Store::Priority(seeds, sequence) => {
158                *sequence += 1;
159                let priority = seed.priority();
160                seeds.push(Prioritized {
161                    priority,
162                    sequence: *sequence,
163                    seed,
164                });
165            }
166        }
167    }
168}
169
170impl Default for InMemFrontier {
171    fn default() -> Self {
172        Self::fifo()
173    }
174}
175
176impl fmt::Debug for InMemFrontier {
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        let store = lock(&self.seeds);
179        let scheduled = lock(&self.scheduled)
180            .as_ref()
181            .map_or(0, |scheduled| scheduled.seeds.len());
182        f.debug_struct("InMemFrontier")
183            .field("order", &store.order())
184            .field("len", &store.len())
185            .field("scheduled", &scheduled)
186            .finish()
187    }
188}
189
190impl Frontier for InMemFrontier {
191    fn push(&self, seed: Seed) -> impl Future<Output = Result<Admission>> + Send {
192        self.insert(seed);
193        ready(Ok(Admission::Queued))
194    }
195
196    /// Hands out a due scheduled seed first, then the queue.
197    fn pop(&self) -> impl Future<Output = Result<NextSeed>> + Send {
198        let due = lock(&self.scheduled)
199            .as_mut()
200            .and_then(|scheduled| scheduled.pop_due(Instant::now()));
201        if let Some(seed) = due {
202            return ready(Ok(NextSeed::Ready(seed)));
203        }
204        let seed = match &mut *lock(&self.seeds) {
205            Store::Fifo(seeds) => seeds.pop_front(),
206            Store::Lifo(seeds) => seeds.pop_back(),
207            Store::Priority(seeds, _) => seeds.pop().map(|prioritized| prioritized.seed),
208        };
209        ready(Ok(seed.map_or(NextSeed::Empty, NextSeed::Ready)))
210    }
211
212    /// Re-queues seeds the crawl stopped before finishing, so a
213    /// [`snapshot`](Frontier::snapshot) taken after the run includes them.
214    fn report(&self, report: Report) -> impl Future<Output = Result<Ack>> + Send {
215        let ack = match report {
216            Report::Cancelled(seed) => {
217                self.insert(seed);
218                Ack::Applied
219            }
220            Report::Retry { .. } => Ack::RetryRefused,
221            _ => Ack::Applied,
222        };
223        ready(Ok(ack))
224    }
225
226    fn size_hint(&self) -> Option<usize> {
227        Some(lock(&self.seeds).len())
228    }
229
230    /// Lists queued and scheduled seeds, the latter without when they are
231    /// due.
232    fn snapshot(&self) -> Option<Vec<Seed>> {
233        let store = lock(&self.seeds);
234        let mut seeds: Vec<Seed> = match &*store {
235            Store::Fifo(seeds) | Store::Lifo(seeds) => seeds.iter().cloned().collect(),
236            Store::Priority(seeds, _) => seeds.iter().map(|entry| entry.seed.clone()).collect(),
237        };
238        if let Some(scheduled) = &*lock(&self.scheduled) {
239            seeds.extend(scheduled.seeds.values().map(|(_, _, seed)| seed.clone()));
240        }
241        Some(seeds)
242    }
243}
244
245impl Schedule for InMemFrontier {
246    fn schedule(&self, seed: Seed, due: Instant) -> impl Future<Output = Result<()>> + Send {
247        lock(&self.scheduled)
248            .get_or_insert_with(Scheduled::default)
249            .insert(seed, due);
250        ready(Ok(()))
251    }
252
253    fn next_due(&self) -> impl Future<Output = Result<Option<Instant>>> + Send {
254        ready(Ok(lock(&self.scheduled)
255            .as_mut()
256            .and_then(Scheduled::next_due)))
257    }
258}
259
260impl Extend<Seed> for InMemFrontier {
261    /// Queues seeds as is, e.g. from a [`snapshot`](Frontier::snapshot),
262    /// without any facade refusing them.
263    fn extend<I: IntoIterator<Item = Seed>>(&mut self, seeds: I) {
264        for seed in seeds {
265            self.insert(seed);
266        }
267    }
268}
269
270impl FromIterator<Seed> for InMemFrontier {
271    /// Collects seeds into a FIFO frontier, e.g. to restore a snapshot.
272    fn from_iter<I: IntoIterator<Item = Seed>>(seeds: I) -> Self {
273        let mut frontier = Self::fifo();
274        frontier.extend(seeds);
275        frontier
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::frontier::test_util::{pop_path, seed};
283
284    #[tokio::test]
285    async fn fifo_and_lifo_order() {
286        let fifo = InMemFrontier::fifo();
287        let lifo = InMemFrontier::lifo();
288        for path in ["a", "b"] {
289            assert!(fifo.push(seed(path)).await.unwrap().is_queued());
290            assert!(lifo.push(seed(path)).await.unwrap().is_queued());
291        }
292        assert_eq!(fifo.size_hint(), Some(2));
293
294        assert_eq!(pop_path(&fifo).await.as_deref(), Some("a"));
295        assert_eq!(pop_path(&lifo).await.as_deref(), Some("b"));
296        assert_eq!(pop_path(&fifo).await.as_deref(), Some("b"));
297        assert!(matches!(fifo.pop().await.unwrap(), NextSeed::Empty));
298    }
299
300    #[tokio::test]
301    async fn priority_order() {
302        let frontier = InMemFrontier::priority();
303        for (path, priority) in [("low", -1), ("first", 0), ("high", 5), ("second", 0)] {
304            frontier
305                .push(seed(path).with_priority(priority))
306                .await
307                .unwrap();
308        }
309        let mut order = Vec::new();
310        while let Some(path) = pop_path(&frontier).await {
311            order.push(path);
312        }
313        assert_eq!(order, ["high", "first", "second", "low"]);
314    }
315}
spire_core/frontier/memory.rs

spire_core/frontier/
memory.rs
