krabka_client_producer/
compression.rs

1//! The `Compression` enum, and the mapping from the producer's choice to a
2//! `RecordBatch` v2 `attributes` value and a
3//! `krabka-compression::CompressionType`.
4
5use std::{fmt, str::FromStr};
6
7use bytes::Bytes;
8
9use crate::error::ProducerError;
10
11#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
12pub enum Compression {
13    #[default]
14    None,
15    Gzip,
16    Snappy,
17    Lz4,
18    Zstd,
19}
20
21impl fmt::Display for Compression {
22    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
23        formatter.write_str(match self {
24            Self::None => "none",
25            Self::Gzip => "gzip",
26            Self::Snappy => "snappy",
27            Self::Lz4 => "lz4",
28            Self::Zstd => "zstd",
29        })
30    }
31}
32
33impl FromStr for Compression {
34    type Err = String;
35
36    fn from_str(value: &str) -> Result<Self, Self::Err> {
37        match value {
38            "none" => Ok(Self::None),
39            "gzip" => Ok(Self::Gzip),
40            "snappy" => Ok(Self::Snappy),
41            "lz4" => Ok(Self::Lz4),
42            "zstd" => Ok(Self::Zstd),
43            _ => Err(format!("unsupported producer compression: {value}")),
44        }
45    }
46}
47
48impl Compression {
49    #[must_use]
50    pub(crate) fn compression_type(self) -> krabka_compression::CompressionType {
51        match self {
52            Self::None => krabka_compression::CompressionType::None,
53            Self::Gzip => krabka_compression::CompressionType::Gzip,
54            Self::Snappy => krabka_compression::CompressionType::Snappy,
55            Self::Lz4 => krabka_compression::CompressionType::Lz4,
56            Self::Zstd => krabka_compression::CompressionType::Zstd,
57        }
58    }
59
60    /// The 3-bit `compression_type` field that goes into the `RecordBatch` v2
61    /// `attributes`, at bits 0..3.
62    #[must_use]
63    pub fn attribute_bits(self) -> i16 {
64        match self {
65            Self::None => 0,
66            Self::Gzip => 1,
67            Self::Snappy => 2,
68            Self::Lz4 => 3,
69            Self::Zstd => 4,
70        }
71    }
72
73    /// Compress the encoded record body. Returns the byte payload that goes
74    /// into the `RecordBatch.records_body` slot.
75    /// # Errors
76    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
77    pub fn compress(self, raw: &[u8]) -> Result<Bytes, ProducerError> {
78        Ok(krabka_compression::compress(self.compression_type(), raw)?)
79    }
80}
81
82/// Default gzip compression level. Kafka's `compression.gzip.level` default
83/// is `CompressionType.GZIP.defaultLevel()`, which is
84/// `Deflater.DEFAULT_COMPRESSION` (-1).
85pub const DEFAULT_PRODUCER_COMPRESSION_GZIP_LEVEL: i32 = -1;
86/// Default lz4 compression level. Kafka's `compression.lz4.level` default is
87/// 9.
88pub const DEFAULT_PRODUCER_COMPRESSION_LZ4_LEVEL: i32 = 9;
89/// Default zstd compression level. Kafka's `compression.zstd.level` default
90/// is 3.
91pub const DEFAULT_PRODUCER_COMPRESSION_ZSTD_LEVEL: i32 = 3;
92
93/// The lowest and highest gzip levels: `Deflater.BEST_SPEED` and
94/// `Deflater.BEST_COMPRESSION`.
95const GZIP_LEVELS: (i32, i32) = (1, 9);
96/// The lowest and highest lz4 levels, from `net.jpountz.lz4.LZ4Constants`.
97const LZ4_LEVELS: (i32, i32) = (1, 17);
98/// The lowest and highest zstd levels: `ZSTD_minCLevel` and `ZSTD_MAX_CLEVEL`.
99const ZSTD_LEVELS: (i32, i32) = (-131_072, 22);
100
101/// Validated compression levels of the producer (KIP-390).
102///
103/// Kafka's `ProducerConfig` defines `compression.gzip.level`,
104/// `compression.lz4.level` and `compression.zstd.level`, and checks each one
105/// with `CompressionType.levelValidator`, whatever `compression.type` is.
106/// `KafkaProducer.configureCompression` then gives the level of the chosen
107/// codec to the codec.
108#[derive(Clone, Copy, Debug, Eq, PartialEq)]
109pub struct CompressionLevels {
110    gzip: i32,
111    lz4: i32,
112    zstd: i32,
113}
114
115impl CompressionLevels {
116    /// Validate the three levels as Kafka's `ProducerConfig` does.
117    ///
118    /// # Errors
119    ///
120    /// Returns Kafka's `ConfigException` message, with the krabka option name,
121    /// for the first level that is out of range: gzip 1 to 9 or -1, lz4 1 to
122    /// 17, zstd -131072 to 22.
123    pub fn new(gzip: i32, lz4: i32, zstd: i32) -> Result<Self, String> {
124        let (gzip_min, gzip_max) = GZIP_LEVELS;
125        if gzip > gzip_max || (gzip < gzip_min && gzip != DEFAULT_PRODUCER_COMPRESSION_GZIP_LEVEL) {
126            return Err(invalid_level(
127                "compression_gzip_level",
128                gzip,
129                &format!(
130                    "Value must be between {gzip_min} and {gzip_max} or equal to {DEFAULT_PRODUCER_COMPRESSION_GZIP_LEVEL}"
131                ),
132            ));
133        }
134        check_range("compression_lz4_level", lz4, LZ4_LEVELS)?;
135        check_range("compression_zstd_level", zstd, ZSTD_LEVELS)?;
136        Ok(Self { gzip, lz4, zstd })
137    }
138
139    #[must_use]
140    pub const fn gzip(self) -> i32 {
141        self.gzip
142    }
143
144    #[must_use]
145    pub const fn lz4(self) -> i32 {
146        self.lz4
147    }
148
149    #[must_use]
150    pub const fn zstd(self) -> i32 {
151        self.zstd
152    }
153
154    /// The level of `compression`, or `None` for a codec with no levels.
155    #[must_use]
156    pub const fn level(self, compression: Compression) -> Option<i32> {
157        match compression {
158            Compression::Gzip => Some(self.gzip),
159            Compression::Lz4 => Some(self.lz4),
160            Compression::Zstd => Some(self.zstd),
161            Compression::None | Compression::Snappy => None,
162        }
163    }
164}
165
166impl Default for CompressionLevels {
167    fn default() -> Self {
168        Self {
169            gzip: DEFAULT_PRODUCER_COMPRESSION_GZIP_LEVEL,
170            lz4: DEFAULT_PRODUCER_COMPRESSION_LZ4_LEVEL,
171            zstd: DEFAULT_PRODUCER_COMPRESSION_ZSTD_LEVEL,
172        }
173    }
174}
175
176/// Kafka's `ConfigDef.Range.ensureValid`.
177fn check_range(name: &str, level: i32, (min, max): (i32, i32)) -> Result<(), String> {
178    if level < min {
179        Err(invalid_level(
180            name,
181            level,
182            &format!("Value must be at least {min}"),
183        ))
184    } else if level > max {
185        Err(invalid_level(
186            name,
187            level,
188            &format!("Value must be no more than {max}"),
189        ))
190    } else {
191        Ok(())
192    }
193}
194
195/// Kafka's `ConfigException(name, value, message)` text.
196fn invalid_level(name: &str, level: i32, message: &str) -> String {
197    format!("Invalid value {level} for configuration {name}: {message}")
198}
199
200#[cfg(test)]
201mod tests {
202    use krabka_units::convert::ByteSizeExt;
203
204    use super::*;
205
206    #[test]
207    fn none_round_trip_is_identity() {
208        let raw = b"hello producer";
209        let out = Compression::None.compress(raw).unwrap();
210        assert2::assert!(out.as_ref() == raw);
211    }
212
213    /// Each row is one set of levels and the result of Kafka's validators:
214    /// `CompressionType.GZIP.levelValidator` for gzip, and
215    /// `ConfigDef.Range.ensureValid` for lz4 and zstd.
216    #[test]
217    fn compression_levels_follow_kafka_validators() {
218        let levels = |gzip, lz4, zstd| CompressionLevels { gzip, lz4, zstd };
219        let gzip_error = |level: i32| {
220            Err(format!(
221                "Invalid value {level} for configuration compression_gzip_level: Value must be \
222                 between 1 and 9 or equal to -1"
223            ))
224        };
225        let cases = [
226            ("defaults", (-1, 9, 3), Ok(levels(-1, 9, 3))),
227            (
228                "lowest levels",
229                (1, 1, -131_072),
230                Ok(levels(1, 1, -131_072)),
231            ),
232            ("highest levels", (9, 17, 22), Ok(levels(9, 17, 22))),
233            ("gzip zero", (0, 9, 3), gzip_error(0)),
234            ("gzip below the default", (-2, 9, 3), gzip_error(-2)),
235            ("gzip above nine", (10, 9, 3), gzip_error(10)),
236            (
237                "lz4 zero",
238                (-1, 0, 3),
239                Err(
240                    "Invalid value 0 for configuration compression_lz4_level: Value must be at \
241                     least 1"
242                        .to_owned(),
243                ),
244            ),
245            (
246                "lz4 above seventeen",
247                (-1, 18, 3),
248                Err(
249                    "Invalid value 18 for configuration compression_lz4_level: Value must be \
250                     no more than 17"
251                        .to_owned(),
252                ),
253            ),
254            (
255                "zstd below the minimum",
256                (-1, 9, -131_073),
257                Err(
258                    "Invalid value -131073 for configuration compression_zstd_level: Value must \
259                     be at least -131072"
260                        .to_owned(),
261                ),
262            ),
263            (
264                "zstd above twenty-two",
265                (-1, 9, 23),
266                Err(
267                    "Invalid value 23 for configuration compression_zstd_level: Value must be \
268                     no more than 22"
269                        .to_owned(),
270                ),
271            ),
272        ];
273        let mut actual = Vec::new();
274        let mut wanted = Vec::new();
275        for (name, (gzip, lz4, zstd), expected) in cases {
276            actual.push((name, CompressionLevels::new(gzip, lz4, zstd)));
277            wanted.push((name, expected));
278        }
279        assert2::assert!(actual == wanted);
280        assert2::assert!(CompressionLevels::default() == levels(-1, 9, 3));
281        assert2::assert!(
282            [
283                Compression::None,
284                Compression::Gzip,
285                Compression::Snappy,
286                Compression::Lz4,
287                Compression::Zstd,
288            ]
289            .map(|compression| levels(5, 12, 19).level(compression))
290                == [None, Some(5), None, Some(12), Some(19)]
291        );
292        let set = levels(5, 12, 19);
293        assert2::assert!([set.gzip(), set.lz4(), set.zstd()] == [5, 12, 19]);
294    }
295
296    #[test]
297    fn attribute_bits_match_kafka_table() {
298        for (_name, compression, want) in [
299            ("none", Compression::None, 0),
300            ("gzip", Compression::Gzip, 1),
301            ("snappy", Compression::Snappy, 2),
302            ("lz4", Compression::Lz4, 3),
303            ("zstd", Compression::Zstd, 4),
304        ] {
305            assert2::assert!(compression.attribute_bits() == want);
306        }
307    }
308
309    #[test]
310    fn gzip_round_trip_via_decoder() {
311        use krabka_compression::CompressionType;
312        let raw = b"the quick brown fox";
313        let compressed = Compression::Gzip.compress(raw).unwrap();
314        let decoded = krabka_compression::decompress(
315            CompressionType::Gzip,
316            &compressed,
317            ByteSizeExt::from_bytes(u64::MAX),
318        )
319        .unwrap();
320        assert2::assert!(decoded.as_ref() == raw);
321    }
322
323    #[test]
324    fn compression_parses_and_displays_canonical_names() {
325        for (name, compression) in [
326            ("none", Compression::None),
327            ("gzip", Compression::Gzip),
328            ("snappy", Compression::Snappy),
329            ("lz4", Compression::Lz4),
330            ("zstd", Compression::Zstd),
331        ] {
332            assert_eq!(name.parse::<Compression>(), Ok(compression));
333            assert_eq!(compression.to_string(), name);
334        }
335        assert_eq!(
336            "brotli".parse::<Compression>(),
337            Err("unsupported producer compression: brotli".to_owned())
338        );
339    }
340}
krabka_client_producer/compression.rs

krabka_client_producer/
compression.rs
