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