b0c103b44f
Implements the --audit-log FILE flag on serve, mcp --bind, and inspect subcommands. Emits per-request NDJSON audit lines with ts, client_ip, tool, fingerprint, duration_ms, status, and diagnostics fields. The AuditLogWriter wraps a BufWriter<File> behind a Mutex and flushes after each line for crash safety. Core changes: - Added pdftract-core/src/audit.rs with AuditRecord schema and AuditLogWriter - Added chrono dependency to pdftract-core/Cargo.toml for timestamp generation - Added crates/pdftract-cli/src/middleware/audit.rs with axum middleware - Integrated AuditState into ServeState, McpServerState, and InspectorState - Added --audit-log flag to Serve, Mcp, and InspectArgs CLI structures - Stdio MCP mode: audit goes to stderr (not stdout, which is JSON-RPC) Acceptance criteria: - pdftract serve --audit-log /var/log/pdftract.ndjson → per-request NDJSON lines appear - Each line is single-line valid JSON (no embedded newlines in values) - client_ip captured from X-Real-IP or X-Forwarded-For header - Stdio MCP audit goes to stderr (with --audit-log /dev/stderr or implicitly) - Concurrent requests: writes don't interleave (Mutex ensures atomic line writes) - Crash mid-request: log line either fully present or fully absent (BufWriter flushes after each write) Closes: pdftract-5boxq
761 lines
26 KiB
Rust
761 lines
26 KiB
Rust
//! HTTP serve mode for pdftract.
|
|
//!
|
|
//! This module implements Phase 6.4's `pdftract serve` subcommand: a long-running
|
|
//! HTTP service for multi-tenant extraction with cache integration.
|
|
//!
|
|
//! # Endpoints
|
|
//!
|
|
//! - `POST /extract` — Extract and return JSON with cache status in response body
|
|
//! - `POST /extract/text` — Extract and return plain text with X-Pdftract-Cache header
|
|
//! - `POST /extract/stream` — Extract and return streaming NDJSON with X-Pdftract-Cache header
|
|
//! - `GET /health` — Health check (always returns 200 OK)
|
|
//!
|
|
//! # Cache headers
|
|
//!
|
|
//! All endpoints return `X-Pdftract-Cache: hit | miss | skipped` header:
|
|
//! - `hit`: Served from cache
|
|
//! - `miss`: Ran extraction; populated cache
|
|
//! - `skipped`: Cache not configured or --no-cache equivalent
|
|
//!
|
|
//! # Concurrency model
|
|
//!
|
|
//! The serve mode uses a two-level concurrency architecture:
|
|
//!
|
|
//! - **tokio**: Per-request concurrency via the async executor. Each HTTP request
|
|
//! is handled asynchronously on tokio's multi-threaded runtime.
|
|
//! - **rayon**: Per-document parallelism within each extraction. PDF pages are
|
|
//! processed in parallel using rayon's work-stealing thread pool.
|
|
//!
|
|
//! The bridge between async (tokio) and sync (rayon) is `tokio::task::spawn_blocking`.
|
|
//! Each POST handler wraps the synchronous extraction call in `spawn_blocking`, which
|
|
//! runs the work on tokio's blocking thread pool (separate from the async reactor).
|
|
//!
|
|
//! This design ensures:
|
|
//! - The async reactor is never blocked by extraction work
|
|
//! - Multiple PDFs can be extracted concurrently (one per request)
|
|
//! - Within each PDF, pages are processed in parallel (rayon)
|
|
//! - Thread pools are sized appropriately (tokio: 512 blocking threads; rayon: num_cpus)
|
|
//!
|
|
//! # Error codes
|
|
//!
|
|
//! - `REQUEST_TOO_LARGE`: Request body exceeds --max-upload-mb limit
|
|
//! - `BAD_REQUEST`: Invalid request parameters or missing file
|
|
//! - `EXTRACTION_ERROR`: PDF parsing or extraction failure
|
|
//! - `INTERNAL_PANIC`: spawn_blocking task panicked (indicates a bug)
|
|
|
|
use anyhow::{Context, Result};
|
|
use axum::{
|
|
body::Body,
|
|
extract::{DefaultBodyLimit, Multipart, State},
|
|
http::{HeaderMap, HeaderValue, StatusCode},
|
|
response::{IntoResponse, Json, Response as AxumResponse},
|
|
routing::{get, post},
|
|
Router,
|
|
};
|
|
use bytes;
|
|
use pdftract_core::audit::AuditLogWriter;
|
|
use pdftract_core::cache;
|
|
use pdftract_core::extract::{extract_pdf, extract_pdf_ndjson, result_to_json};
|
|
use pdftract_core::options::{ExtractionOptions, ReceiptsMode};
|
|
use crate::middleware::{AuditState, audit_middleware};
|
|
use serde::Deserialize;
|
|
use std::path::{Path, PathBuf};
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
use tower_http::limit::RequestBodyLimitLayer;
|
|
|
|
/// Cache state for the HTTP server.
|
|
#[derive(Clone)]
|
|
pub struct CacheState {
|
|
/// Cache directory path
|
|
pub cache_dir: Option<PathBuf>,
|
|
/// Cache size limit in bytes
|
|
pub cache_size_bytes: u64,
|
|
/// Whether cache is disabled
|
|
pub cache_disabled: bool,
|
|
}
|
|
|
|
/// Server state for the HTTP serve mode.
|
|
#[derive(Clone)]
|
|
pub struct ServeState {
|
|
/// Cache configuration
|
|
pub cache: Arc<Mutex<CacheState>>,
|
|
/// Audit log state
|
|
pub audit: AuditState,
|
|
}
|
|
|
|
impl ServeState {
|
|
/// Create a new serve state.
|
|
pub fn new(
|
|
cache_dir: Option<PathBuf>,
|
|
cache_size_bytes: u64,
|
|
cache_disabled: bool,
|
|
audit_writer: Option<AuditLogWriter>,
|
|
) -> Self {
|
|
let cache = CacheState {
|
|
cache_dir,
|
|
cache_size_bytes,
|
|
cache_disabled,
|
|
};
|
|
Self {
|
|
cache: Arc::new(Mutex::new(cache)),
|
|
audit: AuditState::new(audit_writer),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Cache status for response headers and metadata.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum CacheStatus {
|
|
Hit,
|
|
Miss,
|
|
Skipped,
|
|
}
|
|
|
|
impl CacheStatus {
|
|
/// Convert to string for header/metadata.
|
|
pub fn as_str(self) -> &'static str {
|
|
match self {
|
|
CacheStatus::Hit => "hit",
|
|
CacheStatus::Miss => "miss",
|
|
CacheStatus::Skipped => "skipped",
|
|
}
|
|
}
|
|
|
|
/// Create header value.
|
|
pub fn header_value(self) -> HeaderValue {
|
|
HeaderValue::from_static(self.as_str())
|
|
}
|
|
|
|
/// Create from string.
|
|
pub fn from_string(s: &str) -> Self {
|
|
match s {
|
|
"hit" => CacheStatus::Hit,
|
|
"miss" => CacheStatus::Miss,
|
|
"skipped" => CacheStatus::Skipped,
|
|
_ => CacheStatus::Skipped,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Extraction request parameters.
|
|
#[derive(Debug, Deserialize)]
|
|
struct ExtractParams {
|
|
/// Receipts mode (off, lite, svg)
|
|
#[serde(default)]
|
|
receipts: String,
|
|
/// Disable cache for this request
|
|
#[serde(default)]
|
|
no_cache: bool,
|
|
/// Enable full-render path using PDFium
|
|
#[serde(default)]
|
|
full_render: bool,
|
|
}
|
|
|
|
/// Run the HTTP serve mode.
|
|
///
|
|
/// # Arguments
|
|
///
|
|
/// * `bind_addr` — Address to bind (e.g., "127.0.0.1:8080")
|
|
/// * `cache_dir` — Optional cache directory
|
|
/// * `cache_size_bytes` — Cache size limit in bytes
|
|
/// * `cache_disabled` — Whether cache is globally disabled
|
|
/// * `max_upload_mb` — Maximum request body size in MB
|
|
/// * `audit_log` — Optional audit log file path
|
|
pub async fn run(
|
|
bind_addr: String,
|
|
cache_dir: Option<PathBuf>,
|
|
cache_size_bytes: u64,
|
|
cache_disabled: bool,
|
|
max_upload_mb: usize,
|
|
audit_log: Option<PathBuf>,
|
|
) -> Result<()> {
|
|
let cache_dir_for_logging = cache_dir.as_deref();
|
|
|
|
// Create audit log writer if specified
|
|
let audit_writer = if let Some(ref path) = audit_log {
|
|
Some(AuditLogWriter::open(path).context(format!(
|
|
"Failed to open audit log: {}",
|
|
path.display()
|
|
))?)
|
|
} else {
|
|
None
|
|
};
|
|
|
|
let state = ServeState::new(cache_dir.clone(), cache_size_bytes, cache_disabled, audit_writer);
|
|
|
|
let max_body_bytes = max_upload_mb * 1024 * 1024;
|
|
|
|
let app = Router::new()
|
|
.route("/", get(root_handler))
|
|
.route("/extract", post(extract_handler))
|
|
.route("/extract/text", post(extract_text_handler))
|
|
.route("/extract/stream", post(extract_stream_handler))
|
|
.route("/health", get(health_handler))
|
|
.layer(axum::middleware::from_fn_with_state(
|
|
state.audit.clone(),
|
|
audit_middleware,
|
|
))
|
|
.layer(DefaultBodyLimit::max(max_body_bytes))
|
|
.layer(RequestBodyLimitLayer::new(max_body_bytes))
|
|
.with_state(state);
|
|
|
|
let listener = tokio::net::TcpListener::bind(&bind_addr)
|
|
.await
|
|
.context(format!("Failed to bind to {}", bind_addr))?;
|
|
|
|
eprintln!("pdftract serve listening on http://{}", bind_addr);
|
|
if let Some(dir) = cache_dir_for_logging {
|
|
eprintln!(
|
|
"Cache enabled: {} (max {} bytes)",
|
|
dir.display(),
|
|
cache_size_bytes
|
|
);
|
|
} else {
|
|
eprintln!("Cache disabled");
|
|
}
|
|
if let Some(ref path) = audit_log {
|
|
eprintln!("Audit log: {}", path.display());
|
|
}
|
|
|
|
axum::serve(listener, app)
|
|
.await
|
|
.context("HTTP server error")?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Root handler - returns server info.
|
|
async fn root_handler() -> impl IntoResponse {
|
|
Json(serde_json::json!({
|
|
"service": "pdftract",
|
|
"version": env!("CARGO_PKG_VERSION"),
|
|
"endpoints": [
|
|
"POST /extract - Extract PDF and return JSON",
|
|
"POST /extract/text - Extract PDF and return plain text",
|
|
"POST /extract/stream - Extract PDF and return streaming NDJSON",
|
|
"GET /health - Health check"
|
|
]
|
|
}))
|
|
}
|
|
|
|
/// Health check handler.
|
|
async fn health_handler() -> impl IntoResponse {
|
|
Json(serde_json::json!({
|
|
"status": "ok",
|
|
"version": env!("CARGO_PKG_VERSION")
|
|
}))
|
|
}
|
|
|
|
/// Extract handler - returns JSON with cache status in metadata.
|
|
async fn extract_handler(
|
|
State(state): State<ServeState>,
|
|
mut multipart: Multipart,
|
|
) -> Result<impl IntoResponse, AxumError> {
|
|
let (pdf_file, params) = receive_pdf(&mut multipart).await?;
|
|
let options = build_options(¶ms)?;
|
|
|
|
// Get cache configuration
|
|
let cache_state = state.cache.lock().await;
|
|
let cache_dir = cache_state.cache_dir.clone();
|
|
let cache_size_bytes = cache_state.cache_size_bytes;
|
|
let cache_disabled = params.no_cache || cache_state.cache_disabled || cache_dir.is_none();
|
|
drop(cache_state);
|
|
|
|
// Perform extraction with cache integration
|
|
let pdf_file_clone = pdf_file.clone();
|
|
let (result, cache_status, cache_age) = tokio::task::spawn_blocking(move || {
|
|
let cache_dir_ref = cache_dir.as_deref();
|
|
cache::extract_with_cache(
|
|
&pdf_file_clone,
|
|
&options,
|
|
cache_dir_ref,
|
|
cache_disabled,
|
|
Some(cache_size_bytes),
|
|
)
|
|
.map_err(|e| AxumError::Extraction(format!("{:?}", e)))
|
|
})
|
|
.await
|
|
.map_err(|e| {
|
|
// Distinguish between cancellation (task dropped) and panic
|
|
if e.is_cancelled() {
|
|
AxumError::Internal(format!("Task cancelled: {}", e))
|
|
} else {
|
|
// is_panic() true means the task panicked - indicates a bug
|
|
AxumError::InternalPanic(format!("Extraction task panicked: {}", e))
|
|
}
|
|
})?
|
|
.map_err(|e| e)?;
|
|
|
|
// Build JSON response with cache status
|
|
let mut result = result;
|
|
result.metadata.cache_status = Some(cache_status.clone());
|
|
result.metadata.cache_age_seconds = cache_age;
|
|
|
|
let json = result_to_json(&result);
|
|
|
|
let response = AxumResponse::builder()
|
|
.status(StatusCode::OK)
|
|
.header("Content-Type", "application/json")
|
|
.header(
|
|
"X-Pdftract-Cache",
|
|
CacheStatus::from_string(&cache_status).header_value(),
|
|
)
|
|
.body(Body::from(serde_json::to_string(&json).unwrap()))
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?;
|
|
|
|
Ok(response)
|
|
}
|
|
|
|
/// Extract text handler - returns plain text with X-Pdftract-Cache header.
|
|
async fn extract_text_handler(
|
|
State(state): State<ServeState>,
|
|
mut multipart: Multipart,
|
|
) -> Result<impl IntoResponse, AxumError> {
|
|
let (pdf_file, params) = receive_pdf(&mut multipart).await?;
|
|
let options = build_options(¶ms)?;
|
|
|
|
// Get cache configuration
|
|
let cache_state = state.cache.lock().await;
|
|
let cache_dir = cache_state.cache_dir.clone();
|
|
let cache_size_bytes = cache_state.cache_size_bytes;
|
|
let cache_disabled = params.no_cache || cache_state.cache_disabled || cache_dir.is_none();
|
|
drop(cache_state);
|
|
|
|
let (result, cache_status, _cache_age) = tokio::task::spawn_blocking(move || {
|
|
let cache_dir_ref = cache_dir.as_deref();
|
|
cache::extract_with_cache(
|
|
&pdf_file,
|
|
&options,
|
|
cache_dir_ref,
|
|
cache_disabled,
|
|
Some(cache_size_bytes),
|
|
)
|
|
.map_err(|e| AxumError::Extraction(format!("{:?}", e)))
|
|
})
|
|
.await
|
|
.map_err(|e| {
|
|
// Distinguish between cancellation (task dropped) and panic
|
|
if e.is_cancelled() {
|
|
AxumError::Internal(format!("Task cancelled: {}", e))
|
|
} else {
|
|
// is_panic() true means the task panicked - indicates a bug
|
|
AxumError::InternalPanic(format!("Extraction task panicked: {}", e))
|
|
}
|
|
})?
|
|
.map_err(|e| e)?;
|
|
|
|
let mut text = String::new();
|
|
for page in &result.pages {
|
|
for span in &page.spans {
|
|
text.push_str(&span.text);
|
|
text.push('\n');
|
|
}
|
|
}
|
|
|
|
let response = AxumResponse::builder()
|
|
.status(StatusCode::OK)
|
|
.header(
|
|
"X-Pdftract-Cache",
|
|
CacheStatus::from_string(&cache_status).header_value(),
|
|
)
|
|
.body(Body::from(text))
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?;
|
|
|
|
Ok(response)
|
|
}
|
|
|
|
/// Extract stream handler - returns true async streaming NDJSON.
|
|
///
|
|
/// This handler spawns a background task that extracts pages sequentially
|
|
/// and sends them over a channel. The response body is a stream that yields
|
|
/// each page as NDJSON immediately after it's extracted.
|
|
///
|
|
/// Cache status is always "skipped" for streaming since we bypass the cache
|
|
/// to provide true incremental output.
|
|
async fn extract_stream_handler(
|
|
State(state): State<ServeState>,
|
|
mut multipart: Multipart,
|
|
) -> Result<impl IntoResponse, AxumError> {
|
|
use tokio_stream::wrappers::ReceiverStream;
|
|
use tokio_stream::StreamExt;
|
|
|
|
let (pdf_file, params) = receive_pdf(&mut multipart).await?;
|
|
let options = build_options(¶ms)?;
|
|
|
|
// Get cache configuration (for logging only - streaming bypasses cache)
|
|
let cache_state = state.cache.lock().await;
|
|
let _cache_dir = cache_state.cache_dir.clone();
|
|
drop(cache_state);
|
|
|
|
// Create a channel for streaming pages
|
|
let (tx, rx) = tokio::sync::mpsc::channel::<Vec<u8>>(16);
|
|
|
|
// Spawn extraction task in background
|
|
tokio::task::spawn_blocking(move || {
|
|
use pdftract_core::extract::extract_pdf_ndjson;
|
|
|
|
// Clone sender for error handling
|
|
let tx_for_error = tx.clone();
|
|
|
|
// Write to a custom writer that sends to the channel
|
|
struct ChannelWriter {
|
|
tx: tokio::sync::mpsc::Sender<Vec<u8>>,
|
|
};
|
|
|
|
impl std::io::Write for ChannelWriter {
|
|
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
|
|
// Clone the buffer since we need to send it
|
|
self.tx
|
|
.blocking_send(buf.to_vec())
|
|
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))?;
|
|
Ok(buf.len())
|
|
}
|
|
|
|
fn flush(&mut self) -> std::io::Result<()> {
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
let writer = ChannelWriter { tx };
|
|
|
|
// Extract to NDJSON, streaming each page as it's extracted
|
|
if let Err(e) = extract_pdf_ndjson(&pdf_file, &options, writer) {
|
|
// Send error as a JSON line
|
|
let error_json = serde_json::json!({
|
|
"error": format!("{:?}", e)
|
|
});
|
|
if let Ok(json_bytes) = serde_json::to_vec(&error_json) {
|
|
let _ = tx_for_error.blocking_send(json_bytes);
|
|
let _ = tx_for_error.blocking_send(b"\n".to_vec());
|
|
}
|
|
}
|
|
|
|
Ok::<(), AxumError>(())
|
|
});
|
|
|
|
// Create a stream from the receiver
|
|
let stream = ReceiverStream::new(rx).map(|item| Ok::<_, axum::Error>(bytes::Bytes::from(item)));
|
|
|
|
// Return a streaming body
|
|
let body = Body::from_stream(stream);
|
|
|
|
let response = AxumResponse::builder()
|
|
.status(StatusCode::OK)
|
|
.header("X-Pdftract-Cache", CacheStatus::Skipped.header_value())
|
|
.header("Content-Type", "application/x-ndjson")
|
|
.body(body)
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?;
|
|
|
|
Ok(response)
|
|
}
|
|
|
|
/// Receive uploaded PDF file and extraction parameters.
|
|
async fn receive_pdf(multipart: &mut Multipart) -> Result<(PathBuf, ExtractParams), AxumError> {
|
|
let mut pdf_path = None;
|
|
let mut params = ExtractParams {
|
|
receipts: "off".to_string(),
|
|
no_cache: false,
|
|
full_render: false,
|
|
};
|
|
|
|
while let Some(field) = multipart
|
|
.next_field()
|
|
.await
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?
|
|
{
|
|
let name = field.name().unwrap_or("").to_string();
|
|
|
|
if name == "file" || name == "pdf" {
|
|
let data = field
|
|
.bytes()
|
|
.await
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?;
|
|
|
|
// Create a temp file that will persist for the duration of the request
|
|
let temp_dir = std::env::temp_dir();
|
|
let temp_file = temp_dir.join(format!("pdftract-upload-{}.pdf", uuid::Uuid::new_v4()));
|
|
tokio::fs::write(&temp_file, &data)
|
|
.await
|
|
.map_err(|e| AxumError::Internal(format!("{:?}", e)))?;
|
|
pdf_path = Some(temp_file);
|
|
} else if name == "receipts" {
|
|
if let Ok(value) = field.text().await {
|
|
params.receipts = value;
|
|
}
|
|
} else if name == "no_cache" {
|
|
params.no_cache = true;
|
|
} else if name == "full_render" {
|
|
// Check if full_render is requested
|
|
if let Ok(value) = field.text().await {
|
|
params.full_render = value == "true" || value == "1";
|
|
}
|
|
// Checkbox without value also means true
|
|
if params.full_render == false {
|
|
params.full_render = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
let pdf_path =
|
|
pdf_path.ok_or_else(|| AxumError::BadRequest("No PDF file uploaded".to_string()))?;
|
|
|
|
Ok((pdf_path, params))
|
|
}
|
|
|
|
/// Build extraction options from parameters.
|
|
///
|
|
/// Validates that full_render is only used when the feature is available.
|
|
/// If full_render is requested but the feature is not compiled in,
|
|
/// the request still succeeds but falls back to direct compositing.
|
|
fn build_options(params: &ExtractParams) -> Result<ExtractionOptions, AxumError> {
|
|
let receipts_mode = match params.receipts.as_str() {
|
|
"lite" => ReceiptsMode::Lite,
|
|
"svg" => ReceiptsMode::SvgClip,
|
|
_ => ReceiptsMode::Off,
|
|
};
|
|
|
|
// Check if full_render is requested
|
|
if params.full_render {
|
|
// Validate that full_render is available at runtime
|
|
#[cfg(all(feature = "ocr", feature = "full-render"))]
|
|
{
|
|
use pdftract_core::render::pdfium_path::has_full_render;
|
|
if !has_full_render() {
|
|
return Err(AxumError::BadRequest(
|
|
"full_render requested but PDFium is not available at runtime. \
|
|
Ensure the PDFium native library is installed."
|
|
.to_string(),
|
|
));
|
|
}
|
|
}
|
|
|
|
#[cfg(not(all(feature = "ocr", feature = "full-render")))]
|
|
{
|
|
// Feature not compiled in - fall back to direct compositing
|
|
// Log a debug message but don't fail the request
|
|
tracing::debug!(
|
|
"full_render requested but full-render feature not compiled; using direct compositing path"
|
|
);
|
|
}
|
|
}
|
|
|
|
Ok(ExtractionOptions {
|
|
receipts: receipts_mode,
|
|
full_render: params.full_render,
|
|
..Default::default()
|
|
})
|
|
}
|
|
|
|
/// Error types for the HTTP server.
|
|
#[derive(Debug)]
|
|
pub enum AxumError {
|
|
/// Bad request (400) - invalid parameters or missing file
|
|
BadRequest(String),
|
|
/// Extraction error (422) - PDF parsing or extraction failure
|
|
Extraction(String),
|
|
/// Internal error (500) - server-side failure
|
|
Internal(String),
|
|
/// Internal panic (500) - spawn_blocking task panicked (indicates a bug)
|
|
InternalPanic(String),
|
|
}
|
|
|
|
impl IntoResponse for AxumError {
|
|
fn into_response(self) -> AxumResponse {
|
|
let (status, error_code, message) = match self {
|
|
AxumError::BadRequest(msg) => (StatusCode::BAD_REQUEST, "BAD_REQUEST", msg),
|
|
AxumError::Extraction(msg) => {
|
|
(StatusCode::UNPROCESSABLE_ENTITY, "EXTRACTION_ERROR", msg)
|
|
}
|
|
AxumError::Internal(msg) => (StatusCode::INTERNAL_SERVER_ERROR, "INTERNAL_ERROR", msg),
|
|
AxumError::InternalPanic(msg) => {
|
|
(StatusCode::INTERNAL_SERVER_ERROR, "INTERNAL_PANIC", msg)
|
|
}
|
|
};
|
|
|
|
let body = serde_json::json!({
|
|
"error": error_code,
|
|
"message": message,
|
|
});
|
|
|
|
(status, Json(body)).into_response()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::time::Duration;
|
|
|
|
/// Test that the AxumError enum converts to correct status codes and error codes.
|
|
#[test]
|
|
fn test_error_into_response() {
|
|
// Test BadRequest
|
|
let err = AxumError::BadRequest("test".to_string());
|
|
let resp = err.into_response();
|
|
assert_eq!(resp.status(), StatusCode::BAD_REQUEST);
|
|
|
|
// Test Extraction
|
|
let err = AxumError::Extraction("test".to_string());
|
|
let resp = err.into_response();
|
|
assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY);
|
|
|
|
// Test Internal
|
|
let err = AxumError::Internal("test".to_string());
|
|
let resp = err.into_response();
|
|
assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
|
|
|
// Test InternalPanic
|
|
let err = AxumError::InternalPanic("test".to_string());
|
|
let resp = err.into_response();
|
|
assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
|
}
|
|
|
|
/// Test that CacheStatus converts correctly to/from strings.
|
|
#[test]
|
|
fn test_cache_status_conversions() {
|
|
assert_eq!(CacheStatus::Hit.as_str(), "hit");
|
|
assert_eq!(CacheStatus::Miss.as_str(), "miss");
|
|
assert_eq!(CacheStatus::Skipped.as_str(), "skipped");
|
|
|
|
assert_eq!(CacheStatus::from_string("hit"), CacheStatus::Hit);
|
|
assert_eq!(CacheStatus::from_string("miss"), CacheStatus::Miss);
|
|
assert_eq!(CacheStatus::from_string("skipped"), CacheStatus::Skipped);
|
|
assert_eq!(CacheStatus::from_string("invalid"), CacheStatus::Skipped);
|
|
}
|
|
|
|
/// Helper to load a valid test PDF.
|
|
fn load_test_pdf() -> Vec<u8> {
|
|
// Use the existing test fixture from pdftract-libpdftract
|
|
let pdf_path = concat!(
|
|
env!("CARGO_MANIFEST_DIR"),
|
|
"/../pdftract-libpdftract/tests/hello.pdf"
|
|
);
|
|
std::fs::read(pdf_path).expect("Failed to read test PDF")
|
|
}
|
|
|
|
/// Integration test: 8 concurrent requests complete in parallel.
|
|
///
|
|
/// This is the critical test from the plan (line 2146). It verifies that:
|
|
/// - All 8 requests complete (proves no deadlock or serialization)
|
|
/// - Wallclock time is similar to a single request (proves parallelism)
|
|
/// - /health responds quickly during concurrent extractions (proves /health doesn't block)
|
|
#[tokio::test]
|
|
async fn test_concurrent_requests_parallel() {
|
|
use axum::{
|
|
body::Body,
|
|
http::{HeaderMap, HeaderValue, Method, StatusCode},
|
|
};
|
|
use reqwest::multipart::{Form, Part};
|
|
use tokio::time::Instant;
|
|
|
|
// Start the server in the background
|
|
let state = ServeState::new(None, 1024 * 1024 * 1024, true); // No cache
|
|
let app = Router::new()
|
|
.route("/extract", post(extract_handler))
|
|
.route("/health", get(health_handler))
|
|
.with_state(state);
|
|
|
|
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("Failed to bind");
|
|
let addr = listener.local_addr().expect("Failed to get local address");
|
|
let port = addr.port();
|
|
|
|
tokio::spawn(async move {
|
|
axum::serve(listener, app).await.expect("Server error");
|
|
});
|
|
|
|
// Give the server a moment to start
|
|
tokio::time::sleep(Duration::from_millis(100)).await;
|
|
|
|
let base_url = format!("http://127.0.0.1:{}", port);
|
|
let client = reqwest::Client::new();
|
|
let pdf_bytes = load_test_pdf();
|
|
|
|
// First, test that /health responds quickly
|
|
let health_start = Instant::now();
|
|
let health_resp = client
|
|
.get(format!("{}/health", base_url))
|
|
.send()
|
|
.await
|
|
.expect("Health request failed");
|
|
let health_duration = health_start.elapsed();
|
|
|
|
assert_eq!(health_resp.status(), StatusCode::OK);
|
|
assert!(
|
|
health_duration < Duration::from_millis(100),
|
|
"/health should respond in < 100ms, took {:?}",
|
|
health_duration
|
|
);
|
|
|
|
// Now launch 8 concurrent extraction requests
|
|
let mut handles = Vec::new();
|
|
let start = Instant::now();
|
|
|
|
for i in 0..8 {
|
|
let client = client.clone();
|
|
let url = format!("{}/extract", base_url);
|
|
let pdf = pdf_bytes.clone();
|
|
|
|
let handle = tokio::spawn(async move {
|
|
let part = Part::bytes(pdf).file_name(format!("test{}.pdf", i));
|
|
let form = Form::new().part("file", part);
|
|
|
|
let resp = client
|
|
.post(&url)
|
|
.multipart(form)
|
|
.send()
|
|
.await
|
|
.expect("Extraction request failed");
|
|
|
|
(i, resp.status(), client)
|
|
});
|
|
|
|
handles.push(handle);
|
|
}
|
|
|
|
// Wait for all requests to complete
|
|
let mut results = Vec::new();
|
|
for handle in handles {
|
|
let (i, status, _) = handle.await.expect("Task panicked");
|
|
results.push((i, status));
|
|
}
|
|
|
|
let total_duration = start.elapsed();
|
|
|
|
// The critical test: all 8 requests completed (proves no deadlock or serialization)
|
|
// We don't assert OK status because the test PDF might not extract correctly;
|
|
// the important thing is that all requests got a response.
|
|
assert_eq!(results.len(), 8, "All 8 requests should have completed");
|
|
|
|
// The critical assertion: if requests were serialized, total time would be
|
|
// roughly 8x a single request. With parallelism, it should be much less.
|
|
// We use a very loose threshold to account for system load and variability.
|
|
let single_request_estimate = Duration::from_millis(100); // Rough estimate
|
|
let serialized_estimate = single_request_estimate * 8;
|
|
|
|
assert!(
|
|
total_duration < serialized_estimate,
|
|
"Requests appear serialized: completed in {:?}, expected < {:?}",
|
|
total_duration,
|
|
serialized_estimate
|
|
);
|
|
|
|
// Also verify /health still responds quickly during load
|
|
let health_start = Instant::now();
|
|
let health_resp = client
|
|
.get(format!("{}/health", base_url))
|
|
.send()
|
|
.await
|
|
.expect("Health request failed");
|
|
let health_duration = health_start.elapsed();
|
|
|
|
assert_eq!(health_resp.status(), StatusCode::OK);
|
|
assert!(
|
|
health_duration < Duration::from_millis(100),
|
|
"/health should respond in < 100ms during load, took {:?}",
|
|
health_duration
|
|
);
|
|
}
|
|
}
|