extern crate alloc;
#[cfg(not(feature = "std"))]
use alloc::{vec, vec::Vec, boxed::Box};
#[cfg(feature = "std")]
use std::vec;
use crate::cache_padded::CachePadded;
use crate::daemon::{Command, Daemon};
use crate::lossy_queue::{LossyQueue, OneshotAck};
use crate::unsafe_core::{Cache, T1, T2, WorkerSlot};
use ahash::RandomState;
use core::hash::{BuildHasher, Hash};
use crate::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering};
use crate::sync::index_types::{AtomicTick, TickType};
use crate::sync::{Arc, ArcSlice, new_arc_slice};
use crate::config::Config;
#[cfg(any(feature = "loom", loom))]
loom::lazy_static! {
pub static ref GLOBAL_EPOCH: loom::sync::atomic::AtomicUsize = loom::sync::atomic::AtomicUsize::new(1);
}
#[cfg(not(any(feature = "loom", loom)))]
pub static GLOBAL_EPOCH: AtomicUsize = AtomicUsize::new(1);
pub struct WorkerState {
pub local_epoch: CachePadded<AtomicUsize>,
}
impl WorkerState {
pub fn new() -> Self {
Self {
local_epoch: CachePadded::new(AtomicUsize::new(0)),
}
}
}
pub trait DaemonSpawner: Send + Sync {
fn spawn(&self, f: alloc::boxed::Box<dyn FnOnce() + Send + 'static>);
}
#[cfg(feature = "std")]
#[derive(Debug, Clone, Copy, Default)]
pub struct DefaultSpawner;
#[cfg(feature = "std")]
impl DaemonSpawner for DefaultSpawner {
#[inline]
fn spawn(&self, f: alloc::boxed::Box<dyn FnOnce() + Send + 'static>) {
#[cfg(any(feature = "loom", loom))]
loom::thread::spawn(move || f());
#[cfg(not(any(feature = "loom", loom)))]
std::thread::spawn(move || f());
}
}
pub trait TlsProvider: Send + Sync {
fn get_worker_id(&self) -> Option<usize>;
fn with_hit_buf(&self, f: &mut dyn FnMut(&mut ([usize; 64], usize)));
fn with_l1_filter(&self, f: &mut dyn FnMut(&mut ([u8; 4096], usize)));
fn with_last_flush_tick(&self, f: &mut dyn FnMut(&mut TickType));
}
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
use std::sync::Mutex;
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
struct IdAllocator {
free_list: Mutex<Vec<usize>>,
next_id: AtomicUsize,
}
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
static ALLOCATOR: IdAllocator = IdAllocator {
free_list: Mutex::new(Vec::new()),
next_id: AtomicUsize::new(0),
};
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
struct ThreadIdGuard {
id: usize,
}
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
impl Drop for ThreadIdGuard {
fn drop(&mut self) {
if let Ok(mut list) = ALLOCATOR.free_list.lock() {
list.push(self.id);
}
}
}
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
use core::cell::{Cell, RefCell};
#[cfg(all(feature = "std", not(any(feature = "loom", loom))))]
thread_local! {
static WORKER_ID: usize = {
let id = if let Ok(mut list) = ALLOCATOR.free_list.lock() {
list.pop().unwrap_or_else(|| ALLOCATOR.next_id.fetch_add(1, Ordering::Relaxed))
} else {
ALLOCATOR.next_id.fetch_add(1, Ordering::Relaxed)
};
GUARD.with(|g| {
*g.borrow_mut() = Some(ThreadIdGuard { id });
});
id
};
static GUARD: RefCell<Option<ThreadIdGuard>> = const { RefCell::new(None) };
static HIT_BUF: RefCell<([usize; 64], usize)> = const { RefCell::new(([0; 64], 0)) };
static L1_FILTER: RefCell<([u8; 4096], usize)> = const { RefCell::new(([0; 4096], 0)) };
static LAST_FLUSH_TICK: Cell<TickType> = const { Cell::new(0) };
}
#[cfg(any(feature = "loom", loom))]
loom::lazy_static! {
static ref NEXT_THREAD_ID: loom::sync::atomic::AtomicUsize = loom::sync::atomic::AtomicUsize::new(0);
}
#[cfg(any(feature = "loom", loom))]
use core::cell::{Cell, RefCell};
#[cfg(any(feature = "loom", loom))]
loom::thread_local! {
static WORKER_ID: usize = NEXT_THREAD_ID.fetch_add(1, Ordering::Relaxed);
static HIT_BUF: RefCell<([usize; 64], usize)> = RefCell::new(([0; 64], 0));
static L1_FILTER: RefCell<(Box<[u8]>, usize)> = RefCell::new((vec![0u8; 4096].into_boxed_slice(), 0));
static LAST_FLUSH_TICK: Cell<TickType> = Cell::new(0);
}
pub struct DefaultTls;
impl Clone for DefaultTls {
fn clone(&self) -> Self {
Self
}
}
#[cfg(any(feature = "std", feature = "loom", loom))]
impl TlsProvider for DefaultTls {
#[inline(always)]
fn get_worker_id(&self) -> Option<usize> {
Some(WORKER_ID.with(|id| *id))
}
#[inline(always)]
fn with_hit_buf(&self, f: &mut dyn FnMut(&mut ([usize; 64], usize))) {
HIT_BUF.with(|buf| {
f(&mut *buf.borrow_mut());
});
}
#[inline(always)]
fn with_l1_filter(&self, f: &mut dyn FnMut(&mut ([u8; 4096], usize))) {
L1_FILTER.with(|filter| {
#[cfg(any(feature = "loom", loom))]
{
let mut state = filter.borrow_mut();
let mut arr = [0u8; 4096];
arr.copy_from_slice(&state.0);
let mut temp = (arr, state.1);
f(&mut temp);
state.0.copy_from_slice(&temp.0);
state.1 = temp.1;
}
#[cfg(not(any(feature = "loom", loom)))]
{
f(&mut *filter.borrow_mut());
}
});
}
#[inline(always)]
fn with_last_flush_tick(&self, f: &mut dyn FnMut(&mut TickType)) {
LAST_FLUSH_TICK.with(|cell| {
let mut val = cell.get();
f(&mut val);
cell.set(val);
});
}
}
#[cfg(not(any(feature = "std", feature = "loom", loom)))]
impl TlsProvider for DefaultTls {
#[inline(always)]
fn get_worker_id(&self) -> Option<usize> {
None
}
#[inline(always)]
fn with_hit_buf(&self, _f: &mut dyn FnMut(&mut ([usize; 64], usize))) {}
#[inline(always)]
fn with_l1_filter(&self, _f: &mut dyn FnMut(&mut ([u8; 4096], usize))) {}
#[inline(always)]
fn with_last_flush_tick(&self, _f: &mut dyn FnMut(&mut TickType)) {}
}
pub struct DualCacheFF<K, V, S = RandomState, Tls: TlsProvider = DefaultTls> {
pub hasher: S,
pub t1: Arc<T1<K, V>>,
pub t2: Arc<T2<K, V>>,
pub cache: Arc<Cache<K, V>>,
pub cmd_tx: Arc<LossyQueue<Command<K, V>>>,
pub hit_tx: Arc<LossyQueue<[usize; 64]>>,
pub epoch: Arc<AtomicU32>,
pub worker_states: ArcSlice<WorkerState>,
pub miss_buffers: ArcSlice<WorkerSlot<K, V>>,
pub daemon_tick: Arc<AtomicTick>,
pub flush_tick_threshold: TickType,
pub is_cold_start: Arc<AtomicBool>,
pub tls: Tls,
}
impl<K, V, S: Clone, Tls: TlsProvider + Clone> Clone for DualCacheFF<K, V, S, Tls> {
fn clone(&self) -> Self {
Self {
hasher: self.hasher.clone(),
t1: self.t1.clone(),
t2: self.t2.clone(),
cache: self.cache.clone(),
cmd_tx: self.cmd_tx.clone(),
hit_tx: self.hit_tx.clone(),
epoch: self.epoch.clone(),
worker_states: self.worker_states.clone(),
miss_buffers: self.miss_buffers.clone(),
daemon_tick: self.daemon_tick.clone(),
flush_tick_threshold: self.flush_tick_threshold,
is_cold_start: self.is_cold_start.clone(),
tls: self.tls.clone(),
}
}
}
#[cfg(feature = "std")]
impl<K, V> DualCacheFF<K, V, RandomState, DefaultTls>
where
K: Hash + Eq + Send + Sync + Clone + 'static,
V: Send + Sync + Clone + 'static,
{
#[inline]
pub fn new(config: Config) -> Self {
Self::new_with_spawner(config, DefaultSpawner)
}
}
#[cfg(feature = "std")]
impl<K, V, Tls> DualCacheFF<K, V, RandomState, Tls>
where
K: Hash + Eq + Send + Sync + Clone + 'static,
V: Send + Sync + Clone + 'static,
Tls: TlsProvider + 'static,
{
#[inline]
pub fn new_with_tls(config: Config, tls: Tls) -> Self {
Self::new_with_tls_and_spawner(config, tls, DefaultSpawner)
}
}
impl<K, V> DualCacheFF<K, V, RandomState, DefaultTls>
where
K: Hash + Eq + Send + Sync + Clone + 'static,
V: Send + Sync + Clone + 'static,
{
pub fn new_with_spawner<Sp: DaemonSpawner + 'static>(config: Config, spawner: Sp) -> Self {
let (cache, daemon) = Self::new_headless(config);
#[cfg(any(feature = "loom", loom))]
{
let _ = daemon;
let _ = spawner;
}
#[cfg(not(any(feature = "loom", loom)))]
{
spawner.spawn(alloc::boxed::Box::new(move || daemon.run()));
}
cache
}
pub fn new_headless(config: Config) -> (Self, Daemon<K, V, RandomState>) {
Self::new_headless_with_tls(config, DefaultTls)
}
}
impl<K, V, Tls> DualCacheFF<K, V, RandomState, Tls>
where
K: Hash + Eq + Send + Sync + Clone + 'static,
V: Send + Sync + Clone + 'static,
Tls: TlsProvider + 'static,
{
pub fn new_headless_with_tls(
config: Config,
tls: Tls,
) -> (Self, Daemon<K, V, RandomState>) {
let hasher = RandomState::new();
let t1 = Arc::new(T1::new(config.t1_slots));
let t2 = Arc::new(T2::new(config.t2_slots));
let cache = Arc::new(Cache::new(config.capacity));
let cmd_q: Arc<LossyQueue<Command<K, V>>> = Arc::new(LossyQueue::new(8192));
let hit_q: Arc<LossyQueue<[usize; 64]>> = Arc::new(LossyQueue::new(1024));
let epoch = Arc::new(AtomicU32::new(0));
let daemon_tick = Arc::new(AtomicTick::new(0));
let is_cold_start = Arc::new(AtomicBool::new(true));
let mut buffers = Vec::with_capacity(config.threads);
let mut states = Vec::with_capacity(config.threads);
for _ in 0..config.threads {
buffers.push(WorkerSlot::new());
states.push(WorkerState::new());
}
let miss_buffers = new_arc_slice(buffers);
let worker_states = new_arc_slice(states);
let daemon = Daemon::new(
hasher.clone(),
config.capacity,
t1.clone(),
t2.clone(),
cache.clone(),
cmd_q.clone(),
hit_q.clone(),
epoch.clone(),
config.duration,
config.poll_us,
worker_states.clone(),
daemon_tick.clone(),
is_cold_start.clone(),
);
let this = Self {
hasher,
t1,
t2,
cache,
cmd_tx: cmd_q,
hit_tx: hit_q,
epoch,
worker_states,
miss_buffers,
daemon_tick,
flush_tick_threshold: (config.poll_us as TickType).max(1),
is_cold_start,
tls,
};
(this, daemon)
}
pub fn new_with_tls_and_spawner<Sp: DaemonSpawner + 'static>(config: Config, tls: Tls, spawner: Sp) -> Self
{
let (cache, daemon) = Self::new_headless_with_tls(config, tls);
#[cfg(any(feature = "loom", loom))]
{
let _ = daemon;
let _ = spawner;
}
#[cfg(not(any(feature = "loom", loom)))]
{
spawner.spawn(alloc::boxed::Box::new(move || daemon.run()));
}
cache
}
}
impl<K, V, S, Tls: TlsProvider> DualCacheFF<K, V, S, Tls>
where
K: Hash + Eq + Send + Sync + Clone + 'static,
V: Send + Sync + Clone + 'static,
S: BuildHasher + Clone + Send + 'static,
{
pub fn sync(&self) {
self.with_hit_buf(|state| {
if state.1 > 0_usize {
let _ = self.hit_tx.try_send(state.0);
state.1 = 0;
}
});
for slot in self.miss_buffers.iter() {
let buf: &mut crate::unsafe_core::BatchBuf<K, V> = slot.get_mut_safe();
if buf.len() > 0 {
let batch = buf.drain_to_vec();
let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
}
}
let ack = OneshotAck::new();
self.cmd_tx.send_blocking(Command::Sync(ack.clone()));
ack.wait();
}
pub fn get(&self, key: &K) -> Option<V> {
let hash = self.hash(key);
let current_epoch_cache = self.epoch.load(Ordering::Relaxed);
let mut id_opt = None;
let global_epoch = GLOBAL_EPOCH.load(Ordering::Relaxed);
self.with_worker_id(|id| {
if id < self.worker_states.len() {
self.worker_states[id]
.local_epoch
.store(global_epoch, Ordering::Relaxed);
id_opt = Some(id);
}
});
let has_epoch = id_opt.is_some() || {
#[cfg(not(feature = "std"))]
{ true }
#[cfg(feature = "std")]
{ false }
};
let mut res: Option<V> = None;
let mut hit_g_idx: Option<u32> = None;
if has_epoch {
if let Some(node) = self.t1.get_node(hash) {
if node.key == *key
&& (node.expire_at == 0 || node.expire_at >= current_epoch_cache)
{
res = Some(node.value.clone());
hit_g_idx = Some(node.g_idx);
}
}
if res.is_none() {
if let Some(node) = self.t2.get_node(hash) {
if node.key == *key
&& (node.expire_at == 0 || node.expire_at >= current_epoch_cache)
{
res = Some(node.value.clone());
hit_g_idx = Some(node.g_idx);
}
}
}
if res.is_none() {
let tag = (hash >> 48) as u16;
if let Some(global_idx) = self.cache.index_probe(hash, tag) {
if let Some(v) = self
.cache
.node_get_full(global_idx, key, current_epoch_cache)
{
res = Some(v);
hit_g_idx = Some(global_idx as u32);
}
}
}
}
if let Some(id) = id_opt {
self.worker_states[id]
.local_epoch
.store(0, Ordering::Relaxed);
}
if let Some(g_idx) = hit_g_idx {
self.record_hit(g_idx as usize);
}
res
}
pub fn insert(&self, key: K, value: V) {
let hash = self.hash(&key);
let mut id_opt = None;
let is_cold = self.is_cold_start.load(Ordering::Relaxed);
let mut bypass = is_cold;
if !bypass {
let global_epoch = GLOBAL_EPOCH.load(Ordering::Relaxed);
self.with_worker_id(|id| {
if id < self.worker_states.len() {
self.worker_states[id]
.local_epoch
.store(global_epoch, Ordering::Relaxed);
id_opt = Some(id);
}
});
if id_opt.is_some() {
if let Some(node) = self.t1.get_node(hash) {
if node.key == key {
bypass = true;
}
}
if !bypass {
if let Some(node) = self.t2.get_node(hash) {
if node.key == key {
bypass = true;
}
}
}
if !bypass {
let tag = (hash >> 48) as u16;
if let Some(global_idx) = self.cache.index_probe(hash, tag) {
if let Some(node) = self.cache.get_node(global_idx) {
if node.key == key {
bypass = true;
}
}
}
}
}
if let Some(id) = id_opt {
self.worker_states[id]
.local_epoch
.store(0, Ordering::Relaxed);
}
}
let pass = if bypass {
true
} else {
self.with_l1_filter(|state| {
let idx = (hash as usize) & 4095_usize;
let val = state.0[idx];
state.1 += 1;
if state.1 >= 4096_usize {
for x in state.0.iter_mut() {
*x >>= 1;
}
state.1 = 0;
}
if val < 1_u8 {
state.0[idx] = 1;
false
} else {
if val < 2_u8 {
state.0[idx] = 2;
}
true
}
}).unwrap_or(true) };
if !pass {
return;
}
let current_tick = self.daemon_tick.load(Ordering::Relaxed);
let mut should_time_flush = false;
self.with_last_flush_tick(|tick| {
should_time_flush = current_tick.wrapping_sub(*tick) >= self.flush_tick_threshold;
});
let mut option_kv = Some((key, value));
let pushed_to_buf = self.with_worker_id(|id| {
let (k, v) = option_kv.take().unwrap();
if id >= self.miss_buffers.len() {
let _ = self.cmd_tx.try_send(Command::Insert(k, v, hash));
return;
}
let buf: &mut crate::unsafe_core::BatchBuf<K, V> = self.miss_buffers[id].get_mut_safe();
let capacity_flush = buf.push((k, v, hash));
if capacity_flush || (should_time_flush && !buf.is_empty()) {
let batch = buf.drain_to_vec();
let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
self.with_last_flush_tick(|tick| {
*tick = current_tick;
});
}
});
if pushed_to_buf.is_none() {
if let Some((k, v)) = option_kv {
let _ = self.cmd_tx.try_send(Command::Insert(k, v, hash));
}
}
}
pub fn insert_t1(&self, key: K, value: V) {
let hash = self.hash(&key);
let _ = self.cmd_tx.try_send(Command::InsertT1(key, value, hash));
}
pub fn remove(&self, key: &K) {
let hash = self.hash(key);
self.with_worker_id(|id| {
if id < self.miss_buffers.len() {
let buf: &mut crate::unsafe_core::BatchBuf<K, V> = self.miss_buffers[id].get_mut_safe();
if buf.len() > 0 {
let batch = buf.drain_to_vec();
let _ = self.cmd_tx.try_send(Command::BatchInsert(batch));
let tick = self.daemon_tick.load(Ordering::Relaxed);
self.with_last_flush_tick(|c| *c = tick);
}
}
});
self.cmd_tx.send_blocking(Command::Remove(key.clone(), hash));
}
pub fn clear(&self) {
let ack = OneshotAck::new();
self.cmd_tx.send_blocking(Command::Clear(ack.clone()));
ack.wait();
}
#[inline(always)]
fn with_worker_id<F, R>(&self, f: F) -> Option<R>
where
F: FnOnce(usize) -> R,
{
self.tls.get_worker_id().map(f)
}
#[inline(always)]
fn with_hit_buf<F, R>(&self, mut f: F) -> Option<R>
where
F: FnMut(&mut ([usize; 64], usize)) -> R,
{
let mut res = None;
self.tls.with_hit_buf(&mut |buf| {
res = Some(f(buf));
});
res
}
#[inline(always)]
fn with_l1_filter<F, R>(&self, mut f: F) -> Option<R>
where
F: FnMut(&mut ([u8; 4096], usize)) -> R,
{
let mut res = None;
self.tls.with_l1_filter(&mut |filter| {
res = Some(f(filter));
});
res
}
#[inline(always)]
fn with_last_flush_tick<F, R>(&self, mut f: F) -> Option<R>
where
F: FnMut(&mut TickType) -> R,
{
let mut res = None;
self.tls.with_last_flush_tick(&mut |tick| {
res = Some(f(tick));
});
res
}
#[inline(always)]
fn hash(&self, key: &K) -> u64 {
self.hasher.hash_one(key)
}
#[inline(always)]
fn record_hit(&self, global_idx: usize) {
let opt = self.with_hit_buf(|state| {
let idx = state.1;
state.0[idx] = global_idx;
state.1 += 1;
if state.1 == 64_usize {
let _ = self.hit_tx.try_send(state.0);
state.1 = 0;
}
});
if opt.is_none() {
let mut batch = [0usize; 64];
batch[0] = global_idx;
let _ = self.hit_tx.try_send(batch);
}
}
}
impl<K, V, S, Tls: TlsProvider> Drop for DualCacheFF<K, V, S, Tls> {
fn drop(&mut self) {
if Arc::strong_count(&self.cmd_tx) <= 2 {
let _ = self.cmd_tx.try_send(Command::Shutdown);
}
}
}