1use 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
23const 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
32pub struct CuckooCache {
36 data: Box<[u8]>,
38 hashers: Box<[RandomState; D]>,
40 item_size: usize,
42 nitem: usize,
44 max_displace: usize,
46 max_ttl: u32,
48 policy: Policy,
50 started: Instant,
52}
53
54impl CuckooCache {
55 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 #[inline]
108 fn slot_offset(&self, index: usize) -> usize {
109 index * self.item_size
110 }
111
112 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 #[inline]
120 fn slot_is_empty(&self, index: usize) -> bool {
121 self.slot_item(index).expire() == 0
122 }
123
124 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; }
130 let elapsed = (Instant::now() - self.started).as_secs();
131 elapsed > expire
132 }
133
134 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 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 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 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 fn compute_expire(&self, ttl: std::time::Duration) -> u32 {
175 if ttl.is_zero() {
176 return u32::MAX; }
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 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 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 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 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 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 fn find_slot_for_insert(&mut self, key: &[u8], positions: &[usize; D]) -> (usize, bool) {
321 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 for &pos in positions {
333 if self.slot_is_empty(pos) {
334 return (pos, false);
335 }
336 }
337
338 for &pos in positions {
340 if self.slot_is_expired(pos) {
341 self.handle_expired(pos);
342 return (pos, false);
343 }
344 }
345
346 if let Some(freed_idx) = self.try_displace(positions) {
348 return (positions[freed_idx], false);
349 }
350
351 let victim_idx = self.select_victim(positions);
353 let pos = positions[victim_idx];
354 self.evict_at(pos);
355 (pos, false)
356 }
357
358 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 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 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 pub fn clear(&mut self) {
510 self.data.fill(0);
511 }
512
513 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 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 #[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}