subms_lsm_tree/features/
zstd.rs

1//! Zstd block compression for SSTable data blocks.
2//!
3//! Block format mirrors the LZ4 wrapper exactly so the read path can dispatch
4//! by the algo byte if multiple compressors share a file:
5//! ```text
6//! marker:    u8   = 0x5A ('Z')
7//! algo:      u8   = 0x00 (stored) | 0x01 (zstd)
8//! uncomp:    u32  uncompressed byte length
9//! data:      bytes
10//! ```
11//!
12//! Compression level defaults to 3 (zstd's default speed/ratio knee). Use
13//! [`ZstdBlockCompressor::with_level`] to override. Levels outside 1..=22 are clamped to
14//! that range.
15
16use std::io;
17
18const MARKER: u8 = 0x5A;
19const ALGO_STORED: u8 = 0x00;
20const ALGO_ZSTD: u8 = 0x01;
21const HEADER_LEN: usize = 1 + 1 + 4;
22const DEFAULT_LEVEL: i32 = 3;
23const MIN_LEVEL: i32 = 1;
24const MAX_LEVEL: i32 = 22;
25
26pub struct ZstdBlockCompressor {
27    level: i32,
28}
29
30impl ZstdBlockCompressor {
31    pub fn new() -> Self {
32        Self {
33            level: DEFAULT_LEVEL,
34        }
35    }
36
37    pub fn with_level(level: i32) -> Self {
38        Self {
39            level: level.clamp(MIN_LEVEL, MAX_LEVEL),
40        }
41    }
42
43    pub fn level(&self) -> i32 {
44        self.level
45    }
46
47    pub fn compress(&self, block: &[u8]) -> io::Result<Vec<u8>> {
48        let compressed = zstd::bulk::compress(block, self.level)?;
49        let mut out = Vec::with_capacity(HEADER_LEN + compressed.len().max(block.len()));
50        out.push(MARKER);
51        if compressed.len() < block.len() {
52            out.push(ALGO_ZSTD);
53            out.extend_from_slice(&(block.len() as u32).to_be_bytes());
54            out.extend_from_slice(&compressed);
55        } else {
56            out.push(ALGO_STORED);
57            out.extend_from_slice(&(block.len() as u32).to_be_bytes());
58            out.extend_from_slice(block);
59        }
60        Ok(out)
61    }
62
63    pub fn decompress(&self, buf: &[u8]) -> io::Result<Vec<u8>> {
64        if buf.len() < HEADER_LEN {
65            return Err(io::Error::new(
66                io::ErrorKind::InvalidData,
67                "zstd block: too short",
68            ));
69        }
70        if buf[0] != MARKER {
71            return Err(io::Error::new(
72                io::ErrorKind::InvalidData,
73                "zstd block: bad marker",
74            ));
75        }
76        let algo = buf[1];
77        let uncomp_len = u32::from_be_bytes(buf[2..6].try_into().unwrap()) as usize;
78        let payload = &buf[HEADER_LEN..];
79        match algo {
80            ALGO_STORED => {
81                if payload.len() != uncomp_len {
82                    return Err(io::Error::new(
83                        io::ErrorKind::InvalidData,
84                        "zstd block: stored payload size mismatch",
85                    ));
86                }
87                Ok(payload.to_vec())
88            }
89            ALGO_ZSTD => zstd::bulk::decompress(payload, uncomp_len),
90            other => Err(io::Error::new(
91                io::ErrorKind::InvalidData,
92                format!("zstd block: unknown algo byte 0x{other:02x}"),
93            )),
94        }
95    }
96}
97
98impl Default for ZstdBlockCompressor {
99    fn default() -> Self {
100        Self::new()
101    }
102}
103
104#[cfg(test)]
105#[path = "zstd_tests.rs"]
106mod tests;
subms_lsm_tree/features/zstd.rs

subms_lsm_tree/features/
zstd.rs
