Files
papercrate/backend/src/workers/thumbnails.rs
T
2025-10-21 21:24:04 +02:00

735 lines
24 KiB
Rust

use std::{convert::TryInto, io::Cursor, panic, sync::Arc, time::Duration};
use async_trait::async_trait;
use chrono::Utc;
use diesel::{pg::upsert::excluded, prelude::*};
use image::{GenericImageView, ImageFormat, ImageReader};
use pdfium_render::prelude::*;
use serde::Deserialize;
use serde_json::{json, Map, Value};
use tokio::task;
use tracing::{error, info, warn};
use uuid::Uuid;
use crate::{
jobs::JOB_GENERATE_THUMBNAILS,
models::{
Document, DocumentAsset, DocumentAssetObject, DocumentVersion, NewDocumentAsset,
NewDocumentAssetObject,
},
schema::{document_asset_objects, document_assets, document_versions, documents},
state::AppState,
utils::storage_paths::document_asset_object_key,
};
use super::{analyze::determine_thumbnail_support, JobExecution, JobHandler};
const THUMBNAIL_WIDTH: u32 = 512;
const THUMBNAIL_HEIGHT: u32 = 512;
const PREVIEW_WIDTH: u32 = THUMBNAIL_WIDTH * 4;
const PREVIEW_HEIGHT: u32 = THUMBNAIL_HEIGHT * 4;
const THUMBNAIL_ASSET_TYPE: &str = "thumbnail";
const PREVIEW_ASSET_TYPE: &str = "preview";
#[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_preview_and_thumbnail(&initial.document, &bytes) {
Ok(result) => result,
Err(err) => {
return JobExecution::Failed { error: err };
}
};
if let Some(page_count) = generation.page_count {
let state_clone = state.clone();
let document_id = initial.document.id;
let version_id = initial.version.id;
match task::spawn_blocking(move || {
persist_document_page_count(state_clone, document_id, version_id, page_count)
})
.await
{
Ok(Ok(())) => {}
Ok(Err(err)) => {
warn!(
job_id = %job.id,
document_id = %document_id,
version_id = %version_id,
error = %err,
"failed to update document page count metadata; retrying"
);
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err,
};
}
Err(join_err) => {
error!(
job_id = %job.id,
document_id = %document_id,
version_id = %version_id,
error = %join_err,
"page count metadata task panicked"
);
return JobExecution::Retry {
delay: Duration::from_secs(60),
error: format!("metadata panic: {join_err}"),
};
}
}
}
if initial.existing_preview.is_some() {
for object in &initial.existing_preview_objects {
if let Err(err) = state.storage.delete_object(&object.s3_key).await {
warn!(
job_id = %job.id,
error = %err,
s3_key = %object.s3_key,
"failed to delete existing preview object"
);
}
}
}
if initial.existing_thumbnail.is_some() {
for object in &initial.existing_thumbnail_objects {
if let Err(err) = state.storage.delete_object(&object.s3_key).await {
warn!(
job_id = %job.id,
error = %err,
s3_key = %object.s3_key,
"failed to delete existing thumbnail object"
);
}
}
}
let preview_asset_id = Uuid::new_v4();
let thumbnail_asset_id = Uuid::new_v4();
let mut preview_objects: Vec<AssetObjectPersistence> =
Vec::with_capacity(generation.preview.objects.len());
for (index, image) in generation.preview.objects.iter().enumerate() {
if index + 1 > i32::MAX as usize {
return JobExecution::Failed {
error: "too many preview objects".to_string(),
};
}
let ordinal = (index + 1) as i32;
let s3_key = document_asset_object_key(
initial.document.id,
initial.version.version_number,
PREVIEW_ASSET_TYPE,
preview_asset_id,
ordinal,
);
if let Err(err) = state
.storage
.put_object(
&s3_key,
image.image_bytes.clone(),
Some("image/png".into()),
None,
)
.await
{
warn!(job_id = %job.id, error = %err, ordinal, "failed to upload preview; retrying");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
preview_objects.push(AssetObjectPersistence {
ordinal,
s3_key,
width: image.width,
height: image.height,
});
}
let mut thumbnail_objects: Vec<AssetObjectPersistence> =
Vec::with_capacity(generation.thumbnail.objects.len());
for (index, image) in generation.thumbnail.objects.iter().enumerate() {
if index + 1 > i32::MAX as usize {
return JobExecution::Failed {
error: "too many thumbnail objects".to_string(),
};
}
let ordinal = (index + 1) as i32;
let s3_key = document_asset_object_key(
initial.document.id,
initial.version.version_number,
THUMBNAIL_ASSET_TYPE,
thumbnail_asset_id,
ordinal,
);
if let Err(err) = state
.storage
.put_object(
&s3_key,
image.image_bytes.clone(),
Some("image/png".into()),
None,
)
.await
{
warn!(job_id = %job.id, error = %err, ordinal, "failed to upload thumbnail; retrying");
return JobExecution::Retry {
delay: Duration::from_secs(30),
error: err.to_string(),
};
}
thumbnail_objects.push(AssetObjectPersistence {
ordinal,
s3_key,
width: image.width,
height: image.height,
});
}
let asset_persistences = vec![
AssetPersistence {
asset_type: PREVIEW_ASSET_TYPE,
asset_id: preview_asset_id,
objects: preview_objects,
},
AssetPersistence {
asset_type: THUMBNAIL_ASSET_TYPE,
asset_id: thumbnail_asset_id,
objects: thumbnail_objects,
},
];
let state_clone = state.clone();
match task::spawn_blocking(move || {
persist_assets_metadata(state_clone, &initial, &asset_persistences)
})
.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,
existing_thumbnail: Option<DocumentAsset>,
existing_thumbnail_objects: Vec<DocumentAssetObject>,
existing_preview: Option<DocumentAsset>,
existing_preview_objects: Vec<DocumentAssetObject>,
skip: bool,
}
struct GeneratedImage {
image_bytes: Vec<u8>,
width: Option<i32>,
height: Option<i32>,
}
struct GeneratedAsset {
objects: Vec<GeneratedImage>,
}
struct GeneratedAssets {
thumbnail: GeneratedAsset,
preview: GeneratedAsset,
page_count: Option<u32>,
}
struct AssetObjectPersistence {
ordinal: i32,
s3_key: String,
width: Option<i32>,
height: Option<i32>,
}
struct AssetPersistence {
asset_type: &'static str,
asset_id: Uuid,
objects: Vec<AssetObjectPersistence>,
}
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_assets: Vec<DocumentAsset> = document_assets::table
.filter(document_assets::document_version_id.eq(payload.document_version_id))
.filter(document_assets::asset_type.eq_any(vec![
THUMBNAIL_ASSET_TYPE.to_string(),
PREVIEW_ASSET_TYPE.to_string(),
]))
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let mut existing_thumbnail = None;
let mut existing_thumbnail_objects: Vec<DocumentAssetObject> = Vec::new();
let mut existing_preview = None;
let mut existing_preview_objects: Vec<DocumentAssetObject> = Vec::new();
for asset in existing_assets {
match asset.asset_type.as_str() {
THUMBNAIL_ASSET_TYPE => {
existing_thumbnail_objects = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.id))
.order(document_asset_objects::ordinal.asc())
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
existing_thumbnail = Some(asset);
}
PREVIEW_ASSET_TYPE => {
existing_preview_objects = document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.id))
.order(document_asset_objects::ordinal.asc())
.load(&mut conn)
.map_err(|err| format!("{err:?}"))?;
existing_preview = Some(asset);
}
_ => {}
}
}
let (supported, _) = determine_thumbnail_support(&document);
if !supported {
return Err("thumbnail generation not supported for this document".into());
}
let expected_cardinality = expected_asset_cardinality(&document, &version);
let preview_cardinality = existing_preview
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_preview_objects.len() as i32);
let thumbnail_cardinality = existing_thumbnail
.as_ref()
.and_then(|asset| asset.cardinality)
.unwrap_or_else(|| existing_thumbnail_objects.len() as i32);
let needs_regeneration = preview_cardinality < expected_cardinality
|| thumbnail_cardinality < expected_cardinality
|| (existing_preview_objects.len() as i32) < expected_cardinality
|| (existing_thumbnail_objects.len() as i32) < expected_cardinality;
let skip = existing_thumbnail.is_some()
&& existing_preview.is_some()
&& !payload.force
&& !needs_regeneration;
Ok(ThumbnailContext {
document,
version,
existing_thumbnail,
existing_thumbnail_objects,
existing_preview,
existing_preview_objects,
skip,
})
}
fn generate_preview_and_thumbnail(
document: &Document,
bytes: &[u8],
) -> Result<GeneratedAssets, String> {
let is_pdf = document_is_pdf(document);
if is_pdf {
let pdf_assets = generate_pdf_assets(bytes)?;
Ok(GeneratedAssets {
preview: pdf_assets.preview,
thumbnail: pdf_assets.thumbnail,
page_count: Some(pdf_assets.page_count),
})
} else {
let (preview, thumbnail) = generate_image_assets(bytes)?;
Ok(GeneratedAssets {
preview,
thumbnail,
page_count: None,
})
}
}
fn generate_image_assets(bytes: &[u8]) -> Result<(GeneratedAsset, GeneratedAsset), String> {
let reader = ImageReader::new(Cursor::new(bytes))
.with_guessed_format()
.map_err(|err| err.to_string())?;
let image = reader.decode().map_err(|err| err.to_string())?;
let preview_image = if image.width() > PREVIEW_WIDTH || image.height() > PREVIEW_HEIGHT {
image.thumbnail(PREVIEW_WIDTH, PREVIEW_HEIGHT)
} else {
image.clone()
};
let thumbnail_image =
if preview_image.width() > THUMBNAIL_WIDTH || preview_image.height() > THUMBNAIL_HEIGHT {
preview_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
} else {
preview_image.clone()
};
let preview = encode_dynamic_image(preview_image)?;
let thumbnail = encode_dynamic_image(thumbnail_image)?;
Ok((
GeneratedAsset {
objects: vec![preview],
},
GeneratedAsset {
objects: vec![thumbnail],
},
))
}
struct PdfGeneratedAssets {
preview: GeneratedAsset,
thumbnail: GeneratedAsset,
page_count: u32,
}
fn generate_pdf_assets(bytes: &[u8]) -> Result<PdfGeneratedAssets, 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 pages = document.pages();
let total_pages = pages.len() as usize;
let render_config = PdfRenderConfig::new()
.set_target_width(PREVIEW_WIDTH as i32)
.set_maximum_height(PREVIEW_HEIGHT as i32)
.render_form_data(true)
.rotate_if_landscape(PdfPageRenderRotation::None, true);
let mut preview_objects: Vec<GeneratedImage> = Vec::with_capacity(total_pages);
let mut thumbnail_objects: Vec<GeneratedImage> = Vec::with_capacity(total_pages);
for page_index in 0..total_pages {
let page = pages
.get(u16::try_from(page_index).map_err(|_| "page index overflow".to_string())?)
.map_err(|err| format!("load page {page_index}: {err}"))?;
let bitmap = page
.render_with_config(&render_config)
.map_err(|err| format!("render pdf page {page_index}: {err}"))?;
let preview_buffer = bitmap.as_image().to_rgb8();
let preview_image = image::DynamicImage::ImageRgb8(preview_buffer);
let thumbnail_image = if preview_image.width() > THUMBNAIL_WIDTH
|| preview_image.height() > THUMBNAIL_HEIGHT
{
preview_image.thumbnail(THUMBNAIL_WIDTH, THUMBNAIL_HEIGHT)
} else {
preview_image.clone()
};
preview_objects.push(encode_dynamic_image(preview_image)?);
thumbnail_objects.push(encode_dynamic_image(thumbnail_image)?);
}
let page_count: u32 = total_pages
.try_into()
.map_err(|_| "page count exceeds supported range".to_string())?;
Ok(PdfGeneratedAssets {
preview: GeneratedAsset {
objects: preview_objects,
},
thumbnail: GeneratedAsset {
objects: thumbnail_objects,
},
page_count,
})
}
fn encode_dynamic_image(image: image::DynamicImage) -> Result<GeneratedImage, String> {
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())?;
Ok(GeneratedImage {
image_bytes: cursor.into_inner(),
width: Some(width as i32),
height: Some(height as i32),
})
}
fn persist_assets_metadata(
state: Arc<AppState>,
context: &ThumbnailContext,
assets: &[AssetPersistence],
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
if let Some(existing_preview) = &context.existing_preview {
diesel::delete(document_assets::table.filter(document_assets::id.eq(existing_preview.id)))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
if let Some(existing_thumbnail) = &context.existing_thumbnail {
diesel::delete(
document_assets::table.filter(document_assets::id.eq(existing_thumbnail.id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
for asset in assets {
if asset.objects.is_empty() {
return Err(format!(
"asset {} has no generated objects",
asset.asset_type
));
}
let object_count: i32 = asset
.objects
.len()
.try_into()
.map_err(|_| "asset contains too many objects".to_string())?;
let new_asset = NewDocumentAsset {
id: asset.asset_id,
document_version_id: context.version.id,
asset_type: asset.asset_type.to_string(),
mime_type: "image/png".to_string(),
metadata: json!({
"generated_at": Utc::now().to_rfc3339(),
}),
cardinality: Some(object_count),
};
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::mime_type.eq(excluded(document_assets::mime_type)),
document_assets::metadata.eq(excluded(document_assets::metadata)),
document_assets::cardinality.eq(excluded(document_assets::cardinality)),
))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
diesel::delete(
document_asset_objects::table
.filter(document_asset_objects::asset_id.eq(asset.asset_id)),
)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
for object in &asset.objects {
let mut metadata_map = Map::new();
if let Some(width) = object.width {
metadata_map.insert("width".to_string(), Value::from(width));
}
if let Some(height) = object.height {
metadata_map.insert("height".to_string(), Value::from(height));
}
let object_metadata = Value::Object(metadata_map);
let new_object = NewDocumentAssetObject {
id: Uuid::new_v4(),
asset_id: asset.asset_id,
ordinal: object.ordinal,
s3_key: object.s3_key.clone(),
metadata: object_metadata,
};
diesel::insert_into(document_asset_objects::table)
.values(&new_object)
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
}
}
Ok(())
}
fn persist_document_page_count(
state: Arc<AppState>,
document_id: Uuid,
document_version_id: Uuid,
page_count: u32,
) -> Result<(), String> {
let mut conn = state.db().map_err(|err| format!("{err:?}"))?;
let existing_metadata: Value = document_versions::table
.filter(document_versions::id.eq(document_version_id))
.filter(document_versions::document_id.eq(document_id))
.select(document_versions::metadata)
.first(&mut conn)
.map_err(|err| format!("{err:?}"))?;
let updated = match existing_metadata {
Value::Object(mut map) => {
map.insert("page_count".to_string(), Value::from(page_count));
Value::Object(map)
}
_ => {
let mut map = Map::new();
map.insert("page_count".to_string(), Value::from(page_count));
Value::Object(map)
}
};
diesel::update(
document_versions::table
.filter(document_versions::id.eq(document_version_id))
.filter(document_versions::document_id.eq(document_id)),
)
.set(document_versions::metadata.eq(updated))
.execute(&mut conn)
.map_err(|err| format!("{err:?}"))?;
Ok(())
}
fn document_is_pdf(document: &Document) -> bool {
document
.content_type
.as_deref()
.map(|mime| mime.eq_ignore_ascii_case("application/pdf"))
.unwrap_or_else(|| {
document
.original_name
.rsplit('.')
.next()
.map(|ext| ext.eq_ignore_ascii_case("pdf"))
.unwrap_or(false)
})
}
fn expected_asset_cardinality(document: &Document, version: &DocumentVersion) -> i32 {
if let Value::Object(map) = &version.metadata {
if let Some(count) = map.get("page_count").and_then(|v| v.as_i64()) {
if count > 0 {
return count
.min(i64::from(i32::MAX))
.try_into()
.unwrap_or(i32::MAX);
}
}
}
if document_is_pdf(document) {
1
} else {
1
}
}