perf: thumbnail semaphore, WOPI streaming, store_bytes guard, spawn_blocking SHA-256

- Issue #1: Add Semaphore(4) + 50MP resolution guard to thumbnail_service
  Bounds peak RAM from 4.8GB (50 uploads) to 384MB (4 concurrent decodes)
- Issue #3: Migrate WOPI PutFile from Bytes to streaming temp file + SHA-256
  RAM per WOPI PUT: ~100MB → ~256KB regardless of file size
- Issue #3: Add 10MB guard in dedup store_bytes (defense-in-depth)
- Issue #5: Move chunked upload assembly to spawn_blocking (sync I/O)
  Frees Tokio workers during SHA-256 hashing (~130ms for 500MB)
- Clean up unused tokio imports (OpenOptions, BufWriter)
This commit is contained in:
Dionisio
2026-02-24 16:11:52 +01:00
parent 71c2cb5edb
commit 6aa38d0d24
3 changed files with 154 additions and 11 deletions
@@ -155,15 +155,33 @@ impl DedupService {
// ── Core store operations ────────────────────────────────────
/// Maximum payload accepted by `store_bytes`. Anything larger
/// should use `store_from_file` (streaming — constant RAM).
const MAX_STORE_BYTES: usize = 10 * 1024 * 1024; // 10 MB
/// Store content with deduplication (from bytes).
///
/// Uses `SELECT … FOR UPDATE` + `INSERT … ON CONFLICT` for atomic
/// upsert — completely TOCTOU-free.
///
/// **Guard**: rejects payloads >10 MB. Large content must go through
/// `store_from_file` which streams from disk with constant RAM.
pub async fn store_bytes(
&self,
content: &[u8],
content_type: Option<String>,
) -> Result<DedupResultDto, DomainError> {
if content.len() > Self::MAX_STORE_BYTES {
return Err(DomainError::internal_error(
"Dedup",
format!(
"store_bytes called with {} bytes (max {}). Use store_from_file for large content.",
content.len(),
Self::MAX_STORE_BYTES
),
));
}
let size = content.len() as u64;
let hash = Self::hash_bytes(content);
@@ -16,6 +16,7 @@ use image::{ImageFormat, imageops::FilterType};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::fs;
use tokio::sync::Semaphore;
use crate::application::ports::thumbnail_ports::{
ThumbnailPort, ThumbnailSize as PortThumbnailSize, ThumbnailStatsDto,
@@ -69,6 +70,14 @@ struct ThumbnailCacheKey {
size: ThumbnailSize,
}
/// Maximum pixel count before rejecting decode (50 megapixels → ~200 MB RGBA).
/// Images above this are silently skipped — protects against single-image OOM.
const MAX_DECODE_PIXELS: u64 = 50_000_000;
/// Default max concurrent thumbnail decode operations.
/// 4 × 96 MB (6000×4000 RGBA) = 384 MB worst-case peak.
const DEFAULT_MAX_CONCURRENT_DECODES: usize = 4;
/// Thumbnail service for generating and caching image thumbnails
pub struct ThumbnailService {
/// Root path for thumbnail storage
@@ -77,6 +86,10 @@ pub struct ThumbnailService {
cache: moka::future::Cache<ThumbnailCacheKey, Bytes>,
/// Configured maximum cache weight (for stats reporting)
max_cache_bytes: u64,
/// Limits how many images are decoded in parallel to bound RAM usage.
/// Without this, 50 simultaneous uploads would decode 50 bitmaps
/// (~96 MB each for 6000×4000) = 4.8 GB peak.
decode_semaphore: Arc<Semaphore>,
}
impl ThumbnailService {
@@ -105,6 +118,7 @@ impl ThumbnailService {
thumbnails_root,
cache,
max_cache_bytes: max_cache_bytes as u64,
decode_semaphore: Arc::new(Semaphore::new(DEFAULT_MAX_CONCURRENT_DECODES)),
}
}
@@ -217,7 +231,11 @@ impl ThumbnailService {
Ok(bytes)
}
/// Generate a thumbnail from an image file
/// Generate a thumbnail from an image file.
///
/// Concurrency is bounded by `decode_semaphore` to prevent OOM when
/// many images are uploaded simultaneously. Resolution is also
/// capped at `MAX_DECODE_PIXELS` to reject pathologically large images.
async fn generate_thumbnail(
&self,
original_path: &Path,
@@ -226,9 +244,25 @@ impl ThumbnailService {
let path = original_path.to_path_buf();
let max_dim = size.max_dimension();
// Acquire semaphore permit — bounds peak RAM from concurrent decodes
let _permit = self.decode_semaphore.acquire().await
.map_err(|_| ThumbnailError::TaskError("Decode semaphore closed".into()))?;
// Run image processing in blocking thread pool
let result = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ThumbnailError> {
// Load image
// Safety check: read dimensions from headers only (no full decode)
let (w, h) = image::ImageReader::open(&path)
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?
.into_dimensions()
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
if (w as u64) * (h as u64) > MAX_DECODE_PIXELS {
return Err(ThumbnailError::ImageError(format!(
"Image too large for thumbnail: {w}×{h} ({} MP, max {MAX_DECODE_PIXELS})",
w as u64 * h as u64 / 1_000_000
)));
}
// Load image (full decode — now safe within semaphore + resolution guard)
let img = image::open(&path).map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
// Calculate new dimensions preserving aspect ratio
@@ -274,10 +308,31 @@ impl ThumbnailService {
tokio::spawn(async move {
tracing::info!("🖼️ Background thumbnail generation starting: {}", file_id);
// Acquire semaphore permit — bounds peak RAM from concurrent decodes
let _permit = match self.decode_semaphore.acquire().await {
Ok(p) => p,
Err(_) => {
tracing::warn!("Decode semaphore closed, skipping thumbnails for {}", file_id);
return;
}
};
let path = original_path.clone();
// Single spawn_blocking: 1 read + 1 decode + 3 resize + 3 encode
let results = tokio::task::spawn_blocking(move || {
// Safety check: read dimensions from headers only (no full decode)
let (w, h) = image::ImageReader::open(&path)
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?
.into_dimensions()
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
if (w as u64) * (h as u64) > MAX_DECODE_PIXELS {
return Err(ThumbnailError::ImageError(format!(
"Image too large for thumbnail: {w}×{h} ({} MP, max {MAX_DECODE_PIXELS})",
w as u64 * h as u64 / 1_000_000
)));
}
let img = image::open(&path)
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
+79 -9
View File
@@ -11,9 +11,9 @@
use crate::interfaces::middleware::auth::AuthUser;
use axum::{
Router,
body::Bytes,
body::Body,
extract::{Path, Query, State},
http::{HeaderMap, StatusCode},
http::{HeaderMap, StatusCode, Request},
response::{Html, IntoResponse, Response},
routing::{get, post},
};
@@ -153,13 +153,22 @@ async fn get_file(
}
/// POST /wopi/files/{file_id}/contents — PutFile
///
/// **Streaming implementation**: the request body is spooled to a temp file
/// with incremental SHA-256 hashing. Peak RAM usage is ~256 KB regardless
/// of file size (previously buffered the entire body as `Bytes`).
async fn put_file(
Path(file_id): Path<String>,
Query(token_query): Query<WopiTokenQuery>,
headers: HeaderMap,
State(state): State<WopiState>,
body: Bytes,
req: Request<Body>,
) -> Response {
use http_body_util::BodyStream;
use sha2::{Digest, Sha256};
use tokio::io::AsyncWriteExt;
use tokio_stream::StreamExt;
let claims = match state
.token_service
.validate_token(&token_query.access_token)
@@ -197,7 +206,7 @@ async fn put_file(
}
}
// Get file path for update_file
// Get file metadata for the path
let file = match state
.app_state
.applications
@@ -209,14 +218,75 @@ async fn put_file(
Err(_) => return StatusCode::NOT_FOUND.into_response(),
};
// Save the file content using path-based update
match state
// ── Streaming spool: body → temp file + incremental SHA-256 ──
let temp_file = match tempfile::NamedTempFile::new() {
Ok(f) => f,
Err(e) => {
tracing::error!("WOPI PutFile: failed to create temp file: {}", e);
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
};
let temp_path = temp_file.path().to_path_buf();
let mut file_out = match tokio::fs::File::create(&temp_path).await {
Ok(f) => f,
Err(e) => {
tracing::error!("WOPI PutFile: failed to open temp file: {}", e);
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
};
let content_type = file.mime_type.clone();
let mut hasher = Sha256::new();
let mut total_bytes: u64 = 0;
let mut stream = BodyStream::new(req.into_body());
while let Some(frame_result) = stream.next().await {
let frame = match frame_result {
Ok(f) => f,
Err(e) => {
let _ = tokio::fs::remove_file(&temp_path).await;
tracing::error!("WOPI PutFile: body read error: {}", e);
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
};
if let Some(chunk) = frame.data_ref() {
total_bytes += chunk.len() as u64;
hasher.update(chunk);
if let Err(e) = file_out.write_all(chunk).await {
let _ = tokio::fs::remove_file(&temp_path).await;
tracing::error!("WOPI PutFile: temp write error: {}", e);
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
}
}
if let Err(e) = file_out.flush().await {
let _ = tokio::fs::remove_file(&temp_path).await;
tracing::error!("WOPI PutFile: flush error: {}", e);
return StatusCode::INTERNAL_SERVER_ERROR.into_response();
}
drop(file_out);
let hash = hex::encode(hasher.finalize());
// ── Atomic store: temp file → dedup blob + DB metadata update ──
let result = state
.app_state
.applications
.file_upload_service
.update_file(&file.path, &body)
.await
{
.update_file_streaming(
&file.path,
&temp_path,
total_bytes,
&content_type,
Some(hash),
)
.await;
// Clean up temp file (may already be moved by dedup, ignore error)
let _ = tokio::fs::remove_file(&temp_path).await;
match result {
Ok(_) => StatusCode::OK.into_response(),
Err(e) => {
tracing::error!("WOPI PutFile failed: {}", e);