Mangalord/backend/src/analysis/daemon.rs

//! The AI page-analysis worker daemon.
//!
//! A lean sibling of [`crate::crawler::daemon`]: it leases `analyze_page`
//! jobs from the SAME `crawler_jobs` queue (filtered by kind, so it never
//! contends with crawl jobs), dispatches each through an [`AnalyzeDispatcher`]
//! seam, and acks done/failed. No cron, no browser, no session state — it
//! runs whenever `ANALYSIS_ENABLED` is set, independent of the crawler.
//!
//! Per job: skip if already `done` (unless the payload sets `force`),
//! heartbeat the lease while the (slow) vision call runs, isolate panics
//! and an outer timeout, and on terminal failure write a `failed`
//! `page_analysis` row so the page's state is observable.

use std::panic::AssertUnwindSafe;
use std::sync::Arc;
use std::time::Duration;

use async_trait::async_trait;
use futures_util::FutureExt;
use sqlx::PgPool;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use uuid::Uuid;

use crate::analysis::events::{AnalysisEvent, AnalysisEvents};
use crate::analysis::vision::VisionClient;
use crate::crawler::jobs::{self, JobPayload, Lease, KIND_ANALYZE_PAGE};
use crate::repo;
use crate::storage::Storage;

/// Lease window; continuously extended by the per-job heartbeat.
const LEASE_DURATION: Duration = Duration::from_secs(60);
/// Heartbeat cadence — a third of the lease window.
const LEASE_HEARTBEAT: Duration = Duration::from_secs(20);
/// How long to wait before re-probing vision when the readiness gate is
/// closed. Short enough to resume promptly after an autoscaler cold start,
/// long enough not to hammer `/health` while vision is down.
const READINESS_POLL: Duration = Duration::from_secs(2);

/// The unit of work: analyze one page. Implemented by
/// [`RealAnalyzeDispatcher`] in production and stubbed in tests.
#[async_trait]
pub trait AnalyzeDispatcher: Send + Sync {
    async fn dispatch(&self, page_id: Uuid) -> anyhow::Result<()>;
}

/// Probe answering "is the vision backend up and the model loaded?". When a
/// probe is wired in, the worker refuses to *lease* a job until it reports
/// ready — so an idle-stopped vision (the autoscaler's normal state) never
/// burns a job's retries or writes a `failed` row. Implemented by
/// [`HttpVisionReadiness`] in production (a `GET /health`) and stubbed in
/// tests.
#[async_trait]
pub trait VisionReadiness: Send + Sync {
    async fn ready(&self) -> bool;
}

/// Production readiness probe: `GET {health_url}`, ready iff it answers 2xx.
/// Any transport error (connection refused while stopped, 503 while the
/// model loads) is treated as not-ready.
pub struct HttpVisionReadiness {
    pub http: reqwest::Client,
    pub health_url: String,
}

#[async_trait]
impl VisionReadiness for HttpVisionReadiness {
    async fn ready(&self) -> bool {
        match self.http.get(&self.health_url).send().await {
            Ok(resp) => resp.status().is_success(),
            Err(_) => false,
        }
    }
}

pub struct AnalysisDaemonConfig {
    pub dispatcher: Arc<dyn AnalyzeDispatcher>,
    pub workers: usize,
    pub job_timeout: Duration,
    /// Live-event sink (Started/Completed/Failed) for admin SSE.
    pub events: Arc<AnalysisEvents>,
    /// Optional vision readiness gate. `None` ⇒ no gate (today's behavior,
    /// for always-on endpoints). `Some` ⇒ the worker parks instead of
    /// leasing while the probe is not ready.
    pub readiness: Option<Arc<dyn VisionReadiness>>,
}

pub struct AnalysisDaemonHandle {
    cancel: CancellationToken,
    join: JoinSet<()>,
}

impl AnalysisDaemonHandle {
    pub async fn shutdown(mut self) {
        self.cancel.cancel();
        while self.join.join_next().await.is_some() {}
    }
}

/// Spawn the analysis workers. Returns immediately; tasks run in the
/// background until the handle is shut down (or `cancel` fires).
pub fn spawn(
    pool: PgPool,
    cancel: CancellationToken,
    cfg: AnalysisDaemonConfig,
) -> AnalysisDaemonHandle {
    let mut join = JoinSet::new();
    for id in 0..cfg.workers.max(1) {
        let ctx = WorkerContext {
            pool: pool.clone(),
            cancel: cancel.clone(),
            dispatcher: Arc::clone(&cfg.dispatcher),
            job_timeout: cfg.job_timeout,
            events: Arc::clone(&cfg.events),
            readiness: cfg.readiness.clone(),
            id,
        };
        join.spawn(async move { ctx.run().await });
    }
    AnalysisDaemonHandle { cancel, join }
}

struct WorkerContext {
    pool: PgPool,
    cancel: CancellationToken,
    dispatcher: Arc<dyn AnalyzeDispatcher>,
    job_timeout: Duration,
    events: Arc<AnalysisEvents>,
    readiness: Option<Arc<dyn VisionReadiness>>,
    id: usize,
}

impl WorkerContext {
    async fn run(self) {
        // Last observed readiness, so we log only on transitions (not every
        // poll). `None` until the first probe.
        let mut was_ready: Option<bool> = None;
        loop {
            if self.cancel.is_cancelled() {
                tracing::info!(worker = self.id, "analysis worker: shutdown");
                return;
            }
            // Readiness gate: park (without leasing) while vision is down or
            // still loading its model, so the autoscaler's idle-stop never
            // costs a job its retries. Probe *before* the lease — leasing
            // increments `attempts` in SQL and can't be undone.
            if let Some(readiness) = &self.readiness {
                let ready = readiness.ready().await;
                // Log on flip so a typo'd health URL (which parks the worker
                // forever) is diagnosable, without spamming every poll.
                if was_ready != Some(ready) {
                    if ready {
                        tracing::info!(worker = self.id, "analysis worker: vision ready — resuming");
                    } else {
                        tracing::warn!(
                            worker = self.id,
                            "analysis worker: vision not ready — parking until /health is 2xx"
                        );
                    }
                    was_ready = Some(ready);
                }
                if !ready {
                    if self.sleep_or_cancel(READINESS_POLL).await {
                        return;
                    }
                    continue;
                }
            }
            let leases =
                match jobs::lease(&self.pool, Some(KIND_ANALYZE_PAGE), 1, LEASE_DURATION).await {
                    Ok(v) => v,
                    Err(e) => {
                        tracing::warn!(worker = self.id, ?e, "analysis worker: lease failed");
                        if self.sleep_or_cancel(Duration::from_secs(5)).await {
                            return;
                        }
                        continue;
                    }
                };
            let Some(lease) = leases.into_iter().next() else {
                if self.sleep_or_cancel(Duration::from_secs(1)).await {
                    return;
                }
                continue;
            };
            self.process_lease(lease).await;
        }
    }

    /// Sleep `dur` or return `true` if cancelled while waiting.
    async fn sleep_or_cancel(&self, dur: Duration) -> bool {
        tokio::select! {
            _ = tokio::time::sleep(dur) => false,
            _ = self.cancel.cancelled() => true,
        }
    }

    async fn process_lease(&self, lease: Lease) {
        let JobPayload::AnalyzePage { page_id, force } = lease.payload else {
            // Shouldn't happen — we lease only analyze_page — but ack done
            // so a misrouted job doesn't loop forever.
            tracing::warn!(worker = self.id, "analysis worker: non-analyze payload leased");
            let _ = jobs::ack_done(&self.pool, lease.id, lease.lease_generation).await;
            return;
        };

        // Skip-if-done net: a non-forced job for an already-analyzed page is
        // a no-op (e.g. a duplicate enqueue). Re-analysis goes through
        // `force` or a fresh page row.
        if !force {
            if let Ok(Some(row)) = repo::page_analysis::load(&self.pool, page_id).await {
                if row.status == crate::domain::page_analysis::AnalysisStatus::Done {
                    let _ = jobs::ack_done(&self.pool, lease.id, lease.lease_generation).await;
                    return;
                }
            }
        }

        // Resolve the labeled page breadcrumb once so live events carry the
        // manga/chapter/number AND human labels the dashboard's "now
        // analyzing" banner names. A missing breadcrumb (page deleted) just
        // suppresses events — the dispatch still runs and acks normally.
        let breadcrumb = repo::page::locate_labeled(&self.pool, page_id)
            .await
            .ok()
            .flatten();
        if let Some((manga_id, manga_title, chapter_id, chapter_number, page_number)) =
            breadcrumb.clone()
        {
            self.events.publish(AnalysisEvent::Started {
                page_id,
                manga_id,
                manga_title,
                chapter_id,
                chapter_number,
                page_number,
            });
        }

        // Heartbeat the lease while the vision call runs.
        let heartbeat = {
            let hb_pool = self.pool.clone();
            let hb_id = lease.id;
            let hb_gen = lease.lease_generation;
            tokio::spawn(async move {
                loop {
                    tokio::time::sleep(LEASE_HEARTBEAT).await;
                    match jobs::renew(&hb_pool, hb_id, hb_gen, LEASE_DURATION).await {
                        Ok(true) => {}
                        Ok(false) => break,
                        Err(e) => tracing::warn!(lease_id = %hb_id, ?e, "heartbeat renew failed"),
                    }
                }
            })
        };

        let started = std::time::Instant::now();
        let dispatch = AssertUnwindSafe(self.dispatcher.dispatch(page_id)).catch_unwind();
        // Race the vision call against shutdown. Without this, shutdown
        // (SIGTERM / `Supervisors::reload_analysis` toggling analysis off)
        // blocks for up to `job_timeout` (default 600s) on whichever call
        // happens to be in flight. On cancel we drop the dispatch future
        // and `jobs::release` the lease — next tick picks it up without
        // burning an attempt.
        let cancel_result = tokio::select! {
            biased;
            _ = self.cancel.cancelled() => None,
            o = tokio::time::timeout(self.job_timeout, dispatch) => Some(o),
        };
        let elapsed_ms = started.elapsed().as_millis() as i64;
        heartbeat.abort();
        let Some(outcome) = cancel_result else {
            tracing::info!(
                worker = self.id,
                lease_id = %lease.id,
                "analysis worker: cancelled mid-dispatch — releasing lease"
            );
            let _ = jobs::release(&self.pool, lease.id, lease.lease_generation).await;
            // Don't publish Failed (the page didn't actually fail) and
            // don't write a duration row — neither outcome is true. DO
            // publish Cancelled so the dashboard's "now analyzing" banner
            // can clear the page we previously sent Started for. Without
            // this the banner stuck on the cancelled page until a later
            // page kicked off and overwrote it.
            if let Some((manga_id, manga_title, chapter_id, chapter_number, page_number)) =
                breadcrumb
            {
                self.events.publish(AnalysisEvent::Cancelled {
                    page_id,
                    manga_id,
                    manga_title,
                    chapter_id,
                    chapter_number,
                    page_number,
                });
            }
            return;
        };

        // Flatten timeout / panic / dispatch error into one Result + message.
        let result: Result<(), String> = match outcome {
            Err(_elapsed) => Err("analysis dispatch timed out".to_string()),
            Ok(Err(_panic)) => Err("analysis dispatcher panicked".to_string()),
            Ok(Ok(Err(e))) => Err(format!("{e:#}")),
            Ok(Ok(Ok(()))) => Ok(()),
        };

        if let Some((manga_id, manga_title, chapter_id, chapter_number, page_number)) =
            breadcrumb
        {
            let event = if result.is_ok() {
                AnalysisEvent::Completed {
                    page_id,
                    manga_id,
                    manga_title,
                    chapter_id,
                    chapter_number,
                    page_number,
                }
            } else {
                AnalysisEvent::Failed {
                    page_id,
                    manga_id,
                    manga_title,
                    chapter_id,
                    chapter_number,
                    page_number,
                }
            };
            self.events.publish(event);
        }

        // Stamp how long the vision dispatch took — but ONLY when a row
        // was actually written this attempt. Otherwise a non-terminal
        // failure on a force-re-analyze (page already has a `done` row)
        // would overwrite the prior duration with the duration of a
        // failed retry that wrote nothing new.
        let wrote_row = match &result {
            Ok(()) => true,
            Err(_) if lease.attempts >= lease.max_attempts => true, // mark_failed wrote below
            Err(_) => false,
        };

        match result {
            Ok(()) => {
                let _ = jobs::ack_done(&self.pool, lease.id, lease.lease_generation).await;
            }
            Err(msg) => {
                tracing::warn!(
                    worker = self.id,
                    %page_id,
                    error = %msg,
                    "analysis worker: dispatch failed — ack failed"
                );
                let _ = jobs::ack_failed(
                    &self.pool,
                    lease.id,
                    &msg,
                    lease.attempts,
                    lease.max_attempts,
                    lease.lease_generation,
                )
                .await;
                // Terminal failure (retries exhausted) → record a `failed`
                // row so the page's analysis state is observable rather than
                // silently absent.
                if lease.attempts >= lease.max_attempts {
                    let _ = repo::page_analysis::mark_failed(&self.pool, page_id, &msg).await;
                }
            }
        }

        if wrote_row {
            let _ = repo::page_analysis::record_duration(&self.pool, page_id, elapsed_ms).await;
        }
    }
}

/// Production dispatcher: load the page, read its image from storage, call
/// the vision model, and persist the analysis.
pub struct RealAnalyzeDispatcher {
    pub db: PgPool,
    pub storage: Arc<dyn Storage>,
    pub vision: VisionClient,
    pub model: String,
    pub max_image_bytes: usize,
}

#[async_trait]
impl AnalyzeDispatcher for RealAnalyzeDispatcher {
    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 through the byte cap so an oversized blob is rejected as it's
        // read rather than after it's fully buffered into memory (mirrors the
        // OCR dispatcher).
        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)
                })?;
        let analysis = self.vision.analyze(&bytes, &page.content_type).await?;
        repo::page_analysis::persist_analysis(&self.db, page_id, &analysis, &self.model).await?;
        Ok(())
    }
}

/// Stubs for the daemon's integration tests. Public because the tests live
/// in the `tests/` dir (a separate crate).
pub mod test_support {
    use super::*;
    use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};

    /// A readiness probe whose answer can be flipped from the test thread, to
    /// drive the "vision is down, then comes up" transition.
    pub struct ToggleReadiness {
        ready: AtomicBool,
    }

    impl ToggleReadiness {
        pub fn new(ready: bool) -> Arc<Self> {
            Arc::new(Self { ready: AtomicBool::new(ready) })
        }
        pub fn set(&self, ready: bool) {
            self.ready.store(ready, Ordering::Release);
        }
    }

    #[async_trait]
    impl VisionReadiness for ToggleReadiness {
        async fn ready(&self) -> bool {
            self.ready.load(Ordering::Acquire)
        }
    }

    /// Counts dispatch calls and returns a configurable result.
    pub struct CountingDispatcher {
        pub calls: AtomicUsize,
        pub fail: bool,
        pub panic: bool,
    }

    impl CountingDispatcher {
        pub fn ok() -> Arc<Self> {
            Arc::new(Self { calls: AtomicUsize::new(0), fail: false, panic: false })
        }
        pub fn failing() -> Arc<Self> {
            Arc::new(Self { calls: AtomicUsize::new(0), fail: true, panic: false })
        }
        pub fn panicking() -> Arc<Self> {
            Arc::new(Self { calls: AtomicUsize::new(0), fail: false, panic: true })
        }
        pub fn call_count(&self) -> usize {
            self.calls.load(Ordering::Acquire)
        }
    }

    /// A dispatcher that sleeps for `delay` before returning `Ok`. Lets a
    /// test fire shutdown WHILE a dispatch is in flight, to exercise the
    /// cancellation race in `process_lease`.
    pub struct SlowDispatcher {
        pub calls: AtomicUsize,
        pub delay: Duration,
    }

    impl SlowDispatcher {
        pub fn new(delay: Duration) -> Arc<Self> {
            Arc::new(Self { calls: AtomicUsize::new(0), delay })
        }
        pub fn call_count(&self) -> usize {
            self.calls.load(Ordering::Acquire)
        }
    }

    #[async_trait]
    impl AnalyzeDispatcher for SlowDispatcher {
        async fn dispatch(&self, _page_id: Uuid) -> anyhow::Result<()> {
            self.calls.fetch_add(1, Ordering::AcqRel);
            tokio::time::sleep(self.delay).await;
            Ok(())
        }
    }

    #[async_trait]
    impl AnalyzeDispatcher for CountingDispatcher {
        async fn dispatch(&self, _page_id: Uuid) -> anyhow::Result<()> {
            self.calls.fetch_add(1, Ordering::AcqRel);
            if self.panic {
                panic!("intentional analysis dispatcher panic");
            }
            if self.fail {
                anyhow::bail!("intentional analysis dispatch failure");
            }
            Ok(())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use tokio::net::TcpListener;

    /// Bind an ephemeral port that answers every request with `status_line`
    /// (e.g. `"200 OK"`), and return its `/health` URL. The probe under test
    /// only inspects the status code, so a zero-length body is enough.
    async fn serve(status_line: &'static str) -> String {
        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            while let Ok((mut sock, _)) = listener.accept().await {
                let mut buf = [0u8; 1024];
                let _ = sock.read(&mut buf).await;
                let resp = format!(
                    "HTTP/1.1 {status_line}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"
                );
                let _ = sock.write_all(resp.as_bytes()).await;
            }
        });
        format!("http://{addr}/health")
    }

    fn probe(url: String) -> HttpVisionReadiness {
        HttpVisionReadiness {
            http: reqwest::Client::builder()
                .timeout(Duration::from_millis(500))
                .build()
                .unwrap(),
            health_url: url,
        }
    }

    #[tokio::test]
    async fn http_readiness_2xx_is_ready() {
        assert!(probe(serve("200 OK").await).ready().await);
    }

    #[tokio::test]
    async fn http_readiness_non_2xx_is_not_ready() {
        // llama-server answers 503 while the model is still loading.
        assert!(!probe(serve("503 Service Unavailable").await).ready().await);
    }

    #[tokio::test]
    async fn http_readiness_connection_error_is_not_ready() {
        // Nothing listening (vision stopped) → not ready, never an error.
        assert!(!probe("http://127.0.0.1:1/health".to_string()).ready().await);
    }
}