fix: fix w-TinyLFU frequencies updates

foyer-memory/Cargo.toml

@@ -15,7 +15,7 @@ readme = { workspace = true }
 ahash = { workspace = true }
 arc-swap = "1"
 bitflags = "2"
-cmsketch = "0.2.1"
+cmsketch = "0.2.2"
 equivalent = { workspace = true }
 fastrace = { workspace = true }
 foyer-common = { workspace = true }

foyer-memory/src/eviction/fifo.rs

@@ -18,7 +18,7 @@
 use intrusive_collections::{intrusive_adapter, LinkedList, LinkedListAtomicLink};
 use serde::{Deserialize, Serialize};
 
-use super::{Eviction, Op};
+use super::{Context, Eviction, Op};
 use crate::{
     error::Result,
     record::{CacheHint, Record},
@@ -52,6 +52,17 @@
 
 intrusive_adapter! { Adapter<K, V> = Arc<Record<Fifo<K, V>>>: Record<Fifo<K, V>> { ?offset = Record::<Fifo<K, V>>::STATE_OFFSET + offset_of!(FifoState, link) => LinkedListAtomicLink } where K: Key, V: Value }
 
+#[derive(Debug, Clone)]
+pub struct FifoContext;
+
+impl Context for FifoContext {
+    type Config = FifoConfig;
+
+    fn init(_: &Self::Config) -> Self {
+        Self
+    }
+}
+
 pub struct Fifo<K, V>
 where
     K: Key,
@@ -70,8 +81,9 @@
     type Value = V;
     type Hint = FifoHint;
     type State = FifoState;
+    type Context = FifoContext;
 
-    fn new(_capacity: usize, _config: &Self::Config) -> Self
+    fn new(_: usize, _: &Self::Config, _: &Self::Context) -> Self
     where
         Self: Sized,
     {
@@ -156,7 +168,7 @@
             .collect_vec();
         let r = |i: usize| rs[i].clone();
 
-        let mut fifo = TestFifo::new(100, &FifoConfig {});
+        let mut fifo = TestFifo::new(100, &FifoConfig {}, &FifoContext);
 
         // 0, 1, 2, 3
         fifo.push(r(0));

foyer-memory/src/eviction/lfu.rs

@@ -14,15 +14,15 @@
 
 use std::{mem::offset_of, sync::Arc};
 
-use cmsketch::CMSketchU16;
+use cmsketch::CMSketchAtomicU16;
 use foyer_common::{
     code::{Key, Value},
     strict_assert, strict_assert_eq, strict_assert_ne,
 };
 use intrusive_collections::{intrusive_adapter, LinkedList, LinkedListAtomicLink};
 use serde::{Deserialize, Serialize};
 
-use super::{Eviction, Op};
+use super::{Context, Eviction, Op, OpCtx};
 use crate::{
     error::{Error, Result},
     record::{CacheHint, Record},
@@ -105,6 +105,21 @@
 
 intrusive_adapter! { Adapter<K, V> = Arc<Record<Lfu<K, V>>>: Record<Lfu<K, V>> { ?offset = Record::<Lfu<K, V>>::STATE_OFFSET + offset_of!(LfuState, link) => LinkedListAtomicLink } where K: Key, V: Value }
 
+#[derive(Debug, Clone)]
+pub struct LfuContext {
+    frequencies: Arc<CMSketchAtomicU16>,
+}
+
+impl Context for LfuContext {
+    type Config = LfuConfig;
+
+    fn init(config: &Self::Config) -> Self {
+        let frequencies = CMSketchAtomicU16::new(config.cmsketch_eps, config.cmsketch_confidence);
+        let frequencies = Arc::new(frequencies);
+        Self { frequencies }
+    }
+}
+
 /// This implementation is inspired by [Caffeine](https://github.com/ben-manes/caffeine) under Apache License 2.0
 ///
 /// A new and hot entry is kept in `window`.
@@ -133,8 +148,7 @@
     window_weight_capacity: usize,
     protected_weight_capacity: usize,
 
-    // TODO(MrCroxx): use a count-min-sketch impl with atomic u16
-    frequencies: CMSketchU16,
+    frequencies: Arc<CMSketchAtomicU16>,
 
     step: usize,
     decay: usize,
@@ -185,8 +199,9 @@
     type Value = V;
     type Hint = LfuHint;
     type State = LfuState;
+    type Context = LfuContext;
 
-    fn new(capacity: usize, config: &Self::Config) -> Self
+    fn new(capacity: usize, config: &Self::Config, context: &Self::Context) -> Self
     where
         Self: Sized,
     {
@@ -212,7 +227,7 @@
 
         let window_weight_capacity = (capacity as f64 * config.window_capacity_ratio) as usize;
         let protected_weight_capacity = (capacity as f64 * config.protected_capacity_ratio) as usize;
-        let frequencies = CMSketchU16::new(config.cmsketch_eps, config.cmsketch_confidence);
+        let frequencies = context.frequencies.clone();
         let decay = frequencies.width();
 
         Self {
@@ -284,7 +299,6 @@
         record.set_in_eviction(true);
         state.queue = Queue::Window;
         self.increase_queue_weight(Queue::Window, record.weight());
-        self.update_frequencies(record.hash());
         self.window.push_back(record);
 
         // If `window` weight exceeds the capacity, overflow entry from `window` to `probation`.
@@ -365,48 +379,57 @@
     }
 
     fn acquire() -> Op<Self> {
-        Op::mutable(|this: &mut Self, record| {
-            // Update frequency by access.
-            this.update_frequencies(record.hash());
-
-            if !record.is_in_eviction() {
-                return;
-            }
-
-            let state = unsafe { &mut *record.state().get() };
-
-            strict_assert!(state.link.is_linked());
+        Op::mutable_ctx(|this: &mut Self, ctx| {
+            match ctx {
+                OpCtx::Hit { record } => {
+                    // Update frequency by access.
+                    this.update_frequencies(record.hash());
+
+                    if !record.is_in_eviction() {
+                        return;
+                    }
 
-            match state.queue {
-                Queue::None => unreachable!(),
-                Queue::Window => {
-                    // Move to MRU position of `window`.
-                    let r = unsafe { this.window.remove_from_ptr(Arc::as_ptr(record)) };
-                    this.window.push_back(r);
-                }
-                Queue::Probation => {
-                    // Promote to MRU position of `protected`.
-                    let r = unsafe { this.probation.remove_from_ptr(Arc::as_ptr(record)) };
-                    this.decrease_queue_weight(Queue::Probation, record.weight());
-                    state.queue = Queue::Protected;
-                    this.increase_queue_weight(Queue::Protected, record.weight());
-                    this.protected.push_back(r);
-
-                    // If `protected` weight exceeds the capacity, overflow entry from `protected` to `probation`.
-                    while this.protected_weight > this.protected_weight_capacity {
-                        strict_assert!(!this.protected.is_empty());
-                        let r = this.protected.pop_front().unwrap();
-                        let s = unsafe { &mut *r.state().get() };
-                        this.decrease_queue_weight(Queue::Protected, r.weight());
-                        s.queue = Queue::Probation;
-                        this.increase_queue_weight(Queue::Probation, r.weight());
-                        this.probation.push_back(r);
+                    let state = unsafe { &mut *record.state().get() };
+
+                    strict_assert!(state.link.is_linked());
+
+                    match state.queue {
+                        Queue::None => unreachable!(),
+                        Queue::Window => {
+                            // Move to MRU position of `window`.
+                            let r = unsafe { this.window.remove_from_ptr(Arc::as_ptr(record)) };
+                            this.window.push_back(r);
+                        }
+                        Queue::Probation => {
+                            // Promote to MRU position of `protected`.
+                            let r = unsafe { this.probation.remove_from_ptr(Arc::as_ptr(record)) };
+                            this.decrease_queue_weight(Queue::Probation, record.weight());
+                            state.queue = Queue::Protected;
+                            this.increase_queue_weight(Queue::Protected, record.weight());
+                            this.protected.push_back(r);
+
+                            // If `protected` weight exceeds the capacity, overflow entry from `protected` to
+                            // `probation`.
+                            while this.protected_weight > this.protected_weight_capacity {
+                                strict_assert!(!this.protected.is_empty());
+                                let r = this.protected.pop_front().unwrap();
+                                let s = unsafe { &mut *r.state().get() };
+                                this.decrease_queue_weight(Queue::Protected, r.weight());
+                                s.queue = Queue::Probation;
+                                this.increase_queue_weight(Queue::Probation, r.weight());
+                                this.probation.push_back(r);
+                            }
+                        }
+                        Queue::Protected => {
+                            // Move to MRU position of `protected`.
+                            let r = unsafe { this.protected.remove_from_ptr(Arc::as_ptr(record)) };
+                            this.protected.push_back(r);
+                        }
                     }
                 }
-                Queue::Protected => {
-                    // Move to MRU position of `protected`.
-                    let r = unsafe { this.protected.remove_from_ptr(Arc::as_ptr(record)) };
-                    this.protected.push_back(r);
+                OpCtx::Miss { hash } => {
+                    // Update frequency by access.
+                    this.update_frequencies(hash);
                 }
             }
         })
@@ -495,7 +518,8 @@
             cmsketch_eps: 0.01,
             cmsketch_confidence: 0.95,
         };
-        let mut lfu = TestLfu::new(10, &config);
+        let context = LfuContext::init(&config);
+        let mut lfu = TestLfu::new(10, &config, &context);
 
         assert_eq!(lfu.window_weight_capacity, 2);
         assert_eq!(lfu.protected_weight_capacity, 6);

foyer-memory/src/eviction/lru.rs

@@ -21,7 +21,7 @@ use foyer_common::{
 use intrusive_collections::{intrusive_adapter, LinkedList, LinkedListAtomicLink};
 use serde::{Deserialize, Serialize};
 
-use super::{Eviction, Op};
+use super::{Context, Eviction, Op};
 use crate::{
     error::{Error, Result},
     record::{CacheHint, Record},
@@ -90,6 +90,17 @@ pub struct LruState {
 
 intrusive_adapter! { Adapter<K, V> = Arc<Record<Lru<K, V>>>: Record<Lru<K, V>> { ?offset = Record::<Lru<K, V>>::STATE_OFFSET + offset_of!(LruState, link) => LinkedListAtomicLink } where K: Key, V: Value }
 
+#[derive(Debug, Clone)]
+pub struct LruContext;
+
+impl Context for LruContext {
+    type Config = LruConfig;
+
+    fn init(_: &Self::Config) -> Self {
+        Self
+    }
+}
+
 pub struct Lru<K, V>
 where
     K: Key,
@@ -135,8 +146,9 @@ where
     type Value = V;
     type Hint = LruHint;
     type State = LruState;
+    type Context = LruContext;
 
-    fn new(capacity: usize, config: &Self::Config) -> Self
+    fn new(capacity: usize, config: &Self::Config, _: &Self::Context) -> Self
     where
         Self: Sized,
     {
@@ -400,7 +412,7 @@ pub mod tests {
         let config = LruConfig {
             high_priority_pool_ratio: 0.5,
         };
-        let mut lru = TestLru::new(8, &config);
+        let mut lru = TestLru::new(8, &config, &LruContext);
 
         assert_eq!(lru.high_priority_weight_capacity, 4);
 
@@ -477,7 +489,7 @@ pub mod tests {
         let config = LruConfig {
             high_priority_pool_ratio: 0.5,
         };
-        let mut lru = TestLru::new(8, &config);
+        let mut lru = TestLru::new(8, &config, &LruContext);
 
         assert_eq!(lru.high_priority_weight_capacity, 4);