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, CrawlSink, Engine};
24use crate::error::{ApiError, ErrorCode};
25use crate::limits::{Caller, 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    /// 任务存储（监控也落在这里）。
112    pub(crate) fn store(&self) -> Option<&Store> {
113        self.inner.store.as_ref()
114    }
115
116    /// 附件写进的地方：任务存储和镜像（见 `attachments`）。
117    pub(crate) fn disk(&self) -> Option<crate::attachments::Disk> {
118        self.inner
119            .store
120            .clone()
121            .map(|store| crate::attachments::Disk {
122                store,
123                mirror: self.inner.mirror.clone(),
124            })
125    }
126
127    pub async fn get<R>(&self, id: &str, f: impl FnOnce(&Job) -> R) -> Option<R> {
128        self.sweep().await;
129        self.inner.map.lock().await.get(id).map(f)
130    }
131
132    /// 状态和从 `skip` 开始最多 `limit` 个文档。
133    pub(crate) async fn status(&self, id: &str, skip: usize, limit: usize) -> Option<StatusPage> {
134        self.sweep().await;
135        let mut page = {
136            let map = self.inner.map.lock().await;
137            let job = map.get(id)?;
138            let end = skip.saturating_add(limit).min(job.data.len());
139            StatusPage {
140                state: job.state,
141                total: job.total,
142                completed: job.completed,
143                expires_at: job.expires_at(),
144                docs: job.docs(),
145                data: job
146                    .data
147                    .get(skip.min(end)..end)
148                    .unwrap_or_default()
149                    .to_vec(),
150                summary: job.summary.clone(),
151            }
152        };
153        if let Some(store) = &self.inner.store
154            && page.data.is_empty()
155            && skip < page.docs
156        {
157            let key = list_key(id, "pages");
158            page.data = store
159                .blocking(move |s| s.tasks().list_range(key, skip as u64, limit))
160                .await
161                .unwrap_or_else(|e| {
162                    tracing::warn!(error = %e, "reading job pages failed");
163                    Vec::new()
164                })
165                .iter()
166                .filter_map(|v| serde_json::from_slice(v).ok())
167                .collect();
168            for doc in &mut page.data {
169                for a in &mut doc.attachments {
170                    if a.inline_data && a.data.is_none() {
171                        a.data = crate::attachments::load(self.disk(), a.id.clone())
172                            .await
173                            .map(|(data, _)| data);
174                    }
175                }
176            }
177        }
178        Some(page)
179    }
180
181    /// 订阅任务的变化（新文档、状态变化）。任务不存在时返回 `None`；任务过期删掉后接收端随之关闭。
182    pub(crate) async fn subscribe(&self, id: &str) -> Option<tokio::sync::watch::Receiver<()>> {
183        self.sweep().await;
184        self.inner
185            .map
186            .lock()
187            .await
188            .get(id)
189            .map(|j| j.changed.subscribe())
190    }
191
192    pub(crate) async fn subscribe_previews(
193        &self,
194        id: &str,
195    ) -> Option<tokio::sync::broadcast::Receiver<crate::types::MarkdownPreview>> {
196        self.inner
197            .map
198            .lock()
199            .await
200            .get(id)
201            .map(|job| job.previews.subscribe())
202    }
203
204    /// 任务的页面错误和被 robots.txt 拦下的 URL。
205    pub(crate) async fn errors(&self, id: &str) -> Option<(Vec<JobError>, Vec<String>)> {
206        self.sweep().await;
207        let (mut errors, mut robots) = {
208            let map = self.inner.map.lock().await;
209            let job = map.get(id)?;
210            (job.errors.clone(), job.robots_blocked.clone())
211        };
212        if let Some(store) = &self.inner.store {
213            // 页面错误，然后是任务级的错误（配额、失败原因、投不出去的 webhook 事件）
214            let (key, notes) = (list_key(id, "errors"), list_key(id, "notes"));
215            let stored = store
216                .blocking(move |s| {
217                    let t = s.tasks();
218                    let mut v = t.list_items(key, usize::MAX).unwrap_or_default();
219                    v.extend(t.list_items(notes, usize::MAX).unwrap_or_default());
220                    v
221                })
222                .await;
223            for v in stored {
224                if let Ok(e) = serde_json::from_slice::<JobError>(&v) {
225                    if e.code == ErrorCode::RobotsDisallowed.as_str() {
226                        robots.push(e.url.clone());
227                    }
228                    errors.push(e);
229                }
230            }
231        }
232        Some((errors, robots))
233    }
234
235    pub async fn cancel(&self, id: &str) -> bool {
236        let found = {
237            let mut map = self.inner.map.lock().await;
238            match map.get_mut(id) {
239                Some(j) => {
240                    if j.state == JobState::Scraping {
241                        j.state = JobState::Cancelled;
242                        j.touch();
243                    }
244                    true
245                }
246                None => false,
247            }
248        };
249        if found {
250            self.save_state(id).await;
251        }
252        found
253    }
254
255    pub async fn active(&self) -> Vec<String> {
256        let map = self.inner.map.lock().await;
257        map.iter()
258            .filter(|(_, j)| j.state == JobState::Scraping)
259            .map(|(id, _)| id.clone())
260            .collect()
261    }
262
263    pub(crate) async fn owned_by(&self, id: &str, caller: Option<&Caller>) -> bool {
264        self.sweep().await;
265        let map = self.inner.map.lock().await;
266        map.get(id).is_some_and(|job| match (&job.owner, caller) {
267            (Some(owner), Some(caller)) => owner.account == caller.account,
268            (None, Some(caller)) => caller.account.as_ref() == "anonymous",
269            (None, None) => true,
270            _ => false,
271        })
272    }
273
274    pub(crate) async fn active_for(&self, caller: Option<&Caller>) -> Vec<String> {
275        let map = self.inner.map.lock().await;
276        map.iter()
277            .filter(|(_, job)| {
278                job.state == JobState::Scraping
279                    && match (&job.owner, caller) {
280                        (Some(owner), Some(caller)) => owner.account == caller.account,
281                        (None, Some(caller)) => caller.account.as_ref() == "anonymous",
282                        (None, None) => true,
283                        _ => false,
284                    }
285            })
286            .map(|(id, _)| id.clone())
287            .collect()
288    }
289
290    /// 按 (类型, 状态) 统计的任务数，给 `/metrics`。
291    pub(crate) async fn counts(&self) -> BTreeMap<(&'static str, &'static str), u64> {
292        let mut out = BTreeMap::new();
293        for j in self.inner.map.lock().await.values() {
294            *out.entry((j.kind.as_str(), j.state.as_str())).or_default() += 1;
295        }
296        out
297    }
298
299    /// 登记新任务：检查调用方的并发任务数，落盘元数据，发 `started` 事件。
300    async fn register(
301        &self,
302        kind: JobKind,
303        total: usize,
304        start: Start,
305        crawl: Option<CrawlRequest>,
306        batch: Option<BatchSpec>,
307    ) -> Result<(String, Option<Hook>), ApiError> {
308        let id = uuid::Uuid::new_v4().to_string();
309        let job = Job::new(kind, total, start.owner.clone());
310        let meta = Meta {
311            kind,
312            state: JobState::Scraping,
313            created_ms: to_ms(job.created),
314            total,
315            completed: 0,
316            owner: start.owner.as_ref().map(|c| c.account.to_string()),
317            webhook: start.webhook.clone(),
318            crawl,
319            batch,
320            summary: None,
321        };
322        {
323            let mut map = self.inner.map.lock().await;
324            if let Some(c) = &start.owner {
325                if c.limits.max_concurrent_jobs.is_some() {
326                    let active = map
327                        .values()
328                        .filter(|j| {
329                            j.state == JobState::Scraping
330                                && j.owner.as_ref().is_some_and(|o| o.account == c.account)
331                        })
332                        .count();
333                    self.inner.limiter.check_jobs(c, active)?;
334                }
335                self.inner.limiter.check_quota(c)?;
336            }
337            map.insert(id.clone(), job);
338        }
339        if let Some(store) = &self.inner.store {
340            let (key, value) = (meta_key(&id), serde_json::to_vec(&meta).unwrap_or_default());
341            let m = self.inner.mirror.clone();
342            let put = move |s: &Store| mirror::write(s, m.as_deref(), |w| w.kv_put(&key, &value));
343            if let Err(e) = store.blocking(put).await {
344                tracing::warn!(job = %id, error = %e, "saving job failed; it will not survive a restart");
345            }
346        }
347        self.inner.metrics.job_started(kind.as_str());
348        tracing::info!(
349            job = %id,
350            kind = kind.as_str(),
351            account = start.owner.as_ref().map(|c| &*c.account),
352            "job started"
353        );
354        let hook = match (&self.inner.webhooks, start.webhook) {
355            (Some(w), Some(cfg)) => {
356                let hook = w.start(&id, kind.as_str(), cfg, self.inner.store.clone());
357                hook.send(Event::Started, Vec::new(), None).await;
358                Some(hook)
359            }
360            _ => None,
361        };
362        Ok((id, hook))
363    }
364
365    /// 启动批量抓取，立即返回任务 id；每个 URL 的结果和错误陆续写入任务。
366    pub async fn start_batch(
367        &self,
368        engine: Arc<dyn Engine>,
369        urls: Vec<String>,
370        opts: ScrapeOptions,
371        concurrency: usize,
372    ) -> String {
373        self.start_batch_with(engine, urls, opts, concurrency, Start::default())
374            .await
375            .expect("no limits without a caller")
376    }
377
378    pub(crate) async fn start_batch_with(
379        &self,
380        engine: Arc<dyn Engine>,
381        urls: Vec<String>,
382        opts: ScrapeOptions,
383        concurrency: usize,
384        start: Start,
385    ) -> Result<String, ApiError> {
386        let spec = BatchSpec {
387            urls,
388            options: opts,
389            concurrency,
390        };
391        let total = spec.urls.len();
392        let (id, hook) = self
393            .register(
394                JobKind::BatchScrape,
395                total,
396                start,
397                None,
398                self.inner.store.is_some().then(|| spec.clone()),
399            )
400            .await?;
401        let pending = spec.urls.into_iter().enumerate().collect();
402        self.spawn_batch(id.clone(), engine, pending, spec.options, concurrency, hook);
403        Ok(id)
404    }
405
406    fn sink(
407        &self,
408        id: &str,
409        kind: JobKind,
410        hook: Option<Hook>,
411        resume: Option<CrawlResume>,
412    ) -> Arc<JobSink> {
413        Arc::new(JobSink {
414            jobs: self.clone(),
415            id: id.to_string(),
416            kind,
417            hook,
418            resume: StdMutex::new(resume),
419            quota_hit: AtomicBool::new(false),
420            keep_summary: true,
421            store_failed: StdMutex::new(None),
422            asset_hashes: tokio::sync::Mutex::new(None),
423            holds: tokio::sync::Mutex::new(Default::default()),
424        })
425    }
426
427    pub(super) fn spawn_batch(
428        &self,
429        id: String,
430        engine: Arc<dyn Engine>,
431        pending: Vec<(usize, String)>,
432        mut opts: ScrapeOptions,
433        concurrency: usize,
434        hook: Option<Hook>,
435    ) {
436        let sink = self.sink(&id, JobKind::BatchScrape, hook, None);
437        if opts.image_assets.is_some() {
438            opts.asset_fetch_cache = Some(Default::default());
439        }
440        tokio::spawn(async move {
441            if opts.image_assets.is_some() {
442                opts.markdown_progress = sink
443                    .jobs
444                    .inner
445                    .map
446                    .lock()
447                    .await
448                    .get(&id)
449                    .map(|job| crate::types::MarkdownProgress(job.previews.clone()));
450            }
451            if let Some(cache) = &opts.asset_fetch_cache {
452                sink.jobs.seed_asset_cache(&id, cache).await;
453            }
454            let opts = Arc::new(opts);
455            let sem = Arc::new(Semaphore::new(concurrency.max(1)));
456            let mut tasks = tokio::task::JoinSet::new();
457            for (i, url) in pending {
458                let (sink, engine, opts, sem) =
459                    (sink.clone(), engine.clone(), opts.clone(), sem.clone());
460                tasks.spawn(async move {
461                    let _permit = sem.acquire_owned().await.expect("semaphore open");
462                    if sink.stopped().await {
463                        return;
464                    }
465                    let done = i.to_string();
466                    // 发请求之前原子预留；配额不够时这一页不抓，任务其余页面也会被 stopped() 拦住。
467                    if let Err(e) = sink.reserve_page(&url).await {
468                        sink.record_error(done, url, e).await;
469                        return;
470                    }
471                    match engine.scrape(&url, &opts).await {
472                        Ok(doc) => {
473                            sink.settle_page(&url).await;
474                            sink.record_page(done, doc).await;
475                        }
476                        Err(e) => {
477                            sink.release_page(&url).await;
478                            sink.record_error(done, url, e).await;
479                        }
480                    }
481                });
482            }
483            while tasks.join_next().await.is_some() {}
484            sink.finish(Ok(())).await;
485        });
486    }
487
488    /// 启动整站爬取，立即返回任务 id。
489    pub async fn start_crawl(&self, engine: Arc<dyn Engine>, req: CrawlRequest) -> String {
490        self.start_crawl_with(engine, req, Start::default())
491            .await
492            .expect("no limits without a caller")
493    }
494
495    pub(crate) async fn start_crawl_with(
496        &self,
497        engine: Arc<dyn Engine>,
498        req: CrawlRequest,
499        start: Start,
500    ) -> Result<String, ApiError> {
501        let (id, hook) = self
502            .register(
503                JobKind::Crawl,
504                0,
505                start,
506                self.inner.store.is_some().then(|| req.clone()),
507                None,
508            )
509            .await?;
510        self.spawn_crawl(id.clone(), engine, req, hook, None);
511        Ok(id)
512    }
513
514    pub(super) fn spawn_crawl(
515        &self,
516        id: String,
517        engine: Arc<dyn Engine>,
518        mut req: CrawlRequest,
519        hook: Option<Hook>,
520        resume: Option<CrawlResume>,
521    ) {
522        if req.scrape_options.image_assets.is_some() {
523            req.scrape_options.asset_fetch_cache = Some(Default::default());
524        }
525        let mut sink = self.sink(&id, JobKind::Crawl, hook, resume);
526        if let Some(s) = Arc::get_mut(&mut sink) {
527            s.keep_summary = req.query.is_some();
528        }
529        tokio::spawn(async move {
530            if req.scrape_options.image_assets.is_some() {
531                req.scrape_options.markdown_progress = sink
532                    .jobs
533                    .inner
534                    .map
535                    .lock()
536                    .await
537                    .get(&id)
538                    .map(|job| crate::types::MarkdownProgress(job.previews.clone()));
539            }
540            if let Some(cache) = &req.scrape_options.asset_fetch_cache {
541                sink.jobs.seed_asset_cache(&id, cache).await;
542            }
543            let result = engine.crawl(&req, sink.clone()).await;
544            sink.finish(result.map_err(|e| (req.url.clone(), e))).await;
545        });
546    }
547
548    /// Restore bounded URL-to-content entries from successful saved pages.
549    async fn seed_asset_cache(&self, id: &str, cache: &crate::types::AssetFetchCache) {
550        let Some(store) = &self.inner.store else {
551            return;
552        };
553        let key = list_key(id, "pages");
554        let pairs = store
555            .blocking(move |s| -> pagesift_store::Result<Vec<(String, String)>> {
556                let tasks = s.tasks();
557                let mut pairs = std::collections::BTreeMap::new();
558                for start in (0..tasks.list_len(&key)).step_by(256) {
559                    for raw in tasks.list_range(&key, start, 256)? {
560                        let doc: crate::types::Document =
561                            serde_json::from_slice(&raw).map_err(std::io::Error::other)?;
562                        for asset in doc.image_assets {
563                            if asset.status == crate::types::ImageAssetStatus::Ok
564                                && (asset.source.starts_with("http://")
565                                    || asset.source.starts_with("https://"))
566                                && let Some(id) = asset.id
567                            {
568                                pairs.entry(asset.source).or_insert(id);
569                            }
570                            if pairs.len() >= 32 {
571                                return Ok(pairs.into_iter().collect());
572                            }
573                        }
574                    }
575                }
576                Ok(pairs.into_iter().collect())
577            })
578            .await;
579        let Ok(pairs) = pairs else { return };
580        for (url, id) in pairs {
581            let Some((bytes, _)) = crate::attachments::load(self.disk(), id.clone()).await else {
582                continue;
583            };
584            if bytes.len() > 512 * 1024 {
585                continue;
586            }
587            let Some(format) = crate::types::attachment::sniff_image_format(&bytes) else {
588                continue;
589            };
590            let attachment =
591                crate::types::Attachment::new(bytes.to_vec(), format, Some(url.clone()));
592            if attachment.id == id {
593                cache.seed(&url, attachment);
594            }
595        }
596    }
597}
pagesift_api/jobs/mod.rs

pagesift_api/jobs/
mod.rs
