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;
39pub const DEFAULT_BATCH_CONCURRENCY: usize = 8;
40
41/// 服务端超时，防 slowloris 和挂死的请求。
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub struct Timeouts {
44    /// 读完请求头的时限（HTTP/1），超时直接断开连接。只在 [`serve_with`] 里生效。
45    pub header_read: Duration,
46    /// 从请求进来到读完请求体的总时限，不是两次读之间的间隔：每 5 秒滴 1 字节也会超时。
47    pub body: Duration,
48    /// `/v2/parse` 上传文件的请求体时限（同样是总时限）：50 MB 的文件 30 秒传完要 1.7 MB/s，
49    /// 上传单独放宽。
50    pub upload_body: Duration,
51    /// 整个请求（到响应头发出为止）的时限。要比同步 scrape 的 `timeout` 上限大，
52    /// crawl/batch 只是创建任务，很快返回，不受影响。
53    pub request: Duration,
54}
55
56impl Default for Timeouts {
57    fn default() -> Self {
58        Self {
59            header_read: Duration::from_secs(30),
60            body: Duration::from_secs(30),
61            upload_body: Duration::from_secs(120),
62            request: Duration::from_millis(crate::validate::MAX_SCRAPE_TIMEOUT_MS)
63                + Duration::from_secs(60),
64        }
65    }
66}
67
68#[derive(Clone, Default)]
69pub struct Config {
70    /// When empty, authentication is disabled; otherwise requires `Authorization: Bearer <key>`.
71    pub api_keys: Vec<String>,
72    /// Public URL for building absolute resource links.
73    pub public_url: Option<String>,
74    pub max_batch_concurrency: Option<usize>,
75    /// Persistent store for jobs and resume capability.
76    pub store: Option<Store>,
77    /// Per-key custom limits.
78    pub key_limits: HashMap<String, KeyLimits>,
79    /// Default limits applied to other keys or unauthenticated callers.
80    pub default_limits: KeyLimits,
81    pub webhook: WebhookSettings,
82    /// 请求体和整请求超时；请求头超时由 [`serve_with`] 的参数决定。
83    pub timeouts: Timeouts,
84}
85
86impl std::fmt::Debug for Config {
87    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
88        f.debug_struct("Config")
89            .field("api_keys", &self.api_keys.len())
90            .field("public_url", &self.public_url)
91            .field("max_batch_concurrency", &self.max_batch_concurrency)
92            .field("store", &self.store.is_some())
93            .field("key_limits", &self.key_limits.len())
94            .field("default_limits", &self.default_limits)
95            .field("webhook", &self.webhook)
96            .field("timeouts", &self.timeouts)
97            .finish()
98    }
99}
100
101#[derive(Clone)]
102pub struct AppState {
103    pub engine: Arc<dyn Engine>,
104    pub keys: Arc<HashMap<String, Arc<KeyLimits>>>,
105    pub default_limits: Arc<KeyLimits>,
106    pub public_url: Option<String>,
107    pub batch_concurrency: usize,
108    pub jobs: Jobs,
109    pub limiter: Arc<Limiter>,
110    pub webhooks: Webhooks,
111    pub extracts: Arc<Mutex<HashMap<String, ExtractJob>>>,
112    pub metrics: Arc<Metrics>,
113}
114
115pub struct ExtractJob {
116    pub status: &'static str,
117    pub data: Option<serde_json::Value>,
118    pub error: Option<ApiError>,
119    pub created: std::time::SystemTime,
120}
121
122/// Construct complete HTTP router with Firecrawl v1/v2, Jina Reader, and metrics.
123pub fn router(engine: Arc<dyn Engine>, config: Config) -> Router {
124    let timeouts = config.timeouts;
125    let default_limits = Arc::new(config.default_limits);
126    let mut keys: HashMap<String, Arc<KeyLimits>> = config
127        .api_keys
128        .into_iter()
129        .map(|k| (k, default_limits.clone()))
130        .collect();
131    keys.extend(config.key_limits.into_iter().map(|(k, l)| (k, Arc::new(l))));
132
133    let accounts: HashMap<String, Arc<KeyLimits>> = keys
134        .iter()
135        .map(|(k, l)| {
136            (
137                Caller::new(Some(k), l.clone()).account.to_string(),
138                l.clone(),
139            )
140        })
141        .collect();
142
143    let metrics = Arc::new(Metrics::default());
144    let limiter = Arc::new(Limiter::new(config.store.clone()));
145    let webhooks = Webhooks::new(config.webhook, metrics.clone());
146    let jobs = Jobs::new(
147        config.store,
148        limiter.clone(),
149        Some(webhooks.clone()),
150        metrics.clone(),
151    );
152
153    if tokio::runtime::Handle::try_current().is_ok() {
154        let fallback = default_limits.clone();
155        jobs.restore(&engine, &|account| {
156            accounts
157                .get(account)
158                .cloned()
159                .unwrap_or_else(|| fallback.clone())
160        });
161    } else {
162        tracing::warn!("router built outside a tokio runtime: saved jobs are not resumed");
163    }
164
165    let state = AppState {
166        engine: engine.clone(),
167        keys: Arc::new(keys),
168        default_limits,
169        public_url: config
170            .public_url
171            .map(|u| u.trim_end_matches('/').to_string()),
172        batch_concurrency: config
173            .max_batch_concurrency
174            .unwrap_or(DEFAULT_BATCH_CONCURRENCY),
175        jobs,
176        limiter,
177        webhooks,
178        extracts: Arc::default(),
179        metrics,
180    };
181
182    let firecrawl = crate::firecrawl::firecrawl_routes()
183        .route_layer(middleware::from_fn_with_state(state.clone(), auth));
184    let jina = crate::jina::jina_router(engine);
185
186    Router::new()
187        .route("/healthz", get(|| async { "ok" }))
188        .route("/metrics", get(render_metrics))
189        .merge(firecrawl)
190        .with_state(state.clone())
191        .merge(jina)
192        .fallback(|| async { ApiError::new(ErrorCode::NotFound, "no such endpoint") })
193        .layer(middleware::from_fn(uniform_errors))
194        .layer(middleware::from_fn_with_state(timeouts, deadline))
195        .layer(middleware::from_fn_with_state(state.clone(), observe))
196}
197
198/// 用默认超时在 `listener` 上提供服务，`shutdown` 完成后优雅退出。
199pub async fn serve(
200    listener: TcpListener,
201    app: Router,
202    shutdown: impl Future<Output = ()> + Send,
203) -> std::io::Result<()> {
204    serve_with(listener, app, Timeouts::default(), shutdown).await
205}
206
207/// 和 `axum::serve(..).with_graceful_shutdown(..)` 行为相同，另外给 HTTP/1 设了请求头读取超时
208/// （axum::serve 设不了）：收到 `shutdown` 后停止 accept，通知每个连接优雅关闭，等它们都结束再返回。
209pub async fn serve_with(
210    listener: TcpListener,
211    app: Router,
212    timeouts: Timeouts,
213    shutdown: impl Future<Output = ()> + Send,
214) -> std::io::Result<()> {
215    let mut builder = auto::Builder::new(TokioExecutor::new());
216    builder
217        .http1()
218        .timer(TokioTimer::new())
219        .header_read_timeout(timeouts.header_read);
220    let graceful = GracefulShutdown::new();
221    let mut shutdown = pin!(shutdown);
222    loop {
223        let stream = tokio::select! {
224            r = listener.accept() => match r {
225                Ok((stream, _)) => stream,
226                Err(e) => {
227                    // 和 axum::serve 一样：单个连接的错误忽略，其他错误（比如文件描述符用完）歇 1 秒再接
228                    if !matches!(
229                        e.kind(),
230                        std::io::ErrorKind::ConnectionRefused
231                            | std::io::ErrorKind::ConnectionAborted
232                            | std::io::ErrorKind::ConnectionReset
233                    ) {
234                        tracing::error!(error = %e, "accept failed");
235                        tokio::time::sleep(Duration::from_secs(1)).await;
236                    }
237                    continue;
238                }
239            },
240            () = &mut shutdown => break,
241        };
242        // 响应头和分块正文分两次写时，Nagle 会让第二次写等对端的延迟确认（约 40 ms）
243        let _ = stream.set_nodelay(true);
244        let conn = builder
245            .serve_connection_with_upgrades(
246                TokioIo::new(stream),
247                TowerToHyperService::new(app.clone()),
248            )
249            .into_owned();
250        let conn = graceful.watch(conn);
251        tokio::spawn(async move {
252            if let Err(e) = conn.await {
253                tracing::trace!(error = %e, "connection ended with an error");
254            }
255        });
256    }
257    drop(listener);
258    graceful.shutdown().await;
259    Ok(())
260}
261
262/// 请求体加总时限；整个请求加时限。超时都回统一的 JSON 错误（408 TIMEOUT）。
263async fn deadline(State(t): State<Timeouts>, req: Request, next: Next) -> Response {
264    let fired = Arc::new(AtomicBool::new(false));
265    // 上传放宽的请求体时间也加到整请求时限上，解析仍有完整的 `timeout` 可用
266    let (body_limit, request_limit) = if req.uri().path() == "/v2/parse" {
267        (
268            t.upload_body,
269            t.request + t.upload_body.saturating_sub(t.body),
270        )
271    } else {
272        (t.body, t.request)
273    };
274    let req = if req.body().is_end_stream() {
275        req
276    } else {
277        let fired = fired.clone();
278        req.map(|inner| {
279            Body::new(DeadlineBody {
280                inner,
281                sleep: Box::pin(tokio::time::sleep(body_limit)),
282                fired,
283            })
284        })
285    };
286    match tokio::time::timeout(request_limit, next.run(req)).await {
287        // 读请求体超时时提取器给的是 400，这里换成 408
288        Ok(_) if fired.load(Ordering::Relaxed) => ApiError::new(
289            ErrorCode::Timeout,
290            format!(
291                "request body was not received within {} s",
292                body_limit.as_secs_f64()
293            ),
294        )
295        .into_response(),
296        Ok(resp) => resp,
297        Err(_) => ApiError::new(
298            ErrorCode::Timeout,
299            format!(
300                "request did not complete within {} s",
301                request_limit.as_secs_f64()
302            ),
303        )
304        .into_response(),
305    }
306}
307
308/// 到点还没读完就报错的请求体。
309struct DeadlineBody {
310    inner: Body,
311    sleep: Pin<Box<Sleep>>,
312    fired: Arc<AtomicBool>,
313}
314
315impl HttpBody for DeadlineBody {
316    type Data = Bytes;
317    type Error = axum::Error;
318
319    fn poll_frame(
320        self: Pin<&mut Self>,
321        cx: &mut Context<'_>,
322    ) -> Poll<Option<Result<Frame<Bytes>, axum::Error>>> {
323        let this = self.get_mut();
324        if let Poll::Ready(frame) = Pin::new(&mut this.inner).poll_frame(cx) {
325            return Poll::Ready(frame);
326        }
327        if this.sleep.as_mut().poll(cx).is_ready() {
328            this.fired.store(true, Ordering::Relaxed);
329            return Poll::Ready(Some(Err(axum::Error::new("request body timed out"))));
330        }
331        Poll::Pending
332    }
333
334    fn is_end_stream(&self) -> bool {
335        self.inner.is_end_stream()
336    }
337
338    fn size_hint(&self) -> SizeHint {
339        self.inner.size_hint()
340    }
341}
342
343/// axum 自带的拒绝（查询参数、路径参数、JSON 解析失败、405 等）是纯文本或空 body，
344/// 统一换成 `{success: false, code, error}`，状态码和 `Allow` 等响应头不变。
345async fn uniform_errors(req: Request, next: Next) -> Response {
346    let resp = next.run(req).await;
347    let status = resp.status();
348    // 自己的错误都带 ErrorCode 扩展，不用动
349    if !status.is_client_error() || resp.extensions().get::<ErrorCode>().is_some() {
350        return resp;
351    }
352    let plain = resp
353        .headers()
354        .get(header::CONTENT_TYPE)
355        .is_none_or(|v| v.as_bytes().starts_with(b"text/plain"));
356    if !plain {
357        return resp;
358    }
359    let (parts, body) = resp.into_parts();
360    let text = axum::body::to_bytes(body, 64 * 1024)
361        .await
362        .map(|b| String::from_utf8_lossy(&b).trim().to_string())
363        .unwrap_or_default();
364    let code = match status {
365        StatusCode::NOT_FOUND => ErrorCode::NotFound,
366        StatusCode::PAYLOAD_TOO_LARGE => ErrorCode::TooLarge,
367        StatusCode::UNAUTHORIZED => ErrorCode::Unauthorized,
368        _ => ErrorCode::InvalidRequest,
369    };
370    let message = if text.is_empty() {
371        status
372            .canonical_reason()
373            .unwrap_or("bad request")
374            .to_lowercase()
375    } else {
376        text
377    };
378    let mut out = ApiError::new(code, message).into_response();
379    *out.status_mut() = status;
380    for (name, value) in &parts.headers {
381        if name != header::CONTENT_TYPE && name != header::CONTENT_LENGTH {
382            out.headers_mut().append(name, value.clone());
383        }
384    }
385    out
386}
387
388/// Authentication and rate limiting middleware.
389async fn auth(
390    State(s): State<AppState>,
391    headers: HeaderMap,
392    mut req: Request,
393    next: Next,
394) -> Response {
395    let caller = if s.keys.is_empty() {
396        Caller::new(None, s.default_limits.clone())
397    } else {
398        let key = headers
399            .get(header::AUTHORIZATION)
400            .and_then(|v| v.to_str().ok())
401            .and_then(|v| v.strip_prefix("Bearer "))
402            .map(str::trim);
403        match key.and_then(|k| s.keys.get(k).map(|l| (k, l))) {
404            Some((k, l)) => Caller::new(Some(k), l.clone()),
405            None => {
406                return ApiError::new(ErrorCode::Unauthorized, "missing or invalid API key")
407                    .into_response();
408            }
409        }
410    };
411    if req.method() == Method::POST
412        && let Err(e) = s.limiter.check_rate(&caller)
413    {
414        let mut resp = e.into_response();
415        resp.extensions_mut().insert(caller);
416        return resp;
417    }
418    req.extensions_mut().insert(caller.clone());
419    let mut resp = next.run(req).await;
420    resp.extensions_mut().insert(caller);
421    resp
422}
423
424/// Metrics and structured request tracing middleware.
425async fn observe(State(s): State<AppState>, req: Request, next: Next) -> Response {
426    let route = req
427        .extensions()
428        .get::<axum::extract::MatchedPath>()
429        .map(|p| p.as_str().to_string())
430        .unwrap_or_else(|| "unmatched".into());
431    let method = req.method().clone();
432    let started = Instant::now();
433    s.metrics.request_started();
434    let resp = next.run(req).await;
435    let elapsed = started.elapsed();
436    let status = resp.status().as_u16();
437    s.metrics.request_finished(&route, status, elapsed);
438    let code = resp.extensions().get::<ErrorCode>().map(|c| c.as_str());
439    if let Some(code) = code {
440        s.metrics.error(code, "http");
441    }
442    let account = resp.extensions().get::<Caller>().map(|c| c.account.clone());
443    let latency_ms = elapsed.as_secs_f64() * 1000.0;
444    if route.starts_with("/v2/") || route.starts_with("/v1/") || route.starts_with("/r") {
445        tracing::info!(
446            target: "pagesift::http",
447            %method, %route, status, latency_ms, account = account.as_deref(), code,
448            "request"
449        );
450    } else {
451        tracing::debug!(target: "pagesift::http", %method, %route, status, latency_ms, "request");
452    }
453    resp
454}
455
456pub fn body<T>(r: Result<Json<T>, JsonRejection>) -> Result<T, ApiError> {
457    r.map(|Json(v)| v)
458        .map_err(|e| ApiError::invalid(format!("invalid request body: {}", e.body_text())))
459}
460
461pub fn base_url(s: &AppState, headers: &HeaderMap) -> String {
462    if let Some(u) = &s.public_url {
463        return u.clone();
464    }
465    let host = headers
466        .get(header::HOST)
467        .and_then(|v| v.to_str().ok())
468        .unwrap_or("localhost:3002");
469    format!("http://{host}")
470}
471
472pub fn webhook(
473    s: &AppState,
474    v: Option<&serde_json::Value>,
475) -> Result<Option<WebhookConfig>, ApiError> {
476    v.filter(|v| !v.is_null())
477        .map(|v| WebhookConfig::parse(v, s.webhooks.ssrf()))
478        .transpose()
479}
480
481async fn render_metrics(State(s): State<AppState>) -> impl IntoResponse {
482    let out = s.metrics.render(&s.jobs.counts().await);
483    (
484        StatusCode::OK,
485        [(header::CONTENT_TYPE, "text/plain; version=0.0.4")],
486        out,
487    )
488}
pagesift_api/server.rs

pagesift_api/
server.rs
