diff --git a/libdd-profiling/src/pprof/test_utils.rs b/libdd-profiling/src/pprof/test_utils.rs index d315352b65..42a15b30e3 100644 --- a/libdd-profiling/src/pprof/test_utils.rs +++ b/libdd-profiling/src/pprof/test_utils.rs @@ -6,8 +6,8 @@ use libdd_profiling_protobuf::prost_impls::{Profile, Sample}; fn deserialize_compressed_pprof(encoded: &[u8]) -> anyhow::Result { use prost::Message; - // Native zstd uses FFI and is unavailable under Miri, where the default - // profile codec remains uncompressed. + // The zstd bindings use FFI so they don't work under miri. This means the + // buffer isn't compressed, so simply convert to a vec. #[cfg(miri)] let buf = encoded.to_vec(); #[cfg(all(not(miri), not(target_arch = "wasm32")))] diff --git a/libdd-profiling/src/profiles/compressor.rs b/libdd-profiling/src/profiles/compressor.rs index cc90786838..3a340771ed 100644 --- a/libdd-profiling/src/profiles/compressor.rs +++ b/libdd-profiling/src/profiles/compressor.rs @@ -3,6 +3,25 @@ use std::io::{self, BufWriter, Read, Write}; +const DEFAULT_COMPRESSION_LEVEL: i32 = 3; +#[cfg(target_arch = "wasm32")] +const MIN_ZRIP_COMPRESSION_LEVEL: i32 = -7; +#[cfg(target_arch = "wasm32")] +const MAX_ZRIP_COMPRESSION_LEVEL: i32 = 4; + +fn zstd_compression_level(level: i32) -> i32 { + let level = if level == 0 { + DEFAULT_COMPRESSION_LEVEL + } else { + level + }; + + #[cfg(target_arch = "wasm32")] + let level = level.clamp(MIN_ZRIP_COMPRESSION_LEVEL, MAX_ZRIP_COMPRESSION_LEVEL); + + level +} + /// This type wraps a [`Vec`] to provide a [`Write`] interface that has a max /// capacity that won't be exceeded. Additionally, it gracefully handles /// out-of-memory conditions instead of panicking (unfortunately not compatible @@ -119,6 +138,11 @@ impl ProfileCodec for NoopProfileCodec { } } +/// A zstd-compatible profile codec. +/// +/// Native targets accept the range reported by `zstd::compression_level_range()`. +/// WASM clamps levels to zrip's supported range of `-7..=4`. Level `0` selects +/// level `3` on every target. #[allow(unused)] pub struct ZstdProfileCodec; @@ -132,7 +156,10 @@ impl ProfileCodec for ZstdProfileCodec { compression_level: i32, ) -> io::Result { let buffer = SizeRestrictedBuffer::try_new(size_hint, max_capacity)?; - zstd::Encoder::<'static, SizeRestrictedBuffer>::new(buffer, compression_level) + zstd::Encoder::<'static, SizeRestrictedBuffer>::new( + buffer, + zstd_compression_level(compression_level), + ) } fn finish(encoder: Self::Encoder) -> io::Result> { @@ -157,7 +184,8 @@ impl ProfileCodec for ZstdProfileCodec { compression_level: i32, ) -> io::Result { let buffer = SizeRestrictedBuffer::try_new(size_hint, max_capacity)?; - zrip::FrameEncoder::new(buffer, compression_level).map_err(io::Error::other) + zrip::FrameEncoder::new(buffer, zstd_compression_level(compression_level)) + .map_err(io::Error::other) } fn finish(encoder: Self::Encoder) -> io::Result> { @@ -267,7 +295,8 @@ impl Compressor { /// - `size_hint`: beginning capacity for the output buffer. This is a hint for the starting /// size, and the implementation may use something different. /// - `max_capacity`: the maximum size for the output buffer (hard limit). - /// - `compression_level`: must be supported by the target's zstd encoder. + /// - `compression_level`: passed to `C::new_encoder`. For the default [`ZstdProfileCodec`], see + /// its target-specific level documentation. pub fn try_new( size_hint: usize, max_capacity: usize, @@ -299,3 +328,25 @@ impl Write for Compressor { self.encoder.flush() } } + +#[cfg(all(test, not(miri), not(target_arch = "wasm32")))] +mod tests { + use super::*; + + fn compress(level: i32) -> Vec { + let mut compressor = Compressor::::try_new(256, 4096, level).unwrap(); + compressor.write_all(b"hello profile").unwrap(); + compressor.finish().unwrap() + } + + #[test] + fn zero_uses_default_compression_level() { + assert_eq!(compress(0), compress(DEFAULT_COMPRESSION_LEVEL)); + } + + #[test] + fn native_compression_level_is_not_clamped() { + assert_eq!(zstd_compression_level(22), 22); + assert!(!compress(22).is_empty()); + } +} diff --git a/libdd-trace-utils/src/send_with_retry/compression.rs b/libdd-trace-utils/src/send_with_retry/compression.rs index 53f7bcf918..466da0a7b5 100644 --- a/libdd-trace-utils/src/send_with_retry/compression.rs +++ b/libdd-trace-utils/src/send_with_retry/compression.rs @@ -6,11 +6,51 @@ use std::io::Write as _; #[cfg(feature = "compression")] const CONTENT_ENCODING_ZSTD: http::HeaderValue = http::HeaderValue::from_static("zstd"); +#[cfg(feature = "compression")] +const DEFAULT_COMPRESSION_LEVEL: i32 = 3; +#[cfg(all(feature = "compression", target_arch = "wasm32"))] +const MIN_ZRIP_COMPRESSION_LEVEL: i32 = -7; +#[cfg(all(feature = "compression", target_arch = "wasm32"))] +const MAX_ZRIP_COMPRESSION_LEVEL: i32 = 4; + +#[cfg(all(feature = "compression", not(target_arch = "wasm32")))] +type ZstdEncoder = zstd::Encoder<'static, Vec>; +#[cfg(all(feature = "compression", target_arch = "wasm32"))] +type ZstdEncoder = zrip::FrameEncoder>; + +#[cfg(feature = "compression")] +fn zstd_compression_level(level: i32) -> i32 { + let level = if level == 0 { + DEFAULT_COMPRESSION_LEVEL + } else { + level + }; + + #[cfg(all(feature = "compression", target_arch = "wasm32"))] + let level = level.clamp(MIN_ZRIP_COMPRESSION_LEVEL, MAX_ZRIP_COMPRESSION_LEVEL); + + level +} + +#[cfg(all(feature = "compression", not(target_arch = "wasm32")))] +fn new_zstd_encoder(writer: Vec, level: i32) -> std::io::Result { + zstd::Encoder::new(writer, level) +} + +#[cfg(all(feature = "compression", target_arch = "wasm32"))] +fn new_zstd_encoder(writer: Vec, level: i32) -> std::io::Result { + zrip::FrameEncoder::new(writer, level).map_err(std::io::Error::other) +} #[derive(Clone, Copy, Debug)] pub enum CompressionStrategy { None, #[cfg(feature = "compression")] + /// Zstd-compatible compression. + /// + /// Native targets accept the range reported by `zstd::compression_level_range()`. + /// WASM clamps levels to zrip's supported range of `-7..=4`. Level `0` selects + /// level `3` on every target. Zstd { level: i32, }, @@ -23,22 +63,16 @@ pub fn compress(data: Vec, strategy: CompressionStrategy) -> (Vec, Compr CompressionStrategy::None => (data, CompressionStrategy::None), #[cfg(feature = "compression")] CompressionStrategy::Zstd { level } => { + let level = zstd_compression_level(level); + let strategy = CompressionStrategy::Zstd { level }; // Start with an initial buffer // Allocate 1/10th of the original buffer, so we shouldn't add too // much memory usage, and no less than 256 bytes let writer = Vec::with_capacity((data.len() / 10).max(256)); - #[cfg(not(target_arch = "wasm32"))] - let result = zstd::Encoder::new(writer, level).and_then(|mut e| { - e.write_all(&data)?; - Ok((e.finish()?, strategy)) + let result = new_zstd_encoder(writer, level).and_then(|mut encoder| { + encoder.write_all(&data)?; + Ok((encoder.finish()?, strategy)) }); - #[cfg(target_arch = "wasm32")] - let result = zrip::FrameEncoder::new(writer, level) - .map_err(std::io::Error::other) - .and_then(|mut e| { - e.write_all(&data)?; - Ok((e.finish()?, strategy)) - }); result.unwrap_or((data, CompressionStrategy::None)) } } @@ -72,4 +106,30 @@ mod tests { assert!(matches!(strategy, CompressionStrategy::Zstd { level: 1 })); assert_eq!(decompress(&compressed).unwrap(), data); } + + #[test] + fn zero_uses_default_compression_level() { + let data = b"hello zstd".repeat(100); + let (default_compressed, strategy) = + compress(data.clone(), CompressionStrategy::Zstd { level: 0 }); + let (level_three_compressed, _) = compress( + data, + CompressionStrategy::Zstd { + level: DEFAULT_COMPRESSION_LEVEL, + }, + ); + + assert_eq!(default_compressed, level_three_compressed); + assert!(matches!(strategy, CompressionStrategy::Zstd { level: 3 })); + } + + #[test] + fn native_compression_level_is_not_clamped() { + let data = b"hello zstd".repeat(100); + let (compressed, strategy) = + compress(data.clone(), CompressionStrategy::Zstd { level: 22 }); + + assert!(matches!(strategy, CompressionStrategy::Zstd { level: 22 })); + assert_eq!(decompress(&compressed).unwrap(), data); + } } diff --git a/tests/wasm/tests/compression.rs b/tests/wasm/tests/compression.rs index ca04a4b92a..abca3cc384 100644 --- a/tests/wasm/tests/compression.rs +++ b/tests/wasm/tests/compression.rs @@ -15,6 +15,14 @@ mod profiling_compressor; use std::io::{Read, Write}; +fn compress_profile(payload: &[u8], level: i32) -> std::io::Result> { + use profiling_compressor::{Compressor, ZstdProfileCodec}; + + let mut compressor = Compressor::::try_new(256, 4096, level)?; + compressor.write_all(payload)?; + compressor.finish() +} + #[wasm_bindgen_test::wasm_bindgen_test] fn trace_compression_uses_zrip() { use trace_compression::{add_headers, compress, CompressionStrategy}; @@ -28,16 +36,44 @@ fn trace_compression_uses_zrip() { assert_eq!(headers["content-encoding"], "zstd"); } +#[wasm_bindgen_test::wasm_bindgen_test] +fn trace_compression_levels_are_clamped() { + use trace_compression::{compress, CompressionStrategy}; + + let payload = b"hello zstd".repeat(100); + for level in [-7, 4] { + let (compressed, strategy) = compress(payload.clone(), CompressionStrategy::Zstd { level }); + assert!(matches!(strategy, CompressionStrategy::Zstd { level: actual } if actual == level)); + assert_eq!(zrip::decompress(&compressed).unwrap(), payload); + } + for (level, expected) in [(-8, -7), (5, 4), (22, 4)] { + let (compressed, strategy) = compress(payload.clone(), CompressionStrategy::Zstd { level }); + assert!( + matches!(strategy, CompressionStrategy::Zstd { level: actual } if actual == expected) + ); + assert_eq!(zrip::decompress(&compressed).unwrap(), payload); + } +} + +#[wasm_bindgen_test::wasm_bindgen_test] +fn trace_compression_zero_uses_level_three() { + use trace_compression::{compress, CompressionStrategy}; + + let payload = b"hello zstd".repeat(100); + let (default_compressed, strategy) = + compress(payload.clone(), CompressionStrategy::Zstd { level: 0 }); + let (level_three_compressed, _) = compress(payload, CompressionStrategy::Zstd { level: 3 }); + + assert_eq!(default_compressed, level_three_compressed); + assert!(matches!(strategy, CompressionStrategy::Zstd { level: 3 })); +} + #[wasm_bindgen_test::wasm_bindgen_test] fn profiling_codecs_use_zrip() { - use profiling_compressor::{ - Compressor, ObservationCodec, ZstdObservationCodec, ZstdProfileCodec, - }; + use profiling_compressor::{ObservationCodec, ZstdObservationCodec}; let payload = b"hello profile".repeat(100); - let mut compressor = Compressor::::try_new(256, 4096, 1).unwrap(); - compressor.write_all(&payload).unwrap(); - let compressed = compressor.finish().unwrap(); + let compressed = compress_profile(&payload, 1).unwrap(); assert_eq!(zrip::decompress(&compressed).unwrap(), payload); let mut encoder = ZstdObservationCodec::new_encoder(256, 4096).unwrap(); @@ -47,3 +83,28 @@ fn profiling_codecs_use_zrip() { decoder.read_to_end(&mut decoded).unwrap(); assert_eq!(decoded, payload); } + +#[wasm_bindgen_test::wasm_bindgen_test] +fn profiling_compression_levels_are_clamped() { + let payload = b"hello profile".repeat(100); + for level in [-7, 4] { + let compressed = compress_profile(&payload, level).unwrap(); + assert_eq!(zrip::decompress(&compressed).unwrap(), payload); + } + for (level, expected) in [(-8, -7), (5, 4), (22, 4)] { + assert_eq!( + compress_profile(&payload, level).unwrap(), + compress_profile(&payload, expected).unwrap() + ); + } +} + +#[wasm_bindgen_test::wasm_bindgen_test] +fn profiling_compression_zero_uses_level_three() { + let payload = b"hello profile".repeat(100); + + assert_eq!( + compress_profile(&payload, 0).unwrap(), + compress_profile(&payload, 3).unwrap() + ); +}