pagesift_api/jobs/
mod.rs

1//! 异步任务（批量抓取和整站爬取）：进度、结果、错误，过期后清理。
2//!
3//! 给了存储（F-6.4）时任务落盘，服务重启后：已结束的任务照样能查（直到过期），
4//! 进行中的任务从断点继续。存储布局（都在 [`pagesift_store::Tasks`] 里）：
5//!
6//! - KV `job/{id}`：任务元数据（类型、状态、创建时间、原始请求、webhook、调用方）；
7//! - 列表 `job/{id}/pages`：完成的文档，按完成顺序，序号就是分页下标；
8//! - 列表 `job/{id}/errors`：页面错误；
9//! - 列表 `job/{id}/frontier`：爬取入过队的 URL 和深度，按入队顺序；
10//! - 集合 `job/{id}/done`：已完成的 URL（爬取）或下标（批量抓取）。
11//!
12//! 每个页面先写文档再记完成：崩溃时最多重复抓一个页面，不会丢页面。
13//! 不给存储时一切留在内存里，行为和以前一样。
14
15use std::collections::{BTreeMap, HashMap};
16use std::sync::atomic::AtomicBool;
17use std::sync::{Arc, Mutex as StdMutex};
18use std::time::SystemTime;
19
20use pagesift_store::Store;
21use tokio::sync::{Mutex, Semaphore};
22
23use crate::engine::{CrawlResume, Engine};
24use crate::error::{ApiError, ErrorCode};
25use crate::limits::Limiter;
26use crate::metrics::Metrics;
27use crate::types::{CrawlRequest, JobError, ScrapeOptions};
28use crate::webhook::{Event, Hook, Webhooks};
29
30pub mod mirror;
31mod sink;
32pub mod state;
33pub(crate) mod store;
34pub mod time;
35
36pub use state::{JOB_TTL, Job, JobKind, JobState, Start, StatusPage};
37pub use time::rfc3339;
38
39use sink::JobSink;
40use store::{BatchSpec, Meta, delete_job, list_key, meta_key, to_ms};
41
42struct Inner {
43    map: Mutex<HashMap<String, Job>>,
44    store: Option<Store>,
45    limiter: Arc<Limiter>,
46    webhooks: Option<Webhooks>,
47    metrics: Arc<Metrics>,
48    /// 多服务器网关：任务存储的写入同时复制到温备（见 [`mirror`]）。
49    mirror: Option<Arc<mirror::JobMirror>>,
50}
51
52#[derive(Clone)]
53pub struct Jobs {
54    inner: Arc<Inner>,
55}
56
57impl Default for Jobs {
58    /// 只在内存里的任务表：不落盘、不限额、不发 webhook。
59    fn default() -> Self {
60        Self::new(None, Arc::default(), None, Arc::default(), None)
61    }
62}
63
64impl Jobs {
65    pub(crate) fn new(
66        store: Option<Store>,
67        limiter: Arc<Limiter>,
68        webhooks: Option<Webhooks>,
69        metrics: Arc<Metrics>,
70        mirror: Option<Arc<mirror::JobMirror>>,
71    ) -> Self {
72        Self {
73            inner: Arc::new(Inner {
74                map: Mutex::new(HashMap::new()),
75                store,
76                limiter,
77                webhooks,
78                metrics,
79                mirror,
80            }),
81        }
82    }
83
84    /// 清掉过期任务（连同存储里的数据）。
85    async fn sweep(&self) {
86        let now = SystemTime::now();
87        let expired: Vec<String> = {
88            let mut map = self.inner.map.lock().await;
89            let ids: Vec<String> = map
90                .iter()
91                .filter(|(_, j)| j.expires_at() <= now)
92                .map(|(id, _)| id.clone())
93                .collect();
94            for id in &ids {
95                map.remove(id);
96            }
97            ids
98        };
99        if let (Some(store), false) = (&self.inner.store, expired.is_empty()) {
100            let m = self.inner.mirror.clone();
101            store
102                .blocking(move |s| {
103                    for id in expired {
104                        delete_job(s, m.as_deref(), &id);
105                    }
106                })
107                .await;
108        }
109    }
110
111    /// 附件写进的地方：任务存储和镜像（见 `attachments`）。
112    pub(crate) fn disk(&self) -> Option<crate::attachments::Disk> {
113        self.inner
114            .store
115            .clone()
116            .map(|store| crate::attachments::Disk {
117                store,
118                mirror: self.inner.mirror.clone(),
119            })
120    }
121
122    pub async fn get<R>(&self, id: &str, f: impl FnOnce(&Job) -> R) -> Option<R> {
123        self.sweep().await;
124        self.inner.map.lock().await.get(id).map(f)
125    }
126
127    /// 状态和从 `skip` 开始最多 `limit` 个文档。
128    pub(crate) async fn status(&self, id: &str, skip: usize, limit: usize) -> Option<StatusPage> {
129        self.sweep().await;
130        let mut page = {
131            let map = self.inner.map.lock().await;
132            let job = map.get(id)?;
133            let end = skip.saturating_add(limit).min(job.data.len());
134            StatusPage {
135                state: job.state,
136                total: job.total,
137                completed: job.completed,
138                expires_at: job.expires_at(),
139                docs: job.docs(),
140                data: job
141                    .data
142                    .get(skip.min(end)..end)
143                    .unwrap_or_default()
144                    .to_vec(),
145                summary: job.summary.clone(),
146            }
147        };
148        if let Some(store) = &self.inner.store
149            && page.data.is_empty()
150            && skip < page.docs
151        {
152            let key = list_key(id, "pages");
153            page.data = store
154                .blocking(move |s| s.tasks().list_range(key, skip as u64, limit))
155                .await
156                .unwrap_or_else(|e| {
157                    tracing::warn!(error = %e, "reading job pages failed");
158                    Vec::new()
159                })
160                .iter()
161                .filter_map(|v| serde_json::from_slice(v).ok())
162                .collect();
163        }
164        Some(page)
165    }
166
167    /// 订阅任务的变化（新文档、状态变化）。任务不存在时返回 `None`；任务过期删掉后接收端随之关闭。
168    pub(crate) async fn subscribe(&self, id: &str) -> Option<tokio::sync::watch::Receiver<()>> {
169        self.sweep().await;
170        self.inner
171            .map
172            .lock()
173            .await
174            .get(id)
175            .map(|j| j.changed.subscribe())
176    }
177
178    /// 任务的页面错误和被 robots.txt 拦下的 URL。
179    pub(crate) async fn errors(&self, id: &str) -> Option<(Vec<JobError>, Vec<String>)> {
180        self.sweep().await;
181        let (mut errors, mut robots) = {
182            let map = self.inner.map.lock().await;
183            let job = map.get(id)?;
184            (job.errors.clone(), job.robots_blocked.clone())
185        };
186        if let Some(store) = &self.inner.store {
187            // 页面错误，然后是任务级的错误（配额、失败原因、投不出去的 webhook 事件）
188            let (key, notes) = (list_key(id, "errors"), list_key(id, "notes"));
189            let stored = store
190                .blocking(move |s| {
191                    let t = s.tasks();
192                    let mut v = t.list_items(key, usize::MAX).unwrap_or_default();
193                    v.extend(t.list_items(notes, usize::MAX).unwrap_or_default());
194                    v
195                })
196                .await;
197            for v in stored {
198                if let Ok(e) = serde_json::from_slice::<JobError>(&v) {
199                    if e.code == ErrorCode::RobotsDisallowed.as_str() {
200                        robots.push(e.url.clone());
201                    }
202                    errors.push(e);
203                }
204            }
205        }
206        Some((errors, robots))
207    }
208
209    pub async fn cancel(&self, id: &str) -> bool {
210        let found = {
211            let mut map = self.inner.map.lock().await;
212            match map.get_mut(id) {
213                Some(j) => {
214                    if j.state == JobState::Scraping {
215                        j.state = JobState::Cancelled;
216                        j.touch();
217                    }
218                    true
219                }
220                None => false,
221            }
222        };
223        if found {
224            self.save_state(id).await;
225        }
226        found
227    }
228
229    pub async fn active(&self) -> Vec<String> {
230        let map = self.inner.map.lock().await;
231        map.iter()
232            .filter(|(_, j)| j.state == JobState::Scraping)
233            .map(|(id, _)| id.clone())
234            .collect()
235    }
236
237    /// 按 (类型, 状态) 统计的任务数，给 `/metrics`。
238    pub(crate) async fn counts(&self) -> BTreeMap<(&'static str, &'static str), u64> {
239        let mut out = BTreeMap::new();
240        for j in self.inner.map.lock().await.values() {
241            *out.entry((j.kind.as_str(), j.state.as_str())).or_default() += 1;
242        }
243        out
244    }
245
246    /// 登记新任务：检查调用方的并发任务数，落盘元数据，发 `started` 事件。
247    async fn register(
248        &self,
249        kind: JobKind,
250        total: usize,
251        start: Start,
252        crawl: Option<CrawlRequest>,
253        batch: Option<BatchSpec>,
254    ) -> Result<(String, Option<Hook>), ApiError> {
255        let id = uuid::Uuid::new_v4().to_string();
256        let job = Job::new(kind, total, start.owner.clone());
257        let meta = Meta {
258            kind,
259            state: JobState::Scraping,
260            created_ms: to_ms(job.created),
261            total,
262            completed: 0,
263            owner: start.owner.as_ref().map(|c| c.account.to_string()),
264            webhook: start.webhook.clone(),
265            crawl,
266            batch,
267            summary: None,
268        };
269        {
270            let mut map = self.inner.map.lock().await;
271            if let Some(c) = &start.owner {
272                if c.limits.max_concurrent_jobs.is_some() {
273                    let active = map
274                        .values()
275                        .filter(|j| {
276                            j.state == JobState::Scraping
277                                && j.owner.as_ref().is_some_and(|o| o.account == c.account)
278                        })
279                        .count();
280                    self.inner.limiter.check_jobs(c, active)?;
281                }
282                self.inner.limiter.check_quota(c)?;
283            }
284            map.insert(id.clone(), job);
285        }
286        if let Some(store) = &self.inner.store {
287            let (key, value) = (meta_key(&id), serde_json::to_vec(&meta).unwrap_or_default());
288            let m = self.inner.mirror.clone();
289            let put = move |s: &Store| mirror::write(s, m.as_deref(), |w| w.kv_put(&key, &value));
290            if let Err(e) = store.blocking(put).await {
291                tracing::warn!(job = %id, error = %e, "saving job failed; it will not survive a restart");
292            }
293        }
294        self.inner.metrics.job_started(kind.as_str());
295        tracing::info!(
296            job = %id,
297            kind = kind.as_str(),
298            account = start.owner.as_ref().map(|c| &*c.account),
299            "job started"
300        );
301        let hook = match (&self.inner.webhooks, start.webhook) {
302            (Some(w), Some(cfg)) => {
303                let hook = w.start(&id, kind.as_str(), cfg, self.inner.store.clone());
304                hook.send(Event::Started, Vec::new(), None).await;
305                Some(hook)
306            }
307            _ => None,
308        };
309        Ok((id, hook))
310    }
311
312    /// 启动批量抓取，立即返回任务 id；每个 URL 的结果和错误陆续写入任务。
313    pub async fn start_batch(
314        &self,
315        engine: Arc<dyn Engine>,
316        urls: Vec<String>,
317        opts: ScrapeOptions,
318        concurrency: usize,
319    ) -> String {
320        self.start_batch_with(engine, urls, opts, concurrency, Start::default())
321            .await
322            .expect("no limits without a caller")
323    }
324
325    pub(crate) async fn start_batch_with(
326        &self,
327        engine: Arc<dyn Engine>,
328        urls: Vec<String>,
329        opts: ScrapeOptions,
330        concurrency: usize,
331        start: Start,
332    ) -> Result<String, ApiError> {
333        let spec = BatchSpec {
334            urls,
335            options: opts,
336            concurrency,
337        };
338        let total = spec.urls.len();
339        let (id, hook) = self
340            .register(
341                JobKind::BatchScrape,
342                total,
343                start,
344                None,
345                self.inner.store.is_some().then(|| spec.clone()),
346            )
347            .await?;
348        let pending = spec.urls.into_iter().enumerate().collect();
349        self.spawn_batch(id.clone(), engine, pending, spec.options, concurrency, hook);
350        Ok(id)
351    }
352
353    fn sink(
354        &self,
355        id: &str,
356        kind: JobKind,
357        hook: Option<Hook>,
358        resume: Option<CrawlResume>,
359    ) -> Arc<JobSink> {
360        Arc::new(JobSink {
361            jobs: self.clone(),
362            id: id.to_string(),
363            kind,
364            hook,
365            resume: StdMutex::new(resume),
366            quota_hit: AtomicBool::new(false),
367            keep_summary: true,
368            store_failed: StdMutex::new(None),
369        })
370    }
371
372    pub(super) fn spawn_batch(
373        &self,
374        id: String,
375        engine: Arc<dyn Engine>,
376        pending: Vec<(usize, String)>,
377        opts: ScrapeOptions,
378        concurrency: usize,
379        hook: Option<Hook>,
380    ) {
381        let sink = self.sink(&id, JobKind::BatchScrape, hook, None);
382        let opts = Arc::new(opts);
383        tokio::spawn(async move {
384            let sem = Arc::new(Semaphore::new(concurrency.max(1)));
385            let mut tasks = tokio::task::JoinSet::new();
386            for (i, url) in pending {
387                let (sink, engine, opts, sem) =
388                    (sink.clone(), engine.clone(), opts.clone(), sem.clone());
389                tasks.spawn(async move {
390                    let _permit = sem.acquire_owned().await.expect("semaphore open");
391                    if sink.stopped().await {
392                        return;
393                    }
394                    let done = i.to_string();
395                    match engine.scrape(&url, &opts).await {
396                        Ok(doc) => sink.record_page(done, doc).await,
397                        Err(e) => sink.record_error(done, url, e).await,
398                    }
399                });
400            }
401            while tasks.join_next().await.is_some() {}
402            sink.finish(Ok(())).await;
403        });
404    }
405
406    /// 启动整站爬取，立即返回任务 id。
407    pub async fn start_crawl(&self, engine: Arc<dyn Engine>, req: CrawlRequest) -> String {
408        self.start_crawl_with(engine, req, Start::default())
409            .await
410            .expect("no limits without a caller")
411    }
412
413    pub(crate) async fn start_crawl_with(
414        &self,
415        engine: Arc<dyn Engine>,
416        req: CrawlRequest,
417        start: Start,
418    ) -> Result<String, ApiError> {
419        let (id, hook) = self
420            .register(
421                JobKind::Crawl,
422                0,
423                start,
424                self.inner.store.is_some().then(|| req.clone()),
425                None,
426            )
427            .await?;
428        self.spawn_crawl(id.clone(), engine, req, hook, None);
429        Ok(id)
430    }
431
432    pub(super) fn spawn_crawl(
433        &self,
434        id: String,
435        engine: Arc<dyn Engine>,
436        req: CrawlRequest,
437        hook: Option<Hook>,
438        resume: Option<CrawlResume>,
439    ) {
440        let mut sink = self.sink(&id, JobKind::Crawl, hook, resume);
441        if let Some(s) = Arc::get_mut(&mut sink) {
442            s.keep_summary = req.query.is_some();
443        }
444        tokio::spawn(async move {
445            let result = engine.crawl(&req, sink.clone()).await;
446            sink.finish(result.map_err(|e| (req.url.clone(), e))).await;
447        });
448    }
449}
pagesift_api/jobs/mod.rs

pagesift_api/jobs/
mod.rs
