diff --git a/fluxer_media_proxy/src/media_process/mod.rs b/fluxer_media_proxy/src/media_process/mod.rs index 783f33009..f48ff082e 100644 --- a/fluxer_media_proxy/src/media_process/mod.rs +++ b/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::{ diff --git a/fluxer_media_proxy/src/media_process/native_cost.rs b/fluxer_media_proxy/src/media_process/native_cost.rs index 157c25c5d..c6d246d6d 100644 --- a/fluxer_media_proxy/src/media_process/native_cost.rs +++ b/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 { +) -> Result { 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 { +pub fn metadata_cost(input: &[u8], media_limits: &MediaLimits) -> Result { 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())), } } diff --git a/fluxer_media_proxy/src/media_process/tests/metadata.rs b/fluxer_media_proxy/src/media_process/tests/metadata.rs index 3cf2ea6f5..cb361288c 100644 --- a/fluxer_media_proxy/src/media_process/tests/metadata.rs +++ b/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") ); } diff --git a/fluxer_media_proxy/src/media_process/tests/transform.rs b/fluxer_media_proxy/src/media_process/tests/transform.rs index da7460c9d..f3f8e73d8 100644 --- a/fluxer_media_proxy/src/media_process/tests/transform.rs +++ b/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}"); } diff --git a/fluxer_media_proxy/src/server/native_task_executor.rs b/fluxer_media_proxy/src/server/native_task_executor.rs index 6e906e9ba..e744d7c1d 100644 --- a/fluxer_media_proxy/src/server/native_task_executor.rs +++ b/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 where T: Send + 'static, - C: FnOnce() -> anyhow::Result + Send + 'static, + C: FnOnce() -> anyhow::Result + Send + 'static, F: FnOnce() -> anyhow::Result + 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); + } } diff --git a/fluxer_media_proxy/src/server/routes/internal.rs b/fluxer_media_proxy/src/server/routes/internal.rs index 0f4b5105e..f38a41a20 100644 --- a/fluxer_media_proxy/src/server/routes/internal.rs +++ b/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();