Skip to main content

kacrab_protocol/
compression.rs

1//! Record-batch compression codecs.
2//!
3//! The [`Compression`] enum is always available so [`crate::record::RecordBatch`]
4//! can read the compression type from `attributes` regardless of which codec
5//! features are enabled. Actual compress/decompress operations require the
6//! corresponding Cargo feature:
7//!
8//! | Feature  | Codec  | Backend             | Pure Rust | HC mode |
9//! |----------|--------|---------------------|-----------|---------|
10//! | `gzip`   | Gzip   | `flate2`            | yes       | n/a     |
11//! | `snappy` | Snappy | `snap`              | yes       | n/a     |
12//! | `lz4`    | Lz4    | `lz4_flex`          | yes       | no      |
13//! | `lz4-hc` | Lz4    | `lz4` (C-FFI)       | no        | yes     |
14//! | `zstd`   | Zstd   | `zstd` (C-FFI)      | no        | n/a     |
15//!
16//! `compression` is a meta-feature that enables `gzip + snappy + lz4 + zstd`
17//! (pure-Rust LZ4). Opt into `lz4-hc` explicitly if you need
18//! High-Compression mode — see the LZ4 section below.
19//!
20//! ## LZ4 backend selection
21//!
22//! * **`lz4`** — `lz4_flex` block API behind a Kafka-compatible custom frame. Fast mode only; the
23//!   `level` argument is ignored. Pure Rust, no C toolchain at build time, cross-compile clean.
24//! * **`lz4-hc`** — `lz4` crate (FFI to liblz4) behind the same frame format. Levels 3..=12 use HC
25//!   mode; levels 0..=2 fall back to fast mode. Requires a C compiler at build time.
26//!
27//! Both features are independent — enabling `lz4-hc` alone is sufficient
28//! (the C lib handles fast mode too). When **both** are enabled, `lz4-hc`
29//! wins at runtime; `lz4_flex` is linked but unused. Most users should
30//! pick one.
31//!
32//! For high compression ratio in Kafka workloads, prefer **`zstd`** over
33//! `lz4-hc` — modern Kafka clients (Java, librdkafka) default to fast LZ4
34//! when LZ4 is selected, and switch to zstd when ratio matters.
35
36#[cfg(feature = "gzip")]
37pub mod gzip;
38#[cfg(any(feature = "lz4", feature = "lz4-hc"))]
39pub mod lz4;
40#[cfg(feature = "snappy")]
41pub mod snappy;
42#[cfg(feature = "zstd")]
43pub mod zstd;
44
45pub mod error;
46
47pub use self::error::{CompressionError, CompressionErrorKind};
48
49/// Result alias for compression operations.
50pub type Result<T> = core::result::Result<T, CompressionError>;
51
52/// Hard ceiling on the decompressed size of a single payload.
53///
54/// Wire frames are already capped at [`crate::frame::MAX_FRAME_LENGTH`]
55/// (100 MiB), so this allows roughly a 10:1 expansion of the largest possible
56/// compressed batch while stopping a decompression bomb — a tiny payload that
57/// inflates without bound — from exhausting memory and aborting the process
58/// under `panic = "abort"`.
59pub const MAX_DECOMPRESSED_LEN: usize = 1 << 30;
60
61/// Read `reader` to the end, refusing to produce more than `max_len` bytes.
62/// Used by the streaming codecs (gzip, zstd), whose output size is unknown
63/// until decoded.
64#[cfg(any(feature = "gzip", feature = "zstd"))]
65fn read_to_end_bounded<R: std::io::Read>(
66    reader: R,
67    max_len: usize,
68    codec: Compression,
69) -> Result<Vec<u8>> {
70    use std::io::Read as _;
71
72    // Take one byte past the limit so an over-limit stream is distinguishable
73    // from one that ends exactly at it.
74    let limit = u64::try_from(max_len).unwrap_or(u64::MAX).saturating_add(1);
75    let mut output = Vec::new();
76    let _read = reader.take(limit).read_to_end(&mut output).map_err(|e| {
77        CompressionError::new(
78            codec,
79            CompressionErrorKind::DecodeFailed {
80                message: e.to_string(),
81            },
82        )
83    })?;
84    if output.len() > max_len {
85        return Err(CompressionError::new(
86            codec,
87            CompressionErrorKind::DecompressedTooLarge { limit: max_len },
88        ));
89    }
90    Ok(output)
91}
92
93/// Kafka record-batch compression codec, encoded in bits 0–2 of the batch
94/// `attributes` field.
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
96#[repr(i16)]
97#[non_exhaustive]
98pub enum Compression {
99    /// No compression.
100    None = 0,
101    /// Gzip (`flate2`).
102    Gzip = 1,
103    /// Snappy (`snap`).
104    Snappy = 2,
105    /// LZ4 frame format (`lz4_flex`).
106    Lz4 = 3,
107    /// Zstandard (`zstd` crate).
108    Zstd = 4,
109}
110
111impl Compression {
112    /// Decode the codec from a record-batch `attributes` field.
113    /// Only bits 0–2 are examined; higher bits are masked off.
114    pub const fn from_attributes(attributes: i16) -> Result<Self> {
115        match attributes & 0x07 {
116            0 => Ok(Self::None),
117            1 => Ok(Self::Gzip),
118            2 => Ok(Self::Snappy),
119            3 => Ok(Self::Lz4),
120            4 => Ok(Self::Zstd),
121            n => Err(CompressionError::new(
122                Self::None,
123                CompressionErrorKind::UnknownCodec(n),
124            )),
125        }
126    }
127
128    /// Compress `payload` at the codec default level.
129    pub fn compress(self, payload: &[u8]) -> Result<Vec<u8>> {
130        self.compress_with_level(payload, None)
131    }
132
133    /// Compress `payload` at the given level (codec-specific).
134    ///
135    /// * `Gzip` accepts `0..=9`, default `6`.
136    /// * `Zstd` accepts `1..=22`, default `3`.
137    /// * `Snappy` ignores the level.
138    /// * `Lz4` level handling depends on the active backend feature:
139    ///   * with `lz4` only — level ignored, always fast mode.
140    ///   * with `lz4-hc` — level `0..=2` → fast, `3..=12` → HC mode, values above `12` clamp to
141    ///     `12`, negative/zero levels clamp to fast level `1`.
142    pub fn compress_with_level(self, payload: &[u8], level: Option<i32>) -> Result<Vec<u8>> {
143        match self {
144            Self::None => {
145                let _ = level;
146                Ok(payload.to_vec())
147            },
148            #[cfg(feature = "gzip")]
149            Self::Gzip => gzip::compress_with_level(payload, level),
150            #[cfg(feature = "snappy")]
151            Self::Snappy => snappy::compress_with_level(payload, level),
152            #[cfg(any(feature = "lz4", feature = "lz4-hc"))]
153            Self::Lz4 => lz4::compress_with_level(payload, level),
154            #[cfg(feature = "zstd")]
155            Self::Zstd => zstd::compress_with_level(payload, level),
156            #[cfg(not(feature = "gzip"))]
157            Self::Gzip => Err(CompressionError::new(
158                self,
159                CompressionErrorKind::CodecDisabled,
160            )),
161            #[cfg(not(feature = "snappy"))]
162            Self::Snappy => Err(CompressionError::new(
163                self,
164                CompressionErrorKind::CodecDisabled,
165            )),
166            #[cfg(not(any(feature = "lz4", feature = "lz4-hc")))]
167            Self::Lz4 => Err(CompressionError::new(
168                self,
169                CompressionErrorKind::CodecDisabled,
170            )),
171            #[cfg(not(feature = "zstd"))]
172            Self::Zstd => Err(CompressionError::new(
173                self,
174                CompressionErrorKind::CodecDisabled,
175            )),
176        }
177    }
178
179    /// Decompress `payload`. Returns the decompressed payload, or an error if
180    /// the required Cargo feature is not enabled.
181    pub fn decompress(self, payload: &[u8]) -> Result<Vec<u8>> {
182        match self {
183            Self::None => Ok(payload.to_vec()),
184            #[cfg(feature = "gzip")]
185            Self::Gzip => gzip::decompress(payload),
186            #[cfg(feature = "snappy")]
187            Self::Snappy => snappy::decompress(payload),
188            #[cfg(any(feature = "lz4", feature = "lz4-hc"))]
189            Self::Lz4 => lz4::decompress(payload),
190            #[cfg(feature = "zstd")]
191            Self::Zstd => zstd::decompress(payload),
192            #[cfg(not(feature = "gzip"))]
193            Self::Gzip => Err(CompressionError::new(
194                self,
195                CompressionErrorKind::CodecDisabled,
196            )),
197            #[cfg(not(feature = "snappy"))]
198            Self::Snappy => Err(CompressionError::new(
199                self,
200                CompressionErrorKind::CodecDisabled,
201            )),
202            #[cfg(not(any(feature = "lz4", feature = "lz4-hc")))]
203            Self::Lz4 => Err(CompressionError::new(
204                self,
205                CompressionErrorKind::CodecDisabled,
206            )),
207            #[cfg(not(feature = "zstd"))]
208            Self::Zstd => Err(CompressionError::new(
209                self,
210                CompressionErrorKind::CodecDisabled,
211            )),
212        }
213    }
214}