fix(media-proxy): let video work fit small memory budgets (#3392)

fluxer_media_proxy/src/media_process/mod.rs

@@ -37,7 +37,9 @@ pub use metadata::{
     MetadataOptions, PreparedMetadata, finish_metadata, metadata_json_with_options,
     prepare_metadata,
 };
-pub use native_cost::{AV_NATIVE_COST_BYTES, av_native_cost, image_transform_cost, metadata_cost};
+pub use native_cost::{
+    AV_NATIVE_COST_BYTES, NativeCost, av_native_cost, image_transform_cost, metadata_cost,
+};
 pub use nsfw_processing::encode_static_image_for_nsfw;
 pub use transform::transform_image;
 pub use video_thumbnail::{

fluxer_media_proxy/src/media_process/native_cost.rs

@@ -10,6 +10,12 @@ use crate::{
 };
 
 pub const AV_NATIVE_COST_BYTES: usize = 256 << 20;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum NativeCost {
+    Bytes(usize),
+    UpToBudget(usize),
+}
 const ANIMATED_CANVAS_BUFFERS: usize = 8;
 const RGBA_BYTES_PER_PIXEL: usize = 4;
 
@@ -37,7 +43,7 @@ pub fn image_transform_cost(
     input: &[u8],
     options: &ImageOptions,
     media_limits: &MediaLimits,
-) -> Result<usize, MediaError> {
+) -> Result<NativeCost, MediaError> {
     ensure_vips_init()?;
     let sniffed_mime = mime::sniff(input).mime;
     let dims = probe_image_dims(media_limits, input)?;
@@ -50,27 +56,29 @@ pub fn image_transform_cost(
     } else {
         source_pixels * decode_bytes_per_pixel(sniffed_mime) + output_pixels * encode
     };
-    Ok(input.len().saturating_add(pixel_cost))
+    Ok(NativeCost::Bytes(input.len().saturating_add(pixel_cost)))
 }
 
-pub fn av_native_cost(input_len: usize) -> usize {
-    input_len
-        .min(constants::MAX_MEDIA_PROXY_BYTES)
-        .saturating_add(AV_NATIVE_COST_BYTES)
+pub fn av_native_cost(input_len: usize) -> NativeCost {
+    NativeCost::UpToBudget(
+        input_len
+            .min(constants::MAX_MEDIA_PROXY_BYTES)
+            .saturating_add(AV_NATIVE_COST_BYTES),
+    )
 }
 
-pub fn metadata_cost(input: &[u8], media_limits: &MediaLimits) -> Result<usize, MediaError> {
+pub fn metadata_cost(input: &[u8], media_limits: &MediaLimits) -> Result<NativeCost, MediaError> {
     let sniffed_mime = mime::sniff(input).mime;
     match mime::category(sniffed_mime) {
         Some(mime::Category::Image) => {
             ensure_vips_init()?;
             let dims = probe_image_dims(media_limits, input)?;
             let source_pixels = dims.width as usize * dims.height as usize;
-            Ok(input
-                .len()
-                .saturating_add(source_pixels * decode_bytes_per_pixel(sniffed_mime)))
+            Ok(NativeCost::Bytes(input.len().saturating_add(
+                source_pixels * decode_bytes_per_pixel(sniffed_mime),
+            )))
         }
         Some(mime::Category::Video | mime::Category::Audio) => Ok(av_native_cost(input.len())),
-        _ => Ok(input.len()),
+        _ => Ok(NativeCost::Bytes(input.len())),
     }
 }

fluxer_media_proxy/src/media_process/tests/metadata.rs

@@ -3,7 +3,7 @@
 use super::super::nsfw_processing::{VideoNSFWFramesRequest, extract_video_frames_for_nsfw};
 use super::super::placeholder::{encode_thumbhash, optional_thumbhash};
 use super::super::{
-    AV_NATIVE_COST_BYTES, MediaError, MetadataOptions, NSFW_PREVIEW_MAX_DIMENSION,
+    AV_NATIVE_COST_BYTES, MediaError, MetadataOptions, NSFW_PREVIEW_MAX_DIMENSION, NativeCost,
     extract_video_thumbnail, metadata_cost, metadata_json_with_options, probe_av_metadata,
 };
 use super::fixtures::{metadata_value, test_media_limits};
@@ -296,7 +296,7 @@ fn metadata_refuses_an_image_past_the_buffer_bound() {
 fn native_cost_of_a_video_past_the_buffer_bound_leaves_out_the_mapped_tail() {
     let mp4 = mapped_past_the_buffer_bound(&fixture_h264_mp4());
     assert_eq!(
-        constants::MAX_MEDIA_PROXY_BYTES + AV_NATIVE_COST_BYTES,
+        NativeCost::UpToBudget(constants::MAX_MEDIA_PROXY_BYTES + AV_NATIVE_COST_BYTES),
         metadata_cost(&mp4, &test_media_limits()).expect("video cost")
     );
 }

fluxer_media_proxy/src/media_process/tests/transform.rs

@@ -3,7 +3,7 @@
 use super::super::image_probe::load_image;
 use super::super::native_runtime::ensure_vips_init;
 use super::super::transform::source_supports_pages;
-use super::super::{ImageOptions, MediaError, image_transform_cost};
+use super::super::{ImageOptions, MediaError, NativeCost, image_transform_cost};
 use super::fixtures::{animated_mode, metadata_value, transform_image};
 use crate::{
     mime, native,
@@ -188,6 +188,9 @@ fn transform_cost_charges_the_full_source_decode_for_jxl() {
         &super::fixtures::test_media_limits(),
     )
     .unwrap();
+    let NativeCost::Bytes(cost) = cost else {
+        panic!("an image is charged exact bytes")
+    };
     assert!(cost >= 4100 * 4100 * 40, "cost {cost}");
 }
 
@@ -205,5 +208,8 @@ fn transform_cost_charges_svg_at_its_output_size() {
         &super::fixtures::test_media_limits(),
     )
     .unwrap();
+    let NativeCost::Bytes(cost) = cost else {
+        panic!("an image is charged exact bytes")
+    };
     assert!(cost < 1 << 20, "cost {cost}");
 }

fluxer_media_proxy/src/server/native_task_executor.rs

@@ -2,7 +2,7 @@
 
 use crate::{
     byte_budget::ByteBudget,
-    media_process::MediaError,
+    media_process::{MediaError, NativeCost},
     metrics::{
         self,
         transform::{NativeTransformMetrics, TransformMetrics},
@@ -96,7 +96,7 @@ impl NativeTaskExecutor {
     ) -> anyhow::Result<T>
     where
         T: Send + 'static,
-        C: FnOnce() -> anyhow::Result<usize> + Send + 'static,
+        C: FnOnce() -> anyhow::Result<NativeCost> + Send + 'static,
         F: FnOnce() -> anyhow::Result<T> + Send + 'static,
     {
         let admission = match self.native_transforms.try_admit() {
@@ -121,14 +121,17 @@ impl NativeTaskExecutor {
         self.run_task(deadline, move || {
             let _span = span.enter();
             let _permit = permit;
-            let bytes = cost()?;
-            if bytes > capacity {
-                warn!(
-                    bytes,
-                    capacity, "native work refused, larger than the memory budget"
-                );
-                return Err(MediaError::InvalidImageDimensions.into());
-            }
+            let bytes = match cost()? {
+                NativeCost::Bytes(bytes) if bytes > capacity => {
+                    warn!(
+                        bytes,
+                        capacity, "native work refused, larger than the memory budget"
+                    );
+                    return Err(MediaError::InvalidImageDimensions.into());
+                }
+                NativeCost::Bytes(bytes) => bytes,
+                NativeCost::UpToBudget(bytes) => bytes.min(capacity),
+            };
             let Some(_reserved) = memory.try_reserve(bytes) else {
                 warn!(bytes, capacity, "native work refused, memory budget in use");
                 return Err(MediaError::AllocationFailed.into());
@@ -298,7 +301,7 @@ mod tests {
         let error = executor
             .run_native(
                 Some(Instant::now() + Duration::from_millis(30)),
-                || Ok(0),
+                || Ok(NativeCost::Bytes(0)),
                 move || {
                     released.blocking_recv();
                     let _ = finished.blocking_send(());
@@ -351,7 +354,7 @@ mod tests {
         let value = executor
             .run_native(
                 Some(Instant::now() + Duration::from_secs(30)),
-                || Ok(0),
+                || Ok(NativeCost::Bytes(0)),
                 || Ok(7u32),
             )
             .await
@@ -378,7 +381,7 @@ mod tests {
             holder
                 .run_native(
                     None,
-                    || Ok(0),
+                    || Ok(NativeCost::Bytes(0)),
                     move || {
                         released.blocking_recv();
                         Ok(())
@@ -388,7 +391,7 @@ mod tests {
         });
         tokio::time::sleep(Duration::from_millis(50)).await;
         let error = executor
-            .run_native(None, || Ok(0), || Ok(()))
+            .run_native(None, || Ok(NativeCost::Bytes(0)), || Ok(()))
             .await
             .expect_err("the admission queue is full");
         assert_eq!(
@@ -412,7 +415,7 @@ mod tests {
         let executor = executor(&metrics, 2, 4);
         executor.begin_shutdown();
         let error = executor
-            .run_native(None, || Ok(0), || Ok(()))
+            .run_native(None, || Ok(NativeCost::Bytes(0)), || Ok(()))
             .await
             .expect_err("a closed executor admits nothing");
         assert_eq!(
@@ -429,7 +432,7 @@ mod tests {
         let error = executor
             .run_native(
                 None,
-                || Ok(1025),
+                || Ok(NativeCost::Bytes(1025)),
                 || -> anyhow::Result<()> { panic!("refused work must not run") },
             )
             .await
@@ -451,7 +454,7 @@ mod tests {
             holder
                 .run_native(
                     None,
-                    || Ok(1000),
+                    || Ok(NativeCost::Bytes(1000)),
                     move || {
                         let _ = started.blocking_send(());
                         released.blocking_recv();
@@ -462,7 +465,7 @@ mod tests {
         });
         has_started.recv().await;
         let error = executor
-            .run_native(None, || Ok(100), || Ok(()))
+            .run_native(None, || Ok(NativeCost::Bytes(100)), || Ok(()))
             .await
             .expect_err("the budget is in use");
         assert_eq!(
@@ -472,8 +475,19 @@ mod tests {
         drop(release);
         held.await.expect("held task").expect("held work");
         executor
-            .run_native(None, || Ok(1024), || Ok(()))
+            .run_native(None, || Ok(NativeCost::Bytes(1024)), || Ok(()))
             .await
             .expect("the released budget admits a job of the full size");
     }
+
+    #[tokio::test]
+    async fn work_that_reserves_up_to_the_budget_runs_on_a_budget_smaller_than_its_estimate() {
+        let metrics = metrics::Metrics::new();
+        let executor = executor(&metrics, 1, 0);
+        let value = executor
+            .run_native(None, || Ok(NativeCost::UpToBudget(1 << 30)), || Ok(7u32))
+            .await
+            .expect("the job takes the whole budget instead of being refused");
+        assert_eq!(7, value);
+    }
 }

fluxer_media_proxy/src/server/routes/internal.rs

@@ -685,7 +685,7 @@ mod tests {
                 .tasks()
                 .run_native(
                     None,
-                    || Ok(0),
+                    || Ok(media_process::NativeCost::Bytes(0)),
                     move || {
                         let _ = started.blocking_send(());
                         let _ = released.blocking_recv();