Skip to main content

pagesift_api/
jobs.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, Ordering};
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::{BoxFuture, CrawlResume, CrawlSink, Engine};
25use crate::error::{ApiError, ErrorCode};
26use crate::limits::{Caller, KeyLimits, Limiter, quota_error};
27use crate::metrics::{Metrics, PageSource};
28use crate::types::{CrawlRequest, CrawlSummary, Document, JobError, ScrapeOptions};
29use crate::webhook::{Event, Hook, WebhookConfig, Webhooks};
30
31/// Firecrawl 的任务结果保留 24 小时。
32pub const JOB_TTL: Duration = Duration::from_secs(24 * 3600);
33
34#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
35#[serde(rename_all = "lowercase")]
36pub enum JobState {
37    Scraping,
38    Completed,
39    Cancelled,
40    Failed,
41}
42
43impl JobState {
44    pub fn as_str(self) -> &'static str {
45        match self {
46            JobState::Scraping => "scraping",
47            JobState::Completed => "completed",
48            JobState::Cancelled => "cancelled",
49            JobState::Failed => "failed",
50        }
51    }
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
55#[serde(rename_all = "snake_case")]
56pub enum JobKind {
57    Crawl,
58    BatchScrape,
59}
60
61impl JobKind {
62    /// 也是 webhook 事件类型的前缀(`crawl.page`、`batch_scrape.completed`)。
63    pub fn as_str(self) -> &'static str {
64        match self {
65            JobKind::Crawl => "crawl",
66            JobKind::BatchScrape => "batch_scrape",
67        }
68    }
69
70    fn pages(self) -> PageSource {
71        match self {
72            JobKind::Crawl => PageSource::Crawl,
73            JobKind::BatchScrape => PageSource::Batch,
74        }
75    }
76}
77
78#[derive(Debug)]
79pub struct Job {
80    pub kind: JobKind,
81    pub state: JobState,
82    pub total: usize,
83    pub completed: usize,
84    /// 按完成顺序追加,分页用下标。任务落盘时文档在存储里,这里为空。
85    pub data: Vec<Document>,
86    /// 任务落盘时错误在存储里,这里为空。
87    pub errors: Vec<JobError>,
88    pub robots_blocked: Vec<String>,
89    pub created: SystemTime,
90    /// 带查询的爬取结束后的说明。
91    pub summary: Option<CrawlSummary>,
92    /// 存储里已有的文档数。
93    stored_docs: usize,
94    owner: Option<Caller>,
95}
96
97impl Job {
98    fn new(kind: JobKind, total: usize, owner: Option<Caller>) -> Self {
99        Self {
100            kind,
101            state: JobState::Scraping,
102            total,
103            completed: 0,
104            data: Vec::new(),
105            errors: Vec::new(),
106            robots_blocked: Vec::new(),
107            created: SystemTime::now(),
108            summary: None,
109            stored_docs: 0,
110            owner,
111        }
112    }
113
114    pub fn expires_at(&self) -> SystemTime {
115        self.created + JOB_TTL
116    }
117
118    fn docs(&self) -> usize {
119        self.data.len() + self.stored_docs
120    }
121}
122
123/// 任务状态的一页。
124pub(crate) struct StatusPage {
125    pub state: JobState,
126    pub total: usize,
127    pub completed: usize,
128    pub expires_at: SystemTime,
129    /// 文档总数。
130    pub docs: usize,
131    pub data: Vec<Document>,
132    pub summary: Option<CrawlSummary>,
133}
134
135/// 启动任务时的调用方信息。
136#[derive(Default)]
137pub(crate) struct Start {
138    pub owner: Option<Caller>,
139    pub webhook: Option<WebhookConfig>,
140}
141
142/// 落盘的任务元数据。
143#[derive(Serialize, Deserialize)]
144struct Meta {
145    kind: JobKind,
146    state: JobState,
147    created_ms: u64,
148    total: usize,
149    completed: usize,
150    owner: Option<String>,
151    webhook: Option<WebhookConfig>,
152    crawl: Option<CrawlRequest>,
153    batch: Option<BatchSpec>,
154    /// 带查询的爬取结束后的说明;旧版本存的任务没有这个字段。
155    #[serde(default)]
156    summary: Option<CrawlSummary>,
157}
158
159#[derive(Clone, Serialize, Deserialize)]
160struct BatchSpec {
161    urls: Vec<String>,
162    options: ScrapeOptions,
163    concurrency: usize,
164}
165
166fn meta_key(id: &str) -> String {
167    format!("job/{id}")
168}
169
170fn list_key(id: &str, what: &str) -> String {
171    format!("job/{id}/{what}")
172}
173
174fn to_ms(t: SystemTime) -> u64 {
175    t.duration_since(UNIX_EPOCH)
176        .map_or(0, |d| d.as_millis() as u64)
177}
178
179struct Inner {
180    map: Mutex<HashMap<String, Job>>,
181    store: Option<Store>,
182    limiter: Arc<Limiter>,
183    webhooks: Option<Webhooks>,
184    metrics: Arc<Metrics>,
185}
186
187#[derive(Clone)]
188pub struct Jobs {
189    inner: Arc<Inner>,
190}
191
192impl Default for Jobs {
193    /// 只在内存里的任务表:不落盘、不限额、不发 webhook。
194    fn default() -> Self {
195        Self::new(None, Arc::default(), None, Arc::default())
196    }
197}
198
199impl Jobs {
200    pub(crate) fn new(
201        store: Option<Store>,
202        limiter: Arc<Limiter>,
203        webhooks: Option<Webhooks>,
204        metrics: Arc<Metrics>,
205    ) -> Self {
206        Self {
207            inner: Arc::new(Inner {
208                map: Mutex::new(HashMap::new()),
209                store,
210                limiter,
211                webhooks,
212                metrics,
213            }),
214        }
215    }
216
217    /// 清掉过期任务(连同存储里的数据)。
218    async fn sweep(&self) {
219        let now = SystemTime::now();
220        let expired: Vec<String> = {
221            let mut map = self.inner.map.lock().await;
222            let ids: Vec<String> = map
223                .iter()
224                .filter(|(_, j)| j.expires_at() <= now)
225                .map(|(id, _)| id.clone())
226                .collect();
227            for id in &ids {
228                map.remove(id);
229            }
230            ids
231        };
232        if let (Some(store), false) = (&self.inner.store, expired.is_empty()) {
233            store
234                .blocking(move |s| {
235                    for id in expired {
236                        delete_job(s, &id);
237                    }
238                })
239                .await;
240        }
241    }
242
243    pub async fn get<R>(&self, id: &str, f: impl FnOnce(&Job) -> R) -> Option<R> {
244        self.sweep().await;
245        self.inner.map.lock().await.get(id).map(f)
246    }
247
248    /// 状态和从 `skip` 开始最多 `limit` 个文档。
249    pub(crate) async fn status(&self, id: &str, skip: usize, limit: usize) -> Option<StatusPage> {
250        self.sweep().await;
251        let mut page = {
252            let map = self.inner.map.lock().await;
253            let job = map.get(id)?;
254            let end = (skip + limit).min(job.data.len());
255            StatusPage {
256                state: job.state,
257                total: job.total,
258                completed: job.completed,
259                expires_at: job.expires_at(),
260                docs: job.docs(),
261                data: job
262                    .data
263                    .get(skip.min(end)..end)
264                    .unwrap_or_default()
265                    .to_vec(),
266                summary: job.summary.clone(),
267            }
268        };
269        if let Some(store) = &self.inner.store
270            && page.data.is_empty()
271            && skip < page.docs
272        {
273            let key = list_key(id, "pages");
274            page.data = store
275                .blocking(move |s| s.tasks().list_range(key, skip as u64, limit))
276                .await
277                .unwrap_or_else(|e| {
278                    tracing::warn!(error = %e, "reading job pages failed");
279                    Vec::new()
280                })
281                .iter()
282                .filter_map(|v| serde_json::from_slice(v).ok())
283                .collect();
284        }
285        Some(page)
286    }
287
288    /// 任务的页面错误和被 robots.txt 拦下的 URL。
289    pub(crate) async fn errors(&self, id: &str) -> Option<(Vec<JobError>, Vec<String>)> {
290        self.sweep().await;
291        let (mut errors, mut robots) = {
292            let map = self.inner.map.lock().await;
293            let job = map.get(id)?;
294            (job.errors.clone(), job.robots_blocked.clone())
295        };
296        if let Some(store) = &self.inner.store {
297            let key = list_key(id, "errors");
298            let stored = store
299                .blocking(move |s| s.tasks().list_items(key, usize::MAX))
300                .await
301                .unwrap_or_default();
302            for v in stored {
303                if let Ok(e) = serde_json::from_slice::<JobError>(&v) {
304                    if e.code == ErrorCode::RobotsDisallowed.as_str() {
305                        robots.push(e.url.clone());
306                    }
307                    errors.push(e);
308                }
309            }
310        }
311        Some((errors, robots))
312    }
313
314    pub async fn cancel(&self, id: &str) -> bool {
315        let found = {
316            let mut map = self.inner.map.lock().await;
317            match map.get_mut(id) {
318                Some(j) => {
319                    if j.state == JobState::Scraping {
320                        j.state = JobState::Cancelled;
321                    }
322                    true
323                }
324                None => false,
325            }
326        };
327        if found {
328            self.save_state(id).await;
329        }
330        found
331    }
332
333    pub async fn active(&self) -> Vec<String> {
334        let map = self.inner.map.lock().await;
335        map.iter()
336            .filter(|(_, j)| j.state == JobState::Scraping)
337            .map(|(id, _)| id.clone())
338            .collect()
339    }
340
341    /// 按 (类型, 状态) 统计的任务数,给 `/metrics`。
342    pub(crate) async fn counts(&self) -> BTreeMap<(&'static str, &'static str), u64> {
343        let mut out = BTreeMap::new();
344        for j in self.inner.map.lock().await.values() {
345            *out.entry((j.kind.as_str(), j.state.as_str())).or_default() += 1;
346        }
347        out
348    }
349
350    /// 登记新任务:检查调用方的并发任务数,落盘元数据,发 `started` 事件。
351    async fn register(
352        &self,
353        kind: JobKind,
354        total: usize,
355        start: Start,
356        crawl: Option<CrawlRequest>,
357        batch: Option<BatchSpec>,
358    ) -> Result<(String, Option<Hook>), ApiError> {
359        let id = uuid::Uuid::new_v4().to_string();
360        let job = Job::new(kind, total, start.owner.clone());
361        let meta = Meta {
362            kind,
363            state: JobState::Scraping,
364            created_ms: to_ms(job.created),
365            total,
366            completed: 0,
367            owner: start.owner.as_ref().map(|c| c.account.to_string()),
368            webhook: start.webhook.clone(),
369            crawl,
370            batch,
371            summary: None,
372        };
373        {
374            let mut map = self.inner.map.lock().await;
375            if let Some(c) = &start.owner {
376                if c.limits.max_concurrent_jobs.is_some() {
377                    let active = map
378                        .values()
379                        .filter(|j| {
380                            j.state == JobState::Scraping
381                                && j.owner.as_ref().is_some_and(|o| o.account == c.account)
382                        })
383                        .count();
384                    self.inner.limiter.check_jobs(c, active)?;
385                }
386                self.inner.limiter.check_quota(c)?;
387            }
388            map.insert(id.clone(), job);
389        }
390        if let Some(store) = &self.inner.store {
391            let (key, value) = (meta_key(&id), serde_json::to_vec(&meta).unwrap_or_default());
392            if let Err(e) = store.blocking(move |s| s.tasks().kv_put(key, &value)).await {
393                tracing::warn!(job = %id, error = %e, "saving job failed; it will not survive a restart");
394            }
395        }
396        self.inner.metrics.job_started(kind.as_str());
397        tracing::info!(
398            job = %id,
399            kind = kind.as_str(),
400            account = start.owner.as_ref().map(|c| &*c.account),
401            "job started"
402        );
403        let hook = match (&self.inner.webhooks, start.webhook) {
404            (Some(w), Some(cfg)) => {
405                let hook = w.start(&id, kind.as_str(), cfg);
406                hook.send(Event::Started, Vec::new(), None);
407                Some(hook)
408            }
409            _ => None,
410        };
411        Ok((id, hook))
412    }
413
414    /// 启动批量抓取,立即返回任务 id;每个 URL 的结果和错误陆续写入任务。
415    pub async fn start_batch(
416        &self,
417        engine: Arc<dyn Engine>,
418        urls: Vec<String>,
419        opts: ScrapeOptions,
420        concurrency: usize,
421    ) -> String {
422        self.start_batch_with(engine, urls, opts, concurrency, Start::default())
423            .await
424            .expect("no limits without a caller")
425    }
426
427    pub(crate) async fn start_batch_with(
428        &self,
429        engine: Arc<dyn Engine>,
430        urls: Vec<String>,
431        opts: ScrapeOptions,
432        concurrency: usize,
433        start: Start,
434    ) -> Result<String, ApiError> {
435        let spec = BatchSpec {
436            urls,
437            options: opts,
438            concurrency,
439        };
440        let total = spec.urls.len();
441        let (id, hook) = self
442            .register(
443                JobKind::BatchScrape,
444                total,
445                start,
446                None,
447                self.inner.store.is_some().then(|| spec.clone()),
448            )
449            .await?;
450        let pending = spec.urls.into_iter().enumerate().collect();
451        self.spawn_batch(id.clone(), engine, pending, spec.options, concurrency, hook);
452        Ok(id)
453    }
454
455    fn sink(
456        &self,
457        id: &str,
458        kind: JobKind,
459        hook: Option<Hook>,
460        resume: Option<CrawlResume>,
461    ) -> Arc<JobSink> {
462        Arc::new(JobSink {
463            jobs: self.clone(),
464            id: id.to_string(),
465            kind,
466            hook,
467            resume: StdMutex::new(resume),
468            quota_hit: AtomicBool::new(false),
469        })
470    }
471
472    fn spawn_batch(
473        &self,
474        id: String,
475        engine: Arc<dyn Engine>,
476        pending: Vec<(usize, String)>,
477        opts: ScrapeOptions,
478        concurrency: usize,
479        hook: Option<Hook>,
480    ) {
481        let sink = self.sink(&id, JobKind::BatchScrape, hook, None);
482        let opts = Arc::new(opts);
483        tokio::spawn(async move {
484            let sem = Arc::new(Semaphore::new(concurrency.max(1)));
485            let mut tasks = tokio::task::JoinSet::new();
486            for (i, url) in pending {
487                let (sink, engine, opts, sem) =
488                    (sink.clone(), engine.clone(), opts.clone(), sem.clone());
489                tasks.spawn(async move {
490                    let _permit = sem.acquire_owned().await.expect("semaphore open");
491                    if sink.stopped().await {
492                        return;
493                    }
494                    let done = i.to_string();
495                    match engine.scrape(&url, &opts).await {
496                        Ok(doc) => sink.record_page(done, doc).await,
497                        Err(e) => sink.record_error(done, url, e).await,
498                    }
499                });
500            }
501            while tasks.join_next().await.is_some() {}
502            sink.finish(Ok(())).await;
503        });
504    }
505
506    /// 启动整站爬取,立即返回任务 id。
507    pub async fn start_crawl(&self, engine: Arc<dyn Engine>, req: CrawlRequest) -> String {
508        self.start_crawl_with(engine, req, Start::default())
509            .await
510            .expect("no limits without a caller")
511    }
512
513    pub(crate) async fn start_crawl_with(
514        &self,
515        engine: Arc<dyn Engine>,
516        req: CrawlRequest,
517        start: Start,
518    ) -> Result<String, ApiError> {
519        let (id, hook) = self
520            .register(
521                JobKind::Crawl,
522                0,
523                start,
524                self.inner.store.is_some().then(|| req.clone()),
525                None,
526            )
527            .await?;
528        self.spawn_crawl(id.clone(), engine, req, hook, None);
529        Ok(id)
530    }
531
532    fn spawn_crawl(
533        &self,
534        id: String,
535        engine: Arc<dyn Engine>,
536        req: CrawlRequest,
537        hook: Option<Hook>,
538        resume: Option<CrawlResume>,
539    ) {
540        let sink = self.sink(&id, JobKind::Crawl, hook, resume);
541        tokio::spawn(async move {
542            let result = engine.crawl(&req, sink.clone()).await;
543            sink.finish(result.map_err(|e| (req.url.clone(), e))).await;
544        });
545    }
546
547    /// 把任务的当前状态和计数写回存储。
548    async fn save_state(&self, id: &str) {
549        let Some(store) = &self.inner.store else {
550            return;
551        };
552        let snapshot = {
553            let map = self.inner.map.lock().await;
554            map.get(id)
555                .map(|j| (j.state, j.total, j.completed, j.summary.clone()))
556        };
557        let Some((state, total, completed, summary)) = snapshot else {
558            return;
559        };
560        let key = meta_key(id);
561        let r = store
562            .blocking(move |s| -> pagesift_store::Result<()> {
563                let t = s.tasks();
564                let Some(raw) = t.kv_get(&key)? else {
565                    return Ok(());
566                };
567                let Ok(mut meta) = serde_json::from_slice::<Meta>(&raw) else {
568                    return Ok(());
569                };
570                meta.state = state;
571                meta.total = total;
572                meta.completed = completed;
573                meta.summary = summary;
574                t.kv_put(&key, &serde_json::to_vec(&meta).unwrap_or_default())
575            })
576            .await;
577        if let Err(e) = r {
578            tracing::warn!(job = %id, error = %e, "saving job state failed");
579        }
580    }
581
582    /// 服务启动时从存储恢复任务:结束的任务照样能查,进行中的从断点继续,过期的删掉。
583    /// `limits` 按调用方名字找回当前配置的限额。需要在 tokio 运行时里调用。
584    pub(crate) fn restore(
585        &self,
586        engine: &Arc<dyn Engine>,
587        limits: &dyn Fn(&str) -> Arc<KeyLimits>,
588    ) {
589        let Some(store) = &self.inner.store else {
590            return;
591        };
592        let t = store.tasks();
593        let keys = match t.kv_keys("job/") {
594            Ok(k) => k,
595            Err(e) => {
596                tracing::error!(error = %e, "listing saved jobs failed");
597                return;
598            }
599        };
600        let now = SystemTime::now();
601        let (mut finished, mut resumed, mut expired) = (0, 0, 0);
602        for key in keys {
603            let Some(id) = std::str::from_utf8(&key)
604                .ok()
605                .and_then(|k| k.strip_prefix("job/"))
606                .filter(|id| !id.contains('/'))
607                .map(String::from)
608            else {
609                continue;
610            };
611            let Some(meta) = t
612                .kv_get(&key)
613                .ok()
614                .flatten()
615                .and_then(|v| serde_json::from_slice::<Meta>(&v).ok())
616            else {
617                tracing::warn!(job = %id, "dropping an unreadable saved job");
618                delete_job(store, &id);
619                continue;
620            };
621            let created = UNIX_EPOCH + Duration::from_millis(meta.created_ms);
622            if created + JOB_TTL <= now {
623                delete_job(store, &id);
624                expired += 1;
625                continue;
626            }
627            let pages = t.list_len(list_key(&id, "pages")) as usize;
628            let errors = t.list_len(list_key(&id, "errors")) as usize;
629            let owner = meta.owner.as_deref().map(|a| Caller::named(a, limits(a)));
630            let mut job = Job::new(meta.kind, meta.total, owner);
631            job.created = created;
632            job.state = meta.state;
633            job.summary = meta.summary;
634            job.stored_docs = pages;
635            job.completed = pages + errors;
636            job.total = job.total.max(job.completed);
637            if job.state != JobState::Scraping {
638                finished += 1;
639                self.inner
640                    .map
641                    .try_lock()
642                    .expect("not shared yet")
643                    .insert(id, job);
644                continue;
645            }
646            // 进行中的任务:算出还没做完的部分,接着跑
647            let done = list_key(&id, "done");
648            let hook = match (&self.inner.webhooks, meta.webhook) {
649                (Some(w), Some(cfg)) => Some(w.start(&id, meta.kind.as_str(), cfg)),
650                _ => None,
651            };
652            match (meta.kind, meta.crawl, meta.batch) {
653                (JobKind::BatchScrape, _, Some(spec)) => {
654                    let pending: Vec<(usize, String)> = spec
655                        .urls
656                        .into_iter()
657                        .enumerate()
658                        .filter(|(i, _)| !t.set_contains(&done, i.to_string()).unwrap_or(false))
659                        .collect();
660                    job.total = job.completed + pending.len();
661                    self.inner
662                        .map
663                        .try_lock()
664                        .expect("not shared yet")
665                        .insert(id.clone(), job);
666                    tracing::info!(job = %id, kind = "batch_scrape", pending = pending.len(), "resuming job");
667                    self.spawn_batch(
668                        id,
669                        engine.clone(),
670                        pending,
671                        spec.options,
672                        spec.concurrency,
673                        hook,
674                    );
675                }
676                (JobKind::Crawl, Some(req), _) => {
677                    let frontier: Vec<(String, usize)> = t
678                        .list_items(list_key(&id, "frontier"), usize::MAX)
679                        .unwrap_or_default()
680                        .iter()
681                        .filter_map(|v| serde_json::from_slice(v).ok())
682                        .collect();
683                    let resume = (!frontier.is_empty()).then(|| {
684                        let pending: Vec<(String, usize)> = frontier
685                            .iter()
686                            .filter(|(u, _)| !t.set_contains(&done, u).unwrap_or(false))
687                            .cloned()
688                            .collect();
689                        CrawlResume {
690                            enqueued: frontier.into_iter().map(|(u, _)| u).collect(),
691                            pending,
692                            done: job.completed,
693                        }
694                    });
695                    if let Some(r) = &resume {
696                        job.total =
697                            (job.completed + r.pending.len()).min(req.limit.max(job.completed));
698                    }
699                    self.inner
700                        .map
701                        .try_lock()
702                        .expect("not shared yet")
703                        .insert(id.clone(), job);
704                    tracing::info!(
705                        job = %id,
706                        kind = "crawl",
707                        pending = resume.as_ref().map_or(0, |r| r.pending.len()),
708                        "resuming job"
709                    );
710                    self.spawn_crawl(id, engine.clone(), req, hook, resume);
711                }
712                _ => {
713                    tracing::warn!(job = %id, "saved job has no request to resume from");
714                    job.state = JobState::Failed;
715                    self.inner
716                        .map
717                        .try_lock()
718                        .expect("not shared yet")
719                        .insert(id, job);
720                }
721            }
722            resumed += 1;
723        }
724        if finished + resumed + expired > 0 {
725            tracing::info!(finished, resumed, expired, "restored saved jobs");
726        }
727    }
728}
729
730/// 删除一个任务在存储里的全部数据。
731fn delete_job(store: &Store, id: &str) {
732    let t = store.tasks();
733    let r = (|| -> pagesift_store::Result<()> {
734        for what in ["pages", "errors", "frontier"] {
735            t.list_clear(list_key(id, what))?;
736        }
737        t.set_clear(list_key(id, "done"))?;
738        t.kv_delete(meta_key(id))?;
739        Ok(())
740    })();
741    if let Err(e) = r {
742        tracing::warn!(job = %id, error = %e, "deleting an expired job failed");
743    }
744}
745
746fn job_error(url: String, e: &ApiError) -> JobError {
747    JobError {
748        id: uuid::Uuid::new_v4().to_string(),
749        timestamp: rfc3339(SystemTime::now()),
750        url,
751        code: e.code.as_str().to_string(),
752        error: e.message.clone(),
753    }
754}
755
756/// 一个任务的进度回报:批量抓取和整站爬取共用。
757struct JobSink {
758    jobs: Jobs,
759    id: String,
760    kind: JobKind,
761    hook: Option<Hook>,
762    resume: StdMutex<Option<CrawlResume>>,
763    /// 调用方当天的页面配额用完了,任务提前结束。
764    quota_hit: AtomicBool,
765}
766
767impl JobSink {
768    fn store(&self) -> Option<&Store> {
769        self.jobs.inner.store.as_ref()
770    }
771
772    async fn owner(&self) -> Option<Caller> {
773        let map = self.jobs.inner.map.lock().await;
774        map.get(&self.id).and_then(|j| j.owner.clone())
775    }
776
777    async fn running(&self) -> bool {
778        let map = self.jobs.inner.map.lock().await;
779        map.get(&self.id)
780            .is_some_and(|j| j.state == JobState::Scraping)
781    }
782
783    /// 任务被取消,或者调用方的配额用完了。
784    async fn stopped(&self) -> bool {
785        let owner = {
786            let map = self.jobs.inner.map.lock().await;
787            match map.get(&self.id) {
788                Some(j) if j.state == JobState::Scraping => j.owner.clone(),
789                _ => return true,
790            }
791        };
792        if let Some(c) = owner
793            && self.jobs.inner.limiter.remaining_pages(&c) == Some(0)
794        {
795            self.quota_hit.store(true, Ordering::Relaxed);
796            return true;
797        }
798        false
799    }
800
801    /// 记一个完成的页面;`done` 是恢复时用来认出它的键(URL 或批量里的下标)。
802    async fn record_page(&self, done: String, doc: Document) {
803        let inner = &self.jobs.inner;
804        if !self.running().await {
805            return;
806        }
807        let page_doc = self.hook.as_ref().map(|_| doc.clone());
808        let mut keep = Some(doc);
809        if let Some(store) = self.store() {
810            let doc = keep.take().expect("set above");
811            let (pages, done_set) = (list_key(&self.id, "pages"), list_key(&self.id, "done"));
812            let r = store
813                .blocking(move |s| -> pagesift_store::Result<()> {
814                    let t = s.tasks();
815                    t.list_push(pages, &serde_json::to_vec(&doc).unwrap_or_default())?;
816                    t.set_insert(done_set, done)?;
817                    Ok(())
818                })
819                .await;
820            if let Err(e) = r {
821                tracing::warn!(job = %self.id, error = %e, "saving a job page failed");
822            }
823        }
824        let owner = {
825            let mut map = inner.map.lock().await;
826            let Some(job) = map.get_mut(&self.id) else {
827                return;
828            };
829            if job.state == JobState::Cancelled {
830                return;
831            }
832            job.completed += 1;
833            job.total = job.total.max(job.completed);
834            match keep {
835                Some(doc) => job.data.push(doc),
836                None => job.stored_docs += 1,
837            }
838            job.owner.clone()
839        };
840        inner.metrics.page(self.kind.pages());
841        if let Some(c) = owner {
842            inner.limiter.consume_pages(&c, 1).await;
843        }
844        if let (Some(hook), Some(doc)) = (&self.hook, page_doc) {
845            hook.send(Event::Page, vec![doc], None);
846        }
847    }
848
849    async fn record_error(&self, done: String, url: String, err: ApiError) {
850        let inner = &self.jobs.inner;
851        if !self.running().await {
852            return;
853        }
854        inner.metrics.error(err.code.as_str(), "job");
855        let e = job_error(url, &err);
856        let mut keep = Some(e);
857        if let Some(store) = self.store() {
858            let e = keep.take().expect("set above");
859            let (errors, done_set) = (list_key(&self.id, "errors"), list_key(&self.id, "done"));
860            let r = store
861                .blocking(move |s| -> pagesift_store::Result<()> {
862                    let t = s.tasks();
863                    t.list_push(errors, &serde_json::to_vec(&e).unwrap_or_default())?;
864                    t.set_insert(done_set, done)?;
865                    Ok(())
866                })
867                .await;
868            if let Err(e) = r {
869                tracing::warn!(job = %self.id, error = %e, "saving a job error failed");
870            }
871        }
872        let mut map = inner.map.lock().await;
873        let Some(job) = map.get_mut(&self.id) else {
874            return;
875        };
876        if job.state == JobState::Cancelled {
877            return;
878        }
879        job.completed += 1;
880        job.total = job.total.max(job.completed);
881        if let Some(e) = keep {
882            if err.code == ErrorCode::RobotsDisallowed {
883                job.robots_blocked.push(e.url.clone());
884            }
885            job.errors.push(e);
886        }
887    }
888
889    /// 任务跑完:定下最终状态、落盘、发结束事件。
890    async fn finish(&self, result: Result<(), (String, ApiError)>) {
891        let inner = &self.jobs.inner;
892        let quota = self.quota_hit.load(Ordering::Relaxed);
893        let mut extra = Vec::new();
894        if quota && let Some(c) = self.owner().await {
895            extra.push(job_error(String::new(), &quota_error(&c)));
896        }
897        if let Err((url, e)) = &result {
898            extra.push(job_error(url.clone(), e));
899        }
900        if let (Some(store), false) = (self.store(), extra.is_empty()) {
901            let (key, items) = (list_key(&self.id, "errors"), extra.clone());
902            store
903                .blocking(move |s| {
904                    for e in items {
905                        let _ = s
906                            .tasks()
907                            .list_push(&key, &serde_json::to_vec(&e).unwrap_or_default());
908                    }
909                })
910                .await;
911        }
912        let final_state = {
913            let mut map = inner.map.lock().await;
914            let Some(job) = map.get_mut(&self.id) else {
915                return;
916            };
917            if job.state == JobState::Scraping {
918                job.state = if result.is_ok() {
919                    JobState::Completed
920                } else {
921                    JobState::Failed
922                };
923                if self.store().is_none() {
924                    job.errors.extend(extra);
925                }
926            }
927            (job.state, job.completed)
928        };
929        self.jobs.save_state(&self.id).await;
930        let (state, completed) = final_state;
931        tracing::info!(
932            job = %self.id,
933            kind = self.kind.as_str(),
934            status = state.as_str(),
935            pages = completed,
936            quota_exceeded = quota,
937            "job finished"
938        );
939        if let Some(hook) = &self.hook {
940            match (state, &result) {
941                (JobState::Failed, Err((_, e))) => {
942                    hook.send(Event::Failed, Vec::new(), Some(e.to_string()))
943                }
944                (JobState::Completed, _) => hook.send(Event::Completed, Vec::new(), None),
945                _ => {}
946            }
947        }
948    }
949}
950
951impl CrawlSink for JobSink {
952    fn page(&self, doc: Document) -> BoxFuture<'_, ()> {
953        Box::pin(async move {
954            let done = doc.metadata.source_url.clone().unwrap_or_default();
955            self.record_page(done, doc).await
956        })
957    }
958
959    fn error<'a>(&'a self, url: &'a str, err: ApiError) -> BoxFuture<'a, ()> {
960        Box::pin(self.record_error(url.to_string(), url.to_string(), err))
961    }
962
963    fn total(&self, total: usize) -> BoxFuture<'_, ()> {
964        Box::pin(async move {
965            if let Some(job) = self.jobs.inner.map.lock().await.get_mut(&self.id) {
966                job.total = total.max(job.completed);
967            }
968        })
969    }
970
971    fn cancelled(&self) -> BoxFuture<'_, bool> {
972        Box::pin(self.stopped())
973    }
974
975    fn resume(&self) -> Option<CrawlResume> {
976        self.resume.lock().unwrap_or_else(|e| e.into_inner()).take()
977    }
978
979    fn tracks_frontier(&self) -> bool {
980        self.store().is_some()
981    }
982
983    fn summary(&self, summary: CrawlSummary) -> BoxFuture<'_, ()> {
984        Box::pin(async move {
985            if let Some(job) = self.jobs.inner.map.lock().await.get_mut(&self.id) {
986                job.summary = Some(summary);
987            }
988            // 结束时 `finish` 还会再存一次;这里先存,任务结束前进程退出也不丢
989            self.jobs.save_state(&self.id).await;
990        })
991    }
992
993    fn enqueued(&self, urls: Vec<(String, usize)>) -> BoxFuture<'_, ()> {
994        Box::pin(async move {
995            let (Some(store), false) = (self.store(), urls.is_empty()) else {
996                return;
997            };
998            let key = list_key(&self.id, "frontier");
999            let r = store
1000                .blocking(move |s| -> pagesift_store::Result<()> {
1001                    let t = s.tasks();
1002                    for u in urls {
1003                        t.list_push(&key, &serde_json::to_vec(&u).unwrap_or_default())?;
1004                    }
1005                    Ok(())
1006                })
1007                .await;
1008            if let Err(e) = r {
1009                tracing::warn!(job = %self.id, error = %e, "saving the crawl queue failed");
1010            }
1011        })
1012    }
1013}
1014
1015/// UTC 时间格式化成 RFC 3339(秒级精度后带毫秒),不引入日期库。
1016pub fn rfc3339(t: SystemTime) -> String {
1017    let d = t.duration_since(UNIX_EPOCH).unwrap_or_default();
1018    let secs = d.as_secs() as i64;
1019    let (days, rem) = (secs.div_euclid(86_400), secs.rem_euclid(86_400));
1020    // Howard Hinnant 的 civil_from_days
1021    let z = days + 719_468;
1022    let era = z.div_euclid(146_097);
1023    let doe = z - era * 146_097;
1024    let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
1025    let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
1026    let mp = (5 * doy + 2) / 153;
1027    let day = doy - (153 * mp + 2) / 5 + 1;
1028    let month = if mp < 10 { mp + 3 } else { mp - 9 };
1029    let year = yoe + era * 400 + i64::from(month <= 2);
1030    format!(
1031        "{year:04}-{month:02}-{day:02}T{:02}:{:02}:{:02}.{:03}Z",
1032        rem / 3600,
1033        rem % 3600 / 60,
1034        rem % 60,
1035        d.subsec_millis()
1036    )
1037}
1038
1039#[cfg(test)]
1040mod tests {
1041    use super::*;
1042
1043    #[test]
1044    fn formats_utc_dates() {
1045        assert_eq!(rfc3339(UNIX_EPOCH), "1970-01-01T00:00:00.000Z");
1046        let t = UNIX_EPOCH + Duration::from_millis(1_792_540_800_123);
1047        assert_eq!(rfc3339(t), "2026-10-21T00:00:00.123Z");
1048        let leap = UNIX_EPOCH + Duration::from_secs(951_782_400); // 2000-02-29
1049        assert_eq!(rfc3339(leap), "2000-02-29T00:00:00.000Z");
1050    }
1051}