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::{Duration, SystemTime, UNIX_EPOCH};
19
20use pagesift_store::Store;
21use serde::{Deserialize, Serialize};
22use tokio::sync::{Mutex, Semaphore};
23
24use crate::engine::{CrawlResume, Engine};
25use crate::error::{ApiError, ErrorCode};
26use crate::limits::{Caller, Limiter};
27use crate::metrics::{Metrics, PageSource};
28use crate::types::{CrawlRequest, CrawlSummary, Document, JobError, ScrapeOptions};
29use crate::webhook::{Event, Hook, WebhookConfig, Webhooks};
30
31mod sink;
32mod store;
33
34use sink::JobSink;
35use store::{BatchSpec, Meta, delete_job, list_key, meta_key, to_ms};
36
37/// Firecrawl 的任务结果保留 24 小时。
38pub const JOB_TTL: Duration = Duration::from_secs(24 * 3600);
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
41#[serde(rename_all = "lowercase")]
42pub enum JobState {
43    Scraping,
44    Completed,
45    Cancelled,
46    Failed,
47}
48
49impl JobState {
50    pub fn as_str(self) -> &'static str {
51        match self {
52            JobState::Scraping => "scraping",
53            JobState::Completed => "completed",
54            JobState::Cancelled => "cancelled",
55            JobState::Failed => "failed",
56        }
57    }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
61#[serde(rename_all = "snake_case")]
62pub enum JobKind {
63    Crawl,
64    BatchScrape,
65}
66
67impl JobKind {
68    /// 也是 webhook 事件类型的前缀（`crawl.page`、`batch_scrape.completed`）。
69    pub fn as_str(self) -> &'static str {
70        match self {
71            JobKind::Crawl => "crawl",
72            JobKind::BatchScrape => "batch_scrape",
73        }
74    }
75
76    pub(super) fn pages(self) -> PageSource {
77        match self {
78            JobKind::Crawl => PageSource::Crawl,
79            JobKind::BatchScrape => PageSource::Batch,
80        }
81    }
82}
83
84#[derive(Debug)]
85pub struct Job {
86    pub kind: JobKind,
87    pub state: JobState,
88    pub total: usize,
89    pub completed: usize,
90    /// 按完成顺序追加，分页用下标。任务落盘时文档在存储里，这里为空。
91    pub data: Vec<Document>,
92    /// 任务落盘时错误在存储里，这里为空。
93    pub errors: Vec<JobError>,
94    pub robots_blocked: Vec<String>,
95    pub created: SystemTime,
96    /// 带查询的爬取结束后的说明。
97    pub summary: Option<CrawlSummary>,
98    /// 存储里已有的文档数。
99    pub(super) stored_docs: usize,
100    pub(super) owner: Option<Caller>,
101}
102
103impl Job {
104    pub(super) fn new(kind: JobKind, total: usize, owner: Option<Caller>) -> Self {
105        Self {
106            kind,
107            state: JobState::Scraping,
108            total,
109            completed: 0,
110            data: Vec::new(),
111            errors: Vec::new(),
112            robots_blocked: Vec::new(),
113            created: SystemTime::now(),
114            summary: None,
115            stored_docs: 0,
116            owner,
117        }
118    }
119
120    pub fn expires_at(&self) -> SystemTime {
121        self.created + JOB_TTL
122    }
123
124    fn docs(&self) -> usize {
125        self.data.len() + self.stored_docs
126    }
127}
128
129/// 任务状态的一页。
130pub(crate) struct StatusPage {
131    pub state: JobState,
132    pub total: usize,
133    pub completed: usize,
134    pub expires_at: SystemTime,
135    /// 文档总数。
136    pub docs: usize,
137    pub data: Vec<Document>,
138    pub summary: Option<CrawlSummary>,
139}
140
141/// 启动任务时的调用方信息。
142#[derive(Default)]
143pub(crate) struct Start {
144    pub owner: Option<Caller>,
145    pub webhook: Option<WebhookConfig>,
146}
147
148struct Inner {
149    map: Mutex<HashMap<String, Job>>,
150    store: Option<Store>,
151    limiter: Arc<Limiter>,
152    webhooks: Option<Webhooks>,
153    metrics: Arc<Metrics>,
154}
155
156#[derive(Clone)]
157pub struct Jobs {
158    inner: Arc<Inner>,
159}
160
161impl Default for Jobs {
162    /// 只在内存里的任务表：不落盘、不限额、不发 webhook。
163    fn default() -> Self {
164        Self::new(None, Arc::default(), None, Arc::default())
165    }
166}
167
168impl Jobs {
169    pub(crate) fn new(
170        store: Option<Store>,
171        limiter: Arc<Limiter>,
172        webhooks: Option<Webhooks>,
173        metrics: Arc<Metrics>,
174    ) -> Self {
175        Self {
176            inner: Arc::new(Inner {
177                map: Mutex::new(HashMap::new()),
178                store,
179                limiter,
180                webhooks,
181                metrics,
182            }),
183        }
184    }
185
186    /// 清掉过期任务（连同存储里的数据）。
187    async fn sweep(&self) {
188        let now = SystemTime::now();
189        let expired: Vec<String> = {
190            let mut map = self.inner.map.lock().await;
191            let ids: Vec<String> = map
192                .iter()
193                .filter(|(_, j)| j.expires_at() <= now)
194                .map(|(id, _)| id.clone())
195                .collect();
196            for id in &ids {
197                map.remove(id);
198            }
199            ids
200        };
201        if let (Some(store), false) = (&self.inner.store, expired.is_empty()) {
202            store
203                .blocking(move |s| {
204                    for id in expired {
205                        delete_job(s, &id);
206                    }
207                })
208                .await;
209        }
210    }
211
212    pub async fn get<R>(&self, id: &str, f: impl FnOnce(&Job) -> R) -> Option<R> {
213        self.sweep().await;
214        self.inner.map.lock().await.get(id).map(f)
215    }
216
217    /// 状态和从 `skip` 开始最多 `limit` 个文档。
218    pub(crate) async fn status(&self, id: &str, skip: usize, limit: usize) -> Option<StatusPage> {
219        self.sweep().await;
220        let mut page = {
221            let map = self.inner.map.lock().await;
222            let job = map.get(id)?;
223            let end = (skip + limit).min(job.data.len());
224            StatusPage {
225                state: job.state,
226                total: job.total,
227                completed: job.completed,
228                expires_at: job.expires_at(),
229                docs: job.docs(),
230                data: job
231                    .data
232                    .get(skip.min(end)..end)
233                    .unwrap_or_default()
234                    .to_vec(),
235                summary: job.summary.clone(),
236            }
237        };
238        if let Some(store) = &self.inner.store
239            && page.data.is_empty()
240            && skip < page.docs
241        {
242            let key = list_key(id, "pages");
243            page.data = store
244                .blocking(move |s| s.tasks().list_range(key, skip as u64, limit))
245                .await
246                .unwrap_or_else(|e| {
247                    tracing::warn!(error = %e, "reading job pages failed");
248                    Vec::new()
249                })
250                .iter()
251                .filter_map(|v| serde_json::from_slice(v).ok())
252                .collect();
253        }
254        Some(page)
255    }
256
257    /// 任务的页面错误和被 robots.txt 拦下的 URL。
258    pub(crate) async fn errors(&self, id: &str) -> Option<(Vec<JobError>, Vec<String>)> {
259        self.sweep().await;
260        let (mut errors, mut robots) = {
261            let map = self.inner.map.lock().await;
262            let job = map.get(id)?;
263            (job.errors.clone(), job.robots_blocked.clone())
264        };
265        if let Some(store) = &self.inner.store {
266            let key = list_key(id, "errors");
267            let stored = store
268                .blocking(move |s| s.tasks().list_items(key, usize::MAX))
269                .await
270                .unwrap_or_default();
271            for v in stored {
272                if let Ok(e) = serde_json::from_slice::<JobError>(&v) {
273                    if e.code == ErrorCode::RobotsDisallowed.as_str() {
274                        robots.push(e.url.clone());
275                    }
276                    errors.push(e);
277                }
278            }
279        }
280        Some((errors, robots))
281    }
282
283    pub async fn cancel(&self, id: &str) -> bool {
284        let found = {
285            let mut map = self.inner.map.lock().await;
286            match map.get_mut(id) {
287                Some(j) => {
288                    if j.state == JobState::Scraping {
289                        j.state = JobState::Cancelled;
290                    }
291                    true
292                }
293                None => false,
294            }
295        };
296        if found {
297            self.save_state(id).await;
298        }
299        found
300    }
301
302    pub async fn active(&self) -> Vec<String> {
303        let map = self.inner.map.lock().await;
304        map.iter()
305            .filter(|(_, j)| j.state == JobState::Scraping)
306            .map(|(id, _)| id.clone())
307            .collect()
308    }
309
310    /// 按 (类型, 状态) 统计的任务数，给 `/metrics`。
311    pub(crate) async fn counts(&self) -> BTreeMap<(&'static str, &'static str), u64> {
312        let mut out = BTreeMap::new();
313        for j in self.inner.map.lock().await.values() {
314            *out.entry((j.kind.as_str(), j.state.as_str())).or_default() += 1;
315        }
316        out
317    }
318
319    /// 登记新任务：检查调用方的并发任务数，落盘元数据，发 `started` 事件。
320    async fn register(
321        &self,
322        kind: JobKind,
323        total: usize,
324        start: Start,
325        crawl: Option<CrawlRequest>,
326        batch: Option<BatchSpec>,
327    ) -> Result<(String, Option<Hook>), ApiError> {
328        let id = uuid::Uuid::new_v4().to_string();
329        let job = Job::new(kind, total, start.owner.clone());
330        let meta = Meta {
331            kind,
332            state: JobState::Scraping,
333            created_ms: to_ms(job.created),
334            total,
335            completed: 0,
336            owner: start.owner.as_ref().map(|c| c.account.to_string()),
337            webhook: start.webhook.clone(),
338            crawl,
339            batch,
340            summary: None,
341        };
342        {
343            let mut map = self.inner.map.lock().await;
344            if let Some(c) = &start.owner {
345                if c.limits.max_concurrent_jobs.is_some() {
346                    let active = map
347                        .values()
348                        .filter(|j| {
349                            j.state == JobState::Scraping
350                                && j.owner.as_ref().is_some_and(|o| o.account == c.account)
351                        })
352                        .count();
353                    self.inner.limiter.check_jobs(c, active)?;
354                }
355                self.inner.limiter.check_quota(c)?;
356            }
357            map.insert(id.clone(), job);
358        }
359        if let Some(store) = &self.inner.store {
360            let (key, value) = (meta_key(&id), serde_json::to_vec(&meta).unwrap_or_default());
361            if let Err(e) = store.blocking(move |s| s.tasks().kv_put(key, &value)).await {
362                tracing::warn!(job = %id, error = %e, "saving job failed; it will not survive a restart");
363            }
364        }
365        self.inner.metrics.job_started(kind.as_str());
366        tracing::info!(
367            job = %id,
368            kind = kind.as_str(),
369            account = start.owner.as_ref().map(|c| &*c.account),
370            "job started"
371        );
372        let hook = match (&self.inner.webhooks, start.webhook) {
373            (Some(w), Some(cfg)) => {
374                let hook = w.start(&id, kind.as_str(), cfg);
375                hook.send(Event::Started, Vec::new(), None);
376                Some(hook)
377            }
378            _ => None,
379        };
380        Ok((id, hook))
381    }
382
383    /// 启动批量抓取，立即返回任务 id；每个 URL 的结果和错误陆续写入任务。
384    pub async fn start_batch(
385        &self,
386        engine: Arc<dyn Engine>,
387        urls: Vec<String>,
388        opts: ScrapeOptions,
389        concurrency: usize,
390    ) -> String {
391        self.start_batch_with(engine, urls, opts, concurrency, Start::default())
392            .await
393            .expect("no limits without a caller")
394    }
395
396    pub(crate) async fn start_batch_with(
397        &self,
398        engine: Arc<dyn Engine>,
399        urls: Vec<String>,
400        opts: ScrapeOptions,
401        concurrency: usize,
402        start: Start,
403    ) -> Result<String, ApiError> {
404        let spec = BatchSpec {
405            urls,
406            options: opts,
407            concurrency,
408        };
409        let total = spec.urls.len();
410        let (id, hook) = self
411            .register(
412                JobKind::BatchScrape,
413                total,
414                start,
415                None,
416                self.inner.store.is_some().then(|| spec.clone()),
417            )
418            .await?;
419        let pending = spec.urls.into_iter().enumerate().collect();
420        self.spawn_batch(id.clone(), engine, pending, spec.options, concurrency, hook);
421        Ok(id)
422    }
423
424    fn sink(
425        &self,
426        id: &str,
427        kind: JobKind,
428        hook: Option<Hook>,
429        resume: Option<CrawlResume>,
430    ) -> Arc<JobSink> {
431        Arc::new(JobSink {
432            jobs: self.clone(),
433            id: id.to_string(),
434            kind,
435            hook,
436            resume: StdMutex::new(resume),
437            quota_hit: AtomicBool::new(false),
438        })
439    }
440
441    pub(super) fn spawn_batch(
442        &self,
443        id: String,
444        engine: Arc<dyn Engine>,
445        pending: Vec<(usize, String)>,
446        opts: ScrapeOptions,
447        concurrency: usize,
448        hook: Option<Hook>,
449    ) {
450        let sink = self.sink(&id, JobKind::BatchScrape, hook, None);
451        let opts = Arc::new(opts);
452        tokio::spawn(async move {
453            let sem = Arc::new(Semaphore::new(concurrency.max(1)));
454            let mut tasks = tokio::task::JoinSet::new();
455            for (i, url) in pending {
456                let (sink, engine, opts, sem) =
457                    (sink.clone(), engine.clone(), opts.clone(), sem.clone());
458                tasks.spawn(async move {
459                    let _permit = sem.acquire_owned().await.expect("semaphore open");
460                    if sink.stopped().await {
461                        return;
462                    }
463                    let done = i.to_string();
464                    match engine.scrape(&url, &opts).await {
465                        Ok(doc) => sink.record_page(done, doc).await,
466                        Err(e) => sink.record_error(done, url, e).await,
467                    }
468                });
469            }
470            while tasks.join_next().await.is_some() {}
471            sink.finish(Ok(())).await;
472        });
473    }
474
475    /// 启动整站爬取，立即返回任务 id。
476    pub async fn start_crawl(&self, engine: Arc<dyn Engine>, req: CrawlRequest) -> String {
477        self.start_crawl_with(engine, req, Start::default())
478            .await
479            .expect("no limits without a caller")
480    }
481
482    pub(crate) async fn start_crawl_with(
483        &self,
484        engine: Arc<dyn Engine>,
485        req: CrawlRequest,
486        start: Start,
487    ) -> Result<String, ApiError> {
488        let (id, hook) = self
489            .register(
490                JobKind::Crawl,
491                0,
492                start,
493                self.inner.store.is_some().then(|| req.clone()),
494                None,
495            )
496            .await?;
497        self.spawn_crawl(id.clone(), engine, req, hook, None);
498        Ok(id)
499    }
500
501    pub(super) fn spawn_crawl(
502        &self,
503        id: String,
504        engine: Arc<dyn Engine>,
505        req: CrawlRequest,
506        hook: Option<Hook>,
507        resume: Option<CrawlResume>,
508    ) {
509        let sink = self.sink(&id, JobKind::Crawl, hook, resume);
510        tokio::spawn(async move {
511            let result = engine.crawl(&req, sink.clone()).await;
512            sink.finish(result.map_err(|e| (req.url.clone(), e))).await;
513        });
514    }
515}
516
517/// UTC 时间格式化成 RFC 3339（秒级精度后带毫秒），不引入日期库。
518pub fn rfc3339(t: SystemTime) -> String {
519    let d = t.duration_since(UNIX_EPOCH).unwrap_or_default();
520    let secs = d.as_secs() as i64;
521    let (days, rem) = (secs.div_euclid(86_400), secs.rem_euclid(86_400));
522    // Howard Hinnant 的 civil_from_days
523    let z = days + 719_468;
524    let era = z.div_euclid(146_097);
525    let doe = z - era * 146_097;
526    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
527    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
528    let mp = (5 * doy + 2) / 153;
529    let day = doy - (153 * mp + 2) / 5 + 1;
530    let month = if mp < 10 { mp + 3 } else { mp - 9 };
531    let year = yoe + era * 400 + i64::from(month <= 2);
532    format!(
533        "{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{:03}Z",
534        rem / 3600,
535        rem % 3600 / 60,
536        rem % 60,
537        d.subsec_millis()
538    )
539}
540
541#[cfg(test)]
542mod tests {
543    use super::*;
544
545    #[test]
546    fn formats_utc_dates() {
547        assert_eq!(rfc3339(UNIX_EPOCH), "1970-01-01T00:00:00.000Z");
548        let t = UNIX_EPOCH + Duration::from_millis(1_792_540_800_123);
549        assert_eq!(rfc3339(t), "2026-10-21T00:00:00.123Z");
550        let leap = UNIX_EPOCH + Duration::from_secs(951_782_400); // 2000-02-29
551        assert_eq!(rfc3339(leap), "2000-02-29T00:00:00.000Z");
552    }
553}
pagesift_api/jobs/mod.rs

pagesift_api/jobs/
mod.rs
