libdd_trace_utils/send_with_retry/
compression.rs

1// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/
2// SPDX-License-Identifier: Apache-2.0
3
4#[cfg(feature = "compression")]
5use std::io::Write as _;
6
7#[cfg(feature = "compression")]
8const CONTENT_ENCODING_ZSTD: http::HeaderValue = http::HeaderValue::from_static("zstd");
9
10#[derive(Clone, Copy, Debug)]
11pub enum CompressionStrategy {
12    None,
13    #[cfg(feature = "compression")]
14    Zstd {
15        level: i32,
16    },
17}
18
19/// Returns the compressed data, and the actual compression strategy used.
20/// If an error happens during compression, defaults to [`CompressionStrategy::None`]
21pub fn compress(data: Vec<u8>, strategy: CompressionStrategy) -> (Vec<u8>, CompressionStrategy) {
22    match strategy {
23        CompressionStrategy::None => (data, CompressionStrategy::None),
24        #[cfg(feature = "compression")]
25        CompressionStrategy::Zstd { level } => {
26            // Start with an initial buffer
27            // Allocate 1/10th of the original buffer, so we shouldn't add too
28            // much memory usage, and no less than 256 bytes
29            let writer = Vec::with_capacity((data.len() / 10).max(256));
30            #[cfg(not(target_arch = "wasm32"))]
31            let result = zstd::Encoder::new(writer, level).and_then(|mut e| {
32                e.write_all(&data)?;
33                Ok((e.finish()?, strategy))
34            });
35            #[cfg(target_arch = "wasm32")]
36            let result = zrip::FrameEncoder::new(writer, level)
37                .map_err(std::io::Error::other)
38                .and_then(|mut e| {
39                    e.write_all(&data)?;
40                    Ok((e.finish()?, strategy))
41                });
42            result.unwrap_or((data, CompressionStrategy::None))
43        }
44    }
45}
46
47pub fn add_headers(headers: &mut http::HeaderMap, strategy: CompressionStrategy) {
48    match strategy {
49        CompressionStrategy::None => {
50            let _ = headers;
51        }
52        #[cfg(feature = "compression")]
53        CompressionStrategy::Zstd { .. } => {
54            headers.insert(http::header::CONTENT_ENCODING, CONTENT_ENCODING_ZSTD);
55        }
56    }
57}
58
59#[cfg(all(test, feature = "compression", not(target_arch = "wasm32")))]
60mod tests {
61    use super::*;
62
63    fn decompress(data: &[u8]) -> std::io::Result<Vec<u8>> {
64        zstd::decode_all(data)
65    }
66
67    #[test]
68    fn zstd_compression_roundtrips() {
69        let data = b"hello zstd".repeat(100);
70        let (compressed, strategy) = compress(data.clone(), CompressionStrategy::Zstd { level: 1 });
71
72        assert!(matches!(strategy, CompressionStrategy::Zstd { level: 1 }));
73        assert_eq!(decompress(&compressed).unwrap(), data);
74    }
75}
libdd_trace_utils/send_with_retry/compression.rs

libdd_trace_utils/send_with_retry/
compression.rs
