feat: support filter for in-memory cache, rename old to storage filter (#1094)

Author: MrCroxx

examples/hybrid_full.rs

@@ -16,9 +16,9 @@ use std::{hash::BuildHasherDefault, num::NonZeroUsize};
 
 use chrono::Datelike;
 use foyer::{
-    BlockEngineBuilder, DeviceBuilder, FifoPicker, Filter, FsDeviceBuilder, HybridCache, HybridCacheBuilder,
-    HybridCachePolicy, IoEngineBuilder, IopsCounter, LruConfig, PsyncIoEngineBuilder, RecoverMode, RejectAll, Result,
-    RuntimeOptions, Throttle, TokioRuntimeOptions,
+    BlockEngineBuilder, DeviceBuilder, FifoPicker, FsDeviceBuilder, HybridCache, HybridCacheBuilder, HybridCachePolicy,
+    IoEngineBuilder, IopsCounter, LruConfig, PsyncIoEngineBuilder, RecoverMode, RejectAll, Result, RuntimeOptions,
+    StorageFilter, Throttle, TokioRuntimeOptions,
 };
 use tempfile::tempdir;
 
@@ -50,6 +50,7 @@ async fn main() -> anyhow::Result<()> {
         })
         .with_hash_builder(BuildHasherDefault::default())
         .with_weighter(|_key, value: &String| value.len())
+        .with_filter(|_, _| true)
         .storage()
         .with_io_engine(io_engine)
         .with_engine_config(
@@ -62,8 +63,8 @@ async fn main() -> anyhow::Result<()> {
                 .with_buffer_pool_size(256 * 1024 * 1024)
                 .with_clean_block_threshold(4)
                 .with_eviction_pickers(vec![Box::<FifoPicker>::default()])
-                .with_admission_filter(Filter::new())
-                .with_reinsertion_filter(Filter::new().with_condition(RejectAll))
+                .with_admission_filter(StorageFilter::new())
+                .with_reinsertion_filter(StorageFilter::new().with_condition(RejectAll))
                 .with_tombstone_log(false),
         )
         .with_recover_mode(RecoverMode::Quiet)

foyer-memory/src/cache.rs

@@ -36,7 +36,7 @@ use crate::{
         s3fifo::{S3Fifo, S3FifoConfig},
         sieve::{Sieve, SieveConfig},
     },
-    raw::{FetchContext, FetchState, FetchTarget, RawCache, RawCacheConfig, RawCacheEntry, RawFetch, Weighter},
+    raw::{FetchContext, FetchState, FetchTarget, Filter, RawCache, RawCacheConfig, RawCacheEntry, RawFetch, Weighter},
     Piece, Pipe, Result,
 };
 
@@ -427,6 +427,7 @@ where
 
     hash_builder: S,
     weighter: Arc<dyn Weighter<K, V>>,
+    filter: Arc<dyn Filter<K, V>>,
 
     event_listener: Option<Arc<dyn EventListener<Key = K, Value = V>>>,
 
@@ -450,6 +451,7 @@ where
 
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
 
             registry: Box::new(NoopMetricsRegistry),
@@ -501,6 +503,7 @@ where
             eviction_config: self.eviction_config,
             hash_builder,
             weighter: self.weighter,
+            filter: self.filter,
             event_listener: self.event_listener,
             registry: self.registry,
             metrics: self.metrics,
@@ -513,6 +516,21 @@ where
         self
     }
 
+    /// Set the filter for the in-memory cache.
+    ///
+    /// The filter is used to decide whether to admit or reject an entry based on its key and value.
+    ///
+    /// If the filter returns true, the key value can be inserted into the in-memory cache;
+    /// otherwise, the key value cannot be inserted.
+    ///
+    /// To ensure API consistency, the in-memory cache will still return a cache entry,
+    /// but it will not count towards the in-memory cache usage,
+    /// and it will be immediately reclaimed when the cache entry is dropped.
+    pub fn with_filter(mut self, filter: impl Filter<K, V>) -> Self {
+        self.filter = Arc::new(filter);
+        self
+    }
+
     /// Set event listener.
     pub fn with_event_listener(mut self, event_listener: Arc<dyn EventListener<Key = K, Value = V>>) -> Self {
         self.event_listener = Some(event_listener);
@@ -560,6 +578,7 @@ where
                 eviction_config,
                 hash_builder: self.hash_builder,
                 weighter: self.weighter,
+                filter: self.filter,
                 event_listener: self.event_listener,
                 metrics,
             }))),
@@ -569,6 +588,7 @@ where
                 eviction_config,
                 hash_builder: self.hash_builder,
                 weighter: self.weighter,
+                filter: self.filter,
                 event_listener: self.event_listener,
                 metrics,
             }))),
@@ -578,6 +598,7 @@ where
                 eviction_config,
                 hash_builder: self.hash_builder,
                 weighter: self.weighter,
+                filter: self.filter,
                 event_listener: self.event_listener,
                 metrics,
             }))),
@@ -587,6 +608,7 @@ where
                 eviction_config,
                 hash_builder: self.hash_builder,
                 weighter: self.weighter,
+                filter: self.filter,
                 event_listener: self.event_listener,
                 metrics,
             }))),
@@ -596,6 +618,7 @@ where
                 eviction_config,
                 hash_builder: self.hash_builder,
                 weighter: self.weighter,
+                filter: self.filter,
                 event_listener: self.event_listener,
                 metrics,
             }))),

foyer-memory/src/prelude.rs

@@ -19,5 +19,5 @@ pub use crate::{
     error::{Error, Result},
     eviction::{fifo::FifoConfig, lfu::LfuConfig, lru::LruConfig, s3fifo::S3FifoConfig, Eviction, Op},
     pipe::{Piece, Pipe},
-    raw::{FetchContext, FetchState, FetchTarget, Weighter},
+    raw::{FetchContext, FetchState, FetchTarget, Filter, Weighter},
 };

foyer-memory/src/raw.rs

@@ -60,6 +60,19 @@ use crate::{
 pub trait Weighter<K, V>: Fn(&K, &V) -> usize + Send + Sync + 'static {}
 impl<K, V, T> Weighter<K, V> for T where T: Fn(&K, &V) -> usize + Send + Sync + 'static {}
 
+/// The filter for the in-memory cache.
+///
+/// The filter is used to decide whether to admit or reject an entry based on its key and value.
+///
+/// If the filter returns true, the key value can be inserted into the in-memory cache;
+/// otherwise, the key value cannot be inserted.
+///
+/// To ensure API consistency, the in-memory cache will still return a cache entry,
+/// but it will not count towards the in-memory cache usage,
+/// and it will be immediately reclaimed when the cache entry is dropped.
+pub trait Filter<K, V>: Fn(&K, &V) -> bool + Send + Sync + 'static {}
+impl<K, V, T> Filter<K, V> for T where T: Fn(&K, &V) -> bool + Send + Sync + 'static {}
+
 pub struct RawCacheConfig<E, S>
 where
     E: Eviction,
@@ -70,6 +83,7 @@ where
     pub eviction_config: E::Config,
     pub hash_builder: S,
     pub weighter: Arc<dyn Weighter<E::Key, E::Value>>,
+    pub filter: Arc<dyn Filter<E::Key, E::Value>>,
     pub event_listener: Option<Arc<dyn EventListener<Key = E::Key, Value = E::Value>>>,
     pub metrics: Arc<Metrics>,
 }
@@ -398,6 +412,7 @@ where
 
     hash_builder: Arc<S>,
     weighter: Arc<dyn Weighter<E::Key, E::Value>>,
+    filter: Arc<dyn Filter<E::Key, E::Value>>,
 
     metrics: Arc<Metrics>,
     event_listener: Option<Arc<dyn EventListener<Key = E::Key, Value = E::Value>>>,
@@ -491,6 +506,7 @@ where
             capacity: config.capacity,
             hash_builder: Arc::new(config.hash_builder),
             weighter: config.weighter,
+            filter: config.filter,
             metrics: config.metrics,
             event_listener: config.event_listener,
             pipe: ArcSwap::new(Arc::new(pipe)),
@@ -559,10 +575,13 @@ where
         &self,
         key: E::Key,
         value: E::Value,
-        properties: E::Properties,
+        mut properties: E::Properties,
     ) -> RawCacheEntry<E, S, I> {
         let hash = self.inner.hash_builder.hash_one(&key);
         let weight = (self.inner.weighter)(&key, &value);
+        if !(self.inner.filter)(&key, &value) {
+            properties = properties.with_disposable(true);
+        }
         let record = Arc::new(Record::new(Data {
             key,
             value,
@@ -582,6 +601,17 @@ where
     #[cfg_attr(feature = "tracing", fastrace::trace(name = "foyer::memory::raw::insert_inner"))]
     fn insert_inner(&self, record: Arc<Record<E>>) -> RawCacheEntry<E, S, I> {
         if record.properties().disposable().unwrap_or_default() {
+            // Remove the stale record if it exists.
+            self.inner.shards[self.shard(record.hash())]
+                .write()
+                .remove(record.hash(), record.key())
+                .inspect(|r| {
+                    // Deallocate data out of the lock critical section.
+                    if let Some(listener) = self.inner.event_listener.as_ref() {
+                        listener.on_leave(Event::Replace, r.key(), r.value());
+                    }
+                });
+
             // If the record is disposable, we do not insert it into the cache.
             // Instead, we just return it and let it be dropped immediately after the last reference drops.
             record.inc_refs(1);
@@ -1212,6 +1242,7 @@ mod tests {
             eviction_config: FifoConfig::default(),
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         })
@@ -1226,6 +1257,7 @@ mod tests {
             eviction_config: S3FifoConfig::default(),
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         })
@@ -1240,6 +1272,7 @@ mod tests {
             eviction_config: LruConfig::default(),
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         })
@@ -1254,6 +1287,7 @@ mod tests {
             eviction_config: LfuConfig::default(),
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         })
@@ -1268,6 +1302,7 @@ mod tests {
             eviction_config: SieveConfig {},
             hash_builder: Default::default(),
             weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         })
@@ -1306,6 +1341,43 @@ mod tests {
         assert_eq!(fifo.usage(), 0);
         drop(e2a);
         assert_eq!(fifo.usage(), 0);
+
+        let fifo = fifo_cache_for_test();
+        fifo.insert(1, 1);
+        assert_eq!(fifo.usage(), 1);
+        assert_eq!(fifo.get(&1).unwrap().value(), &1);
+        let e2 = fifo.insert_with_properties(1, 100, TestProperties::default().with_disposable(true));
+        assert_eq!(fifo.usage(), 0);
+        drop(e2);
+        assert_eq!(fifo.usage(), 0);
+        assert!(fifo.get(&1).is_none());
+    }
+
+    #[expect(clippy::type_complexity)]
+    #[test_log::test]
+    fn test_insert_filter() {
+        let fifo: RawCache<
+            Fifo<u64, u64, TestProperties>,
+            ModHasher,
+            HashTableIndexer<Fifo<u64, u64, TestProperties>>,
+        > = RawCache::new(RawCacheConfig {
+            capacity: 256,
+            shards: 4,
+            eviction_config: FifoConfig::default(),
+            hash_builder: Default::default(),
+            weighter: Arc::new(|_, _| 1),
+            filter: Arc::new(|k, _| !matches!(*k, 42)),
+            event_listener: None,
+            metrics: Arc::new(Metrics::noop()),
+        });
+
+        fifo.insert(1, 1);
+        fifo.insert(2, 2);
+        fifo.insert(42, 42);
+        assert_eq!(fifo.usage(), 2);
+        assert_eq!(fifo.get(&1).unwrap().value(), &1);
+        assert_eq!(fifo.get(&2).unwrap().value(), &2);
+        assert!(fifo.get(&42).is_none());
     }
 
     #[test]
@@ -1338,6 +1410,7 @@ mod tests {
             eviction_config: FifoConfig::default(),
             hash_builder: Default::default(),
             weighter: Arc::new(|k, v| k.len() + v.len()),
+            filter: Arc::new(|_, _| true),
             event_listener: None,
             metrics: Arc::new(Metrics::noop()),
         });
@@ -1440,6 +1513,7 @@ mod tests {
                 eviction_config: FifoConfig::default(),
                 hash_builder: Default::default(),
                 weighter: Arc::new(|_, _| 1),
+                filter: Arc::new(|_, _| true),
                 event_listener: None,
                 metrics: Arc::new(Metrics::noop()),
             });
@@ -1455,6 +1529,7 @@ mod tests {
                 eviction_config: S3FifoConfig::default(),
                 hash_builder: Default::default(),
                 weighter: Arc::new(|_, _| 1),
+                filter: Arc::new(|_, _| true),
                 event_listener: None,
                 metrics: Arc::new(Metrics::noop()),
             });
@@ -1470,6 +1545,7 @@ mod tests {
                 eviction_config: LruConfig::default(),
                 hash_builder: Default::default(),
                 weighter: Arc::new(|_, _| 1),
+                filter: Arc::new(|_, _| true),
                 event_listener: None,
                 metrics: Arc::new(Metrics::noop()),
             });
@@ -1485,6 +1561,7 @@ mod tests {
                 eviction_config: LfuConfig::default(),
                 hash_builder: Default::default(),
                 weighter: Arc::new(|_, _| 1),
+                filter: Arc::new(|_, _| true),
                 event_listener: None,
                 metrics: Arc::new(Metrics::noop()),
             });
@@ -1500,6 +1577,7 @@ mod tests {
                 eviction_config: SieveConfig {},
                 hash_builder: Default::default(),
                 weighter: Arc::new(|_, _| 1),
+                filter: Arc::new(|_, _| true),
                 event_listener: None,
                 metrics: Arc::new(Metrics::noop()),
             });