diff --git a/backend/src/analysis/daemon.rs b/backend/src/analysis/daemon.rs index a5624df..67568bb 100644 --- a/backend/src/analysis/daemon.rs +++ b/backend/src/analysis/daemon.rs @@ -32,6 +32,10 @@ use crate::storage::Storage; 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. @@ -40,12 +44,45 @@ 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, pub workers: usize, pub job_timeout: Duration, /// Live-event sink (Started/Completed/Failed) for admin SSE. pub events: Arc, + /// 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>, } pub struct AnalysisDaemonHandle { @@ -75,6 +112,7 @@ pub fn spawn( 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 }); @@ -88,6 +126,7 @@ struct WorkerContext { dispatcher: Arc, job_timeout: Duration, events: Arc, + readiness: Option>, id: usize, } @@ -98,6 +137,18 @@ impl WorkerContext { 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 { + if !readiness.ready().await { + 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, @@ -269,7 +320,29 @@ impl AnalyzeDispatcher for RealAnalyzeDispatcher { /// in the `tests/` dir (a separate crate). pub mod test_support { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; + 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 { + 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 { diff --git a/backend/src/app.rs b/backend/src/app.rs index 3fabc32..35beccf 100644 --- a/backend/src/app.rs +++ b/backend/src/app.rs @@ -420,6 +420,24 @@ fn spawn_analysis_daemon( model: cfg.model.clone(), max_image_bytes: cfg.max_image_bytes, }); + // When a readiness URL is configured, gate leasing on it so an + // autoscaler that idle-stops the vision container never lets a job burn + // its retries. A dedicated short-timeout client keeps the probe snappy + // and independent of the (long) per-request analysis timeout. + let readiness: Option> = + match &cfg.vision_health_url { + Some(url) if !url.is_empty() => { + let probe = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(5)) + .build() + .context("build vision readiness http client")?; + Some(Arc::new(crate::analysis::daemon::HttpVisionReadiness { + http: probe, + health_url: url.clone(), + })) + } + _ => None, + }; let handle = crate::analysis::daemon::spawn( db, CancellationToken::new(), @@ -428,6 +446,7 @@ fn spawn_analysis_daemon( workers: cfg.workers, job_timeout: cfg.job_timeout, events, + readiness, }, ); tracing::info!(workers = cfg.workers, model = %cfg.model, "analysis worker daemon started"); diff --git a/backend/src/config.rs b/backend/src/config.rs index 34eb6ed..a62743b 100644 --- a/backend/src/config.rs +++ b/backend/src/config.rs @@ -124,6 +124,13 @@ pub struct AnalysisConfig { pub workers: usize, /// OpenAI-compatible chat/completions URL (`ANALYSIS_VISION_URL`). pub endpoint: String, + /// Optional vision readiness probe URL (`ANALYSIS_VISION_HEALTH_URL`), + /// e.g. `http://mangalord-vision:8000/health`. When set, the worker + /// refuses to lease a job until this answers 2xx — so an autoscaler that + /// idle-stops the vision container never lets a job burn its retries or + /// land a `failed` row. `None` (the default) disables the gate, matching + /// the prior behavior for always-on endpoints. Env-only, like `api_key`. + pub vision_health_url: Option, /// Model id to request (`ANALYSIS_MODEL`). pub model: String, /// Optional bearer token (`ANALYSIS_API_KEY`); local servers usually @@ -190,6 +197,7 @@ impl Default for AnalysisConfig { enabled: false, workers: 1, endpoint: "http://localhost:8000/v1/chat/completions".to_string(), + vision_health_url: None, model: String::new(), api_key: None, request_timeout: Duration::from_secs(120), @@ -221,6 +229,9 @@ impl AnalysisConfig { enabled: env_bool("ANALYSIS_ENABLED", d.enabled), workers: env_usize("ANALYSIS_WORKERS", d.workers).max(1), endpoint: std::env::var("ANALYSIS_VISION_URL").unwrap_or(d.endpoint), + vision_health_url: std::env::var("ANALYSIS_VISION_HEALTH_URL") + .ok() + .filter(|s| !s.is_empty()), model: std::env::var("ANALYSIS_MODEL").unwrap_or(d.model), api_key: std::env::var("ANALYSIS_API_KEY") .ok() diff --git a/backend/src/settings.rs b/backend/src/settings.rs index 1be7d6b..3c211a0 100644 --- a/backend/src/settings.rs +++ b/backend/src/settings.rs @@ -410,6 +410,8 @@ impl AnalysisSettings { grounding_prompt: prompt(&self.grounding_prompt, GROUNDING_PROMPT_DEFAULT), // Env-only secret preserved from the base. api_key: base.api_key.clone(), + // Env-only readiness probe URL preserved from the base. + vision_health_url: base.vision_health_url.clone(), }) } } diff --git a/backend/tests/analysis_readiness.rs b/backend/tests/analysis_readiness.rs new file mode 100644 index 0000000..4bb3909 --- /dev/null +++ b/backend/tests/analysis_readiness.rs @@ -0,0 +1,149 @@ +//! Integration tests for the analysis worker's *vision readiness gate*. +//! +//! When a readiness probe is wired in and reports "not ready" (the vision +//! container is stopped, or running but the model is still loading), the +//! worker must NOT lease any `analyze_page` job. Leasing increments +//! `attempts` in SQL and a down vision then burns all retries and writes a +//! permanent `failed` row — poisoning the queue while the autoscaler is +//! mid-cold-start. So a not-ready gate must leave jobs untouched +//! (`state='pending'`, `attempts=0`, no `page_analysis` row) and resume the +//! instant the probe flips to ready. + +use std::sync::Arc; +use std::time::Duration; + +use mangalord::analysis::daemon::{ + self, + test_support::{CountingDispatcher, ToggleReadiness}, + AnalysisDaemonConfig, +}; +use mangalord::repo::page_analysis; +use sqlx::PgPool; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +async fn seed_page(pool: &PgPool) -> Uuid { + let manga_id: Uuid = + sqlx::query_scalar("INSERT INTO mangas (title) VALUES ('M') RETURNING id") + .fetch_one(pool) + .await + .unwrap(); + let chapter_id: Uuid = + sqlx::query_scalar("INSERT INTO chapters (manga_id, number) VALUES ($1, 1) RETURNING id") + .bind(manga_id) + .fetch_one(pool) + .await + .unwrap(); + sqlx::query_scalar( + "INSERT INTO pages (chapter_id, page_number, storage_key, content_type) \ + VALUES ($1, 1, 'k/1.png', 'image/png') RETURNING id", + ) + .bind(chapter_id) + .fetch_one(pool) + .await + .unwrap() +} + +async fn job_row(pool: &PgPool, page_id: Uuid) -> (String, i32) { + sqlx::query_as( + "SELECT state, attempts FROM crawler_jobs \ + WHERE payload->>'kind' = 'analyze_page' AND payload->>'page_id' = $1", + ) + .bind(page_id.to_string()) + .fetch_one(pool) + .await + .unwrap() +} + +async fn wait_for_state(pool: &PgPool, page_id: Uuid, want: &str) { + let result = tokio::time::timeout(Duration::from_secs(8), async { + loop { + if job_row(pool, page_id).await.0 == want { + return; + } + tokio::time::sleep(Duration::from_millis(50)).await; + } + }) + .await; + assert!(result.is_ok(), "job for {page_id} never reached state {want}"); +} + +fn spawn_with( + pool: &PgPool, + dispatcher: Arc, + readiness: Option>, +) -> daemon::AnalysisDaemonHandle { + daemon::spawn( + pool.clone(), + CancellationToken::new(), + AnalysisDaemonConfig { + dispatcher, + workers: 1, + job_timeout: Duration::from_secs(5), + events: Arc::new(mangalord::analysis::events::AnalysisEvents::new()), + readiness, + }, + ) +} + +#[sqlx::test(migrations = "./migrations")] +async fn not_ready_gate_never_leases_or_fails_the_page(pool: PgPool) { + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let readiness = ToggleReadiness::new(false); + let handle = spawn_with(&pool, dispatcher.clone(), Some(readiness.clone())); + + // Give the worker several poll cycles to (wrongly) lease if the gate is + // broken. + tokio::time::sleep(Duration::from_millis(800)).await; + handle.shutdown().await; + + assert_eq!(dispatcher.call_count(), 0, "must not dispatch while not ready"); + let (state, attempts) = job_row(&pool, page_id).await; + assert_eq!(state, "pending", "job must stay pending while vision is down"); + assert_eq!(attempts, 0, "job must not be leased (attempts must stay 0)"); + assert!( + page_analysis::load(&pool, page_id).await.unwrap().is_none(), + "no failed/any page_analysis row may be written while not ready" + ); +} + +#[sqlx::test(migrations = "./migrations")] +async fn resumes_when_readiness_flips_to_ready(pool: PgPool) { + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let readiness = ToggleReadiness::new(false); + let handle = spawn_with(&pool, dispatcher.clone(), Some(readiness.clone())); + + // Confirm it is parked, then open the gate. + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!(dispatcher.call_count(), 0); + readiness.set(true); + + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + assert_eq!(dispatcher.call_count(), 1, "must process once ready"); +} + +#[sqlx::test(migrations = "./migrations")] +async fn no_gate_processes_as_before(pool: PgPool) { + // readiness = None preserves today's behavior for always-on endpoints. + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let handle = spawn_with(&pool, dispatcher.clone(), None); + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + assert_eq!(dispatcher.call_count(), 1); +} diff --git a/backend/tests/analysis_worker.rs b/backend/tests/analysis_worker.rs index fd2d236..152be4b 100644 --- a/backend/tests/analysis_worker.rs +++ b/backend/tests/analysis_worker.rs @@ -78,6 +78,7 @@ fn spawn_with( events: std::sync::Arc::new( mangalord::analysis::events::AnalysisEvents::new(), ), + readiness: None, }, ); (handle, cancel) @@ -227,6 +228,7 @@ async fn worker_publishes_started_and_completed_events(pool: PgPool) { workers: 1, job_timeout: Duration::from_secs(5), events: events.clone(), + readiness: None, }, ); @@ -277,6 +279,7 @@ async fn worker_publishes_failed_event_on_dispatch_error(pool: PgPool) { workers: 1, job_timeout: Duration::from_secs(5), events: events.clone(), + readiness: None, }, );