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:
@@ -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()))?;
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
Reference in New Issue
Block a user