From 961de20a002e7332414f10045e2f62214d9c501f Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sun, 10 Dec 2023 19:11:36 -0600
Subject: [PATCH] Add timeout and metrics to media processing

---
 src/error.rs    |  4 ++++
 src/generate.rs | 12 ++++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index 649a6d67..b7cae657 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -153,6 +153,9 @@ pub(crate) enum UploadError {
     #[error("Client took too long to send request")]
     AggregateTimeout,
 
+    #[error("Timed out while waiting for media processing")]
+    ProcessTimeout,
+
     #[error("Failed external validation")]
     FailedExternalValidation,
 }
@@ -187,6 +190,7 @@ impl UploadError {
             Self::Range => ErrorCode::RANGE_NOT_SATISFIABLE,
             Self::Limit(_) => ErrorCode::VALIDATE_FILE_SIZE,
             Self::Timeout(_) | Self::AggregateTimeout => ErrorCode::STREAM_TOO_SLOW,
+            Self::ProcessTimeout => ErrorCode::COMMAND_TIMEOUT,
             Self::FailedExternalValidation => ErrorCode::FAILED_EXTERNAL_VALIDATION,
         }
     }
diff --git a/src/generate.rs b/src/generate.rs
index fe262ea4..a7137d84 100644
--- a/src/generate.rs
+++ b/src/generate.rs
@@ -6,12 +6,17 @@ use crate::{
     details::Details,
     error::{Error, UploadError},
     formats::{ImageFormat, InputProcessableFormat, InternalVideoFormat, ProcessableFormat},
+    future::{WithMetrics, WithTimeout},
     repo::{ArcRepo, Hash, VariantAlreadyExists},
     store::Store,
     tmp_file::TmpDir,
 };
 use actix_web::web::Bytes;
-use std::{path::PathBuf, sync::Arc, time::Instant};
+use std::{
+    path::PathBuf,
+    sync::Arc,
+    time::{Duration, Instant},
+};
 use tokio::io::AsyncReadExt;
 use tracing::Instrument;
 
@@ -79,7 +84,10 @@ pub(crate) async fn generate<S: Store + 'static>(
 
         let (details, bytes) = process_map
             .process(hash, thumbnail_path, process_fut)
-            .await?;
+            .with_timeout(Duration::from_secs(config.media.process_timeout * 4))
+            .with_metrics("pict-rs.generate.process")
+            .await
+            .map_err(|_| UploadError::ProcessTimeout)??;
 
         Ok((details, bytes))
     }