dig_blockstore/compression.rs
1//! Zstd compression and bincode serialization for block bodies and headers.
2//!
3//! This module owns all serialization/deserialization logic:
4//! - [`BlockStore::serialize_block`] / [`BlockStore::deserialize_block`]: bincode + zstd (with optional dictionary).
5//! - [`BlockStore::serialize_header`] / [`BlockStore::deserialize_header`]: bincode only (no compression).
6//! - Dictionary training: [`BlockStore::train_dictionary`], [`BlockStore::maybe_train_dictionary`].
7//! - Dictionary loading: [`resolve_zstd_dictionary`], [`load_zstd_dict_from_db`].
8//!
9//! # Requirements
10//!
11//! - [`SER-001`](../docs/requirements/domains/serialization/specs/SER-001.md) — block serialization.
12//! - [`SER-002`](../docs/requirements/domains/serialization/specs/SER-002.md) — header serialization.
13//! - [`SER-005`](../docs/requirements/domains/serialization/specs/SER-005.md) — dictionary training.
14
15use std::sync::Arc;
16
17use dig_block::{L2Block, L2BlockHeader};
18use rand::seq::SliceRandom;
19use rocksdb::{IteratorMode, DB};
20
21use crate::constants::{
22 CF_BLOCKS, CF_METADATA, DICT_TARGET_SIZE, DICT_TRAINING_THRESHOLD, META_ZSTD_DICT,
23};
24use crate::error::BlockStoreError;
25use crate::store::BlockStore;
26
27/// Compression and serialization methods on [`BlockStore`].
28///
29/// These are `impl BlockStore` (not `BlockStoreInner`) because they were originally
30/// defined in the `impl BlockStore` block and callers use `Self::serialize_header()`
31/// which resolves through `BlockStore`. Field access goes through `Deref<Target=BlockStoreInner>`.
32impl BlockStore {
33 pub fn serialize_header(header: &L2BlockHeader) -> Result<Vec<u8>, BlockStoreError> {
34 bincode::serialize(header).map_err(|e| BlockStoreError::Serialization(e.to_string()))
35 }
36
37 /// Deserialize a header from [`CF_HEADERS`] bytes ([`SER-002`](../docs/requirements/domains/serialization/specs/SER-002.md)).
38 ///
39 /// **Read path:** raw bincode only — callers MUST NOT pass zstd-compressed payloads (those belong in [`CF_BLOCKS`]
40 /// via [`Self::deserialize_block`]).
41 pub fn deserialize_header(bytes: &[u8]) -> Result<L2BlockHeader, BlockStoreError> {
42 bincode::deserialize(bytes).map_err(|e| BlockStoreError::Serialization(e.to_string()))
43 }
44
45 /// Serialize then zstd-compress a block for [`CF_BLOCKS`] ([`SER-001`](../docs/requirements/domains/serialization/specs/SER-001.md)).
46 ///
47 /// **Pipeline:** [`bincode::serialize`] → [`zstd::bulk::Compressor::with_dictionary`] when
48 /// [`Self::use_compression_dict`] and a dictionary are present; otherwise [`zstd::encode_all`] (plain zstd).
49 ///
50 /// **Errors:** [`BlockStoreError::Serialization`] from bincode; [`BlockStoreError::Compression`] from zstd.
51 pub fn serialize_block(&self, block: &L2Block) -> Result<Vec<u8>, BlockStoreError> {
52 let raw = bincode::serialize(block)?;
53 if self.use_compression_dict {
54 let dict_guard = self.zstd_dict.read();
55 if let Some(dict) = dict_guard.as_ref() {
56 let mut compressor = zstd::bulk::Compressor::with_dictionary(
57 self.compression_level,
58 dict.as_slice(),
59 )
60 .map_err(|e| BlockStoreError::Compression(e.to_string()))?;
61 return compressor
62 .compress(raw.as_slice())
63 .map_err(BlockStoreError::compression_from_io);
64 }
65 }
66 zstd::encode_all(raw.as_slice(), self.compression_level)
67 .map_err(BlockStoreError::compression_from_io)
68 }
69
70 /// Reverse [`Self::serialize_block`] ([`SER-001`](../docs/requirements/domains/serialization/specs/SER-001.md)).
71 ///
72 /// **Fallback:** Dictionary decompress is attempted first when configured; on failure, plain
73 /// [`zstd::decode_all`] handles **pre-dictionary** payloads written before training ([`SER-005`](../docs/requirements/domains/serialization/specs/SER-005.md)).
74 ///
75 /// **Hash invariance:** Correct payloads MUST yield an [`L2Block`] whose [`L2Block::hash`] matches the original
76 /// pre-serialize block ([`SER-004`](../docs/requirements/domains/serialization/specs/SER-004.md); verified in `tests/ser_004_tests.rs`).
77 ///
78 /// **Errors:** Decompression failures map to [`BlockStoreError::Serialization`] so callers see a single
79 /// “payload unusable” surface for malformed CF_BYTES; bincode structural errors also use [`Serialization`](BlockStoreError::Serialization).
80 pub fn deserialize_block(&self, compressed: &[u8]) -> Result<L2Block, BlockStoreError> {
81 let raw = self.decompress_block_payload(compressed).map_err(|e| {
82 BlockStoreError::Serialization(format!("deserialize_block: decompress failed: {e}"))
83 })?;
84 bincode::deserialize(&raw).map_err(|e| BlockStoreError::Serialization(e.to_string()))
85 }
86
87 /// Decompress a raw zstd frame from [`CF_BLOCKS`] back to bincode bytes.
88 ///
89 /// # Fallback strategy ([`SER-005`])
90 ///
91 /// When dictionary mode is active, this method tries dictionary decompression first.
92 /// If that fails (because the payload was written *before* the dictionary was trained),
93 /// it falls back to plain [`zstd::decode_all`]. This two-phase approach ensures all
94 /// historical blocks remain readable after dictionary training—a critical invariant
95 /// since DIG does not re-encode existing blocks when a dictionary is installed.
96 ///
97 /// # Decompression bomb protection
98 ///
99 /// [`zstd::bulk::Decompressor::decompress`] accepts `max_decompressed_block_bytes` as
100 /// an upper bound on output size, preventing malicious payloads from exhausting memory.
101 /// The plain fallback path ([`zstd::decode_all`]) does not have this cap; future work
102 /// may wrap it similarly.
103 pub(crate) fn decompress_block_payload(&self, compressed: &[u8]) -> std::io::Result<Vec<u8>> {
104 if self.use_compression_dict {
105 if let Some(dict) = self.zstd_dict.read().as_ref() {
106 // Phase 1: attempt dictionary-aware decompression (post-training payloads).
107 let mut decompressor = zstd::bulk::Decompressor::with_dictionary(dict.as_slice())?;
108 return match decompressor.decompress(compressed, self.max_decompressed_block_bytes)
109 {
110 Ok(bytes) => Ok(bytes),
111 // Phase 2: dictionary decompression failed—payload is likely a pre-training
112 // plain zstd frame. Fall back to standard decoding.
113 Err(_) => zstd::decode_all(compressed),
114 };
115 }
116 }
117 // No dictionary configured or available: standard zstd decompression.
118 zstd::decode_all(compressed)
119 }
120
121 /// Full scan of [`CF_BLOCKS`] to count stored block rows.
122 ///
123 /// Iterates every key in the column family; used by [`crate::BlockStore::stats`] and by
124 /// [`Self::maybe_train_dictionary`] to detect when the training threshold is crossed
125 /// ([`SER-005`](../docs/requirements/domains/serialization/specs/SER-005.md)).
126 pub fn block_count(&self) -> Result<u64, BlockStoreError> {
127 let cf = self.cf(CF_BLOCKS)?;
128 let iter = self.db.iterator_cf(cf, IteratorMode::Start);
129 let mut n = 0u64;
130 for item in iter {
131 let (_k, _v) = item?;
132 n = n.saturating_add(1);
133 }
134 Ok(n)
135 }
136
137 /// **[`SER-005`](../docs/requirements/domains/serialization/specs/SER-005.md)** — Reload dictionary bytes from
138 /// [`META_ZSTD_DICT`] into memory after external maintenance (or to align with [`Self::train_dictionary`]
139 /// persistence).
140 ///
141 /// **Startup:** [`Self::open`] already embeds this via [`load_zstd_dict_from_db`]; public callers use
142 /// `init_dictionary` when a **second process** trains the dictionary or metadata is repaired online.
143 pub fn init_dictionary(&self) -> Result<(), BlockStoreError> {
144 let loaded = load_zstd_dict_from_db(&self.db, self.use_compression_dict)?;
145 *self.zstd_dict.write() = loaded;
146 Ok(())
147 }
148
149 /// Collect `sample_count` **uncompressed** bincode block bodies for [`zstd::dict::from_samples`].
150 ///
151 /// **Randomness:** Keys/values are shuffled with [`rand::thread_rng`] so training sees a representative slice of
152 /// the corpus, not a height-ordered prefix ([`SER-005`](../docs/requirements/domains/serialization/specs/SER-005.md) implementation notes).
153 pub(crate) fn sample_block_bodies(
154 &self,
155 sample_count: usize,
156 ) -> Result<Vec<Vec<u8>>, BlockStoreError> {
157 let cf = self.cf(CF_BLOCKS)?;
158 let mut blobs: Vec<Vec<u8>> = Vec::new();
159 let iter = self.db.iterator_cf(cf, IteratorMode::Start);
160 for item in iter {
161 let (_key, value) = item?;
162 blobs.push(value.to_vec());
163 }
164 if blobs.len() < sample_count {
165 return Err(BlockStoreError::Serialization(format!(
166 "dictionary training: need at least {sample_count} blocks in {CF_BLOCKS}, have {}",
167 blobs.len()
168 )));
169 }
170 blobs.shuffle(&mut rand::thread_rng());
171 blobs.truncate(sample_count);
172 let mut samples = Vec::with_capacity(sample_count);
173 for compressed in blobs {
174 let raw = self.decompress_block_payload(&compressed).map_err(|e| {
175 BlockStoreError::Serialization(format!(
176 "dictionary training sample decompress: {e}"
177 ))
178 })?;
179 samples.push(raw);
180 }
181 Ok(samples)
182 }
183
184 /// Train + persist a zstd dictionary; **idempotent** if [`META_ZSTD_DICT`] already contains bytes.
185 pub(crate) fn train_dictionary(&self) -> Result<Vec<u8>, BlockStoreError> {
186 let meta = self.cf(CF_METADATA)?;
187 if let Some(blob) = self.db.get_cf(meta, META_ZSTD_DICT.as_bytes())? {
188 if !blob.is_empty() {
189 return Ok(blob);
190 }
191 }
192 let n = DICT_TRAINING_THRESHOLD as usize;
193 let samples = self.sample_block_bodies(n)?;
194 let refs: Vec<&[u8]> = samples.iter().map(Vec::as_slice).collect();
195 let dict = zstd::dict::from_samples(&refs, DICT_TARGET_SIZE).map_err(|e| {
196 BlockStoreError::Serialization(format!("dictionary training failed: {e}"))
197 })?;
198 self.db
199 .put_cf(meta, META_ZSTD_DICT.as_bytes(), dict.as_slice())?;
200 Ok(dict)
201 }
202
203 /// If dictionary training is enabled, the live dictionary slot is empty, and [`Self::block_count`] is at or above
204 /// [`DICT_TRAINING_THRESHOLD`], train once and install into memory.
205 pub(crate) fn maybe_train_dictionary(&self) -> Result<(), BlockStoreError> {
206 if !self.use_compression_dict {
207 return Ok(());
208 }
209 if self.zstd_dict.read().is_some() {
210 return Ok(());
211 }
212 let meta = self.cf(CF_METADATA)?;
213 if self
214 .db
215 .get_cf(meta, META_ZSTD_DICT.as_bytes())?
216 .filter(|b| !b.is_empty())
217 .is_some()
218 {
219 self.init_dictionary()?;
220 return Ok(());
221 }
222 if self.block_count()? < DICT_TRAINING_THRESHOLD {
223 return Ok(());
224 }
225 let dict = self.train_dictionary()?;
226 *self.zstd_dict.write() = Some(Arc::new(dict));
227 Ok(())
228 }
229}
230
231pub(crate) fn resolve_zstd_dictionary(
232 db: &DB,
233 use_compression_dict: bool,
234 override_bytes: Option<Vec<u8>>,
235) -> Result<Option<Arc<Vec<u8>>>, BlockStoreError> {
236 if let Some(bytes) = override_bytes {
237 return if bytes.is_empty() {
238 Ok(None)
239 } else {
240 Ok(Some(Arc::new(bytes)))
241 };
242 }
243 load_zstd_dict_from_db(db, use_compression_dict)
244}
245
246/// Load the trained zstd dictionary from [`CF_METADATA`] / [`META_ZSTD_DICT`].
247///
248/// Returns `None` when:
249/// - `use_compression_dict` is `false` (feature disabled in config).
250/// - The [`META_ZSTD_DICT`] key does not exist (no training has occurred).
251/// - The stored blob is empty (edge case: metadata key exists but value is zero-length).
252///
253/// The returned `Arc<Vec<u8>>` is shared between the [`BlockStore`] field `zstd_dict`
254/// and all compress/decompress operations, avoiding per-call copies of the ~100 KB dictionary.
255///
256/// # Called by
257///
258/// [`resolve_zstd_dictionary`] (at open time) and [`BlockStore::init_dictionary`] (runtime reload).
259pub(crate) fn load_zstd_dict_from_db(
260 db: &DB,
261 use_compression_dict: bool,
262) -> Result<Option<Arc<Vec<u8>>>, BlockStoreError> {
263 if !use_compression_dict {
264 return Ok(None);
265 }
266 let meta = db
267 .cf_handle(CF_METADATA)
268 .ok_or_else(|| BlockStoreError::Serialization("missing CF_METADATA".into()))?;
269 let Some(blob) = db.get_cf(meta, META_ZSTD_DICT.as_bytes())? else {
270 return Ok(None);
271 };
272 if blob.is_empty() {
273 return Ok(None);
274 }
275 Ok(Some(Arc::new(blob)))
276}