Skip to main content

cuckoo_cache/
cuckoo.rs

1//! Core cuckoo cache implementation.
2//!
3//! Each item slot uses the [`TinyItem`] layout from the keyvalue crate:
4//!
5//! ```text
6//! ┌──────────┬──────┬──────┬──────────┬──────────┐
7//! │  EXPIRE  │ KLEN │ VLEN │   KEY    │  VALUE   │
8//! │  (u32)   │ (u8) │ (u8) │          │          │
9//! │ 4 bytes  │ 1 b  │ 1 b  │          │          │
10//! └──────────┴──────┴──────┴──────────┴──────────┘
11//! ```
12//!
13//! A slot is empty when `expire == 0` (all bytes zeroed).
14//! Integer values are signalled by `vlen == 0`.
15
16use crate::*;
17use ahash::RandomState;
18use clocksource::coarse::Instant;
19use core::hash::{BuildHasher, Hasher};
20use keyvalue::{TinyItem, Value, TINY_ITEM_HDR_SIZE};
21use rand::RngExt;
22
23/// Fixed seeds for the four independent hash functions, analogous to the
24/// initialization vectors used in the C implementation.
25const SEEDS: [[u64; 4]; D] = [
26    [0x3ac5_d673, 0x6d78_39d0, 0x2b58_1cf5, 0x4dd2_be0a],
27    [0x9e37_79b9, 0x517c_c1b7, 0x27d4_eb2f, 0x3c6e_f372],
28    [0xdead_beef, 0xcafe_babe, 0x1234_5678, 0xfeed_face],
29    [0xa076_1d64, 0xe703_7ed1, 0x8ebc_6af0, 0x5899_65cd],
30];
31
32/// A pre-allocated cache that uses cuckoo hashing with D=4 candidate positions
33/// per key. Items are stored inline in fixed-size slots within a contiguous
34/// array.
35pub struct CuckooCache {
36    /// Backing storage: `nitem` slots of `item_size` bytes each.
37    data: Box<[u8]>,
38    /// Four independent hash builders for cuckoo hashing.
39    hashers: Box<[RandomState; D]>,
40    /// Bytes per item slot.
41    item_size: usize,
42    /// Total number of item slots.
43    nitem: usize,
44    /// Maximum displacement chain depth.
45    max_displace: usize,
46    /// Maximum TTL in seconds.
47    max_ttl: u32,
48    /// Eviction policy.
49    policy: Policy,
50    /// Creation time for computing relative expiration timestamps.
51    started: Instant,
52}
53
54impl CuckooCache {
55    /// Returns a new [`Builder`] for configuring a `CuckooCache`.
56    ///
57    /// ```
58    /// use cuckoo_cache::CuckooCache;
59    ///
60    /// let cache = CuckooCache::builder()
61    ///     .nitem(4096)
62    ///     .item_size(64)
63    ///     .build();
64    /// ```
65    pub fn builder() -> Builder {
66        Builder::default()
67    }
68
69    pub(crate) fn from_builder(b: Builder) -> Self {
70        assert!(
71            b.item_size > TINY_ITEM_HDR_SIZE,
72            "item_size must be greater than {} bytes (header overhead)",
73            TINY_ITEM_HDR_SIZE
74        );
75        assert!(b.nitem > 0, "nitem must be positive");
76
77        let total = b
78            .item_size
79            .checked_mul(b.nitem)
80            .expect("total storage size overflow");
81
82        let data = vec![0u8; total].into_boxed_slice();
83        let hashers = Box::new(SEEDS.map(|s| RandomState::with_seeds(s[0], s[1], s[2], s[3])));
84
85        debug!(
86            "cuckoo cache: {} items x {} bytes = {} bytes total",
87            b.nitem, b.item_size, total,
88        );
89
90        Self {
91            data,
92            hashers,
93            item_size: b.item_size,
94            nitem: b.nitem,
95            max_displace: b.max_displace,
96            max_ttl: b.max_ttl,
97            policy: b.policy,
98            started: Instant::now(),
99        }
100    }
101
102    // -----------------------------------------------------------------------
103    // Slot access helpers
104    // -----------------------------------------------------------------------
105
106    /// Byte offset into `self.data` for slot `index`.
107    #[inline]
108    fn slot_offset(&self, index: usize) -> usize {
109        index * self.item_size
110    }
111
112    /// Get a [`TinyItem`] view of a slot.
113    fn slot_item(&self, index: usize) -> TinyItem {
114        let off = self.slot_offset(index);
115        unsafe { TinyItem::from_ptr((self.data.as_ptr() as *mut u8).add(off)) }
116    }
117
118    /// Check whether a slot is empty (expire == 0).
119    #[inline]
120    fn slot_is_empty(&self, index: usize) -> bool {
121        self.slot_item(index).expire() == 0
122    }
123
124    /// Check whether a slot's item has expired.
125    fn slot_is_expired(&self, index: usize) -> bool {
126        let expire = self.slot_item(index).expire();
127        if expire == 0 || expire == u32::MAX {
128            return false; // empty or no-expiry
129        }
130        let elapsed = (Instant::now() - self.started).as_secs();
131        elapsed > expire
132    }
133
134    /// Clear a slot by zeroing all its bytes.
135    fn clear_slot(&mut self, index: usize) {
136        let off = self.slot_offset(index);
137        self.data[off..off + self.item_size].fill(0);
138    }
139
140    /// Copy a slot's contents from `from` to `to` and clear the source.
141    fn move_slot(&mut self, from: usize, to: usize) {
142        let from_off = self.slot_offset(from);
143        let to_off = self.slot_offset(to);
144        let size = self.item_size;
145        self.data.copy_within(from_off..from_off + size, to_off);
146        self.data[from_off..from_off + size].fill(0);
147    }
148
149    /// Write an item into a slot.
150    fn write_slot(&mut self, index: usize, key: &[u8], value: Value, expire: u32) {
151        self.clear_slot(index);
152        let off = self.slot_offset(index);
153        let ptr = unsafe { self.data.as_mut_ptr().add(off) };
154        let mut item = TinyItem::from_ptr(ptr);
155        item.define(key, value, expire);
156    }
157
158    // -----------------------------------------------------------------------
159    // Hashing
160    // -----------------------------------------------------------------------
161
162    /// Compute the D candidate positions for a key.
163    fn positions(&self, key: &[u8]) -> [usize; D] {
164        let mut positions = [0usize; D];
165        for (i, pos) in positions.iter_mut().enumerate() {
166            let mut hasher = self.hashers[i].build_hasher();
167            hasher.write(key);
168            *pos = (hasher.finish() as usize) % self.nitem;
169        }
170        positions
171    }
172
173    /// Compute the expiration timestamp for a given TTL.
174    fn compute_expire(&self, ttl: std::time::Duration) -> u32 {
175        if ttl.is_zero() {
176            return u32::MAX; // no expiry
177        }
178        let secs = std::cmp::min(ttl.as_secs(), self.max_ttl as u64) as u32;
179        let elapsed = (Instant::now() - self.started).as_secs();
180        elapsed.saturating_add(secs)
181    }
182
183    // -----------------------------------------------------------------------
184    // Expiration and eviction helpers
185    // -----------------------------------------------------------------------
186
187    /// Handle an expired item: update metrics and clear the slot.
188    fn handle_expired(&mut self, index: usize) {
189        #[cfg(feature = "metrics")]
190        {
191            metrics::ITEM_EXPIRE.increment();
192            self.decrement_item_metrics(index);
193        }
194        self.clear_slot(index);
195    }
196
197    /// Evict an item to make room: update metrics and clear the slot.
198    fn evict_at(&mut self, index: usize) {
199        #[cfg(feature = "metrics")]
200        {
201            metrics::ITEM_EVICT.increment();
202            self.decrement_item_metrics(index);
203        }
204        self.clear_slot(index);
205    }
206
207    #[cfg(feature = "metrics")]
208    fn decrement_item_metrics(&self, index: usize) {
209        let item = self.slot_item(index);
210        let klen = item.klen() as i64;
211        let vlen = item.header().value_len() as i64;
212        metrics::ITEM_CURRENT.sub(1);
213        metrics::ITEM_KEY_BYTE.sub(klen);
214        metrics::ITEM_VAL_BYTE.sub(vlen);
215        metrics::ITEM_DATA_BYTE.sub(klen + vlen);
216    }
217
218    #[cfg(feature = "metrics")]
219    fn increment_item_metrics(&self, index: usize) {
220        let item = self.slot_item(index);
221        let klen = item.klen() as i64;
222        let vlen = item.header().value_len() as i64;
223        metrics::ITEM_CURRENT.add(1);
224        metrics::ITEM_KEY_BYTE.add(klen);
225        metrics::ITEM_VAL_BYTE.add(vlen);
226        metrics::ITEM_DATA_BYTE.add(klen + vlen);
227    }
228
229    // -----------------------------------------------------------------------
230    // Displacement
231    // -----------------------------------------------------------------------
232
233    /// Attempt to free one of the candidate positions via displacement.
234    /// Returns the index into `candidates` of the freed slot, or `None`.
235    fn try_displace(&mut self, candidates: &[usize; D]) -> Option<usize> {
236        for (idx, &pos) in candidates.iter().enumerate() {
237            if self.displace_from(pos, 0) {
238                return Some(idx);
239            }
240        }
241        None
242    }
243
244    /// Try to free the slot at `pos` by moving its occupant to one of the
245    /// occupant's alternative candidate positions, recursing up to
246    /// `max_displace` levels deep. Returns `true` if `pos` was freed.
247    fn displace_from(&mut self, pos: usize, depth: usize) -> bool {
248        if self.slot_is_empty(pos) {
249            return true;
250        }
251        if self.slot_is_expired(pos) {
252            self.handle_expired(pos);
253            return true;
254        }
255        if depth >= self.max_displace {
256            return false;
257        }
258
259        let key_buf = self.slot_item(pos).key().to_vec();
260        let alts = self.positions(&key_buf);
261
262        for &alt in &alts {
263            if alt == pos {
264                continue;
265            }
266            if self.slot_is_empty(alt) {
267                self.move_slot(pos, alt);
268                #[cfg(feature = "metrics")]
269                metrics::CUCKOO_DISPLACE.increment();
270                return true;
271            }
272            if self.slot_is_expired(alt) {
273                self.handle_expired(alt);
274                self.move_slot(pos, alt);
275                #[cfg(feature = "metrics")]
276                metrics::CUCKOO_DISPLACE.increment();
277                return true;
278            }
279        }
280
281        for &alt in &alts {
282            if alt == pos {
283                continue;
284            }
285            if self.displace_from(alt, depth + 1) {
286                self.move_slot(pos, alt);
287                #[cfg(feature = "metrics")]
288                metrics::CUCKOO_DISPLACE.increment();
289                return true;
290            }
291        }
292
293        false
294    }
295
296    /// Select a victim candidate index for eviction.
297    fn select_victim(&self, candidates: &[usize; D]) -> usize {
298        match self.policy {
299            Policy::Random => rand::rng().random::<u64>() as usize % D,
300            Policy::Expire => {
301                let mut best = 0;
302                let mut best_expire = u32::MAX;
303                for (i, &pos) in candidates.iter().enumerate() {
304                    let expire = self.slot_item(pos).expire();
305                    if expire < best_expire {
306                        best = i;
307                        best_expire = expire;
308                    }
309                }
310                best
311            }
312        }
313    }
314
315    // -----------------------------------------------------------------------
316    // Insertion helper
317    // -----------------------------------------------------------------------
318
319    /// Find a slot for inserting a key. Returns `(slot_index, is_update)`.
320    fn find_slot_for_insert(&mut self, key: &[u8], positions: &[usize; D]) -> (usize, bool) {
321        // Pass 1: existing non-expired key
322        for &pos in positions {
323            if self.slot_is_empty(pos) || self.slot_is_expired(pos) {
324                continue;
325            }
326            if self.slot_item(pos).key() == key {
327                return (pos, true);
328            }
329        }
330
331        // Pass 2: empty slot
332        for &pos in positions {
333            if self.slot_is_empty(pos) {
334                return (pos, false);
335            }
336        }
337
338        // Pass 3: expired slot
339        for &pos in positions {
340            if self.slot_is_expired(pos) {
341                self.handle_expired(pos);
342                return (pos, false);
343            }
344        }
345
346        // Pass 4: displacement
347        if let Some(freed_idx) = self.try_displace(positions) {
348            return (positions[freed_idx], false);
349        }
350
351        // Pass 5: evict
352        let victim_idx = self.select_victim(positions);
353        let pos = positions[victim_idx];
354        self.evict_at(pos);
355        (pos, false)
356    }
357
358    // -----------------------------------------------------------------------
359    // Public API
360    // -----------------------------------------------------------------------
361
362    /// Look up an item by key.
363    ///
364    /// ```
365    /// use cuckoo_cache::CuckooCache;
366    /// use std::time::Duration;
367    ///
368    /// let mut cache = CuckooCache::builder().build();
369    /// assert!(cache.get(b"coffee").is_none());
370    ///
371    /// cache.insert(b"coffee", b"strong", Duration::ZERO).unwrap();
372    /// let item = cache.get(b"coffee").unwrap();
373    /// assert_eq!(item.value(), b"strong");
374    /// ```
375    pub fn get(&mut self, key: &[u8]) -> Option<Item> {
376        #[cfg(feature = "metrics")]
377        metrics::CUCKOO_GET.increment();
378
379        let positions = self.positions(key);
380
381        for &pos in &positions {
382            if self.slot_is_empty(pos) {
383                continue;
384            }
385
386            if self.slot_item(pos).key() == key {
387                if self.slot_is_expired(pos) {
388                    self.handle_expired(pos);
389                    #[cfg(feature = "metrics")]
390                    metrics::CUCKOO_GET_KEY_MISS.increment();
391                    return None;
392                }
393
394                #[cfg(feature = "metrics")]
395                metrics::CUCKOO_GET_KEY_HIT.increment();
396
397                return Some(Item::new(self.slot_item(pos)));
398            }
399        }
400
401        #[cfg(feature = "metrics")]
402        metrics::CUCKOO_GET_KEY_MISS.increment();
403
404        None
405    }
406
407    /// Insert an item into the cache.
408    ///
409    /// ```
410    /// use cuckoo_cache::CuckooCache;
411    /// use std::time::Duration;
412    ///
413    /// let mut cache = CuckooCache::builder().build();
414    /// cache.insert(b"drink", b"coffee", Duration::ZERO).unwrap();
415    ///
416    /// let item = cache.get(b"drink").unwrap();
417    /// assert_eq!(item.value(), b"coffee");
418    ///
419    /// cache.insert(b"drink", b"whisky", Duration::ZERO).unwrap();
420    /// let item = cache.get(b"drink").unwrap();
421    /// assert_eq!(item.value(), b"whisky");
422    /// ```
423    pub fn insert<'a, T: Into<Value<'a>>>(
424        &mut self,
425        key: &[u8],
426        value: T,
427        ttl: std::time::Duration,
428    ) -> Result<(), CuckooCacheError> {
429        let value: Value = value.into();
430
431        let required = TINY_ITEM_HDR_SIZE + key.len() + keyvalue::size_of(&value);
432        if required > self.item_size {
433            #[cfg(feature = "metrics")]
434            metrics::CUCKOO_INSERT_EX.increment();
435            return Err(CuckooCacheError::ItemOversized {
436                size: required,
437                max: self.item_size,
438            });
439        }
440
441        debug_assert!(!key.is_empty(), "empty keys are not supported");
442        debug_assert!(
443            key.len() <= u8::MAX as usize,
444            "key length exceeds maximum of 255"
445        );
446
447        #[cfg(feature = "metrics")]
448        metrics::CUCKOO_INSERT.increment();
449
450        let expire = self.compute_expire(ttl);
451        let positions = self.positions(key);
452        let (pos, is_update) = self.find_slot_for_insert(key, &positions);
453
454        if is_update {
455            #[cfg(feature = "metrics")]
456            {
457                metrics::CUCKOO_UPDATE.increment();
458                self.decrement_item_metrics(pos);
459            }
460        }
461
462        self.write_slot(pos, key, value, expire);
463
464        #[cfg(feature = "metrics")]
465        self.increment_item_metrics(pos);
466
467        Ok(())
468    }
469
470    /// Remove the item with the given key.
471    ///
472    /// ```
473    /// use cuckoo_cache::CuckooCache;
474    /// use std::time::Duration;
475    ///
476    /// let mut cache = CuckooCache::builder().build();
477    /// assert!(!cache.delete(b"coffee"));
478    ///
479    /// cache.insert(b"coffee", b"strong", Duration::ZERO).unwrap();
480    /// assert!(cache.delete(b"coffee"));
481    /// assert!(cache.get(b"coffee").is_none());
482    /// ```
483    pub fn delete(&mut self, key: &[u8]) -> bool {
484        #[cfg(feature = "metrics")]
485        metrics::CUCKOO_DELETE.increment();
486
487        let positions = self.positions(key);
488
489        for &pos in &positions {
490            if self.slot_is_empty(pos) {
491                continue;
492            }
493            if self.slot_item(pos).key() == key {
494                if self.slot_is_expired(pos) {
495                    self.handle_expired(pos);
496                    return false;
497                }
498                #[cfg(feature = "metrics")]
499                self.decrement_item_metrics(pos);
500                self.clear_slot(pos);
501                return true;
502            }
503        }
504
505        false
506    }
507
508    /// Clear all items from the cache.
509    pub fn clear(&mut self) {
510        self.data.fill(0);
511    }
512
513    /// Perform a wrapping addition on a numeric value.
514    pub fn wrapping_add(&mut self, key: &[u8], rhs: u64) -> Result<Item, CuckooCacheError> {
515        let positions = self.positions(key);
516
517        for &pos in &positions {
518            if self.slot_is_empty(pos) || self.slot_is_expired(pos) {
519                continue;
520            }
521            if self.slot_item(pos).key() == key {
522                let off = self.slot_offset(pos);
523                let ptr = unsafe { self.data.as_mut_ptr().add(off) };
524                let mut item = TinyItem::from_ptr(ptr);
525                item.wrapping_add(rhs)
526                    .map_err(|_| CuckooCacheError::NotNumeric)?;
527                return Ok(Item::new(item));
528            }
529        }
530
531        Err(CuckooCacheError::NotFound)
532    }
533
534    /// Perform a saturating subtraction on a numeric value.
535    pub fn saturating_sub(&mut self, key: &[u8], rhs: u64) -> Result<Item, CuckooCacheError> {
536        let positions = self.positions(key);
537
538        for &pos in &positions {
539            if self.slot_is_empty(pos) || self.slot_is_expired(pos) {
540                continue;
541            }
542            if self.slot_item(pos).key() == key {
543                let off = self.slot_offset(pos);
544                let ptr = unsafe { self.data.as_mut_ptr().add(off) };
545                let mut item = TinyItem::from_ptr(ptr);
546                item.saturating_sub(rhs)
547                    .map_err(|_| CuckooCacheError::NotNumeric)?;
548                return Ok(Item::new(item));
549            }
550        }
551
552        Err(CuckooCacheError::NotFound)
553    }
554
555    /// Get a count of live (non-expired) items.
556    ///
557    /// ```
558    /// use cuckoo_cache::CuckooCache;
559    ///
560    /// let cache = CuckooCache::builder().build();
561    /// assert_eq!(cache.items(), 0);
562    /// ```
563    #[cfg(any(test, feature = "debug"))]
564    pub fn items(&self) -> usize {
565        (0..self.nitem)
566            .filter(|&i| !self.slot_is_empty(i) && !self.slot_is_expired(i))
567            .count()
568    }
569}