Skip to main content

cuckoo_clock/
exporter.rs

1use std::{error::Error, fmt::Display, num::TryFromIntError, sync::Mutex};
2#[allow(deprecated)]
3use std::{
4    hash::{BuildHasher, SipHasher},
5    io::{Read, Write},
6};
7
8use crate::{
9    bucket::Bucket,
10    config::{
11        BitCount, CounterConfig, CuckooConfiguration, LruAgingStrategy, LruConfig, TtlConfig,
12    },
13};
14
15/// Error type for import operations.
16#[derive(Debug)]
17pub enum ImportError {
18    /// Error due to underlying IO error.
19    Io(std::io::Error),
20    /// Error due to reading invalid stored configuration.
21    Config(crate::config::ConfigError),
22    /// Error due to reading a zero value for a field that does not allow zeroes.
23    NonZeroError(TryFromIntError),
24    /// Error due to mismatch in hasher names, between the stored and expected hasher.
25    InvalidHasherName {
26        /// Expected name for the hasher (defined by the type of hasher for [`crate::CuckooFilter`]).
27        expected: String,
28        /// Found name for the hasher (value stored in the exported data).
29        found: String,
30    },
31    /// Error due to malformed header in the export.
32    InvalidHasherHeader {
33        /// Found value in the header that did not match the standard value.
34        found: String,
35    },
36}
37
38/// Error type for export operations.
39#[derive(Debug)]
40pub enum ExportError {
41    /// Error due to underlying IO error.
42    Io(std::io::Error),
43}
44
45impl Display for ImportError {
46    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
47        match self {
48            ImportError::Io(error) => write!(f, "IO error while importing: {}", error),
49            ImportError::Config(error) => {
50                write!(f, "config error on imported data: {}", error)
51            }
52            ImportError::NonZeroError(error) => {
53                write!(f, "conversion error in imported data: {}", error)
54            }
55            ImportError::InvalidHasherName { expected, found } => write!(
56                f,
57                "Invalid hasher name found in header: \"{}\". Expected: \"{}\"",
58                found, expected
59            ),
60            ImportError::InvalidHasherHeader { found } => write!(
61                f,
62                "Invalid hasher header. Expected \"{}\" at the end, but found \"{}\"",
63                HEADER_END, found
64            ),
65        }
66    }
67}
68
69impl Error for ImportError {
70    fn cause(&self) -> Option<&dyn Error> {
71        match self {
72            ImportError::Io(error) => Some(error),
73            ImportError::Config(error) => Some(error),
74            ImportError::NonZeroError(error) => Some(error),
75            ImportError::InvalidHasherName { .. } => None,
76            ImportError::InvalidHasherHeader { .. } => None,
77        }
78    }
79}
80
81impl From<std::io::Error> for ImportError {
82    fn from(value: std::io::Error) -> Self {
83        Self::Io(value)
84    }
85}
86
87impl From<crate::config::ConfigError> for ImportError {
88    fn from(value: crate::config::ConfigError) -> Self {
89        Self::Config(value)
90    }
91}
92
93impl From<TryFromIntError> for ImportError {
94    fn from(value: TryFromIntError) -> Self {
95        Self::NonZeroError(value)
96    }
97}
98
99impl Display for ExportError {
100    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
101        match self {
102            ExportError::Io(error) => write!(f, "IO error while exporting: {}", error),
103        }
104    }
105}
106
107impl Error for ExportError {
108    fn cause(&self) -> Option<&dyn Error> {
109        match self {
110            ExportError::Io(error) => Some(error),
111        }
112    }
113}
114
115impl From<std::io::Error> for ExportError {
116    fn from(value: std::io::Error) -> Self {
117        Self::Io(value)
118    }
119}
120
121const HEADER: &str = "cuckoo-clock:";
122const HEADER_END: &str = ":header-end\n";
123
124/// A trait representing objects that can be exported (e.g. persisted into a file).
125///
126/// This trait is needed for [`cuckoo_clock::CuckooFilter::export`].
127pub trait Exportable: Clone {
128    /// Writes the state of this object into the provided writer.
129    fn write_to(&self, writer: impl Write) -> Result<(), ExportError>;
130
131    /// Reads the stored state to create an instance of this object.
132    fn read_from(reader: impl Read) -> Result<Self, ImportError>;
133}
134
135/// A trait representing [`BuildHasher`] instances that can be exported (e.g. persisted into a file).
136///
137/// This trait is needed for [`cuckoo_clock::CuckooFilter::export`].
138pub trait ExportableBuildHasher: Exportable + BuildHasher {
139    /// Unique of this [`ExportableBuildHasher`]. This will be stored in the header of exported
140    /// data.
141    const NAME: &str;
142}
143
144/// Implementation of [`ExportableRandomState`], similar to [`std::hash::RandomState`].
145#[derive(Clone)]
146pub struct ExportableRandomState {
147    k0: u64,
148    k1: u64,
149}
150
151impl ExportableRandomState {
152    pub(crate) const fn new() -> Self {
153        Self { k0: 0, k1: 0 }
154    }
155
156    /// Creates a new random instance.
157    #[must_use]
158    pub fn new_random() -> Self {
159        Self {
160            k0: rand::random(),
161            k1: rand::random(),
162        }
163    }
164}
165
166impl Default for ExportableRandomState {
167    fn default() -> Self {
168        Self::new()
169    }
170}
171
172#[allow(deprecated)]
173impl BuildHasher for ExportableRandomState {
174    type Hasher = SipHasher;
175
176    fn build_hasher(&self) -> Self::Hasher {
177        SipHasher::new_with_keys(self.k0, self.k1)
178    }
179}
180
181impl Exportable for ExportableRandomState {
182    fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
183        let k0 = self.k0.to_be_bytes();
184        let k1 = self.k1.to_be_bytes();
185        writer.write_all(&k0)?;
186        writer.write_all(&k1)?;
187        Ok(())
188    }
189
190    fn read_from(mut reader: impl Read) -> Result<Self, ImportError> {
191        let mut buf = [0u8; 8];
192        reader.read_exact(&mut buf)?;
193        let k0 = u64::from_be_bytes(buf);
194        reader.read_exact(&mut buf)?;
195        let k1 = u64::from_be_bytes(buf);
196        Ok(Self { k0, k1 })
197    }
198}
199
200impl ExportableBuildHasher for ExportableRandomState {
201    const NAME: &str = "cuckoo_clock::exporter::ExportableRandomState";
202}
203
204/// Exporter for [`crate::CuckooFilter`].
205///
206/// This implements [`Read`] to allow reading this data into whatever storage required.
207pub struct CuckooFilterExporter<'a, H: ExportableBuildHasher> {
208    hasher: &'a H,
209    buckets: &'a Vec<Mutex<Bucket>>,
210    config: &'a CuckooConfiguration,
211}
212
213impl<'a, H: ExportableBuildHasher> CuckooFilterExporter<'a, H> {
214    pub(crate) const fn new(
215        hasher: &'a H,
216        buckets: &'a Vec<Mutex<Bucket>>,
217        config: &'a CuckooConfiguration,
218    ) -> Self {
219        Self {
220            hasher,
221            buckets,
222            config,
223        }
224    }
225
226    /// Reads in the entire [`crate::CuckooFilter`] state into the provided [`Write`] instance.
227    ///
228    /// This will lock buckets one by one and can be run concurrently with other
229    /// [`crate::CuckooFilter`] operations, but it may result in a state which
230    /// combines old values for some buckets and newer valus for some other buckets.
231    /// The resulting state should still be valid.
232    pub fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
233        writer.write_all(format!("{}:{}:", HEADER, H::NAME).as_bytes())?;
234        self.hasher.write_to(&mut writer)?;
235        writer.write_all(HEADER_END.as_bytes())?;
236
237        self.config.write_to(&mut writer)?;
238        writer.write_all(b"\n")?;
239
240        for b in self.buckets.iter() {
241            #[expect(clippy::unwrap_used)]
242            let bucket = b.lock().unwrap();
243            bucket.export(&mut writer)?;
244        }
245
246        Ok(())
247    }
248
249    /// Reads the entire [`crate::CuckooFilter`] state into a [`Vec`].
250    ///
251    /// This can be useful if writing to slower [`Write`] interfaces, to ensure less changes are
252    /// made while exporting is in progress.
253    pub fn snapshot(&self) -> Result<Vec<u8>, ExportError> {
254        let mut result = Vec::new();
255        self.write_to(&mut result)?;
256        Ok(result)
257    }
258}
259
260pub(crate) fn read_hasher_from<H: ExportableBuildHasher + BuildHasher>(
261    mut reader: impl Read,
262) -> Result<H, ImportError> {
263    let mut header_prefix = vec![0; HEADER.len() + H::NAME.len() + 2];
264    reader.read_exact(&mut header_prefix)?;
265    if header_prefix != format!("{}:{}:", HEADER, H::NAME).as_bytes() {
266        return Err(ImportError::InvalidHasherName {
267            expected: H::NAME.to_string(),
268            found: String::from_utf8(header_prefix)
269                .map(|s| {
270                    s.trim_start_matches(&format!("{}:", HEADER))
271                        .trim_end_matches(":")
272                        .to_string()
273                })
274                .unwrap_or("<non UTF-8 data>".to_string()),
275        });
276    };
277    let hasher = H::read_from(&mut reader)?;
278    let mut header_suffix = vec![0; HEADER_END.len()];
279    reader.read_exact(&mut header_suffix)?;
280    if header_suffix != HEADER_END.as_bytes() {
281        return Err(ImportError::InvalidHasherHeader {
282            found: String::from_utf8(header_suffix).unwrap_or("<non UTF-8 data>".to_string()),
283        });
284    };
285    Ok(hasher)
286}
287
288impl Exportable for CuckooConfiguration {
289    fn write_to(&self, mut writer: impl std::io::Write) -> Result<(), crate::ExportError> {
290        // 0 as version start, and the rest is lib version
291        writer.write_all(&[0, 0, 2, 4])?;
292        // Since we are reading a value from out field config, we know that it can't ever be higher
293        // than 32 and should fit into u8.
294        #[allow(clippy::expect_used)]
295        // Value mask is calculated as bits ^ 2 - 1, so we get back the bit count with ilog2
296        writer.write_all(
297            &u8::try_from((self.fingerprint_field_config.value_mask() as usize + 1).ilog2())
298                .expect("Fingeprint bits can't be higher than 32")
299                .to_be_bytes(),
300        )?;
301        writer.write_all(&(self.bucket_size as u64).to_be_bytes())?;
302        // Max entries
303        writer.write_all(&((self.bucket_count * self.bucket_size) as u64).to_be_bytes())?;
304        writer.write_all(&(self.max_kicks as u64).to_be_bytes())?;
305        if let Some((lru, _)) = &self.lru_field_config {
306            writer.write_all(&[1])?;
307            lru.write_to(&mut writer)?;
308        }
309        if let Some((ttl, _)) = &self.ttl_field_config {
310            writer.write_all(&[2])?;
311            ttl.write_to(&mut writer)?;
312        }
313        if let Some((counter, _)) = &self.counter_field_config {
314            writer.write_all(&[3])?;
315            counter.write_to(&mut writer)?;
316        }
317        Ok(())
318    }
319
320    fn read_from(mut reader: impl std::io::Read) -> Result<Self, crate::ImportError> {
321        let mut u8_buf = [0u8; 1];
322        let mut u64_buf = [0u8; 8];
323        reader.read_exact(&mut u8_buf)?;
324        let mut version: usize = 0;
325        if u8::from_be_bytes(u8_buf) == 0 {
326            // This is a version header
327            reader.read_exact(&mut u8_buf)?;
328            version += 1_000_000 * usize::from(u8::from_be_bytes(u8_buf));
329            reader.read_exact(&mut u8_buf)?;
330            version += 1_000 * usize::from(u8::from_be_bytes(u8_buf));
331            reader.read_exact(&mut u8_buf)?;
332            version += usize::from(u8::from_be_bytes(u8_buf));
333            reader.read_exact(&mut u8_buf)?;
334        } else {
335            // This is an old version (< 0.2.0) - without version header
336            version = 1_000;
337        }
338        let fp_bits: BitCount = usize::from(u8::from_be_bytes(u8_buf)).try_into()?;
339        reader.read_exact(&mut u64_buf)?;
340        let bucket_size = u64::from_be_bytes(u64_buf);
341        reader.read_exact(&mut u64_buf)?;
342        let max_entries = u64::from_be_bytes(u64_buf);
343        reader.read_exact(&mut u64_buf)?;
344        let max_kicks = u64::from_be_bytes(u64_buf);
345        let mut builder = CuckooConfiguration::builder(max_entries.try_into()?)
346            .fingerprint_bits(fp_bits)
347            .bucket_size(usize::try_from(bucket_size)?.try_into()?)
348            .max_kicks(max_kicks.try_into()?);
349        while let Ok(()) = reader.read_exact(&mut u8_buf) {
350            let conf_type = u8::from_be_bytes(u8_buf);
351            match conf_type {
352                1 => {
353                    builder = builder.with_lru(LruConfig::read_from(&mut reader, version)?);
354                }
355                2 => {
356                    builder = builder.with_ttl(TtlConfig::read_from(&mut reader, version)?);
357                }
358                3 => {
359                    builder = builder.with_counter(CounterConfig::read_from(&mut reader, version)?);
360                }
361                // We can't handle this type, so abort
362                // TODO: return an error?
363                _ => break,
364            }
365        }
366        Ok(builder.build()?)
367    }
368}
369
370/// A trait representing config objects that can be exported (e.g. persisted into a file).
371pub(crate) trait ExportableConfig: Clone {
372    /// Writes the state of this object into the provided writer.
373    fn write_to(&self, writer: impl Write) -> Result<(), ExportError>;
374
375    /// Reads the stored state to create an instance of this object.
376    fn read_from(reader: impl Read, version: usize) -> Result<Self, ImportError>;
377}
378
379impl ExportableConfig for LruConfig {
380    fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
381        writer.write_all(&u8::from(self.counter_bits).to_be_bytes())?;
382        writer.write_all(&u8::from(self.remove_on_zero).to_be_bytes())?;
383        self.aging_strategy.write_to(&mut writer)?;
384        writer.write_all(&self.starting_value.to_be_bytes())?;
385        writer.write_all(&self.increment.to_be_bytes())?;
386        Ok(())
387    }
388
389    fn read_from(mut reader: impl Read, version: usize) -> Result<Self, ImportError> {
390        let mut u8_buf = [0u8; 1];
391        let mut u32_buf = [0u8; 4];
392        reader.read_exact(&mut u8_buf)?;
393        let bits = u8::from_be_bytes(u8_buf);
394        let mut remove_on_zero = false;
395        let mut aging_strategy = LruAgingStrategy::default();
396        let mut increment = 1;
397        let mut starting_value = 1;
398        if version >= 2_000 {
399            reader.read_exact(&mut u8_buf)?;
400            remove_on_zero = u8::from_be_bytes(u8_buf) != 0;
401        }
402        if version >= 2_002 {
403            aging_strategy = LruAgingStrategy::read_from(&mut reader, version)?;
404            reader.read_exact(&mut u32_buf)?;
405            starting_value = u32::from_be_bytes(u32_buf);
406            reader.read_exact(&mut u32_buf)?;
407            increment = u32::from_be_bytes(u32_buf);
408        }
409        Ok(Self {
410            counter_bits: (bits as usize).try_into()?,
411            aging_strategy,
412            starting_value,
413            remove_on_zero,
414            increment,
415        })
416    }
417}
418
419impl ExportableConfig for LruAgingStrategy {
420    fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
421        match self {
422            LruAgingStrategy::Halving => writer.write_all(&1u8.to_be_bytes())?,
423            LruAgingStrategy::Decrement(value) => {
424                writer.write_all(&2u8.to_be_bytes())?;
425                writer.write_all(&value.to_be_bytes())?;
426            }
427        }
428        Ok(())
429    }
430
431    fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
432        let mut u8_buf = [0u8; 1];
433        let mut u32_buf = [0u8; 4];
434        reader.read_exact(&mut u8_buf)?;
435        let strategy = u8::from_be_bytes(u8_buf);
436        match strategy {
437            1 => Ok(Self::Halving),
438            2 => {
439                reader.read_exact(&mut u32_buf)?;
440                let value = u32::from_be_bytes(u32_buf);
441                Ok(Self::Decrement(value))
442            }
443            // TODO: error for this?
444            _ => Ok(Self::default()),
445        }
446    }
447}
448
449impl ExportableConfig for TtlConfig {
450    fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
451        writer.write_all(&self.ttl.get().to_be_bytes())?;
452        writer.write_all(&u8::from(self.ttl_bits).to_be_bytes())?;
453        Ok(())
454    }
455
456    fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
457        let mut u8_buf = [0u8; 1];
458        let mut u32_buf = [0u8; 4];
459        reader.read_exact(&mut u32_buf)?;
460        let ttl = u32::from_be_bytes(u32_buf);
461        reader.read_exact(&mut u8_buf)?;
462        let bits = u8::from_be_bytes(u8_buf);
463        Ok(Self {
464            ttl: ttl.try_into()?,
465            ttl_bits: (bits as usize).try_into()?,
466        })
467    }
468}
469
470impl ExportableConfig for CounterConfig {
471    fn write_to(&self, mut writer: impl Write) -> Result<(), ExportError> {
472        writer.write_all(&u8::from(self.counter_bits).to_be_bytes())?;
473        writer.write_all(&self.change_on_insert.to_be_bytes())?;
474        writer.write_all(&self.change_on_lookup.to_be_bytes())?;
475        Ok(())
476    }
477
478    fn read_from(mut reader: impl Read, _version: usize) -> Result<Self, ImportError> {
479        let mut u8_buf = [0u8; 1];
480        let mut u32_buf = [0u8; 4];
481        reader.read_exact(&mut u8_buf)?;
482        let bits = u8::from_be_bytes(u8_buf);
483        reader.read_exact(&mut u32_buf)?;
484        let change_on_insert = i32::from_be_bytes(u32_buf);
485        reader.read_exact(&mut u32_buf)?;
486        let change_on_lookup = i32::from_be_bytes(u32_buf);
487        Ok(Self {
488            counter_bits: (bits as usize).try_into()?,
489            change_on_insert,
490            change_on_lookup,
491        })
492    }
493}