Mangalord/backend/src/analysis/ocr.rs

//! The in-process OCR analysis backend (the `ocrs` engine).
//!
//! A lightweight alternative to [`crate::analysis::vision`]: instead of a slow
//! local LLM, each page is run through `ocrs` — a pure-Rust detect→recognize
//! OCR pipeline on the `rten` runtime. It extracts **text only** (no tags,
//! scene description or NSFW flags — those stay the vision backend's job), then
//! reuses [`repo::page_analysis::persist_analysis`] so the OCR lines land in
//! `page_ocr_text` and the weighted `search_doc` tsvector exactly as the vision
//! path produces them. That makes the existing text-search surfaces
//! (`/v1/me/page-search` and the tag aggregations) work with no further wiring.
//!
//! The engine is split behind the [`OcrEngine`] trait so the dispatcher is
//! unit-testable without shipping the (multi-MB) `.rten` model files: tests use
//! [`test_support::StubOcrEngine`], production uses [`OcrsEngine`].

use std::sync::Arc;

use async_trait::async_trait;
use sqlx::PgPool;
use uuid::Uuid;

use crate::analysis::daemon::AnalyzeDispatcher;
use crate::domain::page_analysis::{OcrResult, SafetyFlag, VisionAnalysis};
use crate::repo;
use crate::storage::Storage;

/// The `model` label stamped onto `page_analysis` rows written by this backend.
pub const OCR_MODEL_LABEL: &str = "ocrs";

/// Extracts text lines from a decoded-or-encoded page image. The production
/// impl ([`OcrsEngine`]) decodes the bytes itself; the trait takes the raw
/// stored image bytes so the dispatcher stays engine-agnostic.
pub trait OcrEngine: Send + Sync {
    /// Run OCR over one page image (the bytes as stored, e.g. PNG/JPEG/WebP).
    /// Returns the recognized text lines in reading order (top→bottom).
    fn recognize(&self, image: &[u8]) -> anyhow::Result<Vec<String>>;
}

/// Turn the OCR engine's ordered text lines into the [`VisionAnalysis`] shape
/// that [`repo::page_analysis::persist_analysis`] consumes. OCR-only: the tag,
/// scene and safety fields are left empty/default. The line `kind` is left
/// blank — `persist_analysis` maps an empty kind to the neutral mid-weight
/// `OcrKind::Narration` bucket (classifying speech/sfx/… is the deferred
/// vision backend's job).
pub fn lines_to_analysis(lines: Vec<String>) -> VisionAnalysis {
    let ocr_results = lines
        .into_iter()
        .map(|text| OcrResult { text, kind: String::new(), y: None })
        .collect();
    VisionAnalysis {
        ocr_results,
        tagging_results: Vec::new(),
        scene_description: String::new(),
        safety_flag: SafetyFlag::default(),
    }
}

/// Production OCR engine: an `ocrs` detect+recognize pipeline with the two
/// `.rten` models loaded once at startup. Cheap to share across workers — the
/// recognize path borrows `&self`.
pub struct OcrsEngine {
    engine: ocrs::OcrEngine,
    /// Hard cap on decoded pixel count (decompression-bomb guard). See
    /// [`decode_rgb8_within`].
    max_decode_pixels: u64,
}

impl OcrsEngine {
    /// Load the detection + recognition models from disk and build the engine.
    /// Fails (at startup) if either model file is missing or unreadable, so a
    /// misconfigured path is a loud boot error rather than a per-page failure.
    ///
    /// `max_decode_pixels` bounds the decoded image size (see
    /// [`decode_rgb8_within`]) — wired from `AnalysisConfig::ocr_max_decode_pixels`.
    pub fn from_model_paths(
        detection: &str,
        recognition: &str,
        max_decode_pixels: u64,
    ) -> anyhow::Result<Self> {
        use anyhow::Context;
        let detection_model = rten::Model::load_file(detection)
            .with_context(|| format!("load ocrs detection model {detection}"))?;
        let recognition_model = rten::Model::load_file(recognition)
            .with_context(|| format!("load ocrs recognition model {recognition}"))?;
        let engine = ocrs::OcrEngine::new(ocrs::OcrEngineParams {
            detection_model: Some(detection_model),
            recognition_model: Some(recognition_model),
            ..Default::default()
        })
        .context("construct ocrs engine")?;
        Ok(Self { engine, max_decode_pixels })
    }
}

/// Decode encoded image bytes to RGB8 while refusing decompression bombs.
///
/// `image::load_from_memory` allocates the full decoded buffer up front, so a
/// tiny file declaring enormous dimensions (e.g. a 50000×50000 PNG) inflates
/// to billions of bytes and OOM-kills the process. We cap the decoder's
/// allocation at `max_decode_pixels` worth of RGBA (4 bytes/px headroom over
/// the RGB8 result), which the `image` crate checks against the header
/// *before* allocating — so an over-size image fails fast instead of dying.
fn decode_rgb8_within(image: &[u8], max_decode_pixels: u64) -> anyhow::Result<image::RgbImage> {
    use anyhow::Context;
    use std::io::Cursor;

    let mut reader = image::ImageReader::new(Cursor::new(image))
        .with_guessed_format()
        .context("guess page image format for OCR")?;
    let mut limits = image::Limits::default();
    // 4 bytes/px (RGBA) gives headroom over the eventual RGB8 buffer and any
    // single intermediate the decoder allocates per pixel.
    limits.max_alloc = Some(max_decode_pixels.saturating_mul(4));
    reader.limits(limits);
    Ok(reader
        .decode()
        .context("decode page image for OCR")?
        .into_rgb8())
}

impl OcrEngine for OcrsEngine {
    fn recognize(&self, image: &[u8]) -> anyhow::Result<Vec<String>> {
        // Decode to RGB8 so `ImageSource` gets a known channel layout,
        // bounding the decoded size against the decompression-bomb cap.
        let rgb = decode_rgb8_within(image, self.max_decode_pixels)?;
        let source = ocrs::ImageSource::from_bytes(rgb.as_raw(), rgb.dimensions())
            .map_err(|e| anyhow::anyhow!("build OCR image source: {e}"))?;
        let input = self.engine.prepare_input(source)?;
        // detect words → group into lines → recognize each line. Mirrors
        // `OcrEngine::get_text`, but keeps the lines as a Vec instead of
        // joining them, so each becomes its own `page_ocr_text` row.
        let words = self.engine.detect_words(&input)?;
        let line_rects = self.engine.find_text_lines(&input, &words);
        let lines = self
            .engine
            .recognize_text(&input, &line_rects)?
            .into_iter()
            .filter_map(|line| line.map(|l| l.to_string()))
            .map(|s| s.trim().to_string())
            .filter(|s| !s.is_empty())
            .collect();
        Ok(lines)
    }
}

/// Bound on concurrent CPU-bound OCR inferences, given the number of analysis
/// workers and the host's available parallelism.
///
/// Each worker's `dispatch` fires one `spawn_blocking` OCR run, so without a
/// bound `ANALYSIS_WORKERS` blocking tasks can run at once. OCR is fully
/// CPU-bound, so running more than there are cores just thrashes the scheduler
/// (and balloons the blocking pool). Cap at the core count, but never below 1
/// and never above the worker count (more permits than workers is pointless).
pub fn ocr_concurrency_limit(workers: usize, cores: usize) -> usize {
    workers.min(cores.max(1)).max(1)
}

/// Run a CPU-bound OCR closure on the blocking pool while holding an
/// **owned** permit for the entire duration of the work.
///
/// The permit is acquired with `acquire_owned` and moved *into* the
/// blocking task rather than being held by the caller's future. This
/// matters on cancellation: if the dispatcher future is dropped (graceful
/// shutdown), a borrowed permit would be released the instant the future
/// unwinds — but `spawn_blocking` work is not cancellable and keeps
/// running detached, so the concurrency bound (ANALYSIS_WORKERS) would be
/// briefly exceeded. Moving the permit into the task ties the slot's
/// lifetime to the actual CPU work.
async fn run_ocr_blocking<T, F>(
    permits: Arc<tokio::sync::Semaphore>,
    f: F,
) -> anyhow::Result<T>
where
    F: FnOnce() -> T + Send + 'static,
    T: Send + 'static,
{
    let permit = permits
        .acquire_owned()
        .await
        .map_err(|e| anyhow::anyhow!("OCR semaphore closed: {e}"))?;
    tokio::task::spawn_blocking(move || {
        let _permit = permit; // held until f() returns, even if the caller is cancelled
        f()
    })
    .await
    .map_err(|e| anyhow::anyhow!("OCR task join error: {e}"))
}

/// Production dispatcher for the OCR backend: load the page, read its image
/// from storage, run OCR on the blocking pool, and persist the lines. Mirrors
/// [`crate::analysis::daemon::RealAnalyzeDispatcher`] but with no network I/O.
pub struct OcrAnalyzeDispatcher {
    pub db: PgPool,
    pub storage: Arc<dyn Storage>,
    pub engine: Arc<dyn OcrEngine>,
    pub max_image_bytes: usize,
    /// Caps concurrent CPU-bound OCR inferences across all workers (see
    /// [`ocr_concurrency_limit`]). Shared via the `Arc<dyn AnalyzeDispatcher>`,
    /// so one permit pool covers every worker.
    pub ocr_permits: Arc<tokio::sync::Semaphore>,
}

#[async_trait]
impl AnalyzeDispatcher for OcrAnalyzeDispatcher {
    async fn dispatch(&self, page_id: Uuid) -> anyhow::Result<()> {
        let Some(page) = repo::page::find_by_id(&self.db, page_id).await? else {
            // Page was deleted between enqueue and dispatch — nothing to do.
            return Ok(());
        };
        // Stream the blob through the byte cap so the read itself bails once
        // the running total exceeds `max_image_bytes` — a plain `get()` would
        // buffer the whole (possibly huge) image into memory before the cap
        // could reject it.
        let file = self
            .storage
            .get_stream(&page.storage_key)
            .await
            .map_err(|e| anyhow::anyhow!("read page image {}: {e}", page.storage_key))?;
        let bytes =
            crate::crawler::safety::accumulate_capped(file.stream, self.max_image_bytes)
                .await
                .map_err(|e| {
                    anyhow::anyhow!("page image {} over the byte cap: {e}", page.storage_key)
                })?;
        // OCR inference is CPU-bound and synchronous — keep it off the async
        // worker's runtime thread, and gate it behind the shared permit pool so
        // ANALYSIS_WORKERS > cores can't oversubscribe the blocking pool.
        let engine = Arc::clone(&self.engine);
        let lines =
            run_ocr_blocking(Arc::clone(&self.ocr_permits), move || engine.recognize(&bytes))
                .await??;
        let analysis = lines_to_analysis(lines);
        repo::page_analysis::persist_analysis(&self.db, page_id, &analysis, OCR_MODEL_LABEL).await?;
        Ok(())
    }
}

/// Stubs for the OCR dispatcher's integration tests. Public because the tests
/// live in the `tests/` dir (a separate crate).
pub mod test_support {
    use super::*;

    /// An [`OcrEngine`] that returns a fixed set of lines regardless of input,
    /// so the dispatcher's storage→persist path can be tested without models.
    pub struct StubOcrEngine {
        pub lines: Vec<String>,
    }

    impl StubOcrEngine {
        pub fn new(lines: &[&str]) -> Arc<Self> {
            Arc::new(Self { lines: lines.iter().map(|s| s.to_string()).collect() })
        }
    }

    impl OcrEngine for StubOcrEngine {
        fn recognize(&self, _image: &[u8]) -> anyhow::Result<Vec<String>> {
            Ok(self.lines.clone())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn run_ocr_blocking_holds_permit_until_work_completes_despite_cancellation() {
        use std::sync::mpsc;
        use tokio::sync::{Notify, Semaphore};

        let permits = Arc::new(Semaphore::new(1));
        let (release_tx, release_rx) = mpsc::channel::<()>();
        let started = Arc::new(Notify::new());

        // Model the dispatch path: acquire an owned permit and run blocking
        // work that we hold open via a channel.
        let sem = Arc::clone(&permits);
        let started2 = Arc::clone(&started);
        let caller = tokio::spawn(async move {
            run_ocr_blocking(sem, move || {
                // Signal that the blocking task now holds the permit, then
                // block until the test releases us.
                started2.notify_one();
                release_rx.recv().ok();
            })
            .await
            .unwrap();
        });

        // Wait until the blocking task is running and owns the permit.
        started.notified().await;
        assert_eq!(permits.available_permits(), 0, "permit taken by blocking work");

        // Cancel the caller future (simulates graceful shutdown). The
        // blocking task is not cancellable and keeps running detached; with
        // an *owned* permit the slot must stay occupied. A borrowed permit
        // would have been released here, regressing the concurrency bound.
        caller.abort();
        let _ = caller.await;
        assert_eq!(
            permits.available_permits(),
            0,
            "owned permit must remain held by the still-running blocking task"
        );

        // Let the blocking work finish; the permit is then returned.
        release_tx.send(()).unwrap();
        let permit = tokio::time::timeout(std::time::Duration::from_secs(5), permits.acquire())
            .await
            .expect("permit should be released once blocking work completes")
            .unwrap();
        drop(permit);
    }

    #[test]
    fn lines_to_analysis_maps_lines_in_order_and_leaves_rest_empty() {
        let v = lines_to_analysis(vec!["Hello".to_string(), "world!".to_string()]);
        assert_eq!(v.ocr_results.len(), 2);
        assert_eq!(v.ocr_results[0].text, "Hello");
        assert_eq!(v.ocr_results[1].text, "world!");
        // OCR-only: kind blank (→ Narration at persist), no tags/scene/safety.
        assert!(v.ocr_results.iter().all(|r| r.kind.is_empty()));
        assert!(v.tagging_results.is_empty());
        assert_eq!(v.scene_description, "");
        assert!(!v.safety_flag.is_nsfw);
        assert!(v.safety_flag.content_type.is_empty());
    }

    #[test]
    fn lines_to_analysis_handles_no_text() {
        let v = lines_to_analysis(Vec::new());
        assert!(v.ocr_results.is_empty());
    }

    /// A minimal valid PNG of `w`×`h` (single black pixel scaled via IHDR is
    /// not valid; instead encode a real tiny image, then we patch the IHDR
    /// dimensions for the bomb case). For the happy path we just encode a real
    /// small image.
    fn encode_png(w: u32, h: u32) -> Vec<u8> {
        let img = image::RgbImage::new(w, h);
        let mut buf = std::io::Cursor::new(Vec::new());
        image::DynamicImage::ImageRgb8(img)
            .write_to(&mut buf, image::ImageFormat::Png)
            .unwrap();
        buf.into_inner()
    }

    #[test]
    fn decode_rgb8_within_accepts_normal_page() {
        // A perfectly ordinary page decodes fine under a generous cap.
        let png = encode_png(64, 96);
        let rgb = decode_rgb8_within(&png, 100_000_000).unwrap();
        assert_eq!(rgb.dimensions(), (64, 96));
    }

    #[test]
    fn decode_rgb8_within_rejects_oversize_image() {
        // The same image, but the cap is set below its pixel count: the
        // allocation limit must trip rather than the decode succeeding. This
        // is the decompression-bomb guard in miniature — a real bomb declares
        // a huge size in a few bytes; here we shrink the budget instead.
        let png = encode_png(2000, 2000); // 4 MP
        let err = decode_rgb8_within(&png, 1_000).unwrap_err();
        assert!(
            err.to_string().contains("decode page image"),
            "expected a decode error, got: {err}"
        );
    }

    #[test]
    fn ocr_concurrency_limit_caps_at_cores() {
        // More workers than cores → clamp to cores (don't oversubscribe).
        assert_eq!(ocr_concurrency_limit(8, 4), 4);
    }

    #[test]
    fn ocr_concurrency_limit_caps_at_workers() {
        // Fewer workers than cores → only `workers` ever run anyway.
        assert_eq!(ocr_concurrency_limit(2, 16), 2);
    }

    #[test]
    fn ocr_concurrency_limit_never_zero() {
        // Degenerate inputs still yield at least one permit.
        assert_eq!(ocr_concurrency_limit(0, 0), 1);
        assert_eq!(ocr_concurrency_limit(1, 0), 1);
    }
}