Add get_stale and get_stale_while_update for memory-cache

These functions allow to use stale values while performing updates
in the background to minimize the lookup latency.

Author: eaufavor

.bleep

@@ -1 +1 @@
-028a39daa932d3079b73073c39b35404ae06f78d
+bf7bddac9b1ce4e162d4658673a02b70e86610f5

pingora-memory-cache/Cargo.toml

@@ -25,3 +25,6 @@ pingora-error = { version = "0.4.0", path = "../pingora-error" }
 log = { workspace = true }
 parking_lot = "0"
 pingora-timeout = { version = "0.4.0", path = "../pingora-timeout" }
+
+[dev-dependencies]
+once_cell = { workspace = true }

pingora-memory-cache/src/lib.rs

@@ -33,6 +33,9 @@ pub enum CacheStatus {
     Expired,
     /// The key was not initially found but was found after awaiting a lock.
     LockHit,
+    /// The returned value was expired but still returned. The [Duration] is
+    /// how long it has been since its expiration time.
+    Stale(Duration),
 }
 
 impl CacheStatus {
@@ -43,16 +46,25 @@ impl CacheStatus {
             Self::Miss => "miss",
             Self::Expired => "expired",
             Self::LockHit => "lock_hit",
+            Self::Stale(_) => "stale",
         }
     }
 
     /// Returns whether this status represents a cache hit.
     pub fn is_hit(&self) -> bool {
         match self {
-            CacheStatus::Hit | CacheStatus::LockHit => true,
+            CacheStatus::Hit | CacheStatus::LockHit | CacheStatus::Stale(_) => true,
             CacheStatus::Miss | CacheStatus::Expired => false,
         }
     }
+
+    /// Returns the stale duration if any
+    pub fn stale(&self) -> Option<Duration> {
+        match self {
+            CacheStatus::Stale(time) => Some(*time),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Debug, Clone)]
@@ -71,15 +83,21 @@ impl<T: Clone> Node<T> {
     }
 
     fn will_expire_at(&self, time: &Instant) -> bool {
-        match self.expire_on.as_ref() {
-            Some(t) => t <= time,
-            None => false,
-        }
+        self.stale_duration(time).is_some()
     }
 
     fn is_expired(&self) -> bool {
         self.will_expire_at(&Instant::now())
     }
+
+    fn stale_duration(&self, time: &Instant) -> Option<Duration> {
+        let expire_time = self.expire_on?;
+        if &expire_time <= time {
+            Some(time.duration_since(expire_time))
+        } else {
+            None
+        }
+    }
 }
 
 /// A high performant in-memory cache with S3-FIFO + TinyLFU
@@ -107,14 +125,32 @@ impl<K: Hash, T: Clone + Send + Sync + 'static> MemoryCache<K, T> {
             if !n.is_expired() {
                 (Some(n.value), CacheStatus::Hit)
             } else {
-                // TODO: consider returning the staled value
                 (None, CacheStatus::Expired)
             }
         } else {
             (None, CacheStatus::Miss)
         }
     }
 
+    /// Similar to [get], fetch the key and return its value in addition to a
+    /// [CacheStatus] but also return the value even if it is expired. When the
+    /// value is expired, the [Duration] of how long it has been stale will
+    /// also be returned.
+    pub fn get_stale(&self, key: &K) -> (Option<T>, CacheStatus) {
+        let hashed_key = self.hasher.hash_one(key);
+
+        if let Some(n) = self.store.get(&hashed_key) {
+            let stale_duration = n.stale_duration(&Instant::now());
+            if let Some(stale_duration) = stale_duration {
+                (Some(n.value), CacheStatus::Stale(stale_duration))
+            } else {
+                (Some(n.value), CacheStatus::Hit)
+            }
+        } else {
+            (None, CacheStatus::Miss)
+        }
+    }
+
     /// Insert a key and value pair with an optional TTL into the cache.
     ///
     /// An item with zero TTL of zero will not be inserted.
@@ -245,6 +281,20 @@ mod tests {
         assert_eq!(hit, CacheStatus::Expired);
     }
 
+    #[test]
+    fn test_get_stale() {
+        let cache: MemoryCache<i32, i32> = MemoryCache::new(10);
+        let (res, hit) = cache.get(&1);
+        assert_eq!(res, None);
+        assert_eq!(hit, CacheStatus::Miss);
+        cache.put(&1, 2, Some(Duration::from_secs(1)));
+        sleep(Duration::from_millis(1100));
+        let (res, hit) = cache.get_stale(&1);
+        assert_eq!(res.unwrap(), 2);
+        // we slept 1100ms and the ttl is 1000ms
+        assert!(hit.stale().unwrap() >= Duration::from_millis(100));
+    }
+
     #[test]
     fn test_eviction() {
         let cache: MemoryCache<i32, i32> = MemoryCache::new(2);

pingora-memory-cache/src/read_through.rs

@@ -283,6 +283,58 @@ where
             ret
         }
     }
+
+    /// Similar to [get], query the cache for a given value, but also returns the value even if the
+    /// value is expired up to `stale_ttl`. If it is a cache miss or the value is stale more than
+    /// the `stale_ttl`, a lookup will be performed to populate the cache.
+    pub async fn get_stale(
+        &self,
+        key: &K,
+        ttl: Option<Duration>,
+        extra: Option<&S>,
+        stale_ttl: Duration,
+    ) -> (Result<T, Box<Error>>, CacheStatus) {
+        let (result, cache_status) = self.inner.get_stale(key);
+        if let Some(result) = result {
+            let stale_duration = cache_status.stale();
+            if stale_duration.unwrap_or(Duration::ZERO) <= stale_ttl {
+                return (Ok(result), cache_status);
+            }
+        }
+        let (res, status) = self.get(key, ttl, extra).await;
+        (res, status)
+    }
+}
+
+impl<K, T, CB, S> RTCache<K, T, CB, S>
+where
+    K: Hash + Clone + Send + Sync,
+    T: Clone + Send + Sync + 'static,
+    S: Clone + Send + Sync,
+    CB: Lookup<K, T, S> + Sync + Send,
+{
+    /// Similar to [get_stale], but when it returns the stale value, it also initiates a lookup
+    /// in the background in order to refresh the value.
+    ///
+    /// Note that this function requires the [RTCache] to be static, which can be done by wrapping
+    /// it with something like [once_cell::sync::Lazy].
+    pub async fn get_stale_while_update(
+        &'static self,
+        key: &K,
+        ttl: Option<Duration>,
+        extra: Option<&S>,
+        stale_ttl: Duration,
+    ) -> (Result<T, Box<Error>>, CacheStatus) {
+        let (result, cache_status) = self.get_stale(key, ttl, extra, stale_ttl).await;
+        let key = key.clone();
+        let extra = extra.cloned();
+        if cache_status.stale().is_some() {
+            tokio::spawn(async move {
+                let _ = self.get(&key, ttl, extra.as_ref()).await;
+            });
+        }
+        (result, cache_status)
+    }
 }
 
 impl<K, T, CB, S> RTCache<K, T, CB, S>
@@ -686,4 +738,67 @@ mod tests {
             .await
             .unwrap();
     }
+
+    #[tokio::test]
+    async fn test_get_stale() {
+        let ttl = Some(Duration::from_millis(100));
+        let cache: RTCache<i32, i32, TestCB, ExtraOpt> = RTCache::new(10, None, None);
+        let opt = Some(ExtraOpt {
+            error: false,
+            empty: false,
+            delay_for: None,
+            used: Arc::new(AtomicI32::new(0)),
+        });
+        let (res, hit) = cache.get(&1, ttl, opt.as_ref()).await;
+        assert_eq!(res.unwrap(), 1);
+        assert_eq!(hit, CacheStatus::Miss);
+        let (res, hit) = cache.get(&1, ttl, opt.as_ref()).await;
+        assert_eq!(res.unwrap(), 1);
+        assert_eq!(hit, CacheStatus::Hit);
+        tokio::time::sleep(Duration::from_millis(150)).await;
+        let (res, hit) = cache
+            .get_stale(&1, ttl, opt.as_ref(), Duration::from_millis(1000))
+            .await;
+        assert_eq!(res.unwrap(), 1);
+        assert!(hit.stale().is_some());
+
+        let (res, hit) = cache
+            .get_stale(&1, ttl, opt.as_ref(), Duration::from_millis(30))
+            .await;
+        assert_eq!(res.unwrap(), 2);
+        assert_eq!(hit, CacheStatus::Expired);
+    }
+
+    #[tokio::test]
+    async fn test_get_stale_while_update() {
+        use once_cell::sync::Lazy;
+        let ttl = Some(Duration::from_millis(100));
+        static CACHE: Lazy<RTCache<i32, i32, TestCB, ExtraOpt>> =
+            Lazy::new(|| RTCache::new(10, None, None));
+        let opt = Some(ExtraOpt {
+            error: false,
+            empty: false,
+            delay_for: None,
+            used: Arc::new(AtomicI32::new(0)),
+        });
+        let (res, hit) = CACHE.get(&1, ttl, opt.as_ref()).await;
+        assert_eq!(res.unwrap(), 1);
+        assert_eq!(hit, CacheStatus::Miss);
+        let (res, hit) = CACHE.get(&1, ttl, opt.as_ref()).await;
+        assert_eq!(res.unwrap(), 1);
+        assert_eq!(hit, CacheStatus::Hit);
+        tokio::time::sleep(Duration::from_millis(150)).await;
+        let (res, hit) = CACHE
+            .get_stale_while_update(&1, ttl, opt.as_ref(), Duration::from_millis(1000))
+            .await;
+        assert_eq!(res.unwrap(), 1);
+        assert!(hit.stale().is_some());
+
+        // allow the background lookup to finish
+        tokio::time::sleep(Duration::from_millis(10)).await;
+
+        let (res, hit) = CACHE.get(&1, ttl, opt.as_ref()).await;
+        assert_eq!(res.unwrap(), 2);
+        assert_eq!(hit, CacheStatus::Hit);
+    }
 }