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
30mod sink;
31pub mod state;
32mod store;
33pub mod time;
34
35pub use state::{JOB_TTL, Job, JobKind, JobState, Start, StatusPage};
36pub use time::rfc3339;
37
38use sink::JobSink;
39use store::{BatchSpec, Meta, delete_job, list_key, meta_key, to_ms};
40
41struct Inner {
42    map: Mutex<HashMap<String, Job>>,
43    store: Option<Store>,
44    limiter: Arc<Limiter>,
45    webhooks: Option<Webhooks>,
46    metrics: Arc<Metrics>,
47}
48
49#[derive(Clone)]
50pub struct Jobs {
51    inner: Arc<Inner>,
52}
53
54impl Default for Jobs {
55    /// 只在内存里的任务表：不落盘、不限额、不发 webhook。
56    fn default() -> Self {
57        Self::new(None, Arc::default(), None, Arc::default())
58    }
59}
60
61impl Jobs {
62    pub(crate) fn new(
63        store: Option<Store>,
64        limiter: Arc<Limiter>,
65        webhooks: Option<Webhooks>,
66        metrics: Arc<Metrics>,
67    ) -> Self {
68        Self {
69            inner: Arc::new(Inner {
70                map: Mutex::new(HashMap::new()),
71                store,
72                limiter,
73                webhooks,
74                metrics,
75            }),
76        }
77    }
78
79    /// 清掉过期任务（连同存储里的数据）。
80    async fn sweep(&self) {
81        let now = SystemTime::now();
82        let expired: Vec<String> = {
83            let mut map = self.inner.map.lock().await;
84            let ids: Vec<String> = map
85                .iter()
86                .filter(|(_, j)| j.expires_at() <= now)
87                .map(|(id, _)| id.clone())
88                .collect();
89            for id in &ids {
90                map.remove(id);
91            }
92            ids
93        };
94        if let (Some(store), false) = (&self.inner.store, expired.is_empty()) {
95            store
96                .blocking(move |s| {
97                    for id in expired {
98                        delete_job(s, &id);
99                    }
100                })
101                .await;
102        }
103    }
104
105    pub async fn get<R>(&self, id: &str, f: impl FnOnce(&Job) -> R) -> Option<R> {
106        self.sweep().await;
107        self.inner.map.lock().await.get(id).map(f)
108    }
109
110    /// 状态和从 `skip` 开始最多 `limit` 个文档。
111    pub(crate) async fn status(&self, id: &str, skip: usize, limit: usize) -> Option<StatusPage> {
112        self.sweep().await;
113        let mut page = {
114            let map = self.inner.map.lock().await;
115            let job = map.get(id)?;
116            let end = (skip + limit).min(job.data.len());
117            StatusPage {
118                state: job.state,
119                total: job.total,
120                completed: job.completed,
121                expires_at: job.expires_at(),
122                docs: job.docs(),
123                data: job
124                    .data
125                    .get(skip.min(end)..end)
126                    .unwrap_or_default()
127                    .to_vec(),
128                summary: job.summary.clone(),
129            }
130        };
131        if let Some(store) = &self.inner.store
132            && page.data.is_empty()
133            && skip < page.docs
134        {
135            let key = list_key(id, "pages");
136            page.data = store
137                .blocking(move |s| s.tasks().list_range(key, skip as u64, limit))
138                .await
139                .unwrap_or_else(|e| {
140                    tracing::warn!(error = %e, "reading job pages failed");
141                    Vec::new()
142                })
143                .iter()
144                .filter_map(|v| serde_json::from_slice(v).ok())
145                .collect();
146        }
147        Some(page)
148    }
149
150    /// 任务的页面错误和被 robots.txt 拦下的 URL。
151    pub(crate) async fn errors(&self, id: &str) -> Option<(Vec<JobError>, Vec<String>)> {
152        self.sweep().await;
153        let (mut errors, mut robots) = {
154            let map = self.inner.map.lock().await;
155            let job = map.get(id)?;
156            (job.errors.clone(), job.robots_blocked.clone())
157        };
158        if let Some(store) = &self.inner.store {
159            let key = list_key(id, "errors");
160            let stored = store
161                .blocking(move |s| s.tasks().list_items(key, usize::MAX))
162                .await
163                .unwrap_or_default();
164            for v in stored {
165                if let Ok(e) = serde_json::from_slice::<JobError>(&v) {
166                    if e.code == ErrorCode::RobotsDisallowed.as_str() {
167                        robots.push(e.url.clone());
168                    }
169                    errors.push(e);
170                }
171            }
172        }
173        Some((errors, robots))
174    }
175
176    pub async fn cancel(&self, id: &str) -> bool {
177        let found = {
178            let mut map = self.inner.map.lock().await;
179            match map.get_mut(id) {
180                Some(j) => {
181                    if j.state == JobState::Scraping {
182                        j.state = JobState::Cancelled;
183                    }
184                    true
185                }
186                None => false,
187            }
188        };
189        if found {
190            self.save_state(id).await;
191        }
192        found
193    }
194
195    pub async fn active(&self) -> Vec<String> {
196        let map = self.inner.map.lock().await;
197        map.iter()
198            .filter(|(_, j)| j.state == JobState::Scraping)
199            .map(|(id, _)| id.clone())
200            .collect()
201    }
202
203    /// 按 (类型, 状态) 统计的任务数，给 `/metrics`。
204    pub(crate) async fn counts(&self) -> BTreeMap<(&'static str, &'static str), u64> {
205        let mut out = BTreeMap::new();
206        for j in self.inner.map.lock().await.values() {
207            *out.entry((j.kind.as_str(), j.state.as_str())).or_default() += 1;
208        }
209        out
210    }
211
212    /// 登记新任务：检查调用方的并发任务数，落盘元数据，发 `started` 事件。
213    async fn register(
214        &self,
215        kind: JobKind,
216        total: usize,
217        start: Start,
218        crawl: Option<CrawlRequest>,
219        batch: Option<BatchSpec>,
220    ) -> Result<(String, Option<Hook>), ApiError> {
221        let id = uuid::Uuid::new_v4().to_string();
222        let job = Job::new(kind, total, start.owner.clone());
223        let meta = Meta {
224            kind,
225            state: JobState::Scraping,
226            created_ms: to_ms(job.created),
227            total,
228            completed: 0,
229            owner: start.owner.as_ref().map(|c| c.account.to_string()),
230            webhook: start.webhook.clone(),
231            crawl,
232            batch,
233            summary: None,
234        };
235        {
236            let mut map = self.inner.map.lock().await;
237            if let Some(c) = &start.owner {
238                if c.limits.max_concurrent_jobs.is_some() {
239                    let active = map
240                        .values()
241                        .filter(|j| {
242                            j.state == JobState::Scraping
243                                && j.owner.as_ref().is_some_and(|o| o.account == c.account)
244                        })
245                        .count();
246                    self.inner.limiter.check_jobs(c, active)?;
247                }
248                self.inner.limiter.check_quota(c)?;
249            }
250            map.insert(id.clone(), job);
251        }
252        if let Some(store) = &self.inner.store {
253            let (key, value) = (meta_key(&id), serde_json::to_vec(&meta).unwrap_or_default());
254            if let Err(e) = store.blocking(move |s| s.tasks().kv_put(key, &value)).await {
255                tracing::warn!(job = %id, error = %e, "saving job failed; it will not survive a restart");
256            }
257        }
258        self.inner.metrics.job_started(kind.as_str());
259        tracing::info!(
260            job = %id,
261            kind = kind.as_str(),
262            account = start.owner.as_ref().map(|c| &*c.account),
263            "job started"
264        );
265        let hook = match (&self.inner.webhooks, start.webhook) {
266            (Some(w), Some(cfg)) => {
267                let hook = w.start(&id, kind.as_str(), cfg);
268                hook.send(Event::Started, Vec::new(), None);
269                Some(hook)
270            }
271            _ => None,
272        };
273        Ok((id, hook))
274    }
275
276    /// 启动批量抓取，立即返回任务 id；每个 URL 的结果和错误陆续写入任务。
277    pub async fn start_batch(
278        &self,
279        engine: Arc<dyn Engine>,
280        urls: Vec<String>,
281        opts: ScrapeOptions,
282        concurrency: usize,
283    ) -> String {
284        self.start_batch_with(engine, urls, opts, concurrency, Start::default())
285            .await
286            .expect("no limits without a caller")
287    }
288
289    pub(crate) async fn start_batch_with(
290        &self,
291        engine: Arc<dyn Engine>,
292        urls: Vec<String>,
293        opts: ScrapeOptions,
294        concurrency: usize,
295        start: Start,
296    ) -> Result<String, ApiError> {
297        let spec = BatchSpec {
298            urls,
299            options: opts,
300            concurrency,
301        };
302        let total = spec.urls.len();
303        let (id, hook) = self
304            .register(
305                JobKind::BatchScrape,
306                total,
307                start,
308                None,
309                self.inner.store.is_some().then(|| spec.clone()),
310            )
311            .await?;
312        let pending = spec.urls.into_iter().enumerate().collect();
313        self.spawn_batch(id.clone(), engine, pending, spec.options, concurrency, hook);
314        Ok(id)
315    }
316
317    fn sink(
318        &self,
319        id: &str,
320        kind: JobKind,
321        hook: Option<Hook>,
322        resume: Option<CrawlResume>,
323    ) -> Arc<JobSink> {
324        Arc::new(JobSink {
325            jobs: self.clone(),
326            id: id.to_string(),
327            kind,
328            hook,
329            resume: StdMutex::new(resume),
330            quota_hit: AtomicBool::new(false),
331        })
332    }
333
334    pub(super) fn spawn_batch(
335        &self,
336        id: String,
337        engine: Arc<dyn Engine>,
338        pending: Vec<(usize, String)>,
339        opts: ScrapeOptions,
340        concurrency: usize,
341        hook: Option<Hook>,
342    ) {
343        let sink = self.sink(&id, JobKind::BatchScrape, hook, None);
344        let opts = Arc::new(opts);
345        tokio::spawn(async move {
346            let sem = Arc::new(Semaphore::new(concurrency.max(1)));
347            let mut tasks = tokio::task::JoinSet::new();
348            for (i, url) in pending {
349                let (sink, engine, opts, sem) =
350                    (sink.clone(), engine.clone(), opts.clone(), sem.clone());
351                tasks.spawn(async move {
352                    let _permit = sem.acquire_owned().await.expect("semaphore open");
353                    if sink.stopped().await {
354                        return;
355                    }
356                    let done = i.to_string();
357                    match engine.scrape(&url, &opts).await {
358                        Ok(doc) => sink.record_page(done, doc).await,
359                        Err(e) => sink.record_error(done, url, e).await,
360                    }
361                });
362            }
363            while tasks.join_next().await.is_some() {}
364            sink.finish(Ok(())).await;
365        });
366    }
367
368    /// 启动整站爬取，立即返回任务 id。
369    pub async fn start_crawl(&self, engine: Arc<dyn Engine>, req: CrawlRequest) -> String {
370        self.start_crawl_with(engine, req, Start::default())
371            .await
372            .expect("no limits without a caller")
373    }
374
375    pub(crate) async fn start_crawl_with(
376        &self,
377        engine: Arc<dyn Engine>,
378        req: CrawlRequest,
379        start: Start,
380    ) -> Result<String, ApiError> {
381        let (id, hook) = self
382            .register(
383                JobKind::Crawl,
384                0,
385                start,
386                self.inner.store.is_some().then(|| req.clone()),
387                None,
388            )
389            .await?;
390        self.spawn_crawl(id.clone(), engine, req, hook, None);
391        Ok(id)
392    }
393
394    pub(super) fn spawn_crawl(
395        &self,
396        id: String,
397        engine: Arc<dyn Engine>,
398        req: CrawlRequest,
399        hook: Option<Hook>,
400        resume: Option<CrawlResume>,
401    ) {
402        let sink = self.sink(&id, JobKind::Crawl, hook, resume);
403        tokio::spawn(async move {
404            let result = engine.crawl(&req, sink.clone()).await;
405            sink.finish(result.map_err(|e| (req.url.clone(), e))).await;
406        });
407    }
408}
pagesift_api/jobs/mod.rs

pagesift_api/jobs/
mod.rs
