Skip to main content

trueno_db/kv/
compressed.rs

1//! Compressed KV Store wrapper (GH-5)
2//!
3//! Provides transparent LZ4/ZSTD compression for any `KvStore` backend.
4//! Ideal for reducing memory footprint of LLM KV caches.
5//!
6//! Compression algorithm is shared via `batuta_common::compression`.
7
8use crate::kv::KvStore;
9use crate::Result;
10
11pub use batuta_common::compression::Compression;
12
13/// Compressed KV store wrapper
14///
15/// Wraps any `KvStore` implementation and transparently compresses/decompresses values.
16///
17/// # Example
18///
19/// ```rust,no_run
20/// use trueno_db::kv::{CompressedKvStore, Compression, MemoryKvStore, KvStore};
21///
22/// # async fn example() -> trueno_db::Result<()> {
23/// let inner = MemoryKvStore::new();
24/// let store = CompressedKvStore::new(inner, Compression::Lz4);
25///
26/// // Values are transparently compressed
27/// store.set("key", vec![0u8; 10000]).await?;
28/// let value = store.get("key").await?;
29/// # Ok(())
30/// # }
31/// ```
32#[derive(Debug)]
33pub struct CompressedKvStore<S: KvStore> {
34    inner: S,
35    compression: Compression,
36}
37
38impl<S: KvStore> CompressedKvStore<S> {
39    /// Create a new compressed KV store wrapping the given store
40    #[must_use]
41    pub const fn new(inner: S, compression: Compression) -> Self {
42        Self { inner, compression }
43    }
44
45    /// Get reference to inner store (for inspection/testing)
46    #[must_use]
47    pub const fn inner(&self) -> &S {
48        &self.inner
49    }
50
51    /// Get compression algorithm
52    #[must_use]
53    pub const fn compression(&self) -> Compression {
54        self.compression
55    }
56}
57
58impl<S: KvStore> KvStore for CompressedKvStore<S> {
59    async fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
60        match self.inner.get(key).await? {
61            Some(compressed) => {
62                let decompressed = self.compression.decompress(&compressed)?;
63                Ok(Some(decompressed))
64            }
65            None => Ok(None),
66        }
67    }
68
69    async fn set(&self, key: &str, value: Vec<u8>) -> Result<()> {
70        let compressed = self.compression.compress(&value)?;
71        self.inner.set(key, compressed).await
72    }
73
74    async fn delete(&self, key: &str) -> Result<()> {
75        self.inner.delete(key).await
76    }
77
78    async fn exists(&self, key: &str) -> Result<bool> {
79        self.inner.exists(key).await
80    }
81}
82
83#[cfg(test)]
84mod tests {
85    use super::*;
86
87    #[test]
88    fn test_compression_as_str() {
89        assert_eq!(Compression::Lz4.as_str(), "lz4");
90        assert_eq!(Compression::Zstd.as_str(), "zstd");
91    }
92
93    #[test]
94    fn test_compression_default() {
95        assert_eq!(Compression::default(), Compression::Lz4);
96    }
97
98    #[test]
99    fn test_lz4_compress_decompress() {
100        let data = b"hello world hello world hello world".to_vec();
101        let compressed = Compression::Lz4.compress(&data).unwrap();
102        let decompressed = Compression::Lz4.decompress(&compressed).unwrap();
103        assert_eq!(decompressed, data);
104    }
105
106    #[test]
107    fn test_zstd_compress_decompress() {
108        let data = b"hello world hello world hello world".to_vec();
109        let compressed = Compression::Zstd.compress(&data).unwrap();
110        let decompressed = Compression::Zstd.decompress(&compressed).unwrap();
111        assert_eq!(decompressed, data);
112    }
113
114    #[test]
115    fn test_empty_data_compression() {
116        let empty: Vec<u8> = vec![];
117
118        let lz4_compressed = Compression::Lz4.compress(&empty).unwrap();
119        assert!(lz4_compressed.is_empty());
120        let lz4_decompressed = Compression::Lz4.decompress(&lz4_compressed).unwrap();
121        assert!(lz4_decompressed.is_empty());
122
123        let zstd_compressed = Compression::Zstd.compress(&empty).unwrap();
124        assert!(zstd_compressed.is_empty());
125        let zstd_decompressed = Compression::Zstd.decompress(&zstd_compressed).unwrap();
126        assert!(zstd_decompressed.is_empty());
127    }
128
129    #[test]
130    fn test_lz4_compresses_repeated_data() {
131        let data = vec![0u8; 10000];
132        let compressed = Compression::Lz4.compress(&data).unwrap();
133        // LZ4 should achieve >10x compression on zeros
134        assert!(compressed.len() < data.len() / 10);
135    }
136
137    #[test]
138    fn test_zstd_compresses_repeated_data() {
139        let data = vec![0u8; 10000];
140        let compressed = Compression::Zstd.compress(&data).unwrap();
141        // ZSTD should achieve >10x compression on zeros
142        assert!(compressed.len() < data.len() / 10);
143    }
144}