pagesift_api/
server.rs

1//! HTTP server routing: Firecrawl v1/v2 compatibility layers, Jina Reader adapter,
2//! health checks, and Prometheus metrics.
3
4use std::collections::HashMap;
5use std::future::Future;
6use std::pin::{Pin, pin};
7use std::sync::Arc;
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::task::{Context, Poll};
10use std::time::{Duration, Instant};
11
12use axum::body::{Body, Bytes, HttpBody};
13use axum::extract::rejection::JsonRejection;
14use axum::extract::{Request, State};
15use axum::http::{HeaderMap, Method, StatusCode, header};
16use axum::middleware::{self, Next};
17use axum::response::{IntoResponse, Response};
18use axum::routing::get;
19use axum::{Json, Router};
20use http_body::{Frame, SizeHint};
21use hyper_util::rt::{TokioExecutor, TokioIo, TokioTimer};
22use hyper_util::server::conn::auto;
23use hyper_util::server::graceful::GracefulShutdown;
24use hyper_util::service::TowerToHyperService;
25use pagesift_store::Store;
26use tokio::net::TcpListener;
27use tokio::sync::Mutex;
28use tokio::time::Sleep;
29
30use crate::engine::Engine;
31use crate::error::{ApiError, ErrorCode};
32use crate::jobs::Jobs;
33use crate::limits::{Caller, KeyLimits, Limiter};
34use crate::metrics::Metrics;
35use crate::webhook::{WebhookConfig, WebhookSettings, Webhooks};
36
37/// Maximum documents per page when querying job status.
38pub const PAGE_SIZE: usize = 100;
39/// 批量抓取的并发：引擎的统一准入给出值时用它（[`Engine::admitted_concurrency`]），
40/// `max_batch_concurrency` 配置是用户上限。两者都没有时用这个先验。
41/// [策略] 先验包 v0：0.81 的现值 8（准入接好以前的行为不变）。
42pub const DEFAULT_BATCH_CONCURRENCY: usize = 8;
43
44/// 服务端超时，防 slowloris 和挂死的请求。
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct Timeouts {
47    /// 读完请求头的时限（HTTP/1），超时直接断开连接。只在 [`serve_with`] 里生效。
48    pub header_read: Duration,
49    /// 从请求进来到读完请求体的总时限，不是两次读之间的间隔：每 5 秒滴 1 字节也会超时。
50    pub body: Duration,
51    /// `/v2/parse` 上传文件的请求体时限（同样是总时限）：50 MB 的文件 30 秒传完要 1.7 MB/s，
52    /// 上传单独放宽。
53    pub upload_body: Duration,
54    /// 整个请求（到响应头发出为止）的时限。要比同步 scrape 的 `timeout` 上限大，
55    /// crawl/batch 只是创建任务，很快返回，不受影响。
56    pub request: Duration,
57}
58
59impl Default for Timeouts {
60    fn default() -> Self {
61        Self {
62            header_read: Duration::from_secs(30),
63            body: Duration::from_secs(30),
64            upload_body: Duration::from_secs(120),
65            request: Duration::from_millis(crate::validate::MAX_SCRAPE_TIMEOUT_MS)
66                + Duration::from_secs(60),
67        }
68    }
69}
70
71#[derive(Clone, Default)]
72pub struct Config {
73    /// When empty, authentication is disabled; otherwise requires `Authorization: Bearer <key>`.
74    pub api_keys: Vec<String>,
75    /// Public URL for building absolute resource links.
76    pub public_url: Option<String>,
77    pub max_batch_concurrency: Option<usize>,
78    /// Persistent store for jobs and resume capability.
79    pub store: Option<Store>,
80    /// Per-key custom limits.
81    pub key_limits: HashMap<String, KeyLimits>,
82    /// Default limits applied to other keys or unauthenticated callers.
83    pub default_limits: KeyLimits,
84    pub webhook: WebhookSettings,
85    /// 请求体和整请求超时；请求头超时由 [`serve_with`] 的参数决定。
86    pub timeouts: Timeouts,
87    /// 多服务器网关：配额和每分钟请求数记在账本里（追加记录，崩溃重放、温备回放后一致）。
88    /// 单主机不设，照旧用本进程的令牌桶和存储里的当天用量。
89    pub ledger: Option<Arc<pagesift_cluster::Ledger>>,
90    /// 多服务器网关：任务存储的写入复制到温备（[`crate::JobMirror`]）。单主机不设。
91    pub job_mirror: Option<Arc<crate::JobMirror>>,
92    /// 追加在 `/metrics` 后面的指标（多服务器网关：按工人加标签的统计）。
93    pub extra_metrics: Option<ExtraMetrics>,
94}
95
96/// 生成额外指标文本（Prometheus 文本格式）的回调。
97pub type ExtraMetrics = Arc<dyn Fn() -> String + Send + Sync>;
98
99impl std::fmt::Debug for Config {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        f.debug_struct("Config")
102            .field("api_keys", &self.api_keys.len())
103            .field("public_url", &self.public_url)
104            .field("max_batch_concurrency", &self.max_batch_concurrency)
105            .field("store", &self.store.is_some())
106            .field("key_limits", &self.key_limits.len())
107            .field("default_limits", &self.default_limits)
108            .field("webhook", &self.webhook)
109            .field("timeouts", &self.timeouts)
110            .field("ledger", &self.ledger.is_some())
111            .field("job_mirror", &self.job_mirror.is_some())
112            .field("extra_metrics", &self.extra_metrics.is_some())
113            .finish()
114    }
115}
116
117#[derive(Clone)]
118pub struct AppState {
119    pub engine: Arc<dyn Engine>,
120    pub keys: Arc<HashMap<String, Arc<KeyLimits>>>,
121    pub default_limits: Arc<KeyLimits>,
122    pub public_url: Option<String>,
123    /// 用户配置的批量并发上限（`max_batch_concurrency`），没配是 `None`。
124    pub batch_concurrency: Option<usize>,
125    pub jobs: Jobs,
126    pub limiter: Arc<Limiter>,
127    pub webhooks: Webhooks,
128    pub extracts: Arc<Mutex<HashMap<String, ExtractJob>>>,
129    pub metrics: Arc<Metrics>,
130    pub extra_metrics: Option<ExtraMetrics>,
131}
132
133impl AppState {
134    /// 一个批量任务的并发：`min(请求的 maxConcurrency, 准入值, 用户上限)`。
135    /// 准入值来自引擎的统一准入（[`Engine::admitted_concurrency`]）；引擎还不给时，用户配了上限
136    /// 就用上限（和 0.81 一样），都没有时用先验 [`DEFAULT_BATCH_CONCURRENCY`]。
137    pub fn batch_concurrency_for(&self, requested: Option<usize>) -> usize {
138        batch_concurrency(
139            requested,
140            self.engine.admitted_concurrency(),
141            self.batch_concurrency,
142        )
143    }
144}
145
146fn batch_concurrency(
147    requested: Option<usize>,
148    admitted: Option<usize>,
149    cap: Option<usize>,
150) -> usize {
151    let admitted = admitted.unwrap_or(cap.unwrap_or(DEFAULT_BATCH_CONCURRENCY));
152    requested
153        .unwrap_or(admitted)
154        .min(admitted)
155        .min(cap.unwrap_or(usize::MAX))
156        .max(1)
157}
158
159pub struct ExtractJob {
160    pub status: &'static str,
161    pub data: Option<serde_json::Value>,
162    pub error: Option<ApiError>,
163    pub created: std::time::SystemTime,
164}
165
166/// Construct complete HTTP router with Firecrawl v1/v2, Jina Reader, and metrics.
167pub fn router(engine: Arc<dyn Engine>, config: Config) -> Router {
168    let timeouts = config.timeouts;
169    let default_limits = Arc::new(config.default_limits);
170    let mut keys: HashMap<String, Arc<KeyLimits>> = config
171        .api_keys
172        .into_iter()
173        .map(|k| (k, default_limits.clone()))
174        .collect();
175    keys.extend(config.key_limits.into_iter().map(|(k, l)| (k, Arc::new(l))));
176
177    let accounts: HashMap<String, Arc<KeyLimits>> = keys
178        .iter()
179        .map(|(k, l)| {
180            (
181                Caller::new(Some(k), l.clone()).account.to_string(),
182                l.clone(),
183            )
184        })
185        .collect();
186
187    let metrics = Arc::new(Metrics::default());
188    let limiter = Arc::new(match config.ledger {
189        Some(ledger) => {
190            // 所有账户的限额事先登记进账本（不在每个请求上登记）；没有 key 的调用方是 anonymous
191            let mut all = accounts.clone();
192            all.entry("anonymous".into())
193                .or_insert_with(|| default_limits.clone());
194            Limiter::with_ledger(ledger, &all)
195        }
196        None => Limiter::new(config.store.clone()),
197    });
198    // 附件跟任务结果一样写进存储，重启后链接仍然有效（见 `attachments`）
199    if let Some(s) = &config.store
200        && tokio::runtime::Handle::try_current().is_ok()
201    {
202        crate::attachments::sweep_expired(crate::attachments::Disk {
203            store: s.clone(),
204            mirror: config.job_mirror.clone(),
205        });
206    }
207    let mut webhooks = Webhooks::new(config.webhook, metrics.clone());
208    webhooks.mirror = config.job_mirror.clone();
209    let jobs = Jobs::new(
210        config.store,
211        limiter.clone(),
212        Some(webhooks.clone()),
213        metrics.clone(),
214        config.job_mirror,
215    );
216
217    if tokio::runtime::Handle::try_current().is_ok() {
218        let fallback = default_limits.clone();
219        jobs.restore(&engine, &|account| {
220            accounts
221                .get(account)
222                .cloned()
223                .unwrap_or_else(|| fallback.clone())
224        });
225    } else {
226        tracing::warn!("router built outside a tokio runtime: saved jobs are not resumed");
227    }
228
229    let state = AppState {
230        engine: engine.clone(),
231        keys: Arc::new(keys),
232        default_limits,
233        public_url: config
234            .public_url
235            .map(|u| u.trim_end_matches('/').to_string()),
236        batch_concurrency: config.max_batch_concurrency,
237        jobs,
238        limiter,
239        webhooks,
240        extracts: Arc::default(),
241        metrics,
242        extra_metrics: config.extra_metrics,
243    };
244
245    // Jina Reader 路由也会抓取页面，和 Firecrawl 路由一起放在鉴权、限流和配额之后
246    let api = crate::firecrawl::firecrawl_routes()
247        .merge(crate::jina::jina_routes())
248        .route_layer(middleware::from_fn_with_state(state.clone(), auth));
249
250    Router::new()
251        .route("/healthz", get(|| async { "ok" }))
252        .route("/metrics", get(render_metrics))
253        .merge(api)
254        .with_state(state.clone())
255        .fallback(|| async { ApiError::new(ErrorCode::NotFound, "no such endpoint") })
256        .layer(middleware::from_fn(uniform_errors))
257        .layer(middleware::from_fn_with_state(timeouts, deadline))
258        .layer(middleware::from_fn_with_state(state.clone(), observe))
259}
260
261/// 用默认超时在 `listener` 上提供服务，`shutdown` 完成后优雅退出。
262pub async fn serve(
263    listener: TcpListener,
264    app: Router,
265    shutdown: impl Future<Output = ()> + Send,
266) -> std::io::Result<()> {
267    serve_with(listener, app, Timeouts::default(), shutdown).await
268}
269
270/// accept 因资源耗尽出错后，没有连接结束时最多等多久再试。
271/// [策略] 0.81 的 1 秒作为兜底上限（axum::serve 同值）：资源由别的进程释放时收不到通知，只能定时重试；
272/// 本进程的连接结束会提前叫醒（见 [`serve_with`]）。
273const ACCEPT_RETRY_CAP: Duration = Duration::from_secs(1);
274
275/// 和 `axum::serve(..).with_graceful_shutdown(..)` 行为相同，另外给 HTTP/1 设了请求头读取超时
276/// （axum::serve 设不了）：收到 `shutdown` 后停止 accept，通知每个连接优雅关闭，等它们都结束再返回。
277pub async fn serve_with(
278    listener: TcpListener,
279    app: Router,
280    timeouts: Timeouts,
281    shutdown: impl Future<Output = ()> + Send,
282) -> std::io::Result<()> {
283    let mut builder = auto::Builder::new(TokioExecutor::new());
284    builder
285        .http1()
286        .timer(TokioTimer::new())
287        .header_read_timeout(timeouts.header_read);
288    let graceful = GracefulShutdown::new();
289    let mut shutdown = pin!(shutdown);
290    // 有连接结束（释放了文件描述符）时叫醒因 accept 出错而在等的循环
291    let released = Arc::new(tokio::sync::Notify::new());
292    loop {
293        let stream = tokio::select! {
294            r = listener.accept() => match r {
295                Ok((stream, _)) => stream,
296                Err(e) => {
297                    // 和 axum::serve 一样：单个连接的错误忽略。其他错误多半是文件描述符或内核缓冲
298                    // 用完（EMFILE、ENFILE、ENOBUFS），要等资源释放：事件驱动，任一连接结束就重试
299                    // （它释放了描述符）；没有连接结束时最多等 ACCEPT_RETRY_CAP 再试（别的进程释放的
300                    // 资源我们收不到通知）。取代 0.81 固定歇 1 秒：描述符一释放就接，不白等。
301                    if !matches!(
302                        e.kind(),
303                        std::io::ErrorKind::ConnectionRefused
304                            | std::io::ErrorKind::ConnectionAborted
305                            | std::io::ErrorKind::ConnectionReset
306                    ) {
307                        tracing::error!(error = %e, "accept failed");
308                        tokio::select! {
309                            () = released.notified() => {}
310                            () = tokio::time::sleep(ACCEPT_RETRY_CAP) => {}
311                            () = &mut shutdown => break,
312                        }
313                    }
314                    continue;
315                }
316            },
317            () = &mut shutdown => break,
318        };
319        // 响应头和分块正文分两次写时，Nagle 会让第二次写等对端的延迟确认（约 40 ms）
320        let _ = stream.set_nodelay(true);
321        let conn = builder
322            .serve_connection_with_upgrades(
323                TokioIo::new(stream),
324                TowerToHyperService::new(app.clone()),
325            )
326            .into_owned();
327        let conn = graceful.watch(conn);
328        let released = released.clone();
329        tokio::spawn(async move {
330            if let Err(e) = conn.await {
331                tracing::trace!(error = %e, "connection ended with an error");
332            }
333            released.notify_one();
334        });
335    }
336    drop(listener);
337    graceful.shutdown().await;
338    Ok(())
339}
340
341/// 请求体加总时限；整个请求加时限。超时都回统一的 JSON 错误（408 TIMEOUT）。
342async fn deadline(State(t): State<Timeouts>, req: Request, next: Next) -> Response {
343    let fired = Arc::new(AtomicBool::new(false));
344    // 上传放宽的请求体时间也加到整请求时限上，解析仍有完整的 `timeout` 可用
345    let (body_limit, request_limit) = if req.uri().path() == "/v2/parse" {
346        (
347            t.upload_body,
348            t.request + t.upload_body.saturating_sub(t.body),
349        )
350    } else {
351        (t.body, t.request)
352    };
353    let req = if req.body().is_end_stream() {
354        req
355    } else {
356        let fired = fired.clone();
357        req.map(|inner| {
358            Body::new(DeadlineBody {
359                inner,
360                sleep: Box::pin(tokio::time::sleep(body_limit)),
361                fired,
362            })
363        })
364    };
365    match tokio::time::timeout(request_limit, next.run(req)).await {
366        // 读请求体超时时提取器给的是 400，这里换成 408
367        Ok(_) if fired.load(Ordering::Relaxed) => ApiError::new(
368            ErrorCode::Timeout,
369            format!(
370                "request body was not received within {} s",
371                body_limit.as_secs_f64()
372            ),
373        )
374        .into_response(),
375        Ok(resp) => resp,
376        Err(_) => ApiError::new(
377            ErrorCode::Timeout,
378            format!(
379                "request did not complete within {} s",
380                request_limit.as_secs_f64()
381            ),
382        )
383        .into_response(),
384    }
385}
386
387/// 到点还没读完就报错的请求体。
388struct DeadlineBody {
389    inner: Body,
390    sleep: Pin<Box<Sleep>>,
391    fired: Arc<AtomicBool>,
392}
393
394impl HttpBody for DeadlineBody {
395    type Data = Bytes;
396    type Error = axum::Error;
397
398    fn poll_frame(
399        self: Pin<&mut Self>,
400        cx: &mut Context<'_>,
401    ) -> Poll<Option<Result<Frame<Bytes>, axum::Error>>> {
402        let this = self.get_mut();
403        if let Poll::Ready(frame) = Pin::new(&mut this.inner).poll_frame(cx) {
404            return Poll::Ready(frame);
405        }
406        if this.sleep.as_mut().poll(cx).is_ready() {
407            this.fired.store(true, Ordering::Relaxed);
408            return Poll::Ready(Some(Err(axum::Error::new("request body timed out"))));
409        }
410        Poll::Pending
411    }
412
413    fn is_end_stream(&self) -> bool {
414        self.inner.is_end_stream()
415    }
416
417    fn size_hint(&self) -> SizeHint {
418        self.inner.size_hint()
419    }
420}
421
422/// axum 自带的拒绝（查询参数、路径参数、JSON 解析失败、405 等）是纯文本或空 body，
423/// 统一换成 `{success: false, code, error}`，状态码和 `Allow` 等响应头不变。
424async fn uniform_errors(req: Request, next: Next) -> Response {
425    let resp = next.run(req).await;
426    let status = resp.status();
427    // 自己的错误都带 ErrorCode 扩展，不用动
428    if !status.is_client_error() || resp.extensions().get::<ErrorCode>().is_some() {
429        return resp;
430    }
431    let plain = resp
432        .headers()
433        .get(header::CONTENT_TYPE)
434        .is_none_or(|v| v.as_bytes().starts_with(b"text/plain"));
435    if !plain {
436        return resp;
437    }
438    let (parts, body) = resp.into_parts();
439    let text = axum::body::to_bytes(body, 64 * 1024)
440        .await
441        .map(|b| String::from_utf8_lossy(&b).trim().to_string())
442        .unwrap_or_default();
443    let code = match status {
444        StatusCode::NOT_FOUND => ErrorCode::NotFound,
445        StatusCode::PAYLOAD_TOO_LARGE => ErrorCode::TooLarge,
446        StatusCode::UNAUTHORIZED => ErrorCode::Unauthorized,
447        _ => ErrorCode::InvalidRequest,
448    };
449    let message = if text.is_empty() {
450        status
451            .canonical_reason()
452            .unwrap_or("bad request")
453            .to_lowercase()
454    } else {
455        text
456    };
457    let mut out = ApiError::new(code, message).into_response();
458    *out.status_mut() = status;
459    for (name, value) in &parts.headers {
460        if name != header::CONTENT_TYPE && name != header::CONTENT_LENGTH {
461            out.headers_mut().append(name, value.clone());
462        }
463    }
464    out
465}
466
467/// Authentication and rate limiting middleware.
468async fn auth(
469    State(s): State<AppState>,
470    headers: HeaderMap,
471    mut req: Request,
472    next: Next,
473) -> Response {
474    let caller = if s.keys.is_empty() {
475        Caller::new(None, s.default_limits.clone())
476    } else {
477        let key = headers
478            .get(header::AUTHORIZATION)
479            .and_then(|v| v.to_str().ok())
480            .and_then(|v| v.strip_prefix("Bearer "))
481            .map(str::trim);
482        match key.and_then(|k| s.keys.get(k).map(|l| (k, l))) {
483            Some((k, l)) => Caller::new(Some(k), l.clone()),
484            None => {
485                return ApiError::new(ErrorCode::Unauthorized, "missing or invalid API key")
486                    .into_response();
487            }
488        }
489    };
490    // 计入每分钟请求数的是会抓取的请求：POST，以及 Jina 的 `GET /r/{url}`；轮询类 GET 不算
491    if (req.method() == Method::POST || req.uri().path().starts_with("/r/"))
492        && let Err(e) = s.limiter.check_rate(&caller)
493    {
494        let mut resp = e.into_response();
495        resp.extensions_mut().insert(caller);
496        return resp;
497    }
498    req.extensions_mut().insert(caller.clone());
499    let mut resp = next.run(req).await;
500    resp.extensions_mut().insert(caller);
501    resp
502}
503
504/// Metrics and structured request tracing middleware.
505async fn observe(State(s): State<AppState>, req: Request, next: Next) -> Response {
506    let route = req
507        .extensions()
508        .get::<axum::extract::MatchedPath>()
509        .map(|p| p.as_str().to_string())
510        .unwrap_or_else(|| "unmatched".into());
511    let method = req.method().clone();
512    let started = Instant::now();
513    s.metrics.request_started();
514    let resp = next.run(req).await;
515    let elapsed = started.elapsed();
516    let status = resp.status().as_u16();
517    s.metrics.request_finished(&route, status, elapsed);
518    let code = resp.extensions().get::<ErrorCode>().map(|c| c.as_str());
519    if let Some(code) = code {
520        s.metrics.error(code, "http");
521    }
522    let account = resp.extensions().get::<Caller>().map(|c| c.account.clone());
523    let latency_ms = elapsed.as_secs_f64() * 1000.0;
524    if route.starts_with("/v2/") || route.starts_with("/v1/") || route.starts_with("/r") {
525        tracing::info!(
526            target: "pagesift::http",
527            %method, %route, status, latency_ms, account = account.as_deref(), code,
528            "request"
529        );
530    } else {
531        tracing::debug!(target: "pagesift::http", %method, %route, status, latency_ms, "request");
532    }
533    resp
534}
535
536pub fn body<T>(r: Result<Json<T>, JsonRejection>) -> Result<T, ApiError> {
537    r.map(|Json(v)| v)
538        .map_err(|e| ApiError::invalid(format!("invalid request body: {}", e.body_text())))
539}
540
541pub fn base_url(s: &AppState, headers: &HeaderMap) -> String {
542    if let Some(u) = &s.public_url {
543        return u.clone();
544    }
545    let host = headers
546        .get(header::HOST)
547        .and_then(|v| v.to_str().ok())
548        .unwrap_or("localhost:3002");
549    format!("http://{host}")
550}
551
552pub fn webhook(
553    s: &AppState,
554    v: Option<&serde_json::Value>,
555) -> Result<Option<WebhookConfig>, ApiError> {
556    v.filter(|v| !v.is_null())
557        .map(|v| WebhookConfig::parse(v, s.webhooks.ssrf()))
558        .transpose()
559}
560
561async fn render_metrics(State(s): State<AppState>) -> impl IntoResponse {
562    let mut out = s.metrics.render(&s.jobs.counts().await);
563    if let Some(extra) = &s.extra_metrics {
564        out.push_str(&extra());
565    }
566    (
567        StatusCode::OK,
568        [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
569        out,
570    )
571}
572
573#[cfg(test)]
574mod batch_concurrency_tests {
575    use super::batch_concurrency;
576
577    /// 准入还没接（None）时和 0.81 完全一样：没配置是 min(请求, 8)，配置了是 min(请求, 配置)；
578    /// 准入给值以后按准入，用户配置仍是上限。
579    #[test]
580    fn admission_decides_and_the_user_caps() {
581        assert_eq!(batch_concurrency(None, None, None), 8);
582        assert_eq!(batch_concurrency(Some(20), None, None), 8);
583        assert_eq!(batch_concurrency(Some(3), None, None), 3);
584        assert_eq!(batch_concurrency(None, None, Some(32)), 32);
585        assert_eq!(batch_concurrency(Some(64), None, Some(32)), 32);
586        assert_eq!(batch_concurrency(None, Some(50), None), 50);
587        assert_eq!(batch_concurrency(None, Some(50), Some(32)), 32);
588        assert_eq!(batch_concurrency(Some(5), Some(50), None), 5);
589        assert_eq!(batch_concurrency(Some(0), Some(0), None), 1);
590    }
591}
pagesift_api/server.rs

pagesift_api/
server.rs
