use std::collections::HashMap;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{SyncSender, sync_channel};
use std::sync::{Arc, Condvar, Mutex, MutexGuard};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};
use rapid_structure_rs::{
ALLOWED_DOWNLOAD_HOSTS, DetectionPolicy, DownloadError, FormulaPolicy, ImageInput,
ModelFileState, ModelRole, ModelSet, OcrOutput, OcrRequest, OrtRuntimeFingerprint,
OutputPolicy, PreprocessPolicy, ProviderPreference, RapidOcrError, RecognitionPolicy,
StagePlan, TablePolicy, TableTextPolicy, WordOutputMode, available_disk_bytes,
clear_verification_cache, ort_runtime_fingerprint, ort_runtime_version, peak_memory_source,
peak_working_set_bytes, verification_stats,
};
use serde_json::{Value, json};
use super::admit::QueueReservation;
use super::download::{
self, DOWNLOAD_QUEUE_CAPACITY, DownloadCommand, DownloadJob, DownloaderFactory,
};
use super::engine::{BackendProvider, EngineFactory, OcrBackend};
use super::error::{ErrorBody, ServeError};
use super::evaluate::EvalRoot;
use super::export::{self, ExportError, ExportFormat, ExportRequest};
use super::flowlog::{self, FlowSink, JobIdentity, RequestFlow};
use super::jobs::{
CancelOutcome, DownloadProgress, JobFailure, JobIdGenerator, JobKind, JobQueue,
JobState as JobLifecycle, JobStore, JobStoreLimits, Millis,
};
use super::limits::ServeLimits;
use super::model_plan::{
BlockingFile, ModelPlan, ModelReport, ModelSnapshot, PendingDownload, Pipeline, PlanFileStatus,
PlanReverification, VerifiedPlanFile, blocked_names, corrupt_names, missing_names,
source_label,
};
use super::queue::{DualQueueScheduler, QueueClass, ScheduledJob, SchedulerConfig};
use super::results::{Outcome, ResultStore, SerializeError, Succeeded, serialize_bounded};
use super::security::{LocalOrigin, ServeToken};
use super::state::{
EngineState, EngineStateMachine, OcrAdmission, ProviderStatus, ServeConfigPlan, ServiceState,
};
const WORKER_POLL: Duration = Duration::from_millis(200);
const SWEEP_INTERVAL: Duration = Duration::from_millis(500);
pub(super) const READ_CHUNK_BYTES: usize = 64 * 1024;
pub(super) const REDACTED_MODEL_DIR: &str = "<redacted>";
pub(super) fn monotonic_ms() -> Millis {
static START: std::sync::OnceLock<Instant> = std::sync::OnceLock::new();
let start = START.get_or_init(Instant::now);
start.elapsed().as_millis() as Millis
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(super) struct OcrRouting {
pub formula: bool,
pub table: bool,
pub disabled_reason: Option<String>,
pub table_disabled_reason: Option<String>,
}
impl OcrRouting {
pub const fn formula_enabled() -> Self {
Self {
formula: true,
table: false,
disabled_reason: None,
table_disabled_reason: None,
}
}
pub fn text_only(reason: impl Into<String>) -> Self {
Self {
formula: false,
table: false,
disabled_reason: Some(reason.into()),
table_disabled_reason: None,
}
}
pub fn with_table(mut self, enabled: bool) -> Self {
self.table = enabled;
self
}
pub fn class_for(&self, requested: Option<&str>) -> Result<QueueClass, ServeError> {
match requested {
None | Some("text") => Ok(QueueClass::Text),
Some("formula") if self.formula => Ok(QueueClass::Formula),
Some(_) => Err(ServeError::BadRequest),
}
}
pub fn table_requested(&self, requested: Option<&str>) -> Result<bool, ServeError> {
match requested {
None | Some("0") => Ok(false),
Some("1") if self.table => Ok(true),
Some("1") => Err(ServeError::BadRequest),
Some(_) => Err(ServeError::BadRequest),
}
}
}
pub(super) fn routing_for(detector: Option<&Path>, table_model: Option<&Path>) -> OcrRouting {
let mut routing = match detector {
Some(_) => OcrRouting::formula_enabled(),
None => OcrRouting::text_only(
"formula routing is not enabled on this server: no page formula detector is \
configured, and without one the formula pipeline can only handle caller-declared \
regions (there are none over HTTP). Pass --formula-detector <ONNX> (for example \
pix2text-mfd-1.5.onnx) to enable the formula queue; ordinary OCR is unaffected \
(docs/05 §4.2, §10.8)",
),
};
if table_model.is_none() {
routing.table_disabled_reason = Some(
"table routing is not enabled on this server: no table structure model is configured. \
Pass --table-model <ONNX> (slanet-plus.onnx, declared by models/table/manifest.json) \
to enable `?table=1`; ordinary OCR and the formula queue are unaffected \
(docs/11 §4.2)"
.to_string(),
);
}
routing.with_table(table_model.is_some())
}
impl OcrRouting {
pub fn from_plan(plan: &ModelPlan) -> Self {
routing_for(
plan.formula_detector().map(|spec| spec.path.as_path()),
plan.table_model().map(|spec| spec.path.as_path()),
)
}
}
pub(super) type FreeSpaceFactory = Arc<dyn Fn(&Path) -> Result<u64, RapidOcrError> + Send + Sync>;
pub(super) fn real_free_space() -> FreeSpaceFactory {
Arc::new(|directory: &Path| available_disk_bytes(directory))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(super) enum EngineLoad {
Ready,
BlockedModelsMissing,
Failed,
}
pub(super) struct ServeContext {
pub limits: ServeLimits,
pub plan: ServeConfigPlan,
pub model_plan: ModelPlan,
pub snapshot: ModelSnapshot,
pub token: ServeToken,
pub local: LocalOrigin,
pub page: String,
pub nonce: String,
pub allow_download: bool,
pub allow_download_hosts: Vec<String>,
pub allow_provider_fallback: bool,
pub eval_root: Option<EvalRoot>,
pub flow: FlowSink,
pub engine_factory: EngineFactory,
pub downloader: DownloaderFactory,
pub free_space: FreeSpaceFactory,
pub post_verify: Option<Box<dyn Fn() + Send + Sync>>,
}
struct RetainedOcr {
bytes: Arc<[u8]>,
params: OcrJobParams,
}
#[derive(Debug, Clone, Copy)]
pub(super) struct OcrJobParams {
pub(super) max_side: Option<u32>,
pub(super) table: bool,
pub(super) words: WordOutputMode,
}
struct JobState {
store: JobStore,
scheduler: DualQueueScheduler,
results: ResultStore,
ids: JobIdGenerator,
originals: HashMap<String, RetainedOcr>,
reserved_text: usize,
reserved_formula: usize,
}
impl JobState {
fn reserve(&mut self, class: QueueClass) -> bool {
let used = self.scheduler.queued_len(class) + self.reserved(class);
if used >= self.scheduler.config().capacity(class) {
return false;
}
*self.reserved_mut(class) += 1;
true
}
fn release(&mut self, class: QueueClass) {
let slot = self.reserved_mut(class);
*slot = slot.saturating_sub(1);
}
fn reserved(&self, class: QueueClass) -> usize {
match class {
QueueClass::Text => self.reserved_text,
QueueClass::Formula => self.reserved_formula,
}
}
fn reserved_mut(&mut self, class: QueueClass) -> &mut usize {
match class {
QueueClass::Text => &mut self.reserved_text,
QueueClass::Formula => &mut self.reserved_formula,
}
}
fn sync_originals(&mut self) {
for id in self.store.take_released_originals() {
self.originals.remove(&id);
}
}
fn sync_positions(&mut self) {
let queued: Vec<(String, QueueClass)> = self
.store
.views(monotonic_ms())
.into_iter()
.filter(|view| view.state == JobLifecycle::Queued)
.filter_map(|view| view.queue.class().map(|class| (view.id, class)))
.collect();
for (id, class) in queued {
let position = self.scheduler.position_of(class, &id);
let _ = self.store.set_position(&id, position);
}
}
fn take_next(&mut self) -> Option<ScheduledJob> {
loop {
let scheduled = self.scheduler.take_next()?;
if self.store.start(&scheduled.id, monotonic_ms()).is_ok() {
self.sync_positions();
return Some(scheduled);
}
}
}
}
pub(super) struct ServeShared {
service: ServiceState,
limits: ServeLimits,
plan: Mutex<ServeConfigPlan>,
model_plan: ModelPlan,
token: ServeToken,
local: LocalOrigin,
page: String,
nonce: String,
allow_download: bool,
allow_download_hosts: Vec<String>,
allow_provider_fallback: bool,
eval_root: Option<EvalRoot>,
flow: FlowSink,
engine_state: Mutex<EngineStateMachine>,
engine: Mutex<Option<Box<dyn OcrBackend>>>,
engine_load: Mutex<()>,
engine_load_ms: Mutex<Option<u64>>,
provider_switch: AtomicBool,
evaluation: AtomicBool,
jobs: Mutex<JobState>,
queue_signal: Condvar,
download_tx: SyncSender<DownloadCommand>,
engine_factory: EngineFactory,
downloader: DownloaderFactory,
free_space: FreeSpaceFactory,
post_verify: Option<Box<dyn Fn() + Send + Sync>>,
shutting_down: AtomicBool,
}
impl ServeShared {
pub fn service_state(&self) -> ServiceState {
self.service
}
pub fn engine_state(&self) -> EngineState {
lock(&self.engine_state).state().clone()
}
pub fn plan_snapshot(&self) -> ServeConfigPlan {
lock(&self.plan).clone()
}
pub fn requested_label(&self) -> String {
lock(&self.plan).requested_label()
}
pub fn provider_status(&self) -> ProviderStatus {
self.engine_state().provider_status(&self.requested_label())
}
pub fn token(&self) -> &ServeToken {
&self.token
}
pub fn local(&self) -> &LocalOrigin {
&self.local
}
pub fn page(&self) -> &str {
&self.page
}
pub fn nonce(&self) -> &str {
&self.nonce
}
pub fn limits(&self) -> ServeLimits {
self.limits
}
pub fn flow(&self) -> &FlowSink {
&self.flow
}
pub fn min_side_len(&self) -> u32 {
self.plan_snapshot().engine.global.min_side_len as u32
}
pub fn models(&self) -> ModelReport {
self.model_plan.report()
}
pub fn routing(&self) -> OcrRouting {
OcrRouting::from_plan(&self.model_plan)
}
fn text_blocking(&self) -> Vec<BlockingFile> {
self.model_plan.plan_blocking(Pipeline::Text)
}
pub fn formula_policy(&self) -> Option<FormulaPolicy> {
let recognizer = self.model_plan.plan_file(ModelRole::FormulaRecognizer)?;
let detector = self.model_plan.plan_file(ModelRole::FormulaDetector);
Some(FormulaPolicy {
enabled: true,
model_path: Some(recognizer.path.clone()),
expected_model_sha256: recognizer.declared_sha256.clone(),
detector_path: detector.map(|file| file.path.clone()),
expected_detector_sha256: detector.and_then(|file| file.declared_sha256.clone()),
..FormulaPolicy::default()
})
}
pub fn table_policy(&self) -> Option<TablePolicy> {
let structure = self.model_plan.plan_file(ModelRole::TableStructure)?;
Some(TablePolicy {
enabled: true,
model_path: Some(structure.path.clone()),
expected_model_sha256: structure.declared_sha256.clone(),
text_policy: TableTextPolicy::Escape,
..TablePolicy::default()
})
}
pub fn formula_models_ready(&self) -> bool {
self.model_plan.plan_blocking(Pipeline::Formula).is_empty()
}
pub fn formula_blocked_body(&self) -> Body {
Body::json(409, render_error_body(&self.formula_missing_failure()))
}
pub fn formula_blocked_reason(&self) -> (&'static str, String) {
let body = self.formula_missing_failure();
(body.code, body.detail.to_string())
}
fn formula_missing_failure(&self) -> ErrorBody {
let (_, corrupt) = self.formula_blocking_lists();
ErrorBody {
code: if corrupt.is_empty() {
"models_missing"
} else {
"models_corrupt"
},
message: "the formula model set is incomplete".to_string(),
detail: self.formula_detail(),
}
}
fn formula_blocking_lists(&self) -> (Vec<String>, Vec<String>) {
let blocking = self.model_plan.plan_blocking(Pipeline::Formula);
(missing_names(&blocking), corrupt_names(&blocking))
}
pub fn formula_detector_status(&self) -> Option<PlanFileStatus> {
self.model_plan
.plan_file(ModelRole::FormulaDetector)
.map(|file| file.status())
}
pub fn formula_detail(&self) -> Value {
let blocking = self.model_plan.plan_blocking(Pipeline::Formula);
let missing = missing_names(&blocking);
let corrupt = corrupt_names(&blocking);
let blocked = blocked_names(&blocking);
json!({
"scope": "formula",
"missing": missing,
"corrupt": corrupt,
"blocked": blocked,
"in_plan": self.model_plan.formula_in_plan(),
"source": source_label(self.model_plan.source()),
"model_dir": REDACTED_MODEL_DIR,
"detector": self.formula_detector_detail(),
})
}
fn formula_detector_detail(&self) -> Value {
let planned = self.model_plan.plan_file(ModelRole::FormulaDetector);
let status = self.formula_detector_status();
json!({
"configured": self.model_plan.formula_in_plan(),
"file": status.as_ref().map(|status| status.name.clone()),
"sha256": planned.and_then(|file| file.declared_sha256.clone()),
"state": status.as_ref().map(|status| status.state.as_str()),
})
}
pub fn table_models_ready(&self) -> bool {
self.model_plan.plan_blocking(Pipeline::Table).is_empty()
}
pub fn table_blocked_body(&self) -> Body {
Body::json(409, render_error_body(&self.table_missing_failure()))
}
pub fn table_blocked_reason(&self) -> (&'static str, String) {
let body = self.table_missing_failure();
(body.code, body.detail.to_string())
}
fn table_missing_failure(&self) -> ErrorBody {
let (_, corrupt) = self.table_blocking_lists();
ErrorBody {
code: if corrupt.is_empty() {
"models_missing"
} else {
"models_corrupt"
},
message: "the table structure model is missing or corrupt".to_string(),
detail: self.table_detail(),
}
}
fn table_blocking_lists(&self) -> (Vec<String>, Vec<String>) {
let blocking = self.model_plan.plan_blocking(Pipeline::Table);
(missing_names(&blocking), corrupt_names(&blocking))
}
pub fn table_model_status(&self) -> Option<PlanFileStatus> {
self.model_plan
.plan_file(ModelRole::TableStructure)
.map(|file| file.status())
}
pub fn table_detail(&self) -> Value {
let blocking = self.model_plan.plan_blocking(Pipeline::Table);
let missing = missing_names(&blocking);
let corrupt = corrupt_names(&blocking);
let blocked = blocked_names(&blocking);
json!({
"scope": "table",
"missing": missing,
"corrupt": corrupt,
"blocked": blocked,
"in_plan": self.model_plan.table_in_plan(),
"source": source_label(self.model_plan.source()),
"model_dir": REDACTED_MODEL_DIR,
"structure": self.table_model_detail(),
})
}
fn table_model_detail(&self) -> Value {
let planned = self.model_plan.plan_file(ModelRole::TableStructure);
let status = self.table_model_status();
json!({
"configured": self.model_plan.table_in_plan(),
"file": status.as_ref().map(|status| status.name.clone()),
"sha256": planned.and_then(|file| file.declared_sha256.clone()),
"state": status.as_ref().map(|status| status.state.as_str()),
})
}
pub fn is_shutting_down(&self) -> bool {
self.shutting_down.load(Ordering::SeqCst)
}
pub fn shutdown(&self) {
self.shutting_down.store(true, Ordering::SeqCst);
self.queue_signal.notify_all();
}
pub fn engine_state_name(&self) -> &'static str {
self.engine_state().name()
}
pub fn engine_load_ms(&self) -> Option<u64> {
*lock(&self.engine_load_ms)
}
pub(super) fn engine_failure_reason(&self) -> Option<String> {
match self.engine_state() {
EngineState::Failed { reason } => Some(reason),
_ => None,
}
}
pub(super) fn models_missing_failure(&self) -> ErrorBody {
let error = self.models_missing_error();
ErrorBody {
code: error.code(),
message: error.message(),
detail: self.models_missing_detail(),
}
}
pub(super) fn models_missing_outcome(&self) -> Outcome {
Outcome::Failed(409, self.models_missing_failure())
}
pub fn status_json(&self) -> Value {
let engine_state = self.engine_state();
let provider = engine_state.provider_status(&self.requested_label());
let (queues, retention) = {
let state = lock(&self.jobs);
let scheduler = state.scheduler.config();
let class_json = |class: QueueClass| {
json!({
"used": state.scheduler.queued_len(class),
"capacity": scheduler.capacity(class),
"consecutive_quota": scheduler.consecutive_quota(class),
"served_in_round": state.scheduler.served_in_round(class),
"wait_bound": scheduler.wait_bound(class),
})
};
(
json!({
"text": class_json(QueueClass::Text),
"formula": class_json(QueueClass::Formula),
"round_len": scheduler.round_len(),
"idle": state.scheduler.is_empty(),
}),
json!({
"jobs": state.store.len(),
"terminal": state.store.terminal_count(),
"retained_bytes": state.store.retained_bytes(),
"tombstones": state.store.tombstone_len(),
"results": state.results.len(),
"result_bytes": state.results.bytes(),
"retained_originals": state.originals.len(),
"job_ttl_ms": state.store.limits().ttl_ms,
}),
)
};
json!({
"state": self.service,
"engine": engine_state,
"provider": provider,
"ort": {
"version": ort_runtime_version(),
"fingerprint": public_fingerprint(&ort_runtime_fingerprint()),
},
"memory": {
"peak_working_set_bytes": peak_working_set_bytes(),
"source": peak_memory_source(),
},
"queues": queues,
"retention": retention,
"formula": self.formula_status_json(&self.models()),
"table": self.table_status_json(),
"limits": {
"max_body_bytes": self.limits.max_body_bytes,
"max_result_bytes": self.limits.max_result_bytes,
"max_export_bytes": self.limits.max_export_bytes,
"max_download_bytes": self.limits.max_download_bytes,
"max_retained": self.limits.max_retained,
"max_retained_bytes": self.limits.max_retained_bytes,
"max_tombstones": self.limits.max_tombstones,
"max_eval_cases": self.limits.max_eval_cases,
},
"model_dir": REDACTED_MODEL_DIR,
"source": source_label(self.model_plan.source()),
"downloads_allowed": self.allow_download,
"download_hosts": self.effective_download_hosts(),
"engine_load_ms": self.engine_load_ms(),
"allow_provider_fallback": self.allow_provider_fallback,
"max_side_len": self.plan_snapshot().engine.global.max_side_len,
})
}
pub fn effective_download_hosts(&self) -> Vec<String> {
let mut hosts: Vec<String> = ALLOWED_DOWNLOAD_HOSTS
.iter()
.map(|host| (*host).to_string())
.collect();
hosts.extend(self.allow_download_hosts.iter().cloned());
hosts
}
pub fn models_json(&self) -> Value {
let report = self.models();
let blocking = self.text_blocking();
json!({
"model_dir": REDACTED_MODEL_DIR,
"source": source_label(self.model_plan.source()),
"downloads_allowed": self.allow_download,
"complete": blocking.is_empty(),
"missing": missing_names(&blocking),
"corrupt": corrupt_names(&blocking),
"blocked": blocked_names(&blocking),
"formula": self.formula_status_json(&report),
"table": self.table_status_json(),
"pipelines": self.pipelines_json(None, self.engine_state_name()),
"verification": verification_json(report.cold_this_call()),
"sets": report
.statuses()
.iter()
.map(|status| set_status_json(status, self.model_plan.set_by_id(&status.set_id)))
.collect::<Vec<_>>(),
})
}
pub fn formula_status_json(&self, report: &ModelReport) -> Value {
let blocking = report.formula_inventory_blocking();
let routing = self.routing();
json!({
"complete": report.formula_inventory_complete(),
"missing": missing_names(&blocking),
"corrupt": corrupt_names(&blocking),
"blocked": blocked_names(&blocking),
"routing": routing.formula,
"disabled_reason": routing.disabled_reason,
"required_roles": ["formula_recognizer"],
"detector": self.formula_detector_detail(),
})
}
pub fn table_status_json(&self) -> Value {
let routing = self.routing();
let blocking = self.model_plan.plan_blocking(Pipeline::Table);
let in_plan = self.model_plan.table_in_plan();
json!({
"complete": in_plan && blocking.is_empty(),
"routing": routing.table,
"disabled_reason": routing.table_disabled_reason,
"in_plan": in_plan,
"missing": missing_names(&blocking),
"corrupt": corrupt_names(&blocking),
"blocked": blocked_names(&blocking),
"required_roles": ["table_structure"],
"structure": self.table_model_detail(),
})
}
fn pipelines_json(&self, verified: Option<&PlanReverification>, text_outcome: &str) -> Value {
let pipeline_files = |pipeline: Pipeline| -> Vec<(String, ModelFileState)> {
match verified {
Some(report) => report
.files_in(pipeline)
.iter()
.map(|file| (file.name.clone(), file.state.clone()))
.collect(),
None => self
.model_plan
.plan_files_in(pipeline)
.into_iter()
.map(|file| (file.name.clone(), file.state()))
.collect(),
}
};
let plan_detail = |pipeline: Pipeline| -> Vec<Value> {
match verified {
Some(report) => report
.files_in(pipeline)
.iter()
.map(|file| {
json!({
"name": file.name,
"role": file.role.as_str(),
"pipeline": file.pipeline.as_str(),
"state": file.state.as_str(),
"declared_sha256": file.declared_sha256,
"has_declared_digest": file.declared_sha256.is_some(),
})
})
.collect(),
None => self
.model_plan
.plan_files_in(pipeline)
.into_iter()
.map(|file| {
json!({
"name": file.name,
"role": file.role.as_str(),
"pipeline": file.pipeline.as_str(),
"state": file.state().as_str(),
"declared_sha256": file.declared_sha256,
"has_declared_digest": file.has_declared_digest(),
})
})
.collect(),
}
};
let verdict = |files: &[(String, ModelFileState)]| {
let names: Vec<String> = files.iter().map(|(name, _)| name.clone()).collect();
let missing: Vec<String> = files
.iter()
.filter(|(_, state)| matches!(state, ModelFileState::Missing))
.map(|(name, _)| name.clone())
.collect();
let corrupt: Vec<String> = files
.iter()
.filter(|(_, state)| matches!(state, ModelFileState::Corrupt { .. }))
.map(|(name, _)| name.clone())
.collect();
let mut blocked = missing.clone();
blocked.extend(corrupt.iter().cloned());
(names, missing, corrupt, blocked)
};
let text_files = pipeline_files(Pipeline::Text);
let (text_names, text_missing, text_corrupt, text_blocked) = verdict(&text_files);
let formula_files = pipeline_files(Pipeline::Formula);
let (formula_names, formula_missing, formula_corrupt, formula_blocked) =
verdict(&formula_files);
let table_files = pipeline_files(Pipeline::Table);
let (table_names, table_missing, table_corrupt, table_blocked) = verdict(&table_files);
let routing = self.routing();
let formula_in_plan = self.model_plan.formula_in_plan();
let table_in_plan = self.model_plan.table_in_plan();
let complete = |pipeline: Pipeline, blocked: &[String]| match verified {
Some(_) => blocked.is_empty(),
None => self.model_plan.plan_complete(pipeline),
};
let text_complete = complete(Pipeline::Text, &text_blocked);
let formula_complete = formula_in_plan && complete(Pipeline::Formula, &formula_blocked);
let table_complete = table_in_plan && complete(Pipeline::Table, &table_blocked);
json!({
"text": {
"pipeline": Pipeline::Text.as_str(),
"outcome": text_outcome,
"in_plan": true,
"complete": text_complete,
"files": text_names,
"missing": text_missing,
"corrupt": text_corrupt,
"blocked": text_blocked,
"plan": plan_detail(Pipeline::Text),
},
"formula": {
"pipeline": Pipeline::Formula.as_str(),
"routing": routing.formula,
"disabled_reason": routing.disabled_reason,
"in_plan": formula_in_plan,
"complete": formula_complete,
"files": formula_names,
"missing": formula_missing,
"corrupt": formula_corrupt,
"blocked": formula_blocked,
"plan": plan_detail(Pipeline::Formula),
"detector": self.formula_detector_detail(),
},
"table": {
"pipeline": Pipeline::Table.as_str(),
"routing": routing.table,
"disabled_reason": routing.table_disabled_reason,
"in_plan": table_in_plan,
"complete": table_complete,
"files": table_names,
"missing": table_missing,
"corrupt": table_corrupt,
"blocked": table_blocked,
"plan": plan_detail(Pipeline::Table),
"structure": self.table_model_detail(),
},
})
}
pub fn models_missing_detail(&self) -> Value {
let blocking = self.text_blocking();
json!({
"missing": missing_names(&blocking),
"corrupt": corrupt_names(&blocking),
"blocked": blocked_names(&blocking),
"source": source_label(self.model_plan.source()),
"model_dir": REDACTED_MODEL_DIR,
})
}
pub fn models_missing_error(&self) -> ServeError {
if self.text_blocking().iter().all(|file| !file.corrupt) {
ServeError::ModelsMissing
} else {
ServeError::ModelsCorrupt
}
}
pub(super) fn error_body(&self, error: &ServeError) -> Body {
match error {
ServeError::ModelsMissing | ServeError::ModelsCorrupt => {
Body::error_with_detail(&self.models_missing_error(), self.models_missing_detail())
}
other => Body::error(other),
}
}
pub fn reserve_queue_slot(self: &Arc<Self>, class: QueueClass) -> Option<QueueReservation> {
{
let mut state = lock(&self.jobs);
if !state.reserve(class) {
return None;
}
}
let shared = Arc::clone(self);
Some(QueueReservation::new(move || {
lock(&shared.jobs).release(class);
}))
}
fn has_queued_work(&self) -> bool {
!lock(&self.jobs).scheduler.is_empty()
}
pub fn submit_ocr(
&self,
bytes: Vec<u8>,
class: QueueClass,
params: OcrJobParams,
reservation: Option<QueueReservation>,
flow: &RequestFlow,
) -> Result<Value, ServeError> {
let decision = match self.admit_ocr() {
Ok(decision) => decision,
Err(error) => {
flow.job_rejected(
class.name(),
error.status_code(),
error.code(),
&error.detail_text(),
);
return Err(error);
}
};
let original_bytes = bytes.len() as u64;
let mut state = lock(&self.jobs);
let id = state.ids.generate();
let position = state.scheduler.enqueue(class, id.clone())?;
if let Err(error) = state.store.insert(
id.clone(),
JobKind::Ocr,
JobQueue::from_class(class),
original_bytes,
monotonic_ms(),
flow.id(),
) {
state.scheduler.remove(class, &id);
return Err(error);
}
state.originals.insert(
id.clone(),
RetainedOcr {
bytes: Arc::from(bytes.into_boxed_slice()),
params,
},
);
state.sync_originals();
state.sync_positions();
drop(state);
drop(reservation);
self.queue_signal.notify_all();
flowlog::job_admitted(
&self.flow,
JobIdentity {
request: flow.id(),
job: &id,
queue: class.name(),
},
position,
&decision,
);
Ok(json!({
"job_id": id,
"kind": JobKind::Ocr.name(),
"queue": class.name(),
"position": position,
"state": JobLifecycle::Queued.name(),
}))
}
fn admission_reason(admission: &OcrAdmission) -> String {
match admission {
OcrAdmission::Run => "run".to_string(),
OcrAdmission::Queue { waiting_for } => format!("queue waiting_for={waiting_for}"),
OcrAdmission::Unavailable { reason } => format!("unavailable reason={reason}"),
OcrAdmission::ModelsMissing { missing } => {
format!("models_missing missing={}", missing.join(","))
}
}
}
fn admit_ocr(&self) -> Result<String, ServeError> {
let admission = self.engine_state().ocr_admission();
if matches!(admission, OcrAdmission::ModelsMissing { .. }) {
self.admit_with_models_on_disk()?;
return Ok("models_on_disk".to_string());
}
let decision = Self::admission_reason(&admission);
ServeError::from_ocr_admission(admission)?;
Ok(decision)
}
fn admit_with_models_on_disk(&self) -> Result<(), ServeError> {
let blocking = blocked_names(&self.text_blocking());
if blocking.is_empty() {
let mut machine = lock(&self.engine_state);
let _ = machine.begin_loading();
return Ok(());
}
{
let mut machine = lock(&self.engine_state);
let _ = machine.models_still_missing(blocking);
}
Err(self.models_missing_error())
}
pub fn submit_download(&self, set_id: &str, flow: &RequestFlow) -> Result<Value, ServeError> {
if !self.allow_download {
return Err(ServeError::DownloadsDisabled);
}
if self.model_plan.set_by_id(set_id).is_none() {
return Err(ServeError::ModelSetNotFound {
set_id: set_id.to_string(),
known: self.model_plan.set_ids(),
});
}
let pending = self
.models()
.pending_download(set_id)
.unwrap_or(PendingDownload {
files: 0,
bytes: Some(0),
});
let budget_bytes = self.limits.download_budget().total_bytes();
if let Some(total) = pending.bytes
&& total > budget_bytes
{
return Err(ServeError::Download(DownloadError::TooLarge {
limit_bytes: budget_bytes,
observed_bytes: Some(total),
}));
}
let required = pending.bytes.unwrap_or(budget_bytes);
let available = (self.free_space)(self.model_plan.model_dir())?;
if available < required {
return Err(ServeError::InsufficientDiskSpace {
required_bytes: required,
available_bytes: available,
});
}
let id = {
let mut state = lock(&self.jobs);
let id = state.ids.generate();
state.store.insert(
id.clone(),
JobKind::ModelDownload,
JobQueue::Download,
0,
monotonic_ms(),
flow.id(),
)?;
state.store.set_download_progress(
&id,
DownloadProgress::planned(pending.files, pending.bytes),
)?;
id
};
match self.download_tx.try_send(DownloadCommand {
job_id: id.clone(),
set_id: set_id.to_string(),
}) {
Ok(()) => Ok(json!({
"job_id": id,
"kind": JobKind::ModelDownload.name(),
"queue": JobQueue::Download.name(),
"state": JobLifecycle::Queued.name(),
"download": {
"files_done": 0,
"files_total": pending.files,
"bytes_done": 0,
"bytes_total": pending.bytes,
"current_file": Value::Null,
},
})),
Err(error) => {
{
let mut state = lock(&self.jobs);
let _ = state.store.cancel(&id, monotonic_ms());
}
Err(match error {
std::sync::mpsc::TrySendError::Full(_) => ServeError::Busy,
std::sync::mpsc::TrySendError::Disconnected(_) => ServeError::Internal,
})
}
}
}
pub fn job_view(&self, id: &str) -> Result<Value, ServeError> {
let mut state = lock(&self.jobs);
state.sync_positions();
let view = state.store.view(id, monotonic_ms())?;
Ok(serde_json::to_value(view).expect("JobView serialization cannot fail"))
}
pub fn job_result(&self, id: &str) -> Result<Body, ServeError> {
let snapshot = self.job_snapshot(id)?;
match snapshot.state {
JobLifecycle::Succeeded => {
let Some(output) = snapshot.output else {
return Err(ServeError::Internal);
};
let value = export::result_json(&output).map_err(ServeError::from)?;
match serialize_bounded(&value, self.limits.max_result_bytes) {
Ok(bytes) => Ok(Body::json(200, bytes)),
Err(SerializeError::TooLarge) => Err(ServeError::ResultTooLarge),
Err(SerializeError::Internal(reason)) => {
eprintln!("serve: cannot serialize the result of {id}: {reason}");
Err(ServeError::Internal)
}
}
}
JobLifecycle::Failed => match snapshot.failure {
Some((status, body)) => Ok(Body::json(status, render_error_body(&body))),
None => Err(ServeError::Internal),
},
JobLifecycle::Queued | JobLifecycle::Running | JobLifecycle::Cancelled => {
Err(ServeError::JobNotFinished)
}
}
}
pub fn annotated_png(&self, id: &str) -> Result<Body, ServeError> {
let snapshot = self.job_snapshot(id)?;
if snapshot.state != JobLifecycle::Succeeded {
return Err(ServeError::JobNotFinished);
}
let Some(output) = snapshot.output else {
return Err(ServeError::Internal);
};
let Some(original) = snapshot.original else {
return Err(ServeError::OriginalEvicted);
};
let png = export::annotated_png(original, &output)
.map_err(|error| export_failure(error, 0, ""))?;
Ok(Body::typed(200, "image/png", png))
}
pub fn export(&self, id: &str, format: ExportFormat) -> Result<Body, ServeError> {
let snapshot = self.job_snapshot(id)?;
if snapshot.state != JobLifecycle::Succeeded {
return Err(ServeError::JobNotFinished);
}
let Some(output) = snapshot.output else {
return Err(ServeError::Internal);
};
let limit = self.limits.max_export_bytes;
let annotated = format!("/api/jobs/{id}/annotated.png");
let document = match format {
ExportFormat::Json => export::json_document(&output, limit, snapshot.serialized_bytes),
ExportFormat::Markdown => export::markdown_document(&ExportRequest {
output: &output,
png: &[],
title: &format!("ocr-{id}"),
limit_bytes: limit,
}),
ExportFormat::Html => {
let Some(original) = snapshot.original else {
return Err(ServeError::OriginalEvicted);
};
let png = export::annotated_png(original, &output)
.map_err(|error| export_failure(error, limit, &annotated))?;
export::html_document(&ExportRequest {
output: &output,
png: &png,
title: &format!("ocr-{id}"),
limit_bytes: limit,
})
}
};
let bytes = document.map_err(|error| export_failure(error, limit, &annotated))?;
Ok(Body::typed(200, format.content_type(), bytes))
}
fn job_snapshot(&self, id: &str) -> Result<JobSnapshot, ServeError> {
let state = lock(&self.jobs);
let record = state.store.record(id)?;
let outcome = state.results.get(id);
let (output, serialized_bytes, failure) = match outcome {
Some(Outcome::Succeeded(succeeded)) => (
Some(Arc::clone(&succeeded.output)),
succeeded.serialized_bytes,
None,
),
Some(Outcome::Failed(status, body)) => (None, 0, Some((*status, body.clone()))),
None => (None, 0, None),
};
Ok(JobSnapshot {
state: record.state,
output,
serialized_bytes,
failure,
original: state
.originals
.get(id)
.map(|retained| Arc::clone(&retained.bytes)),
})
}
pub fn cancel_job(&self, id: &str) -> Result<Value, ServeError> {
let mut state = lock(&self.jobs);
let outcome = state.store.cancel(id, monotonic_ms())?;
if outcome == CancelOutcome::Cancelled {
let queue = state.store.record(id)?.queue;
if let Some(class) = queue.class() {
state.scheduler.remove(class, id);
}
}
state.sync_originals();
state.sync_positions();
let view = state.store.view(id, monotonic_ms())?;
Ok(serde_json::to_value(view).expect("JobView serialization cannot fail"))
}
pub(super) fn begin_download(&self, id: &str) -> Result<(), ServeError> {
lock(&self.jobs).store.start(id, monotonic_ms())
}
pub(super) fn abandon_download(&self, id: &str) {
let mut state = lock(&self.jobs);
let _ = state.store.cancel(id, monotonic_ms());
}
pub(super) fn finish_download_failed(&self, id: &str, error: ServeError) {
let mut state = lock(&self.jobs);
let _ = state
.store
.fail_classified(id, JobFailure::from(&error), monotonic_ms());
self.log_terminal(&state, id, "failed", 0, 0);
}
pub(super) fn finish_download_ok(&self, id: &str) {
let mut state = lock(&self.jobs);
let _ = state.store.succeed(id, 0, monotonic_ms());
self.log_terminal(&state, id, "succeeded", 0, 0);
}
pub(super) fn finish_download_cancelled(&self, id: &str) {
let mut state = lock(&self.jobs);
if !state.store.is_cancel_requested(id) {
let _ = state.store.cancel(id, monotonic_ms());
}
let _ = state.store.finish_cancelled(id, monotonic_ms());
self.log_terminal(&state, id, "cancelled", 0, 0);
}
pub(super) fn set_download_progress(&self, id: &str, progress: DownloadProgress) {
let mut state = lock(&self.jobs);
let _ = state.store.set_download_progress(id, progress);
}
pub(super) fn is_cancel_requested(&self, id: &str) -> bool {
lock(&self.jobs).store.is_cancel_requested(id)
}
pub(super) fn download_job(&self, set_id: &str) -> Option<DownloadJob<'_>> {
let set = self.model_plan.set_by_id(set_id)?;
let (connect_timeout, read_timeout) = download::default_timeouts();
Some(DownloadJob {
set,
root: self.model_plan.model_dir(),
budget_bytes: self.limits.max_download_bytes,
connect_timeout,
read_timeout,
allowed_hosts: self.effective_download_hosts(),
})
}
pub(super) fn pending_download(&self, set_id: &str) -> Option<PendingDownload> {
self.models().pending_download(set_id)
}
pub(super) fn known_set_ids(&self) -> Vec<String> {
self.model_plan.set_ids()
}
fn retained_ocr(&self, id: &str) -> Option<(Arc<[u8]>, OcrJobParams)> {
lock(&self.jobs)
.originals
.get(id)
.map(|retained| (Arc::clone(&retained.bytes), retained.params))
}
fn finish(&self, job_id: &str, outcome: Outcome, backend_ms: u64) {
let now = monotonic_ms();
let mut state = lock(&self.jobs);
let bytes = outcome.bytes();
let failure = match &outcome {
Outcome::Failed(status, body) => Some(JobFailure::from_body(*status, body)),
Outcome::Succeeded(_) => None,
};
if !state.results.insert(job_id, outcome) {
let _ = state.store.fail(job_id, "the result store is full", now);
state.sync_originals();
self.log_terminal(&state, job_id, "failed", 0, backend_ms);
return;
}
match failure {
None => {
let _ = state.store.succeed(job_id, bytes, now);
}
Some(failure) => {
let _ = state.store.fail_classified(job_id, failure, now);
}
}
state.sync_originals();
let terminal = state
.store
.record(job_id)
.map(|record| record.state.name())
.unwrap_or("succeeded");
self.log_terminal(&state, job_id, terminal, bytes, backend_ms);
}
fn log_running(&self, scheduled: &ScheduledJob) {
if !self.flow.enabled() {
return;
}
let state = lock(&self.jobs);
let Ok(record) = state.store.record(&scheduled.id) else {
return;
};
let wait_ms = record
.started_ms
.unwrap_or(record.queued_ms)
.saturating_sub(record.queued_ms);
flowlog::job_running(
&self.flow,
JobIdentity {
request: record.request_id,
job: &record.id,
queue: record.queue.name(),
},
wait_ms,
);
}
fn log_terminal(
&self,
state: &JobState,
job_id: &str,
terminal: &str,
result_bytes: u64,
backend_ms: u64,
) {
if !self.flow.enabled() {
return;
}
let Ok(record) = state.store.record(job_id) else {
return;
};
let failure = record
.failure
.as_ref()
.map(|failure| (failure.status, failure.code, failure.detail.to_string()));
flowlog::job_terminal(
&self.flow,
JobIdentity {
request: record.request_id,
job: job_id,
queue: record.queue.name(),
},
terminal,
result_bytes,
backend_ms,
failure,
);
}
fn create_session(
&self,
plan: &ServeConfigPlan,
) -> (Result<Box<dyn OcrBackend>, RapidOcrError>, u64) {
let started = Instant::now();
let mut config = plan.engine.clone();
let result = match self.model_plan.pin_engine_paths(&mut config) {
Ok(()) => (self.engine_factory)(&config),
Err(error) => Err(RapidOcrError::ModelResolve(error.to_string())),
};
let elapsed_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
(result, elapsed_ms)
}
pub(super) fn ensure_engine_loaded(&self, force: bool) -> EngineLoad {
let plan = self.plan_snapshot();
if !force
&& matches!(self.engine_state(), EngineState::Ready { .. })
&& lock(&self.engine).is_some()
{
return EngineLoad::Ready;
}
let _loading = lock(&self.engine_load);
if !force
&& matches!(self.engine_state(), EngineState::Ready { .. })
&& lock(&self.engine).is_some()
{
return EngineLoad::Ready;
}
self.load_engine(&plan)
}
fn load_engine(&self, plan: &ServeConfigPlan) -> EngineLoad {
let blocking = blocked_names(&self.text_blocking());
{
let mut machine = lock(&self.engine_state);
if matches!(machine.state(), EngineState::Loading) {
if !blocking.is_empty() {
let _ = machine.load_failed(missing_models_reason(&blocking));
return EngineLoad::Failed;
}
} else if blocking.is_empty() {
if machine.begin_loading().is_err() {
return EngineLoad::Failed;
}
} else if matches!(machine.state(), EngineState::BlockedModelsMissing { .. }) {
let _ = machine.models_still_missing(blocking);
return EngineLoad::BlockedModelsMissing;
} else {
if machine.begin_loading().is_err() {
return EngineLoad::Failed;
}
let _ = machine.load_failed(missing_models_reason(&blocking));
return EngineLoad::Failed;
}
}
let (result, elapsed_ms) = self.create_session(plan);
let mut provider: Option<BackendProvider> = None;
let mut reason: Option<String> = None;
{
let mut engine = lock(&self.engine);
*engine = None;
match result {
Ok(backend) => {
provider = Some(backend.provider());
*engine = Some(backend);
}
Err(error) => reason = Some(error.to_string()),
}
}
*lock(&self.engine_load_ms) = Some(elapsed_ms);
let mut machine = lock(&self.engine_state);
match (provider, reason) {
(Some(provider), _) => {
let _ = machine.load_succeeded(
plan.requested_label(),
provider.selected_ep,
provider.fallback_to_cpu,
);
EngineLoad::Ready
}
(None, Some(reason)) => {
let _ = machine.load_failed(reason);
EngineLoad::Failed
}
(None, None) => {
let _ = machine.load_failed("the engine factory returned no backend");
EngineLoad::Failed
}
}
}
pub fn reload_engine(&self, provider: Option<ProviderPreference>) -> Result<Value, ServeError> {
let started = Instant::now();
let Some(requested) = provider else {
let outcome = self.ensure_engine_loaded(true);
return Ok(self.engine_payload(outcome, started, None, None));
};
self.apply_provider(requested, started)
}
pub fn reverify_models(&self) -> Result<Value, ServeError> {
let started = Instant::now();
clear_verification_cache();
let report = self.model_plan.reverify();
let status = self.models();
if let Some(hook) = &self.post_verify {
hook();
}
let outcome = self.ensure_engine_loaded(true);
Ok(self.engine_response_with_report(outcome, started, None, None, &report, &status))
}
pub fn begin_provider_switch(self: &Arc<Self>) -> Option<ProviderSwitchGuard> {
self.provider_switch
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.ok()?;
Some(ProviderSwitchGuard {
shared: Arc::clone(self),
})
}
pub fn begin_evaluation(self: &Arc<Self>) -> Option<EvaluationGuard> {
self.evaluation
.compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
.ok()?;
Some(EvaluationGuard {
shared: Arc::clone(self),
})
}
pub fn max_eval_cases(&self) -> usize {
self.limits.max_eval_cases
}
pub fn eval_root(&self) -> Option<&EvalRoot> {
self.eval_root.as_ref()
}
fn apply_provider(
&self,
requested: ProviderPreference,
started: Instant,
) -> Result<Value, ServeError> {
let old_plan = self.plan_snapshot();
let new_plan = old_plan
.clone()
.with_provider(requested, self.allow_provider_fallback)
.map_err(|error| ServeError::ProviderRejected {
provider: rapid_structure_rs::format_provider_preference(requested),
reason: error.to_string(),
})?;
let _loading = lock(&self.engine_load);
{
let mut machine = lock(&self.engine_state);
if machine.begin_rebuild().is_err() {
drop(machine);
*lock(&self.plan) = new_plan.clone();
let outcome = self.load_engine(&new_plan);
return Ok(self.engine_payload(outcome, started, None, None));
}
}
*lock(&self.plan) = new_plan.clone();
let mut engine = lock(&self.engine);
*engine = None;
let (created, session_ms) = self.create_session(&new_plan);
*lock(&self.engine_load_ms) = Some(session_ms);
match created {
Ok(backend) => {
let provider = backend.provider();
*engine = Some(backend);
drop(engine);
{
let mut machine = lock(&self.engine_state);
let _ = machine.load_succeeded(
new_plan.requested_label(),
provider.selected_ep,
provider.fallback_to_cpu,
);
}
Ok(self.engine_payload(EngineLoad::Ready, started, None, None))
}
Err(new_error) => {
let new_reason = new_error.to_string();
let (restored, rollback_ms) = self.create_session(&old_plan);
*lock(&self.engine_load_ms) = Some(rollback_ms);
match restored {
Ok(backend) => {
let provider = backend.provider();
*engine = Some(backend);
drop(engine);
*lock(&self.plan) = old_plan.clone();
{
let mut machine = lock(&self.engine_state);
let _ = machine.load_succeeded(
old_plan.requested_label(),
provider.selected_ep,
provider.fallback_to_cpu,
);
}
Ok(self.engine_payload(
EngineLoad::Ready,
started,
Some(new_reason),
Some(rollback_ms),
))
}
Err(restore_error) => {
let reason = format!(
"switching the execution provider to {} failed ({new_reason}); \
restoring the previous engine ({}) also failed ({restore_error})",
new_plan.requested_label(),
old_plan.requested_label()
);
drop(engine);
{
let mut machine = lock(&self.engine_state);
let _ = machine.load_failed(reason.clone());
}
Ok(self.engine_payload(
EngineLoad::Failed,
started,
Some(reason),
Some(rollback_ms),
))
}
}
}
}
}
fn engine_payload(
&self,
outcome: EngineLoad,
started: Instant,
error: Option<String>,
rollback_ms: Option<u64>,
) -> Value {
let report = self.models();
self.engine_response(outcome, started, error, rollback_ms, &report)
}
fn engine_response(
&self,
outcome: EngineLoad,
started: Instant,
error: Option<String>,
rollback_ms: Option<u64>,
report: &ModelReport,
) -> Value {
let base = self.engine_payload_json(outcome, started, error, rollback_ms);
let Value::Object(mut object) = base else {
return base;
};
object.insert(
"verification".to_string(),
verification_json(report.cold_this_call()),
);
Value::Object(object)
}
fn engine_response_with_report(
&self,
outcome: EngineLoad,
started: Instant,
error: Option<String>,
rollback_ms: Option<u64>,
verified: &PlanReverification,
report: &ModelReport,
) -> Value {
let mut base = self.engine_response(outcome, started, error, rollback_ms, report);
let Value::Object(object) = &mut base else {
return base;
};
let outcome_label = match object.get("outcome") {
Some(Value::String(label)) => label.clone(),
_ => "failed".to_string(),
};
object.insert("missing".to_string(), verified_names(verified, false));
object.insert("corrupt".to_string(), verified_names(verified, true));
object.insert("files".to_string(), plan_files_json(verified));
object.insert(
"computed".to_string(),
Value::from(verified.digests_computed()),
);
object.insert(
"content_changed".to_string(),
Value::Array(
verified
.content_changed()
.iter()
.map(|name| Value::from(name.clone()))
.collect(),
),
);
object.insert(
"pipelines".to_string(),
self.pipelines_json(Some(verified), &outcome_label),
);
if let Some(Value::Object(verification)) = object.get_mut("verification") {
verification.insert(
"digests_computed".to_string(),
Value::from(verified.digests_computed()),
);
}
base
}
fn engine_payload_json(
&self,
outcome: EngineLoad,
started: Instant,
error: Option<String>,
rollback_ms: Option<u64>,
) -> Value {
let provider = self.provider_status();
let outcome_label = outcome_label(&error, outcome);
let blocking = self.text_blocking();
let elapsed_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
json!({
"outcome": outcome_label,
"engine": self.engine_state(),
"provider": provider,
"requested": provider.requested,
"selected_ep": provider.selected_ep,
"fallback_to_cpu": provider.fallback_to_cpu,
"missing": missing_names(&blocking),
"corrupt": corrupt_names(&blocking),
"source": source_label(self.model_plan.source()),
"model_dir": REDACTED_MODEL_DIR,
"load_ms": elapsed_ms,
"rollback_ms": rollback_ms,
"error": error,
})
}
}
fn outcome_label(error: &Option<String>, outcome: EngineLoad) -> &'static str {
match (error, outcome) {
(Some(_), EngineLoad::Ready) => "rolled_back",
_ => match outcome {
EngineLoad::Ready => "ready",
EngineLoad::BlockedModelsMissing => "blocked_models_missing",
EngineLoad::Failed => "failed",
},
}
}
pub(super) struct ProviderSwitchGuard {
shared: Arc<ServeShared>,
}
impl Drop for ProviderSwitchGuard {
fn drop(&mut self) {
self.shared.provider_switch.store(false, Ordering::SeqCst);
}
}
pub(super) struct EvaluationGuard {
shared: Arc<ServeShared>,
}
impl Drop for EvaluationGuard {
fn drop(&mut self) {
self.shared.evaluation.store(false, Ordering::SeqCst);
}
}
fn missing_models_reason(blocking: &[String]) -> String {
format!(
"the model files are missing or corrupt, so no session can be created: {}",
blocking.join(", ")
)
}
struct JobSnapshot {
state: JobLifecycle,
output: Option<Arc<OcrOutput>>,
serialized_bytes: u64,
failure: Option<(u16, ErrorBody)>,
original: Option<Arc<[u8]>>,
}
fn export_failure(error: ExportError, limit_bytes: u64, annotated: &str) -> ServeError {
match error {
ExportError::TooLarge { observed_bytes } => ServeError::ExportTooLarge {
limit_bytes,
observed_bytes,
annotated: annotated.to_string(),
},
ExportError::Decode(error) => ServeError::from(error),
ExportError::Render(reason) | ExportError::Internal(reason) => {
eprintln!("serve: export/annotation failed: {reason}");
ServeError::Internal
}
}
}
#[derive(Debug, Clone)]
pub(super) struct Body {
pub status: u16,
pub content_type: &'static str,
pub bytes: Vec<u8>,
}
impl Body {
pub fn json(status: u16, bytes: Vec<u8>) -> Self {
Self {
status,
content_type: "application/json; charset=utf-8",
bytes,
}
}
pub fn html(bytes: Vec<u8>) -> Self {
Self {
status: 200,
content_type: "text/html; charset=utf-8",
bytes,
}
}
pub fn typed(status: u16, content_type: &'static str, bytes: Vec<u8>) -> Self {
Self {
status,
content_type,
bytes,
}
}
pub fn error(error: &ServeError) -> Self {
Self::json(error.status_code(), error.render_body().into_bytes())
}
pub fn error_with_detail(error: &ServeError, extra: Value) -> Self {
let mut body = error.body();
let Value::Object(extra) = extra else {
return Self::json(error.status_code(), render_error_body(&body));
};
let mut detail = match body.detail {
Value::Object(map) => map,
_ => serde_json::Map::new(),
};
for (key, value) in extra {
detail.insert(key, value);
}
body.detail = Value::Object(detail);
Self::json(error.status_code(), render_error_body(&body))
}
}
pub(super) struct ServeRuntime {
shared: Arc<ServeShared>,
workers: Vec<JoinHandle<()>>,
}
impl ServeRuntime {
pub fn new(context: ServeContext) -> std::io::Result<Self> {
let service = ServiceState::Starting
.listening_succeeded()
.map_err(|error| std::io::Error::other(error.to_string()))?;
let (machine, backend) = load_engine(&context);
let (download_tx, download_rx) = sync_channel(DOWNLOAD_QUEUE_CAPACITY);
let shared = Arc::new(ServeShared {
service,
limits: context.limits,
plan: Mutex::new(context.plan),
model_plan: context.model_plan,
token: context.token,
local: context.local,
page: context.page,
nonce: context.nonce,
allow_download: context.allow_download,
allow_download_hosts: context.allow_download_hosts,
allow_provider_fallback: context.allow_provider_fallback,
eval_root: context.eval_root,
flow: context.flow,
engine_state: Mutex::new(machine),
engine: Mutex::new(backend),
engine_load: Mutex::new(()),
engine_load_ms: Mutex::new(None),
provider_switch: AtomicBool::new(false),
evaluation: AtomicBool::new(false),
jobs: Mutex::new(JobState {
store: JobStore::new(JobStoreLimits::from_limits(&context.limits)),
scheduler: DualQueueScheduler::new(SchedulerConfig::from_limits(&context.limits)),
results: ResultStore::new(
context.limits.max_retained,
context.limits.max_retained_bytes,
),
ids: JobIdGenerator::default(),
originals: HashMap::new(),
reserved_text: 0,
reserved_formula: 0,
}),
queue_signal: Condvar::new(),
download_tx,
engine_factory: context.engine_factory,
downloader: context.downloader,
free_space: context.free_space,
post_verify: context.post_verify,
shutting_down: AtomicBool::new(false),
});
let mut workers = Vec::with_capacity(3);
{
let runtime = Arc::clone(&shared);
workers.push(spawn("serve-ocr", move || ocr_worker(runtime))?);
}
{
let runtime = Arc::clone(&shared);
let downloader = Arc::clone(&shared.downloader);
workers.push(spawn("serve-download", move || {
download::worker(runtime, download_rx, downloader)
})?);
}
{
let runtime = Arc::clone(&shared);
workers.push(spawn("serve-sweeper", move || sweeper(runtime))?);
}
Ok(Self { shared, workers })
}
pub fn shared(&self) -> &Arc<ServeShared> {
&self.shared
}
pub fn stop(&mut self) {
self.shared.shutdown();
for worker in self.workers.drain(..) {
let _ = worker.join();
}
}
}
impl Drop for ServeRuntime {
fn drop(&mut self) {
self.stop();
}
}
fn load_engine(context: &ServeContext) -> (EngineStateMachine, Option<Box<dyn OcrBackend>>) {
let mut machine = EngineStateMachine::start(context.snapshot.readiness());
if !context.snapshot.blocking_names().is_empty() {
return (machine, None);
}
let mut engine_config = context.plan.engine.clone();
if let Err(error) = context.model_plan.pin_engine_paths(&mut engine_config) {
let _ = machine.load_failed(error.to_string());
return (machine, None);
}
match (context.engine_factory)(&engine_config) {
Ok(backend) => {
let BackendProvider {
selected_ep,
fallback_to_cpu,
} = backend.provider();
let _ = machine.load_succeeded(
context.plan.requested_label(),
selected_ep,
fallback_to_cpu,
);
(machine, Some(backend))
}
Err(error) => {
let _ = machine.load_failed(error.to_string());
(machine, None)
}
}
}
fn ocr_worker(runtime: Arc<ServeShared>) {
while !runtime.is_shutting_down() {
if !runtime.has_queued_work() {
wait_for_work(&runtime);
continue;
}
match runtime.ensure_engine_loaded(false) {
EngineLoad::Ready => {}
EngineLoad::BlockedModelsMissing => {
fail_next_queued(&runtime, runtime.models_missing_outcome());
continue;
}
EngineLoad::Failed => {
let reason = runtime
.engine_failure_reason()
.unwrap_or_else(|| "the OCR engine could not be created".to_string());
fail_next_queued(&runtime, failure(503, "engine_unavailable", reason));
continue;
}
}
let Some(scheduled) = next_scheduled(&runtime) else {
continue;
};
runtime.log_running(&scheduled);
let Some((bytes, params)) = runtime.retained_ocr(&scheduled.id) else {
runtime.finish(
&scheduled.id,
failure(500, "internal", "the queued image is no longer available"),
0,
);
continue;
};
let backend_started = Instant::now();
let outcome = recognize(&runtime, &scheduled.id, bytes, scheduled.class, params);
let backend_ms = u64::try_from(backend_started.elapsed().as_millis()).unwrap_or(u64::MAX);
runtime.finish(&scheduled.id, outcome, backend_ms);
}
}
fn fail_next_queued(runtime: &ServeShared, outcome: Outcome) {
let Some(scheduled) = next_scheduled(runtime) else {
return;
};
runtime.finish(&scheduled.id, outcome, 0);
}
fn next_scheduled(runtime: &ServeShared) -> Option<ScheduledJob> {
let mut state = lock(&runtime.jobs);
loop {
if let Some(scheduled) = state.take_next() {
return Some(scheduled);
}
if runtime.is_shutting_down() {
return None;
}
let (guard, timeout) = runtime
.queue_signal
.wait_timeout(state, WORKER_POLL)
.unwrap_or_else(|error| error.into_inner());
state = guard;
if timeout.timed_out() && runtime.is_shutting_down() {
return None;
}
}
}
fn wait_for_work(runtime: &ServeShared) {
let state = lock(&runtime.jobs);
if !state.scheduler.is_empty() || runtime.is_shutting_down() {
return;
}
let _ = runtime.queue_signal.wait_timeout(state, WORKER_POLL);
}
fn recognize(
runtime: &ServeShared,
job_id: &str,
bytes: Arc<[u8]>,
class: QueueClass,
params: OcrJobParams,
) -> Outcome {
let formula = match class {
QueueClass::Text => FormulaPolicy::default(),
QueueClass::Formula => match runtime.formula_policy() {
Some(policy) => policy,
None => return Outcome::Failed(409, runtime.formula_missing_failure()),
},
};
let table_policy = if params.table {
match runtime.table_policy() {
Some(policy) => policy,
None => return Outcome::Failed(409, runtime.table_missing_failure()),
}
} else {
TablePolicy::default()
};
let request = text_request(bytes, params.max_side, formula, table_policy, params.words);
let output = match recognize_with(runtime, request) {
Ok(output) => output,
Err(error) => {
if error.status_code() == 500 {
eprintln!(
"serve: recognition failed for job {job_id}: {}",
error.message()
);
}
return Outcome::Failed(error.status_code(), error.body());
}
};
serialize_output(output, runtime.limits.max_result_bytes)
}
pub(super) fn text_request(
bytes: Arc<[u8]>,
max_side: Option<u32>,
formula: FormulaPolicy,
table: TablePolicy,
words: WordOutputMode,
) -> OcrRequest {
OcrRequest {
input: ImageInput::Encoded(bytes),
roi: None,
scale_hint: None,
stages: StagePlan::default(),
preprocess: PreprocessPolicy {
max_side,
..PreprocessPolicy::default()
},
detection: DetectionPolicy::default(),
recognition: RecognitionPolicy { words },
output: OutputPolicy::default(),
formula,
table,
}
}
pub(super) fn recognize_with(
runtime: &ServeShared,
request: OcrRequest,
) -> Result<OcrOutput, ServeError> {
let mut engine = lock(&runtime.engine);
let Some(backend) = engine.as_mut() else {
return Err(ServeError::EngineUnavailable {
reason: "the OCR engine is not loaded".to_string(),
});
};
backend.recognize(request).map_err(ServeError::from)
}
fn serialize_output(output: OcrOutput, limit: u64) -> Outcome {
let value = match export::result_json(&output) {
Ok(value) => value,
Err(error) => {
let serve_error = ServeError::from(error);
return Outcome::Failed(serve_error.status_code(), serve_error.body());
}
};
match serialize_bounded(&value, limit) {
Ok(bytes) => {
let serialized_bytes = bytes.len() as u64;
Outcome::Succeeded(Succeeded::new(Arc::new(output), serialized_bytes))
}
Err(SerializeError::TooLarge) => {
let error = ServeError::ResultTooLarge;
Outcome::Failed(error.status_code(), error.body())
}
Err(SerializeError::Internal(reason)) => failure(500, "internal", reason),
}
}
fn failure(status: u16, code: &'static str, message: impl Into<String>) -> Outcome {
Outcome::Failed(
status,
ErrorBody {
code,
message: message.into(),
detail: Value::Null,
},
)
}
fn sweeper(runtime: Arc<ServeShared>) {
while !runtime.is_shutting_down() {
thread::sleep(SWEEP_INTERVAL);
let now = monotonic_ms();
let mut state = lock(&runtime.jobs);
let report = state.store.tick(now);
if report.expired_jobs > 0
|| report.evicted_for_count > 0
|| report.evicted_for_bytes > 0
|| report.expired_tombstones > 0
{
eprintln!(
"serve: retention tick expired_jobs={} evicted_for_count={} evicted_for_bytes={} \
expired_tombstones={}",
report.expired_jobs,
report.evicted_for_count,
report.evicted_for_bytes,
report.expired_tombstones
);
}
if !state.store.is_empty() {
let live: std::collections::HashSet<String> = state
.store
.views(now)
.into_iter()
.map(|view| view.id)
.collect();
let dropped_results = state.results.retain_only(|id| live.contains(id));
state.sync_originals();
let before = state.originals.len();
state.originals.retain(|id, _| live.contains(id));
let dropped_originals = before - state.originals.len();
if dropped_results > 0 || dropped_originals > 0 {
eprintln!(
"serve: dropped {dropped_results} result(s) and {dropped_originals} retained \
original(s)"
);
}
}
}
}
fn spawn(name: &str, body: impl FnOnce() + Send + 'static) -> std::io::Result<JoinHandle<()>> {
thread::Builder::new().name(name.to_string()).spawn(body)
}
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(|error| error.into_inner())
}
fn public_fingerprint(fingerprint: &OrtRuntimeFingerprint) -> Value {
let module = &fingerprint.runtime_module;
json!({
"file": file_name(&module.path),
"size_bytes": module.size_bytes,
"sha256": module.sha256,
"source": fingerprint.runtime_source,
"complete": fingerprint.is_complete(),
"providers": fingerprint.provider_dlls.iter().map(|dll| json!({
"name": dll.name,
"size_bytes": dll.size_bytes,
"loaded": dll.loaded,
})).collect::<Vec<_>>(),
})
}
fn file_name(path: &str) -> String {
std::path::Path::new(path)
.file_name()
.map(|name| name.to_string_lossy().into_owned())
.unwrap_or_else(|| path.to_string())
}
fn set_status_json(
status: &rapid_structure_rs::ModelSetStatus,
declared: Option<&ModelSet>,
) -> Value {
json!({
"id": status.set_id,
"family": declared.map(|set| set.family.clone()),
"version": declared.map(|set| set.version.clone()),
"complete": status.complete,
"download_bytes_total": status.download_bytes_total,
"files": status.files.iter().map(|(file, state)| {
let mut value = json!({
"name": file.name,
"role": file.role.as_str(),
"state": state.as_str(),
"size_bytes": file.size_bytes,
"sha256": if file.has_hash() { Value::String(file.sha256.clone()) } else { Value::Null },
"source_url": if file.has_source_url() { Value::String(file.source_url.clone()) } else { Value::Null },
});
if let rapid_structure_rs::ModelFileState::Corrupt { expected, actual } = state
&& let Value::Object(object) = &mut value
{
object.insert("expected".to_string(), Value::String(expected.clone()));
object.insert("actual".to_string(), Value::String(actual.clone()));
}
value
}).collect::<Vec<_>>(),
})
}
fn render_error_body(body: &ErrorBody) -> Vec<u8> {
serde_json::to_vec(body).expect("ErrorBody serialization cannot fail")
}
fn plan_file_json(file: &VerifiedPlanFile) -> Value {
json!({
"name": file.name,
"role": file.role.as_str(),
"pipeline": file.pipeline.as_str(),
"state": file.state.as_str(),
"declared_sha256": file.declared_sha256,
"has_declared_digest": file.declared_sha256.is_some(),
"sha256": file.sha256,
"cause": file.cause.as_str(),
"digest_computed_this_call": file.sha256.is_some(),
})
}
fn plan_files_json(verified: &PlanReverification) -> Value {
Value::Array(verified.files().iter().map(plan_file_json).collect())
}
fn verified_names(verified: &PlanReverification, corrupt: bool) -> Value {
Value::Array(
verified
.blocking_in(Pipeline::Text)
.iter()
.filter(|file| matches!(file.state, ModelFileState::Corrupt { .. }) == corrupt)
.map(|file| Value::from(file.name.clone()))
.collect(),
)
}
fn verification_json(cold_this_call: usize) -> Value {
let stats = verification_stats();
json!({
"identity": "path + size + mtime + SHA-256 of the first and last 64 KiB",
"partial_window_bytes": rapid_structure_rs::PARTIAL_DIGEST_WINDOW_BYTES,
"cold_this_call": cold_this_call,
"cold_verifications": stats.cold_verifications,
"cache_hits": stats.cache_hits,
"partial_reads": stats.partial_reads,
"partial_mismatches": stats.partial_mismatches,
"entries": stats.entries,
"last_cold_ms": stats.last_cold_ms(),
"last_cold_bytes": stats.last_cold_bytes,
"guarantee": "content is verified at first use and re-verified whenever the file's \
identity changes (path, size, mtime, or the first/last 64 KiB); the \
service does not claim that the on-disk content is trusted at all times",
"force_check": "POST /api/models/reverify (or --reverify-models at startup) is the \
deterministic way to force a fresh full check",
"residual_blind_spot": "a same-size, same-mtime edit that also keeps the first and last \
64 KiB byte-identical is still not detected; the partial digest \
is a heuristic that narrows the window, not a security boundary, \
because an attacker who can write the file can also preserve its \
head and tail",
})
}
#[cfg(test)]
mod tests {
use std::path::Path;
use super::{OcrRouting, REDACTED_MODEL_DIR, monotonic_ms, public_fingerprint};
use crate::serve::error::ServeError;
use crate::serve::queue::QueueClass;
use crate::serve::state::{EngineStateMachine, ModelReadiness};
#[test]
fn the_clock_is_monotonic_and_the_redaction_is_a_constant() {
let first = monotonic_ms();
let second = monotonic_ms();
assert!(second >= first, "{first} -> {second}");
assert_eq!(REDACTED_MODEL_DIR, "<redacted>");
}
#[test]
fn the_routing_refuses_to_silently_use_the_formula_queue() {
let disabled = OcrRouting::text_only("test: no detector");
assert_eq!(disabled.class_for(None).expect("text"), QueueClass::Text);
assert_eq!(
disabled.class_for(Some("text")).expect("text"),
QueueClass::Text
);
assert!(matches!(
disabled.class_for(Some("formula")),
Err(ServeError::BadRequest)
));
assert!(matches!(
disabled.class_for(Some("bogus")),
Err(ServeError::BadRequest)
));
assert_eq!(
disabled.disabled_reason.as_deref(),
Some("test: no detector"),
"the page needs a textual reason, not just a disabled control"
);
let enabled = OcrRouting::formula_enabled();
assert_eq!(
enabled.class_for(Some("formula")).expect("formula"),
QueueClass::Formula
);
assert_eq!(enabled.disabled_reason, None);
assert_eq!(enabled.class_for(None).expect("text"), QueueClass::Text);
}
#[test]
fn the_formula_routing_is_decided_by_the_detector_alone() {
let enabled = super::routing_for(Some(Path::new("D:\\m\\pix2text-mfd-1.5.onnx")), None);
assert!(enabled.formula);
assert_eq!(enabled.disabled_reason, None);
let disabled = super::routing_for(None, None);
assert!(!disabled.formula);
let reason = disabled.disabled_reason.expect("a textual reason");
assert!(reason.contains("--formula-detector"), "{reason}");
assert!(reason.contains("ordinary OCR is unaffected"), "{reason}");
}
#[test]
fn the_table_routing_is_decided_by_the_structure_model_alone() {
let detector = Path::new("D:\\m\\pix2text-mfd-1.5.onnx");
let structure = Path::new("D:\\m\\slanet-plus.onnx");
let neither = super::routing_for(None, None);
assert!(!neither.table);
assert!(!neither.formula);
let reason = neither
.table_disabled_reason
.clone()
.expect("a textual reason");
assert!(reason.contains("--table-model"), "{reason}");
assert!(reason.contains("ordinary OCR"), "{reason}");
assert_eq!(
neither
.table_requested(Some("1"))
.unwrap_err()
.status_code(),
400
);
assert!(!neither.table_requested(None).expect("absent means off"));
assert!(!neither.table_requested(Some("0")).expect("explicit off"));
let table_only = super::routing_for(None, Some(structure));
assert!(table_only.table);
assert_eq!(table_only.table_disabled_reason, None);
assert!(!table_only.formula);
assert!(table_only.table_requested(Some("1")).expect("enabled"));
let formula_only = super::routing_for(Some(detector), None);
assert!(formula_only.formula);
assert!(!formula_only.table);
let both = super::routing_for(Some(detector), Some(structure));
assert!(both.formula);
assert!(both.table);
assert_eq!(both.table_disabled_reason, None);
}
#[test]
fn the_table_switch_rejects_anything_but_one_and_zero() {
let enabled = super::routing_for(None, Some(Path::new("D:\\m\\slanet-plus.onnx")));
for raw in ["2", "true", "yes", "", "TRUE", "1 ", "-1"] {
assert!(
enabled.table_requested(Some(raw)).is_err(),
"`?table={raw}` must be rejected"
);
}
}
#[test]
fn the_public_fingerprint_never_carries_a_directory() {
let fingerprint = rapid_structure_rs::ort_runtime_fingerprint();
let public = public_fingerprint(&fingerprint);
let text = public.to_string();
let raw = &fingerprint.runtime_module.path;
if let Some(parent) = Path::new(raw).parent()
&& !parent.as_os_str().is_empty()
{
let directory = parent.to_string_lossy().into_owned();
assert!(
!text.contains(&directory),
"the module directory must not appear: {text}"
);
}
assert!(public.get("file").is_some());
}
#[test]
fn readiness_treats_missing_and_corrupt_the_same_way() {
let machine = EngineStateMachine::start(ModelReadiness::Incomplete {
missing: vec!["a.onnx".to_string(), "b.onnx".to_string()],
});
assert_eq!(machine.state().name(), "blocked_models_missing");
}
}