diff --git a/backend/Cargo.lock b/backend/Cargo.lock index 605e293..609e6cd 100644 --- a/backend/Cargo.lock +++ b/backend/Cargo.lock @@ -1517,7 +1517,7 @@ checksum = "c41e0c4fef86961ac6d6f8a82609f55f31b05e4fce149ac5710e439df7619ba4" [[package]] name = "mangalord" -version = "0.67.0" +version = "0.68.0" dependencies = [ "anyhow", "argon2", diff --git a/backend/Cargo.toml b/backend/Cargo.toml index 1bcb25e..a424043 100644 --- a/backend/Cargo.toml +++ b/backend/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "mangalord" -version = "0.67.0" +version = "0.68.0" edition = "2021" default-run = "mangalord" diff --git a/backend/src/analysis/daemon.rs b/backend/src/analysis/daemon.rs new file mode 100644 index 0000000..87a00a9 --- /dev/null +++ b/backend/src/analysis/daemon.rs @@ -0,0 +1,281 @@ +//! 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::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); + +/// 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<()>; +} + +pub struct AnalysisDaemonConfig { + pub dispatcher: Arc, + pub workers: usize, + pub job_timeout: Duration, +} + +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, + id, + }; + join.spawn(async move { ctx.run().await }); + } + AnalysisDaemonHandle { cancel, join } +} + +struct WorkerContext { + pool: PgPool, + cancel: CancellationToken, + dispatcher: Arc, + job_timeout: Duration, + id: usize, +} + +impl WorkerContext { + async fn run(self) { + loop { + if self.cancel.is_cancelled() { + tracing::info!(worker = self.id, "analysis worker: shutdown"); + return; + } + 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).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).await; + return; + } + } + } + + // Heartbeat the lease while the vision call runs. + let heartbeat = { + let hb_pool = self.pool.clone(); + let hb_id = lease.id; + tokio::spawn(async move { + loop { + tokio::time::sleep(LEASE_HEARTBEAT).await; + match jobs::renew(&hb_pool, hb_id, LEASE_DURATION).await { + Ok(true) => {} + Ok(false) => break, + Err(e) => tracing::warn!(lease_id = %hb_id, ?e, "heartbeat renew failed"), + } + } + }) + }; + + let dispatch = AssertUnwindSafe(self.dispatcher.dispatch(page_id)).catch_unwind(); + let outcome = tokio::time::timeout(self.job_timeout, dispatch).await; + heartbeat.abort(); + + // 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(()), + }; + + match result { + Ok(()) => { + let _ = jobs::ack_done(&self.pool, lease.id).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, + ) + .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; + } + } + } + } +} + +/// 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, + 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(()); + }; + let bytes = self + .storage + .get(&page.storage_key) + .await + .map_err(|e| anyhow::anyhow!("read page image {}: {e}", page.storage_key))?; + if bytes.len() > self.max_image_bytes { + anyhow::bail!( + "page image {} is {} bytes, over the {} cap", + page.storage_key, + bytes.len(), + self.max_image_bytes + ); + } + 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::{AtomicUsize, Ordering}; + + /// 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 { + Arc::new(Self { calls: AtomicUsize::new(0), fail: false, panic: false }) + } + pub fn failing() -> Arc { + Arc::new(Self { calls: AtomicUsize::new(0), fail: true, panic: false }) + } + pub fn panicking() -> Arc { + Arc::new(Self { calls: AtomicUsize::new(0), fail: false, panic: true }) + } + pub fn call_count(&self) -> usize { + self.calls.load(Ordering::Acquire) + } + } + + #[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(()) + } + } +} diff --git a/backend/src/analysis/mod.rs b/backend/src/analysis/mod.rs index 3cd9028..53899a8 100644 --- a/backend/src/analysis/mod.rs +++ b/backend/src/analysis/mod.rs @@ -7,5 +7,6 @@ //! * [`vision`] — the HTTP client: downscale → request → parse → sanitize. //! * [`daemon`] — the job-leasing worker loop (added in the worker phase). +pub mod daemon; pub mod prompt; pub mod vision; diff --git a/backend/src/app.rs b/backend/src/app.rs index b2f3856..d76036e 100644 --- a/backend/src/app.rs +++ b/backend/src/app.rs @@ -91,6 +91,9 @@ pub struct CrawlerControl { pub struct AppHandle { pub router: Router, pub daemon: Option, + /// AI content-analysis worker daemon; `None` when `ANALYSIS_ENABLED` + /// is off. Awaited alongside the crawler daemon on shutdown. + pub analysis_daemon: Option, } impl AppHandle { @@ -98,6 +101,9 @@ impl AppHandle { if let Some(d) = self.daemon { d.shutdown().await; } + if let Some(a) = self.analysis_daemon { + a.shutdown().await; + } } } @@ -131,6 +137,41 @@ pub async fn build(config: Config) -> anyhow::Result { (None, None, None) }; + // AI content-analysis worker. Independent of the crawler daemon (works + // for uploads with the crawler off), gated on ANALYSIS_ENABLED. Uses a + // plain reqwest client — no cookie jar / proxy, unlike the crawler's. + let analysis_daemon = if config.analysis.enabled { + let http = reqwest::Client::builder() + .timeout(config.analysis.request_timeout) + .build() + .context("build analysis http client")?; + let vision = crate::analysis::vision::VisionClient::new(http, &config.analysis); + let dispatcher = Arc::new(crate::analysis::daemon::RealAnalyzeDispatcher { + db: db.clone(), + storage: Arc::clone(&storage), + vision, + model: config.analysis.model.clone(), + max_image_bytes: config.analysis.max_image_bytes, + }); + let handle = crate::analysis::daemon::spawn( + db.clone(), + tokio_util::sync::CancellationToken::new(), + crate::analysis::daemon::AnalysisDaemonConfig { + dispatcher, + workers: config.analysis.workers, + job_timeout: config.analysis.job_timeout, + }, + ); + tracing::info!( + workers = config.analysis.workers, + model = %config.analysis.model, + "analysis worker daemon started" + ); + Some(handle) + } else { + None + }; + let auth_limiter = Arc::new(AuthRateLimiter::new(config.auth.rate_limit)); let state = AppState { db, @@ -144,7 +185,7 @@ pub async fn build(config: Config) -> anyhow::Result { analysis_enabled: config.analysis.enabled, }; let router = router(state).layer(cors_layer(&config.cors_allowed_origins)); - Ok(AppHandle { router, daemon }) + Ok(AppHandle { router, daemon, analysis_daemon }) } /// Bundle returned by [`spawn_crawler_daemon`]. The handle owns the diff --git a/backend/src/main.rs b/backend/src/main.rs index 19464dd..20fc3c4 100644 --- a/backend/src/main.rs +++ b/backend/src/main.rs @@ -21,7 +21,11 @@ async fn main() -> anyhow::Result<()> { let config = mangalord::config::Config::from_env()?; let addr: SocketAddr = config.bind_address.parse()?; - let mangalord::app::AppHandle { router, daemon } = mangalord::app::build(config).await?; + let mangalord::app::AppHandle { + router, + daemon, + analysis_daemon, + } = mangalord::app::build(config).await?; tracing::info!(%addr, "mangalord listening"); let listener = tokio::net::TcpListener::bind(addr).await?; @@ -43,6 +47,17 @@ async fn main() -> anyhow::Result<()> { ); } } + if let Some(a) = analysis_daemon { + if tokio::time::timeout(CRAWLER_SHUTDOWN_TIMEOUT, a.shutdown()) + .await + .is_err() + { + tracing::warn!( + timeout_s = CRAWLER_SHUTDOWN_TIMEOUT.as_secs(), + "analysis daemon shutdown exceeded timeout; abandoning" + ); + } + } Ok(()) } diff --git a/backend/src/repo/page.rs b/backend/src/repo/page.rs index b4d3292..fa524f5 100644 --- a/backend/src/repo/page.rs +++ b/backend/src/repo/page.rs @@ -29,6 +29,19 @@ pub async fn create<'e, E: PgExecutor<'e>>( Ok(row) } +/// Fetch a single page by id. `None` when it doesn't exist (e.g. deleted +/// between an analysis job's enqueue and its dispatch). +pub async fn find_by_id(pool: &PgPool, id: Uuid) -> AppResult> { + let row = sqlx::query_as::<_, Page>( + "SELECT id, chapter_id, page_number, storage_key, content_type \ + FROM pages WHERE id = $1", + ) + .bind(id) + .fetch_optional(pool) + .await?; + Ok(row) +} + pub async fn list_for_chapter(pool: &PgPool, chapter_id: Uuid) -> AppResult> { let rows = sqlx::query_as::<_, Page>( r#" diff --git a/backend/tests/analysis_worker.rs b/backend/tests/analysis_worker.rs new file mode 100644 index 0000000..34d5c07 --- /dev/null +++ b/backend/tests/analysis_worker.rs @@ -0,0 +1,244 @@ +//! Integration tests for the analysis worker daemon: it leases only +//! `analyze_page` jobs, acks done on success, skips already-done pages +//! (unless forced), and on terminal failure writes a `failed` row. + +use std::sync::Arc; +use std::time::Duration; + +use mangalord::analysis::daemon::{self, test_support::CountingDispatcher, AnalysisDaemonConfig}; +use mangalord::domain::page_analysis::{ + AnalysisStatus, SafetyFlag, VisionAnalysis, +}; +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_state(pool: &PgPool, page_id: Uuid) -> Option { + sqlx::query_scalar( + "SELECT state FROM crawler_jobs \ + WHERE payload->>'kind' = 'analyze_page' AND payload->>'page_id' = $1", + ) + .bind(page_id.to_string()) + .fetch_optional(pool) + .await + .unwrap() +} + +/// Poll until the page's analyze_page job reaches `want`, or fail after 5s. +async fn wait_for_state(pool: &PgPool, page_id: Uuid, want: &str) { + let deadline = Duration::from_secs(5); + let result = tokio::time::timeout(deadline, async { + loop { + if job_state(pool, page_id).await.as_deref() == Some(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, +) -> (daemon::AnalysisDaemonHandle, CancellationToken) { + let cancel = CancellationToken::new(); + let handle = daemon::spawn( + pool.clone(), + cancel.clone(), + AnalysisDaemonConfig { + dispatcher, + workers: 1, + job_timeout: Duration::from_secs(5), + }, + ); + (handle, cancel) +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_dispatches_and_acks_done(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 (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + + assert_eq!(dispatcher.call_count(), 1); +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_skips_already_done_page_unless_forced(pool: PgPool) { + let page_id = seed_page(&pool).await; + // Pre-mark the page analyzed. + page_analysis::persist_analysis( + &pool, + page_id, + &VisionAnalysis { + ocr_results: vec![], + tagging_results: vec![], + scene_description: String::new(), + safety_flag: SafetyFlag::default(), + }, + "m", + ) + .await + .unwrap(); + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + + assert_eq!( + dispatcher.call_count(), + 0, + "a non-forced job for a done page must skip dispatch" + ); +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_reanalyzes_done_page_when_forced(pool: PgPool) { + let page_id = seed_page(&pool).await; + page_analysis::persist_analysis( + &pool, + page_id, + &VisionAnalysis { + ocr_results: vec![], + tagging_results: vec![], + scene_description: String::new(), + safety_flag: SafetyFlag::default(), + }, + "m", + ) + .await + .unwrap(); + page_analysis::enqueue_for_page(&pool, page_id, true) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + + assert_eq!(dispatcher.call_count(), 1, "force must re-dispatch"); +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_marks_failed_row_on_terminal_failure(pool: PgPool) { + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + // Make the first failure terminal (no backoff wait in the test). + sqlx::query( + "UPDATE crawler_jobs SET max_attempts = 1 \ + WHERE payload->>'page_id' = $1", + ) + .bind(page_id.to_string()) + .execute(&pool) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::failing(); + let (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "dead").await; + handle.shutdown().await; + + let row = page_analysis::load(&pool, page_id).await.unwrap().unwrap(); + assert_eq!(row.status, AnalysisStatus::Failed); + assert!(row.error.is_some()); +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_isolates_dispatcher_panics(pool: PgPool) { + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + sqlx::query( + "UPDATE crawler_jobs SET max_attempts = 1 WHERE payload->>'page_id' = $1", + ) + .bind(page_id.to_string()) + .execute(&pool) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::panicking(); + let (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "dead").await; + handle.shutdown().await; + + // The worker survived the panic and recorded a failed row. + let row = page_analysis::load(&pool, page_id).await.unwrap().unwrap(); + assert_eq!(row.status, AnalysisStatus::Failed); +} + +#[sqlx::test(migrations = "./migrations")] +async fn worker_ignores_non_analyze_jobs(pool: PgPool) { + // A crawl job must be left untouched by the analysis worker. + use mangalord::crawler::jobs::{self, JobPayload}; + let crawl = jobs::enqueue( + &pool, + &JobPayload::SyncManga { + source_id: "s".into(), + source_manga_key: "k".into(), + }, + ) + .await + .unwrap(); + let crawl_id = match crawl { + jobs::EnqueueResult::Inserted(id) => id, + jobs::EnqueueResult::Skipped => panic!("expected insert"), + }; + + let page_id = seed_page(&pool).await; + page_analysis::enqueue_for_page(&pool, page_id, false) + .await + .unwrap(); + + let dispatcher = CountingDispatcher::ok(); + let (handle, _c) = spawn_with(&pool, dispatcher.clone()); + wait_for_state(&pool, page_id, "done").await; + handle.shutdown().await; + + let crawl_state: String = + sqlx::query_scalar("SELECT state FROM crawler_jobs WHERE id = $1") + .bind(crawl_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(crawl_state, "pending", "crawl job must be left for the crawler"); +} diff --git a/frontend/package.json b/frontend/package.json index 0890197..08833b1 100644 --- a/frontend/package.json +++ b/frontend/package.json @@ -1,6 +1,6 @@ { "name": "mangalord-frontend", - "version": "0.67.0", + "version": "0.68.0", "private": true, "type": "module", "scripts": {