millipede_core/sitemap/
tandem.rs

1use std::{sync::Arc, time::Duration};
2
3use crate::{
4    request::Request,
5    storage::{
6        AddOptions, BatchAddHandle, Lease, LeaseId, QueueOpInfo, ReclaimOptions, RequestQueue,
7        RequestSource, StorageError, StorageResult,
8    },
9};
10
11use super::SitemapRequestList;
12
13/// A request queue that lazily feeds sitemap entries into another queue.
14///
15/// The wrapper lets the crawler consume already-queued work first, then drains the sitemap in
16/// bounded batches whenever the inner queue becomes empty.
17///
18/// # Example
19///
20/// ```no_run
21/// use std::sync::Arc;
22/// use millipede_core::{
23///     sitemap::{RequestQueueWithSitemap, SitemapRequestList},
24///     storage::RequestQueue,
25/// };
26///
27/// # fn wrap(inner: Arc<dyn RequestQueue>, list: SitemapRequestList) {
28/// let queue: Arc<dyn RequestQueue> =
29///     Arc::new(RequestQueueWithSitemap::new(inner, list).batch_size(64));
30/// # drop(queue);
31/// # }
32/// ```
33pub struct RequestQueueWithSitemap {
34    inner: Arc<dyn RequestQueue>,
35    list: SitemapRequestList,
36    batch: usize,
37    drain: tokio::sync::Mutex<DrainState>,
38}
39
40#[derive(Default)]
41struct DrainState {
42    pending_add: Option<Request>,
43}
44
45impl RequestQueueWithSitemap {
46    /// Wraps `inner` with a sitemap source, draining at most 32 entries per pass.
47    pub fn new(inner: Arc<dyn RequestQueue>, list: SitemapRequestList) -> Self {
48        Self {
49            inner,
50            list,
51            batch: 32,
52            drain: tokio::sync::Mutex::new(DrainState::default()),
53        }
54    }
55
56    /// Sets the maximum number of sitemap entries drained per pass.
57    pub fn batch_size(mut self, n: usize) -> Self {
58        self.batch = n.max(1);
59        self
60    }
61
62    async fn persist_after_batch(&self) {
63        if let Err(error) = self.list.persist().await {
64            tracing::warn!(%error, "sitemap tandem checkpoint failed");
65        }
66    }
67}
68
69#[async_trait::async_trait]
70impl RequestQueue for RequestQueueWithSitemap {
71    async fn add(&self, req: Request, opts: AddOptions) -> StorageResult<QueueOpInfo> {
72        self.inner.add(req, opts).await
73    }
74
75    async fn add_batch(
76        &self,
77        reqs: Vec<RequestSource>,
78        opts: AddOptions,
79    ) -> StorageResult<BatchAddHandle> {
80        self.inner.add_batch(reqs, opts).await
81    }
82
83    async fn fetch_next(&self) -> StorageResult<Option<Lease>> {
84        let mut drain = self.drain.lock().await;
85        loop {
86            if !self.inner.is_empty().await? {
87                return self.inner.fetch_next().await;
88            }
89
90            if drain.pending_add.is_none() && self.list.is_finished().await {
91                return self.inner.fetch_next().await;
92            }
93
94            for _ in 0..self.batch {
95                if drain.pending_add.is_none() {
96                    drain.pending_add = match self.list.fetch_next_for_tandem().await {
97                        Ok(request) => request,
98                        Err(error) => {
99                            self.persist_after_batch().await;
100                            return Err(StorageError::Backend(anyhow::Error::new(error)));
101                        }
102                    };
103                    if drain.pending_add.is_none() {
104                        break;
105                    }
106                }
107
108                let request = drain
109                    .pending_add
110                    .as_ref()
111                    .expect("pending sitemap request exists")
112                    .clone();
113                let _ = self.inner.add(request, AddOptions::default()).await?;
114                drain.pending_add = None;
115            }
116            self.persist_after_batch().await;
117
118            if drain.pending_add.is_none() && self.list.is_finished().await {
119                return self.inner.fetch_next().await;
120            }
121        }
122    }
123
124    async fn mark_handled(&self, lease: Lease) -> StorageResult<()> {
125        self.inner.mark_handled(lease).await
126    }
127
128    async fn reclaim(&self, lease: Lease, opts: ReclaimOptions) -> StorageResult<()> {
129        self.inner.reclaim(lease, opts).await
130    }
131
132    async fn renew(&self, lease_id: &LeaseId, extend_by: Duration) -> StorageResult<()> {
133        self.inner.renew(lease_id, extend_by).await
134    }
135
136    async fn abandon(&self, lease: Lease) -> StorageResult<()> {
137        self.inner.abandon(lease).await
138    }
139
140    /// Returns `false` while sitemap entries remain undrained, preventing premature engine
141    /// termination; once the sitemap is drained, delegates to the inner queue.
142    async fn is_empty(&self) -> StorageResult<bool> {
143        let drain = self.drain.lock().await;
144        if drain.pending_add.is_some() {
145            return Ok(false);
146        }
147        if !self.list.is_finished().await {
148            return Ok(false);
149        }
150        self.inner.is_empty().await
151    }
152
153    /// Returns `false` while sitemap entries remain undrained, preventing premature engine
154    /// termination; once the sitemap is drained, delegates to the inner queue.
155    async fn is_finished(&self) -> StorageResult<bool> {
156        let drain = self.drain.lock().await;
157        if drain.pending_add.is_some() {
158            return Ok(false);
159        }
160        if !self.list.is_finished().await {
161            return Ok(false);
162        }
163        self.inner.is_finished().await
164    }
165
166    async fn handled_count(&self) -> StorageResult<u64> {
167        self.inner.handled_count().await
168    }
169
170    /// Returns only the inner queue's pending count. Undrained sitemap entries are not included
171    /// because their count is unknown until the sitemap is streamed.
172    async fn pending_count(&self) -> StorageResult<u64> {
173        self.inner.pending_count().await
174    }
175}
millipede_core/sitemap/tandem.rs

millipede_core/sitemap/
tandem.rs
