S3-FIFO caching

crates/utils/src/cache.rs

@@ -0,0 +1,238 @@
+/*
+ * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
+ */
+
+use std::{
+    hash::Hash,
+    net::IpAddr,
+    sync::Arc,
+    time::{Duration, Instant},
+};
+
+use quick_cache::{
+    sync::{DefaultLifecycle, PlaceholderGuard},
+    Equivalent, Weighter,
+};
+
+use crate::config::Config;
+
+pub struct Cache<K: Eq + Hash + CacheItemWeight, V: Clone + CacheItemWeight>(
+    quick_cache::sync::Cache<K, V, CacheItemWeighter>,
+);
+pub struct CacheWithTtl<K: Eq + Hash + CacheItemWeight, V: Clone + CacheItemWeight>(
+    quick_cache::sync::Cache<K, TtlEntry<V>, CacheItemWeighter>,
+);
+
+#[derive(Clone)]
+pub struct TtlEntry<V: Clone + CacheItemWeight> {
+    value: V,
+    expires: Instant,
+}
+
+impl<K: Eq + Hash + CacheItemWeight, V: Clone + CacheItemWeight> Cache<K, V> {
+    pub fn from_config(config: &mut Config, key: &str) -> Self {
+        Self::new(
+            config
+                .property_or_default((key, "capacity"), "1024")
+                .unwrap_or(100),
+            config
+                .property_or_default((key, "size"), "10485760")
+                .unwrap_or(10485760),
+        )
+    }
+
+    pub fn new(estimated_items_capacity: usize, weight_capacity: u64) -> Self {
+        Self(quick_cache::sync::Cache::with_weighter(
+            estimated_items_capacity,
+            weight_capacity,
+            CacheItemWeighter,
+        ))
+    }
+
+    #[inline(always)]
+    pub fn get<Q>(&self, key: &Q) -> Option<V>
+    where
+        Q: Hash + Equivalent<K> + ?Sized,
+    {
+        self.0.get(key)
+    }
+
+    #[inline(always)]
+    pub async fn get_value_or_guard_async<'a, Q>(
+        &'a self,
+        key: &Q,
+    ) -> Result<
+        V,
+        PlaceholderGuard<'a, K, V, CacheItemWeighter, ahash::RandomState, DefaultLifecycle<K, V>>,
+    >
+    where
+        Q: Hash + Equivalent<K> + ToOwned<Owned = K> + ?Sized,
+    {
+        self.0.get_value_or_guard_async(key).await
+    }
+
+    #[inline(always)]
+    pub fn insert(&self, key: K, value: V) {
+        self.0.insert(key, value);
+    }
+
+    #[inline(always)]
+    pub fn remove(&self, key: &K) {
+        self.0.remove(key);
+    }
+
+    #[inline(always)]
+    pub fn clear(&self) {
+        self.0.clear();
+    }
+}
+
+impl<K: Eq + Hash + CacheItemWeight, V: Clone + CacheItemWeight> CacheWithTtl<K, V> {
+    pub fn from_config(config: &mut Config, key: &str) -> Self {
+        Self::new(
+            config
+                .property_or_default((key, "capacity"), "1024")
+                .unwrap_or(100),
+            config
+                .property_or_default((key, "size"), "10485760")
+                .unwrap_or(10485760),
+        )
+    }
+
+    pub fn new(estimated_items_capacity: usize, weight_capacity: u64) -> Self {
+        Self(quick_cache::sync::Cache::with_weighter(
+            estimated_items_capacity,
+            weight_capacity,
+            CacheItemWeighter,
+        ))
+    }
+
+    #[inline(always)]
+    pub fn get<Q>(&self, key: &Q) -> Option<V>
+    where
+        Q: Hash + Equivalent<K> + ?Sized,
+    {
+        self.0.get(key).and_then(|v| {
+            if v.expires > Instant::now() {
+                Some(v.value)
+            } else {
+                None
+            }
+        })
+    }
+
+    #[inline(always)]
+    pub async fn get_value_or_guard_async<'a, Q>(
+        &'a self,
+        key: &Q,
+    ) -> Result<
+        V,
+        PlaceholderGuard<
+            'a,
+            K,
+            TtlEntry<V>,
+            CacheItemWeighter,
+            ahash::RandomState,
+            DefaultLifecycle<K, TtlEntry<V>>,
+        >,
+    >
+    where
+        Q: Hash + Equivalent<K> + ToOwned<Owned = K> + ?Sized,
+    {
+        match self.0.get_value_or_guard_async(key).await {
+            Ok(value) => {
+                if value.expires > Instant::now() {
+                    Ok(value.value)
+                } else {
+                    self.0.remove(key);
+                    self.0.get_value_or_guard_async(key).await.map(|v| v.value)
+                }
+            }
+            Err(err) => Err(err),
+        }
+    }
+
+    #[inline(always)]
+    pub fn insert(&self, key: K, value: V, expires: Duration) {
+        self.0.insert(key, TtlEntry::new(value, expires));
+    }
+
+    #[inline(always)]
+    pub fn remove(&self, key: &K) {
+        self.0.remove(key);
+    }
+
+    #[inline(always)]
+    pub fn clear(&self) {
+        self.0.clear();
+    }
+}
+
+#[derive(Clone)]
+pub struct CacheItemWeighter;
+
+impl<K: CacheItemWeight, V: CacheItemWeight> Weighter<K, V> for CacheItemWeighter {
+    fn weight(&self, key: &K, val: &V) -> u64 {
+        key.weight() + val.weight()
+    }
+}
+
+pub trait CacheItemWeight {
+    fn weight(&self) -> u64;
+}
+
+impl<T: Clone + CacheItemWeight> CacheItemWeight for TtlEntry<T> {
+    fn weight(&self) -> u64 {
+        self.value.weight() + 8
+    }
+}
+
+impl<T: Clone + CacheItemWeight> CacheItemWeight for Option<T> {
+    fn weight(&self) -> u64 {
+        match self {
+            Some(v) => v.weight(),
+            None => 1,
+        }
+    }
+}
+
+impl<T: CacheItemWeight> CacheItemWeight for Arc<T> {
+    fn weight(&self) -> u64 {
+        self.as_ref().weight()
+    }
+}
+
+impl CacheItemWeight for u64 {
+    fn weight(&self) -> u64 {
+        std::mem::size_of::<u64>() as u64
+    }
+}
+
+impl CacheItemWeight for String {
+    fn weight(&self) -> u64 {
+        self.len() as u64
+    }
+}
+
+impl CacheItemWeight for u32 {
+    fn weight(&self) -> u64 {
+        std::mem::size_of::<u32>() as u64
+    }
+}
+
+impl CacheItemWeight for Vec<IpAddr> {
+    fn weight(&self) -> u64 {
+        (self.len() * std::mem::size_of::<IpAddr>()) as u64
+    }
+}
+
+impl<T: Clone + CacheItemWeight> TtlEntry<T> {
+    pub fn new(value: T, expires: Duration) -> Self {
+        Self {
+            value,
+            expires: Instant::now() + expires,
+        }
+    }
+}

crates/utils/src/lib.rs

@@ -6,10 +6,10 @@
 
 use std::sync::Arc;
 
+pub mod cache;
 pub mod codec;
 pub mod config;
 pub mod glob;
-pub mod lru_cache;
 pub mod map;
 pub mod snowflake;
 pub mod url_params;

crates/utils/src/lru_cache.rs

@@ -1,41 +0,0 @@
-/*
- * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
- *
- * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
- */
-
-use std::{borrow::Borrow, hash::Hash};
-
-use parking_lot::Mutex;
-
-pub type LruCache<K, V> = Mutex<lru_cache::LruCache<K, V, ahash::RandomState>>;
-
-pub trait LruCached<K, V>: Sized {
-    fn with_capacity(capacity: usize) -> Self;
-    fn get<Q>(&self, name: &Q) -> Option<V>
-    where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized;
-    fn insert(&self, name: K, value: V) -> Option<V>;
-}
-
-impl<K: Hash + Eq, V: Clone> LruCached<K, V> for LruCache<K, V> {
-    fn with_capacity(capacity: usize) -> Self {
-        Mutex::new(lru_cache::LruCache::with_hasher(
-            capacity,
-            ahash::RandomState::new(),
-        ))
-    }
-
-    fn get<Q>(&self, name: &Q) -> Option<V>
-    where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
-    {
-        self.lock().get_mut(name).map(|entry| entry.clone())
-    }
-
-    fn insert(&self, name: K, item: V) -> Option<V> {
-        self.lock().insert(name, item)
-    }
-}

crates/utils/src/map/mod.rs

@@ -6,5 +6,4 @@
 
 pub mod bitmap;
 pub mod mutex_map;
-pub mod ttl_dashmap;
 pub mod vec_map;

crates/utils/src/map/ttl_dashmap.rs

@@ -1,64 +0,0 @@
-/*
- * SPDX-FileCopyrightText: 2020 Stalwart Labs Ltd <hello@stalw.art>
- *
- * SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
- */
-
-use std::{borrow::Borrow, hash::Hash, time::Instant};
-
-use dashmap::DashMap;
-
-pub type TtlDashMap<K, V> = DashMap<K, LruItem<V>, ahash::RandomState>;
-pub type ADashMap<K, V> = DashMap<K, V, ahash::RandomState>;
-
-#[derive(Debug, Clone)]
-pub struct LruItem<V> {
-    pub item: V,
-    valid_until: Instant,
-}
-
-pub trait TtlMap<K, V>: Sized {
-    fn with_capacity(capacity: usize, shard_amount: usize) -> Self;
-    fn get_with_ttl<Q>(&self, name: &Q) -> Option<V>
-    where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized;
-    fn insert_with_ttl(&self, name: K, value: V, valid_until: Instant) -> V;
-    fn cleanup(&self);
-}
-
-impl<K: Hash + Eq, V: Clone> TtlMap<K, V> for TtlDashMap<K, V> {
-    fn with_capacity(capacity: usize, shard_amount: usize) -> Self {
-        DashMap::with_capacity_and_hasher_and_shard_amount(
-            capacity,
-            ahash::RandomState::new(),
-            shard_amount,
-        )
-    }
-
-    fn get_with_ttl<Q>(&self, name: &Q) -> Option<V>
-    where
-        K: Borrow<Q>,
-        Q: Hash + Eq + ?Sized,
-    {
-        match self.get(name) {
-            Some(entry) if entry.valid_until >= Instant::now() => entry.item.clone().into(),
-            _ => None,
-        }
-    }
-
-    fn insert_with_ttl(&self, name: K, item: V, valid_until: Instant) -> V {
-        self.insert(
-            name,
-            LruItem {
-                item: item.clone(),
-                valid_until,
-            },
-        );
-        item
-    }
-
-    fn cleanup(&self) {
-        self.retain(|_, entry| entry.valid_until >= Instant::now());
-    }
-}