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