2

backend/src/workers/analyze.rs

@@ -0,0 +1,165 @@
+use std::{collections::HashSet, sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use diesel::prelude::*;
+use serde::Deserialize;
+use serde_json::{json, Map, Value};
+use tokio::task;
+use tracing::{error, warn};
+use uuid::Uuid;
+
+use crate::{
+    jobs::{enqueue_job, JOB_ANALYZE_DOCUMENT, JOB_GENERATE_THUMBNAILS},
+    models::{Document, DocumentVersion},
+    schema::{document_versions, documents},
+    state::AppState,
+};
+
+use super::{JobExecution, JobHandler};
+
+#[derive(Debug, Deserialize)]
+struct AnalyzePayload {
+    document_id: Uuid,
+    document_version_id: Uuid,
+    #[serde(default)]
+    force: bool,
+}
+
+pub struct AnalyzeDocumentJob;
+
+impl AnalyzeDocumentJob {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl JobHandler for AnalyzeDocumentJob {
+    fn job_type(&self) -> &'static str {
+        JOB_ANALYZE_DOCUMENT
+    }
+
+    async fn handle(&self, state: Arc<AppState>, job: crate::models::Job) -> JobExecution {
+        let payload: AnalyzePayload = match serde_json::from_value(job.payload.clone()) {
+            Ok(payload) => payload,
+            Err(err) => {
+                return JobExecution::Failed {
+                    error: format!("invalid analyze payload: {err}"),
+                }
+            }
+        };
+
+        let state_clone = state.clone();
+        match task::spawn_blocking(move || analyze_document(state_clone, payload)).await {
+            Ok(Ok(execution)) => execution,
+            Ok(Err(err)) => {
+                warn!(job_id = %job.id, error = %err, "analyze job will retry");
+                JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err,
+                }
+            }
+            Err(join_err) => {
+                error!(job_id = %job.id, error = %join_err, "analyze task panicked");
+                JobExecution::Retry {
+                    delay: Duration::from_secs(60),
+                    error: format!("worker panicked: {join_err}"),
+                }
+            }
+        }
+    }
+}
+
+fn analyze_document(state: Arc<AppState>, payload: AnalyzePayload) -> Result<JobExecution, String> {
+    let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
+
+    let version: DocumentVersion = document_versions::table
+        .find(payload.document_version_id)
+        .first(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    if version.document_id != payload.document_id {
+        return Err("document/version mismatch".into());
+    }
+
+    let document: Document = documents::table
+        .find(payload.document_id)
+        .first(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    let (supported, reason) = determine_thumbnail_support(&document);
+
+    let mut summary_map = match version.operations_summary {
+        Value::Object(map) => map,
+        _ => Map::new(),
+    };
+    summary_map.insert("thumbnail_supported".to_string(), Value::Bool(supported));
+    if let Some(reason) = reason {
+        summary_map.insert("thumbnail_reason".to_string(), Value::String(reason));
+    } else {
+        summary_map.remove("thumbnail_reason");
+    }
+
+    diesel::update(document_versions::table.find(version.id))
+        .set(document_versions::operations_summary.eq(Value::Object(summary_map)))
+        .execute(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    if supported {
+        let enqueue_result = enqueue_job(
+            &mut conn,
+            JOB_GENERATE_THUMBNAILS,
+            json!({
+                "document_id": payload.document_id,
+                "document_version_id": payload.document_version_id,
+                "force": payload.force,
+            }),
+            None,
+        );
+
+        if let Err(err) = enqueue_result {
+            return Err(err.to_string());
+        }
+    }
+
+    Ok(JobExecution::Success)
+}
+
+pub(crate) fn determine_thumbnail_support(document: &Document) -> (bool, Option<String>) {
+    let supported_mimes: HashSet<&'static str> = [
+        "image/jpeg",
+        "image/png",
+        "image/gif",
+        "image/tiff",
+        "image/bmp",
+        "image/webp",
+        "application/pdf",
+    ]
+    .into_iter()
+    .collect();
+
+    if let Some(ref content_type) = document.content_type {
+        if supported_mimes.contains(content_type.as_str()) {
+            return (true, None);
+        }
+    }
+
+    if let Some(ext) = document
+        .original_name
+        .rsplit('.')
+        .next()
+        .map(|ext| ext.to_ascii_lowercase())
+    {
+        let supported_exts = [
+            "jpg", "jpeg", "png", "gif", "tif", "tiff", "bmp", "webp", "pdf",
+        ];
+        if supported_exts.contains(&ext.as_str()) {
+            return (true, None);
+        }
+    }
+
+    (
+        false,
+        Some("content type not supported for thumbnails".into()),
+    )
+}

backend/src/workers/mod.rs

@@ -0,0 +1,131 @@
+use std::{collections::HashMap, sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use tokio::time::sleep;
+use tracing::{error, info, warn};
+
+use crate::{
+    jobs::{mark_job_failed, mark_job_succeeded, reserve_job, retry_job_after, JobQueueError},
+    models::Job,
+    state::AppState,
+};
+
+pub mod analyze;
+pub mod thumbnails;
+
+#[derive(Debug)]
+pub enum JobExecution {
+    Success,
+    Retry { delay: Duration, error: String },
+    Failed { error: String },
+}
+
+#[async_trait]
+pub trait JobHandler: Send + Sync {
+    fn job_type(&self) -> &'static str;
+    async fn handle(&self, state: Arc<AppState>, job: Job) -> JobExecution;
+}
+
+pub struct Worker {
+    state: Arc<AppState>,
+    handlers: HashMap<&'static str, Arc<dyn JobHandler>>,
+    poll_interval: Duration,
+}
+
+impl Worker {
+    pub fn new(
+        state: Arc<AppState>,
+        handlers: Vec<Arc<dyn JobHandler>>,
+        poll_interval: Duration,
+    ) -> Self {
+        let map = handlers
+            .into_iter()
+            .map(|handler| (handler.job_type(), handler))
+            .collect();
+        Self {
+            state,
+            handlers: map,
+            poll_interval,
+        }
+    }
+
+    pub async fn run(&self) {
+        info!("worker started");
+        loop {
+            match self.tick().await {
+                Ok(true) => {}
+                Ok(false) => sleep(self.poll_interval).await,
+                Err(err) => {
+                    error!(error = %err, "worker tick failed");
+                    sleep(self.poll_interval).await;
+                }
+            }
+        }
+    }
+
+    async fn tick(&self) -> Result<bool, JobQueueError> {
+        let job_types: Vec<&str> = self.handlers.keys().copied().collect();
+        if job_types.is_empty() {
+            return Ok(false);
+        }
+
+        let mut conn = match self.state.db() {
+            Ok(conn) => conn,
+            Err(err) => {
+                error!(?err, "failed to obtain database connection in worker");
+                return Ok(false);
+            }
+        };
+
+        let job_opt = reserve_job(&mut conn, &job_types)?;
+        drop(conn);
+
+        if let Some(job) = job_opt {
+            if let Some(handler) = self.handlers.get(job.job_type.as_str()) {
+                let result = handler.handle(self.state.clone(), job.clone()).await;
+                match result {
+                    JobExecution::Success => {
+                        if let Ok(mut conn) = self.state.db() {
+                            mark_job_succeeded(&mut conn, job.id)?;
+                        } else {
+                            error!("failed to mark job succeeded due to pool error");
+                        }
+                    }
+                    JobExecution::Retry { delay, error } => {
+                        warn!(job_id = %job.id, job_type = %job.job_type, %error, "job will retry");
+                        if let Ok(mut conn) = self.state.db() {
+                            retry_job_after(&mut conn, job.id, delay, &error)?;
+                        } else {
+                            error!("failed to requeue job for retry due to pool error");
+                        }
+                    }
+                    JobExecution::Failed { error } => {
+                        error!(job_id = %job.id, job_type = %job.job_type, %error, "job failed");
+                        if let Ok(mut conn) = self.state.db() {
+                            mark_job_failed(&mut conn, job.id, &error)?;
+                        } else {
+                            error!("failed to mark job failed due to pool error");
+                        }
+                    }
+                }
+            } else {
+                error!(job_type = %job.job_type, "no handler registered for job type");
+                if let Ok(mut conn) = self.state.db() {
+                    mark_job_failed(&mut conn, job.id, "no handler registered")?;
+                } else {
+                    error!("failed to mark job failed for missing handler due to pool error");
+                }
+            }
+            Ok(true)
+        } else {
+            Ok(false)
+        }
+    }
+}
+
+pub fn default_handlers() -> Vec<Arc<dyn JobHandler>> {
+    vec![
+        Arc::new(analyze::AnalyzeDocumentJob::new()),
+        Arc::new(thumbnails::GenerateThumbnailsJob::new()),
+    ]
+}

backend/src/workers/thumbnails.rs

@@ -0,0 +1,321 @@
+use std::{io::Cursor, panic, sync::Arc, time::Duration};
+
+use async_trait::async_trait;
+use chrono::Utc;
+use diesel::{pg::upsert::excluded, prelude::*};
+use image::{
+    codecs::png::PngEncoder, ColorType, GenericImageView, ImageEncoder, ImageFormat, ImageReader,
+};
+use pdfium_render::prelude::*;
+use serde::Deserialize;
+use serde_json::json;
+use tokio::task;
+use tracing::{error, info, warn};
+use uuid::Uuid;
+
+use crate::{
+    jobs::JOB_GENERATE_THUMBNAILS,
+    models::{Document, DocumentAsset, DocumentVersion, NewDocumentAsset},
+    schema::{document_assets, document_versions, documents},
+    state::AppState,
+};
+
+use super::{analyze::determine_thumbnail_support, JobExecution, JobHandler};
+
+const THUMBNAIL_WIDTH: u32 = 512;
+const THUMBNAIL_HEIGHT: u32 = 512;
+const THUMBNAIL_ASSET_TYPE: &str = "thumbnail";
+
+#[derive(Debug, Deserialize)]
+struct ThumbnailPayload {
+    document_id: Uuid,
+    document_version_id: Uuid,
+    #[serde(default)]
+    force: bool,
+}
+
+pub struct GenerateThumbnailsJob;
+
+impl GenerateThumbnailsJob {
+    pub fn new() -> Self {
+        Self
+    }
+}
+
+#[async_trait]
+impl JobHandler for GenerateThumbnailsJob {
+    fn job_type(&self) -> &'static str {
+        JOB_GENERATE_THUMBNAILS
+    }
+
+    async fn handle(&self, state: Arc<AppState>, job: crate::models::Job) -> JobExecution {
+        let payload: ThumbnailPayload = match serde_json::from_value(job.payload.clone()) {
+            Ok(p) => p,
+            Err(err) => {
+                return JobExecution::Failed {
+                    error: format!("invalid thumbnail payload: {err}"),
+                }
+            }
+        };
+
+        let state_clone = state.clone();
+        let initial =
+            match task::spawn_blocking(move || load_thumbnail_context(state_clone, &payload)).await
+            {
+                Ok(Ok(ctx)) => ctx,
+                Ok(Err(err)) => {
+                    warn!(job_id = %job.id, error = %err, "thumbnail job will retry");
+                    return JobExecution::Retry {
+                        delay: Duration::from_secs(30),
+                        error: err,
+                    };
+                }
+                Err(join_err) => {
+                    error!(job_id = %job.id, error = %join_err, "thumbnail task panicked");
+                    return JobExecution::Retry {
+                        delay: Duration::from_secs(60),
+                        error: format!("worker panicked: {join_err}"),
+                    };
+                }
+            };
+
+        if initial.skip {
+            info!(job_id = %job.id, "thumbnails already exist; skipping");
+            return JobExecution::Success;
+        }
+
+        let bytes = match state.storage.get_object(&initial.version.s3_key).await {
+            Ok(bytes) => bytes,
+            Err(err) => {
+                warn!(job_id = %job.id, error = %err, "thumbnail fetch failed; will retry");
+                return JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err.to_string(),
+                };
+            }
+        };
+
+        let generation = match generate_thumbnail(&initial.document, &bytes) {
+            Ok(result) => result,
+            Err(err) => {
+                return JobExecution::Failed { error: err };
+            }
+        };
+
+        if let Err(err) = state
+            .storage
+            .put_object(
+                &generation.s3_key,
+                generation.image_bytes.clone(),
+                Some("image/png".into()),
+            )
+            .await
+        {
+            warn!(job_id = %job.id, error = %err, "failed to upload thumbnail; retrying");
+            return JobExecution::Retry {
+                delay: Duration::from_secs(30),
+                error: err.to_string(),
+            };
+        }
+
+        let state_clone = state.clone();
+        match task::spawn_blocking(move || {
+            persist_thumbnail_metadata(state_clone, &initial, &generation)
+        })
+        .await
+        {
+            Ok(Ok(())) => {}
+            Ok(Err(err)) => {
+                warn!(job_id = %job.id, error = %err, "failed to persist thumbnail metadata; retrying");
+                return JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: err,
+                };
+            }
+            Err(join_err) => {
+                error!(job_id = %job.id, error = %join_err, "thumbnail metadata update panicked");
+                return JobExecution::Retry {
+                    delay: Duration::from_secs(30),
+                    error: format!("metadata update panic: {join_err}"),
+                };
+            }
+        }
+
+        JobExecution::Success
+    }
+}
+
+struct ThumbnailContext {
+    document: Document,
+    version: DocumentVersion,
+    skip: bool,
+}
+
+struct GeneratedThumbnail {
+    image_bytes: Vec<u8>,
+    width: Option<i32>,
+    height: Option<i32>,
+    s3_key: String,
+}
+
+fn load_thumbnail_context(
+    state: Arc<AppState>,
+    payload: &ThumbnailPayload,
+) -> Result<ThumbnailContext, String> {
+    let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
+
+    let version: DocumentVersion = document_versions::table
+        .find(payload.document_version_id)
+        .first(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    if version.document_id != payload.document_id {
+        return Err("document/version mismatch".into());
+    }
+
+    let document: Document = documents::table
+        .find(payload.document_id)
+        .first(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    let existing: Option<DocumentAsset> = document_assets::table
+        .filter(document_assets::document_version_id.eq(payload.document_version_id))
+        .filter(document_assets::asset_type.eq(THUMBNAIL_ASSET_TYPE))
+        .first(&mut conn)
+        .optional()
+        .map_err(|err| format!("{err:?}"))?;
+
+    let (supported, _) = determine_thumbnail_support(&document);
+    if !supported {
+        return Err("thumbnail generation not supported for this document".into());
+    }
+
+    let skip = existing.is_some() && !payload.force;
+
+    Ok(ThumbnailContext {
+        document,
+        version,
+        skip,
+    })
+}
+
+fn generate_thumbnail(document: &Document, bytes: &[u8]) -> Result<GeneratedThumbnail, String> {
+    let is_pdf = document
+        .content_type
+        .as_deref()
+        .map(|mime| mime == "application/pdf")
+        .unwrap_or_else(|| {
+            document
+                .original_name
+                .rsplit('.')
+                .next()
+                .map(|ext| ext.eq_ignore_ascii_case("pdf"))
+                .unwrap_or(false)
+        });
+
+    let (png_bytes, width, height) = if is_pdf {
+        generate_pdf_thumbnail(bytes)?
+    } else {
+        generate_image_thumbnail(bytes)?
+    };
+
+    let s3_key = format!("thumbnails/{}/{}.png", document.id, Uuid::new_v4());
+
+    Ok(GeneratedThumbnail {
+        image_bytes: png_bytes,
+        width,
+        height,
+        s3_key,
+    })
+}
+
+fn generate_image_thumbnail(bytes: &[u8]) -> Result<(Vec<u8>, Option<i32>, Option<i32>), String> {
+    let reader = ImageReader::new(Cursor::new(bytes))
+        .with_guessed_format()
+        .map_err(|err| err.to_string())?;
+    let mut image = reader.decode().map_err(|err| err.to_string())?;
+    if image.width() > THUMBNAIL_WIDTH || image.height() > THUMBNAIL_HEIGHT {
+        image = image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT);
+    }
+
+    let (width, height) = image.dimensions();
+    let mut cursor = Cursor::new(Vec::new());
+    image
+        .write_to(&mut cursor, ImageFormat::Png)
+        .map_err(|err| err.to_string())?;
+    let buffer = cursor.into_inner();
+    Ok((buffer, Some(width as i32), Some(height as i32)))
+}
+
+fn generate_pdf_thumbnail(bytes: &[u8]) -> Result<(Vec<u8>, Option<i32>, Option<i32>), String> {
+    let pdfium = panic::catch_unwind(|| Pdfium::default())
+        .map_err(|_| "failed to initialize PDFium".to_string())?;
+
+    let document = pdfium
+        .load_pdf_from_byte_slice(bytes, None)
+        .map_err(|err| format!("load pdf: {err}"))?;
+
+    let page = document
+        .pages()
+        .get(0)
+        .map_err(|err| format!("load first page: {err}"))?;
+
+    let render_config = PdfRenderConfig::new()
+        .set_target_width(THUMBNAIL_WIDTH as i32)
+        .set_maximum_height(THUMBNAIL_HEIGHT as i32)
+        .render_form_data(true)
+        .rotate_if_landscape(PdfPageRenderRotation::None, true);
+
+    let bitmap = page
+        .render_with_config(&render_config)
+        .map_err(|err| format!("render pdf page: {err}"))?;
+
+    let image = bitmap.as_image().to_rgb8();
+    let (width, height) = image.dimensions();
+    let mut cursor = Cursor::new(Vec::new());
+    PngEncoder::new(&mut cursor)
+        .write_image(image.as_raw(), width, height, ColorType::Rgb8.into())
+        .map_err(|err| format!("encode pdf thumbnail: {err}"))?;
+
+    Ok((cursor.into_inner(), Some(width as i32), Some(height as i32)))
+}
+
+fn persist_thumbnail_metadata(
+    state: Arc<AppState>,
+    context: &ThumbnailContext,
+    generated: &GeneratedThumbnail,
+) -> Result<(), String> {
+    let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
+
+    let new_asset = NewDocumentAsset {
+        id: Uuid::new_v4(),
+        document_version_id: context.version.id,
+        asset_type: THUMBNAIL_ASSET_TYPE.to_string(),
+        s3_key: generated.s3_key.clone(),
+        mime_type: "image/png".to_string(),
+        width: generated.width,
+        height: generated.height,
+        metadata: json!({
+            "generated_at": Utc::now().to_rfc3339(),
+        }),
+    };
+
+    diesel::insert_into(document_assets::table)
+        .values(&new_asset)
+        .on_conflict((
+            document_assets::document_version_id,
+            document_assets::asset_type,
+        ))
+        .do_update()
+        .set((
+            document_assets::s3_key.eq(excluded(document_assets::s3_key)),
+            document_assets::mime_type.eq(excluded(document_assets::mime_type)),
+            document_assets::width.eq(excluded(document_assets::width)),
+            document_assets::height.eq(excluded(document_assets::height)),
+            document_assets::metadata.eq(excluded(document_assets::metadata)),
+        ))
+        .execute(&mut conn)
+        .map_err(|err| format!("{err:?}"))?;
+
+    Ok(())
+}