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 ────────────────────────────────────
|
// ── 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).
|
/// Store content with deduplication (from bytes).
|
||||||
///
|
///
|
||||||
/// Uses `SELECT … FOR UPDATE` + `INSERT … ON CONFLICT` for atomic
|
/// Uses `SELECT … FOR UPDATE` + `INSERT … ON CONFLICT` for atomic
|
||||||
/// upsert — completely TOCTOU-free.
|
/// 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(
|
pub async fn store_bytes(
|
||||||
&self,
|
&self,
|
||||||
content: &[u8],
|
content: &[u8],
|
||||||
content_type: Option<String>,
|
content_type: Option<String>,
|
||||||
) -> Result<DedupResultDto, DomainError> {
|
) -> 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 size = content.len() as u64;
|
||||||
let hash = Self::hash_bytes(content);
|
let hash = Self::hash_bytes(content);
|
||||||
|
|
||||||
|
|||||||
@@ -16,6 +16,7 @@ use image::{ImageFormat, imageops::FilterType};
|
|||||||
use std::path::{Path, PathBuf};
|
use std::path::{Path, PathBuf};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use tokio::fs;
|
use tokio::fs;
|
||||||
|
use tokio::sync::Semaphore;
|
||||||
|
|
||||||
use crate::application::ports::thumbnail_ports::{
|
use crate::application::ports::thumbnail_ports::{
|
||||||
ThumbnailPort, ThumbnailSize as PortThumbnailSize, ThumbnailStatsDto,
|
ThumbnailPort, ThumbnailSize as PortThumbnailSize, ThumbnailStatsDto,
|
||||||
@@ -69,6 +70,14 @@ struct ThumbnailCacheKey {
|
|||||||
size: ThumbnailSize,
|
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
|
/// Thumbnail service for generating and caching image thumbnails
|
||||||
pub struct ThumbnailService {
|
pub struct ThumbnailService {
|
||||||
/// Root path for thumbnail storage
|
/// Root path for thumbnail storage
|
||||||
@@ -77,6 +86,10 @@ pub struct ThumbnailService {
|
|||||||
cache: moka::future::Cache<ThumbnailCacheKey, Bytes>,
|
cache: moka::future::Cache<ThumbnailCacheKey, Bytes>,
|
||||||
/// Configured maximum cache weight (for stats reporting)
|
/// Configured maximum cache weight (for stats reporting)
|
||||||
max_cache_bytes: u64,
|
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 {
|
impl ThumbnailService {
|
||||||
@@ -105,6 +118,7 @@ impl ThumbnailService {
|
|||||||
thumbnails_root,
|
thumbnails_root,
|
||||||
cache,
|
cache,
|
||||||
max_cache_bytes: max_cache_bytes as u64,
|
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)
|
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(
|
async fn generate_thumbnail(
|
||||||
&self,
|
&self,
|
||||||
original_path: &Path,
|
original_path: &Path,
|
||||||
@@ -226,9 +244,25 @@ impl ThumbnailService {
|
|||||||
let path = original_path.to_path_buf();
|
let path = original_path.to_path_buf();
|
||||||
let max_dim = size.max_dimension();
|
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
|
// Run image processing in blocking thread pool
|
||||||
let result = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, ThumbnailError> {
|
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()))?;
|
let img = image::open(&path).map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
|
||||||
|
|
||||||
// Calculate new dimensions preserving aspect ratio
|
// Calculate new dimensions preserving aspect ratio
|
||||||
@@ -274,10 +308,31 @@ impl ThumbnailService {
|
|||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
tracing::info!("🖼️ Background thumbnail generation starting: {}", file_id);
|
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();
|
let path = original_path.clone();
|
||||||
|
|
||||||
// Single spawn_blocking: 1 read + 1 decode + 3 resize + 3 encode
|
// Single spawn_blocking: 1 read + 1 decode + 3 resize + 3 encode
|
||||||
let results = tokio::task::spawn_blocking(move || {
|
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)
|
let img = image::open(&path)
|
||||||
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
|
.map_err(|e| ThumbnailError::ImageError(e.to_string()))?;
|
||||||
|
|
||||||
|
|||||||
@@ -11,9 +11,9 @@
|
|||||||
use crate::interfaces::middleware::auth::AuthUser;
|
use crate::interfaces::middleware::auth::AuthUser;
|
||||||
use axum::{
|
use axum::{
|
||||||
Router,
|
Router,
|
||||||
body::Bytes,
|
body::Body,
|
||||||
extract::{Path, Query, State},
|
extract::{Path, Query, State},
|
||||||
http::{HeaderMap, StatusCode},
|
http::{HeaderMap, StatusCode, Request},
|
||||||
response::{Html, IntoResponse, Response},
|
response::{Html, IntoResponse, Response},
|
||||||
routing::{get, post},
|
routing::{get, post},
|
||||||
};
|
};
|
||||||
@@ -153,13 +153,22 @@ async fn get_file(
|
|||||||
}
|
}
|
||||||
|
|
||||||
/// POST /wopi/files/{file_id}/contents — PutFile
|
/// 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(
|
async fn put_file(
|
||||||
Path(file_id): Path<String>,
|
Path(file_id): Path<String>,
|
||||||
Query(token_query): Query<WopiTokenQuery>,
|
Query(token_query): Query<WopiTokenQuery>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
State(state): State<WopiState>,
|
State(state): State<WopiState>,
|
||||||
body: Bytes,
|
req: Request<Body>,
|
||||||
) -> Response {
|
) -> Response {
|
||||||
|
use http_body_util::BodyStream;
|
||||||
|
use sha2::{Digest, Sha256};
|
||||||
|
use tokio::io::AsyncWriteExt;
|
||||||
|
use tokio_stream::StreamExt;
|
||||||
|
|
||||||
let claims = match state
|
let claims = match state
|
||||||
.token_service
|
.token_service
|
||||||
.validate_token(&token_query.access_token)
|
.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
|
let file = match state
|
||||||
.app_state
|
.app_state
|
||||||
.applications
|
.applications
|
||||||
@@ -209,14 +218,75 @@ async fn put_file(
|
|||||||
Err(_) => return StatusCode::NOT_FOUND.into_response(),
|
Err(_) => return StatusCode::NOT_FOUND.into_response(),
|
||||||
};
|
};
|
||||||
|
|
||||||
// Save the file content using path-based update
|
// ── Streaming spool: body → temp file + incremental SHA-256 ──
|
||||||
match state
|
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
|
.app_state
|
||||||
.applications
|
.applications
|
||||||
.file_upload_service
|
.file_upload_service
|
||||||
.update_file(&file.path, &body)
|
.update_file_streaming(
|
||||||
.await
|
&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(),
|
Ok(_) => StatusCode::OK.into_response(),
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
tracing::error!("WOPI PutFile failed: {}", e);
|
tracing::error!("WOPI PutFile failed: {}", e);
|
||||||
|
|||||||
Reference in New Issue
Block a user