Skip to main content

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            zstd::Encoder::new(writer, level)
31                .and_then(|mut e| {
32                    e.write_all(&data)?;
33                    Ok((e.finish()?, strategy))
34                })
35                .unwrap_or((data, CompressionStrategy::None))
36        }
37    }
38}
39
40pub fn add_headers(headers: &mut http::HeaderMap, strategy: CompressionStrategy) {
41    match strategy {
42        CompressionStrategy::None => {
43            let _ = headers;
44        }
45        #[cfg(feature = "compression")]
46        CompressionStrategy::Zstd { .. } => {
47            headers.insert(http::header::CONTENT_ENCODING, CONTENT_ENCODING_ZSTD);
48        }
49    }
50}