1#![cfg_attr(not(feature = "std"), no_std)]
2extern crate alloc;
3#[cfg(not(feature = "std"))]
4use alloc::{vec, vec::Vec, boxed::Box};
5#[cfg(feature = "std")]
6use std::vec;
7use crate::cache_padded::CachePadded;
8use crate::daemon::{Command, Daemon};
9use crate::lossy_queue::{LossyQueue, OneshotAck};
10use crate::unsafe_core::{Cache, T1, T2, WorkerSlot};
11use ahash::RandomState;
12use core::hash::{BuildHasher, Hash};
13use crate::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering};
14use crate::sync::index_types::{AtomicTick, TickType};
15use crate::sync::{Arc, ArcSlice, new_arc_slice};
16use crate::config::Config;
17
18#[cfg(any(feature = "loom", loom))]
24loom::lazy_static! {
25 pub static ref GLOBAL_EPOCH: loom::sync::atomic::AtomicUsize = loom::sync::atomic::AtomicUsize::new(1);
26}
27
28#[cfg(not(any(feature = "loom", loom)))]
29pub static GLOBAL_EPOCH: AtomicUsize = AtomicUsize::new(1);
30
31pub struct WorkerState {
34 pub local_epoch: CachePadded<AtomicUsize>,
35}
36
37impl WorkerState {
38 pub fn new() -> Self {
39 Self {
40 local_epoch: CachePadded::new(AtomicUsize::new(0)),
41 }
42 }
43}
44
45pub trait DaemonSpawner: Send + Sync {
52 fn spawn(&self, f: alloc::boxed::Box<dyn FnOnce() + Send + 'static>);
54}
55
56#[cfg(feature = "std")]
58#[derive(Debug, Clone, Copy, Default)]
59pub struct DefaultSpawner;
60
61#[cfg(feature = "std")]
62impl DaemonSpawner for DefaultSpawner {
63 #[inline]
64 fn spawn(&self, f: alloc::boxed::Box<dyn FnOnce() + Send + 'static>) {
65 #[cfg(any(feature = "loom", loom))]
66 loom::thread::spawn(move || f());
67 #[cfg(not(any(feature = "loom", loom)))]
68 std::thread::spawn(move || f());
69 }
70}
71
72pub trait TlsProvider: Send + Sync {
77 fn get_worker_id(&self) -> Option<usize>;
81
82 fn with_hit_buf(&self, f: &mut dyn FnMut(&mut ([usize; 64], usize)));
87
88 fn with_l1_filter(&self, f: &mut dyn FnMut(&mut ([u8; 4096], usize)));
93
94 fn with_last_flush_tick(&self, f: &mut dyn FnMut(&mut TickType));
99}
100
101#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
108use std::sync::Mutex;
109
110#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
111struct IdAllocator {
112 free_list: Mutex<Vec<usize>>,
113 next_id: AtomicUsize,
114}
115
116#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
117static ALLOCATOR: IdAllocator = IdAllocator {
118 free_list: Mutex::new(Vec::new()),
119 next_id: AtomicUsize::new(0),
120};
121
122#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
123struct ThreadIdGuard {
124 id: usize,
125}
126
127#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
128impl Drop for ThreadIdGuard {
129 fn drop(&mut self) {
130 if let Ok(mut list) = ALLOCATOR.free_list.lock() {
131 list.push(self.id);
132 }
133 }
134}
135
136#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
137use core::cell::{Cell, RefCell};
138
139#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
140thread_local! {
141 static WORKER_ID: usize = {
142 let id = if let Ok(mut list) = ALLOCATOR.free_list.lock() {
143 list.pop().unwrap_or_else(|| ALLOCATOR.next_id.fetch_add(1, Ordering::Relaxed))
144 } else {
145 ALLOCATOR.next_id.fetch_add(1, Ordering::Relaxed)
146 };
147
148 GUARD.with(|g| {
149 *g.borrow_mut() = Some(ThreadIdGuard { id });
150 });
151 id
152 };
153
154 static GUARD: RefCell<Option<ThreadIdGuard>> = const { RefCell::new(None) };
155
156 static HIT_BUF: RefCell<([usize; 64], usize)> = const { RefCell::new(([0; 64], 0)) };
159
160 static L1_FILTER: RefCell<([u8; 4096], usize)> = const { RefCell::new(([0; 4096], 0)) };
163
164 static LAST_FLUSH_TICK: Cell<TickType> = const { Cell::new(0) };
168}
169
170#[cfg(any(feature = "loom", loom))]
171loom::lazy_static! {
172 static ref NEXT_THREAD_ID: loom::sync::atomic::AtomicUsize = loom::sync::atomic::AtomicUsize::new(0);
173}
174
175#[cfg(any(feature = "loom", loom))]
176use core::cell::{Cell, RefCell};
177
178#[cfg(any(feature = "loom", loom))]
179loom::thread_local! {
180 static WORKER_ID: usize = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
181
182 static HIT_BUF: RefCell<([usize; 64], usize)> = RefCell::new(([0; 64], 0));
185
186 static L1_FILTER: RefCell<(Box<[u8]>, usize)> = RefCell::new((vec![0u8; 4096].into_boxed_slice(), 0));
190
191 static LAST_FLUSH_TICK: Cell<TickType> = Cell::new(0);
195}
196
197pub struct DefaultTls;
199
200impl Clone for DefaultTls {
201 fn clone(&self) -> Self {
202 Self
203 }
204}
205
206#[cfg(any(feature = "std", feature = "loom", loom))]
207impl TlsProvider for DefaultTls {
208 #[inline(always)]
209 fn get_worker_id(&self) -> Option<usize> {
210 Some(WORKER_ID.with(|id| *id))
211 }
212
213 #[inline(always)]
214 fn with_hit_buf(&self, f: &mut dyn FnMut(&mut ([usize; 64], usize))) {
215 HIT_BUF.with(|buf| {
216 f(&mut *buf.borrow_mut());
217 });
218 }
219
220 #[inline(always)]
221 fn with_l1_filter(&self, f: &mut dyn FnMut(&mut ([u8; 4096], usize))) {
222 L1_FILTER.with(|filter| {
223 #[cfg(any(feature = "loom", loom))]
224 {
225 let mut state = filter.borrow_mut();
226 let mut arr = [0u8; 4096];
227 arr.copy_from_slice(&state.0);
228 let mut temp = (arr, state.1);
229 f(&mut temp);
230 state.0.copy_from_slice(&temp.0);
231 state.1 = temp.1;
232 }
233 #[cfg(not(any(feature = "loom", loom)))]
234 {
235 f(&mut *filter.borrow_mut());
236 }
237 });
238 }
239
240 #[inline(always)]
241 fn with_last_flush_tick(&self, f: &mut dyn FnMut(&mut TickType)) {
242 LAST_FLUSH_TICK.with(|cell| {
243 let mut val = cell.get();
244 f(&mut val);
245 cell.set(val);
246 });
247 }
248}
249
250#[cfg(not(any(feature = "std", feature = "loom", loom)))]
251impl TlsProvider for DefaultTls {
252 #[inline(always)]
253 fn get_worker_id(&self) -> Option<usize> {
254 None
255 }
256
257 #[inline(always)]
258 fn with_hit_buf(&self, _f: &mut dyn FnMut(&mut ([usize; 64], usize))) {}
259
260 #[inline(always)]
261 fn with_l1_filter(&self, _f: &mut dyn FnMut(&mut ([u8; 4096], usize))) {}
262
263 #[inline(always)]
264 fn with_last_flush_tick(&self, _f: &mut dyn FnMut(&mut TickType)) {}
265}
266
267pub struct DualCacheFF<K, V, S = RandomState, Tls: TlsProvider = DefaultTls> {
270 pub hasher: S,
271 pub t1: Arc<T1<K, V>>,
272 pub t2: Arc<T2<K, V>>,
273 pub cache: Arc<Cache<K, V>>,
274 pub cmd_tx: Arc<LossyQueue<Command<K, V>>>,
275 pub hit_tx: Arc<LossyQueue<[usize; 64]>>,
276 pub epoch: Arc<AtomicU32>,
277 pub worker_states: ArcSlice<WorkerState>,
279 pub miss_buffers: ArcSlice<WorkerSlot<K, V>>,
281 pub daemon_tick: Arc<AtomicTick>,
284 pub flush_tick_threshold: TickType,
286 pub is_cold_start: Arc<AtomicBool>,
288 pub tls: Tls,
290}
291
292impl<K, V, S: Clone, Tls: TlsProvider + Clone> Clone for DualCacheFF<K, V, S, Tls> {
293 fn clone(&self) -> Self {
294 Self {
295 hasher: self.hasher.clone(),
296 t1: self.t1.clone(),
297 t2: self.t2.clone(),
298 cache: self.cache.clone(),
299 cmd_tx: self.cmd_tx.clone(),
300 hit_tx: self.hit_tx.clone(),
301 epoch: self.epoch.clone(),
302 worker_states: self.worker_states.clone(),
303 miss_buffers: self.miss_buffers.clone(),
304 daemon_tick: self.daemon_tick.clone(),
305 flush_tick_threshold: self.flush_tick_threshold,
306 is_cold_start: self.is_cold_start.clone(),
307 tls: self.tls.clone(),
308 }
309 }
310}
311
312#[cfg(feature = "std")]
315impl<K, V> DualCacheFF<K, V, RandomState, DefaultTls>
316where
317 K: Hash + Eq + Send + Sync + Clone + 'static,
318 V: Send + Sync + Clone + 'static,
319{
320 #[inline]
324 pub fn new(config: Config) -> Self {
325 Self::new_with_spawner(config, DefaultSpawner)
326 }
327}
328
329#[cfg(feature = "std")]
330impl<K, V, Tls> DualCacheFF<K, V, RandomState, Tls>
331where
332 K: Hash + Eq + Send + Sync + Clone + 'static,
333 V: Send + Sync + Clone + 'static,
334 Tls: TlsProvider + 'static,
335{
336 #[inline]
338 pub fn new_with_tls(config: Config, tls: Tls) -> Self {
339 Self::new_with_tls_and_spawner(config, tls, DefaultSpawner)
340 }
341}
342
343impl<K, V> DualCacheFF<K, V, RandomState, DefaultTls>
346where
347 K: Hash + Eq + Send + Sync + Clone + 'static,
348 V: Send + Sync + Clone + 'static,
349{
350 pub fn new_with_spawner<Sp: DaemonSpawner + 'static>(config: Config, spawner: Sp) -> Self {
352 let (cache, daemon) = Self::new_headless(config);
353 #[cfg(any(feature = "loom", loom))]
354 {
355 let _ = daemon;
356 let _ = spawner;
357 }
358 #[cfg(not(any(feature = "loom", loom)))]
359 {
360 spawner.spawn(alloc::boxed::Box::new(move || daemon.run()));
361 }
362 cache
363 }
364
365 pub fn new_headless(config: Config) -> (Self, Daemon<K, V, RandomState>) {
378 Self::new_headless_with_tls(config, DefaultTls)
379 }
380}
381
382impl<K, V, Tls> DualCacheFF<K, V, RandomState, Tls>
383where
384 K: Hash + Eq + Send + Sync + Clone + 'static,
385 V: Send + Sync + Clone + 'static,
386 Tls: TlsProvider + 'static,
387{
388 pub fn new_headless_with_tls(
390 config: Config,
391 tls: Tls,
392 ) -> (Self, Daemon<K, V, RandomState>) {
393 let hasher = RandomState::new();
394 let t1 = Arc::new(T1::new(config.t1_slots));
395 let t2 = Arc::new(T2::new(config.t2_slots));
396 let cache = Arc::new(Cache::new(config.capacity));
397 let cmd_q: Arc<LossyQueue<Command<K, V>>> = Arc::new(LossyQueue::new(8192));
398 let hit_q: Arc<LossyQueue<[usize; 64]>> = Arc::new(LossyQueue::new(1024));
399 let epoch = Arc::new(AtomicU32::new(0));
400 let daemon_tick = Arc::new(AtomicTick::new(0));
401 let is_cold_start = Arc::new(AtomicBool::new(true));
402
403 let mut buffers = Vec::with_capacity(config.threads);
404 let mut states = Vec::with_capacity(config.threads);
405 for _ in 0..config.threads {
406 buffers.push(WorkerSlot::new());
407 states.push(WorkerState::new());
408 }
409 let miss_buffers = new_arc_slice(buffers);
410 let worker_states = new_arc_slice(states);
411
412 let daemon = Daemon::new(
413 hasher.clone(),
414 config.capacity,
415 t1.clone(),
416 t2.clone(),
417 cache.clone(),
418 cmd_q.clone(),
419 hit_q.clone(),
420 epoch.clone(),
421 config.duration,
422 config.poll_us,
423 worker_states.clone(),
424 daemon_tick.clone(),
425 is_cold_start.clone(),
426 );
427
428 let this = Self {
429 hasher,
430 t1,
431 t2,
432 cache,
433 cmd_tx: cmd_q,
434 hit_tx: hit_q,
435 epoch,
436 worker_states,
437 miss_buffers,
438 daemon_tick,
439 flush_tick_threshold: (config.poll_us as TickType).max(1),
440 is_cold_start,
441 tls,
442 };
443
444 (this, daemon)
445 }
446
447 pub fn new_with_tls_and_spawner<Sp: DaemonSpawner + 'static>(config: Config, tls: Tls, spawner: Sp) -> Self
449 {
450 let (cache, daemon) = Self::new_headless_with_tls(config, tls);
451 #[cfg(any(feature = "loom", loom))]
452 {
453 let _ = daemon;
454 let _ = spawner;
455 }
456 #[cfg(not(any(feature = "loom", loom)))]
457 {
458 spawner.spawn(alloc::boxed::Box::new(move || daemon.run()));
459 }
460 cache
461 }
462
463}
464
465impl<K, V, S, Tls: TlsProvider> DualCacheFF<K, V, S, Tls>
468where
469 K: Hash + Eq + Send + Sync + Clone + 'static,
470 V: Send + Sync + Clone + 'static,
471 S: BuildHasher + Clone + Send + 'static,
472{
473 pub fn sync(&self) {
477 self.with_hit_buf(|state| {
479 if state.1 > 0_usize {
480 let _ = self.hit_tx.try_send(state.0);
481 state.1 = 0;
482 }
483 });
484
485 for slot in self.miss_buffers.iter() {
487 let buf: &mut crate::unsafe_core::BatchBuf<K, V> = slot.get_mut_safe();
488 if buf.len() > 0 {
489 let batch = buf.drain_to_vec();
490 let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
491 }
492 }
493
494 let ack = OneshotAck::new();
496 self.cmd_tx.send_blocking(Command::Sync(ack.clone()));
497 ack.wait();
498 }
499
500 pub fn get(&self, key: &K) -> Option<V> {
505 let hash = self.hash(key);
506 let current_epoch_cache = self.epoch.load(Ordering::Relaxed);
507
508 let mut id_opt = None;
510 let global_epoch = GLOBAL_EPOCH.load(Ordering::Relaxed);
511 self.with_worker_id(|id| {
512 if id < self.worker_states.len() {
513 self.worker_states[id]
514 .local_epoch
515 .store(global_epoch, Ordering::Relaxed);
516 id_opt = Some(id);
517 }
518 });
519
520 let has_epoch = id_opt.is_some() || {
521 #[cfg(not(feature = "std"))]
522 { true }
523 #[cfg(feature = "std")]
524 { false }
525 };
526
527 let mut res: Option<V> = None;
528 let mut hit_g_idx: Option<u32> = None;
529
530 if has_epoch {
531 if let Some(node) = self.t1.get_node(hash) {
533 if node.key == *key
534 && (node.expire_at == 0 || node.expire_at >= current_epoch_cache)
535 {
536 res = Some(node.value.clone());
537 hit_g_idx = Some(node.g_idx);
538 }
539 }
540
541 if res.is_none() {
543 if let Some(node) = self.t2.get_node(hash) {
544 if node.key == *key
545 && (node.expire_at == 0 || node.expire_at >= current_epoch_cache)
546 {
547 res = Some(node.value.clone());
548 hit_g_idx = Some(node.g_idx);
549 }
550 }
551 }
552
553 if res.is_none() {
555 let tag = (hash >> 48) as u16;
556 if let Some(global_idx) = self.cache.index_probe(hash, tag) {
557 if let Some(v) = self
558 .cache
559 .node_get_full(global_idx, key, current_epoch_cache)
560 {
561 res = Some(v);
562 hit_g_idx = Some(global_idx as u32);
563 }
564 }
565 }
566 }
567
568 if let Some(id) = id_opt {
570 self.worker_states[id]
571 .local_epoch
572 .store(0, Ordering::Relaxed);
573 }
574
575 if let Some(g_idx) = hit_g_idx {
576 self.record_hit(g_idx as usize);
577 }
578
579 res
580 }
581
582 pub fn insert(&self, key: K, value: V) {
596 let hash = self.hash(&key);
597
598 let mut id_opt = None;
599 let is_cold = self.is_cold_start.load(Ordering::Relaxed);
600 let mut bypass = is_cold;
601
602 if !bypass {
603 let global_epoch = GLOBAL_EPOCH.load(Ordering::Relaxed);
606 self.with_worker_id(|id| {
607 if id < self.worker_states.len() {
608 self.worker_states[id]
609 .local_epoch
610 .store(global_epoch, Ordering::Relaxed);
611 id_opt = Some(id);
612 }
613 });
614
615 if id_opt.is_some() {
616 if let Some(node) = self.t1.get_node(hash) {
618 if node.key == key {
619 bypass = true;
620 }
621 }
622
623 if !bypass {
625 if let Some(node) = self.t2.get_node(hash) {
626 if node.key == key {
627 bypass = true;
628 }
629 }
630 }
631
632 if !bypass {
634 let tag = (hash >> 48) as u16;
635 if let Some(global_idx) = self.cache.index_probe(hash, tag) {
636 if let Some(node) = self.cache.get_node(global_idx) {
637 if node.key == key {
638 bypass = true;
639 }
640 }
641 }
642 }
643 }
644
645 if let Some(id) = id_opt {
647 self.worker_states[id]
648 .local_epoch
649 .store(0, Ordering::Relaxed);
650 }
651 }
652
653 let pass = if bypass {
654 true
655 } else {
656 self.with_l1_filter(|state| {
658 let idx = (hash as usize) & 4095_usize;
659 let val = state.0[idx];
660
661 state.1 += 1;
662 if state.1 >= 4096_usize {
663 for x in state.0.iter_mut() {
664 *x >>= 1;
665 }
666 state.1 = 0;
667 }
668
669 if val < 1_u8 {
670 state.0[idx] = 1;
671 false
672 } else {
673 if val < 2_u8 {
674 state.0[idx] = 2;
675 }
676 true
677 }
678 }).unwrap_or(true) };
680
681 if !pass {
682 return;
683 }
684
685 let current_tick = self.daemon_tick.load(Ordering::Relaxed);
687 let mut should_time_flush = false;
688 self.with_last_flush_tick(|tick| {
689 should_time_flush = current_tick.wrapping_sub(*tick) >= self.flush_tick_threshold;
690 });
691
692 let mut option_kv = Some((key, value));
694 let pushed_to_buf = self.with_worker_id(|id| {
695 let (k, v) = option_kv.take().unwrap();
696 if id >= self.miss_buffers.len() {
697 let _ = self.cmd_tx.try_send(Command::Insert(k, v, hash));
699 return;
700 }
701
702 let buf: &mut crate::unsafe_core::BatchBuf<K, V> = self.miss_buffers[id].get_mut_safe();
704 let capacity_flush = buf.push((k, v, hash));
705
706 if capacity_flush || (should_time_flush && !buf.is_empty()) {
707 let batch = buf.drain_to_vec();
708 let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
709 self.with_last_flush_tick(|tick| {
710 *tick = current_tick;
711 });
712 }
713 });
714
715 if pushed_to_buf.is_none() {
716 if let Some((k, v)) = option_kv {
717 let _ = self.cmd_tx.try_send(Command::Insert(k, v, hash));
718 }
719 }
720 }
721
722 pub fn remove(&self, key: &K) {
724 let hash = self.hash(key);
725
726 self.with_worker_id(|id| {
728 if id < self.miss_buffers.len() {
729 let buf: &mut crate::unsafe_core::BatchBuf<K, V> = self.miss_buffers[id].get_mut_safe();
730 if buf.len() > 0 {
731 let batch = buf.drain_to_vec();
732 let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
733 let tick = self.daemon_tick.load(Ordering::Relaxed);
734 self.with_last_flush_tick(|c| *c = tick);
735 }
736 }
737 });
738
739 self.cmd_tx.send_blocking(Command::Remove(key.clone(), hash));
740 }
741
742 pub fn clear(&self) {
744 let ack = OneshotAck::new();
745 self.cmd_tx.send_blocking(Command::Clear(ack.clone()));
746 ack.wait();
747 }
748
749 #[inline(always)]
752 fn with_worker_id<F, R>(&self, f: F) -> Option<R>
753 where
754 F: FnOnce(usize) -> R,
755 {
756 self.tls.get_worker_id().map(f)
757 }
758
759 #[inline(always)]
760 fn with_hit_buf<F, R>(&self, mut f: F) -> Option<R>
761 where
762 F: FnMut(&mut ([usize; 64], usize)) -> R,
763 {
764 let mut res = None;
765 self.tls.with_hit_buf(&mut |buf| {
766 res = Some(f(buf));
767 });
768 res
769 }
770
771 #[inline(always)]
772 fn with_l1_filter<F, R>(&self, mut f: F) -> Option<R>
773 where
774 F: FnMut(&mut ([u8; 4096], usize)) -> R,
775 {
776 let mut res = None;
777 self.tls.with_l1_filter(&mut |filter| {
778 res = Some(f(filter));
779 });
780 res
781 }
782
783 #[inline(always)]
784 fn with_last_flush_tick<F, R>(&self, mut f: F) -> Option<R>
785 where
786 F: FnMut(&mut TickType) -> R,
787 {
788 let mut res = None;
789 self.tls.with_last_flush_tick(&mut |tick| {
790 res = Some(f(tick));
791 });
792 res
793 }
794
795 #[inline(always)]
796 fn hash(&self, key: &K) -> u64 {
797 self.hasher.hash_one(key)
798 }
799
800 #[inline(always)]
805 fn record_hit(&self, global_idx: usize) {
806 let opt = self.with_hit_buf(|state| {
807 let idx = state.1;
808 state.0[idx] = global_idx;
809 state.1 += 1;
810 if state.1 == 64_usize {
811 let _ = self.hit_tx.try_send(state.0);
812 state.1 = 0;
813 }
814 });
815
816 if opt.is_none() {
817 let mut batch = [0usize; 64];
818 batch[0] = global_idx;
819 let _ = self.hit_tx.try_send(batch);
820 }
821 }
822}
823
824impl<K, V, S, Tls: TlsProvider> Drop for DualCacheFF<K, V, S, Tls> {
825 fn drop(&mut self) {
826 if Arc::strong_count(&self.cmd_tx) <= 2 {
827 let _ = self.cmd_tx.try_send(Command::Shutdown);
828 }
829 }
830}
831