kafka_protocol/compression/
zstd.rs

1use crate::protocol::buf::{ByteBuf, ByteBufMut};
2use anyhow::{Context, Result};
3use bytes::{Buf, BufMut, Bytes, BytesMut};
4
5use super::{Compressor, Decompressor};
6
7/// Gzip compression algorithm. See [Kafka's broker configuration](https://kafka.apache.org/documentation/#brokerconfigs_compression.type)
8/// for more information.
9pub struct Zstd;
10
11const COMPRESSION_LEVEL: i32 = 3;
12
13impl<B: ByteBufMut> Compressor<B> for Zstd {
14    type BufMut = BytesMut;
15    fn compress<R, F>(buf: &mut B, f: F) -> Result<R>
16    where
17        F: FnOnce(&mut Self::BufMut) -> Result<R>,
18    {
19        // Write uncompressed bytes into a temporary buffer
20        let mut tmp = BytesMut::new();
21        let res = f(&mut tmp)?;
22
23        // Compress directly into the target buffer
24        zstd::stream::copy_encode(tmp.reader(), buf.writer(), COMPRESSION_LEVEL)
25            .context("Failed to compress zstd")?;
26        Ok(res)
27    }
28}
29
30impl<B: ByteBuf> Decompressor<B> for Zstd {
31    type Buf = Bytes;
32
33    fn decompress<R, F>(buf: &mut B, f: F) -> Result<R>
34    where
35        F: FnOnce(&mut Self::Buf) -> Result<R>,
36    {
37        let mut tmp = BytesMut::new().writer();
38        // Allocate a temporary buffer to hold the uncompressed bytes
39        let buf = buf.copy_to_bytes(buf.remaining());
40        zstd::stream::copy_decode(buf.reader(), &mut tmp).context("Failed to decompress zstd")?;
41
42        f(&mut tmp.into_inner().into())
43    }
44}
45
46#[cfg(test)]
47mod test {
48    use crate::compression::Zstd;
49    use crate::compression::{Compressor, Decompressor};
50    use anyhow::Result;
51    use bytes::BytesMut;
52    use std::fmt::Write;
53    use std::str;
54    use zstd::zstd_safe::WriteBuf;
55
56    #[test]
57    fn test_zstd() {
58        let mut compressed = BytesMut::new();
59        Zstd::compress(&mut compressed, |buf| -> Result<()> {
60            buf.write_str("hello zstd")?;
61            Ok(())
62        })
63        .unwrap();
64
65        Zstd::decompress(&mut compressed, |buf| -> Result<()> {
66            let decompressed_str = str::from_utf8(buf.as_slice())?;
67            assert_eq!(decompressed_str, "hello zstd");
68            Ok(())
69        })
70        .unwrap();
71    }
72}
kafka_protocol/compression/zstd.rs

kafka_protocol/compression/
zstd.rs
