From dd73b569efcef24eef2215231fe7e02e16df3b66 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Petr=20Bene=C5=A1?= Date: Sun, 8 Sep 2024 12:58:53 +0000 Subject: [PATCH] add Cached::cache_try_get_or_set_with --- examples/kitchen_sink.rs | 14 ++++++++- examples/kitchen_sink_proc_macro.rs | 14 ++++++++- src/lib.rs | 7 +++++ src/stores/expiring_value_cache.rs | 38 +++++++++++++++++++++++ src/stores/mod.rs | 13 +++++++- src/stores/sized.rs | 29 ++++++++++++++++-- src/stores/timed.rs | 47 +++++++++++++++++++++++++++++ src/stores/timed_sized.rs | 43 ++++++++++++++++++++++++++ src/stores/unbound.rs | 33 ++++++++++++++++++++ 9 files changed, 233 insertions(+), 5 deletions(-) diff --git a/examples/kitchen_sink.rs b/examples/kitchen_sink.rs index 3dffb48..32eaaf9 100644 --- a/examples/kitchen_sink.rs +++ b/examples/kitchen_sink.rs @@ -2,7 +2,7 @@ extern crate cached; use std::cmp::Eq; -use std::collections::HashMap; +use std::collections::{hash_map::Entry, HashMap}; use std::hash::Hash; use std::thread::sleep; @@ -83,6 +83,18 @@ impl Cached for MyCache { fn cache_get_or_set_with V>(&mut self, k: K, f: F) -> &mut V { self.store.entry(k).or_insert_with(f) } + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + let v = match self.store.entry(k) { + Entry::Occupied(occupied) => occupied.into_mut(), + Entry::Vacant(vacant) => vacant.insert(f()?), + }; + + Ok(v) + } fn cache_set(&mut self, k: K, v: V) -> Option { self.store.insert(k, v) } diff --git a/examples/kitchen_sink_proc_macro.rs b/examples/kitchen_sink_proc_macro.rs index 106f76d..cbab61d 100644 --- a/examples/kitchen_sink_proc_macro.rs +++ b/examples/kitchen_sink_proc_macro.rs @@ -2,7 +2,7 @@ use cached::proc_macro::cached; use cached::Return; use cached::{Cached, SizedCache, UnboundCache}; use std::cmp::Eq; -use std::collections::HashMap; +use std::collections::{hash_map::Entry, HashMap}; use std::hash::Hash; use std::thread::{sleep, spawn}; use std::time::Duration; @@ -100,6 +100,18 @@ impl Cached for MyCache { fn cache_get_or_set_with V>(&mut self, k: K, f: F) -> &mut V { self.store.entry(k).or_insert_with(f) } + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + let v = match self.store.entry(k) { + Entry::Occupied(occupied) => occupied.into_mut(), + Entry::Vacant(vacant) => vacant.insert(f()?), + }; + + Ok(v) + } fn cache_set(&mut self, k: K, v: V) -> Option { self.store.insert(k, v) } diff --git a/src/lib.rs b/src/lib.rs index deca0b4..163f322 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -318,6 +318,13 @@ pub trait Cached { /// Get or insert a key, value pair fn cache_get_or_set_with V>(&mut self, k: K, f: F) -> &mut V; + /// Get or insert a key, value pair with error handling + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E>; + /// Remove a cached value /// /// ```rust diff --git a/src/stores/expiring_value_cache.rs b/src/stores/expiring_value_cache.rs index 34b550c..42615ad 100644 --- a/src/stores/expiring_value_cache.rs +++ b/src/stores/expiring_value_cache.rs @@ -116,6 +116,21 @@ impl Cached for ExpiringValueCache Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + let (was_present, was_valid, v) = self + .store + .try_get_or_set_with_if(k, f, |v| !v.is_expired())?; + if was_present && was_valid { + self.hits += 1; + } else { + self.misses += 1; + } + Ok(v) + } fn cache_set(&mut self, k: K, v: V) -> Option { self.store.cache_set(k, v) } @@ -278,6 +293,29 @@ mod tests { assert_eq!(c.cache_misses(), Some(1)); } + #[test] + fn expiring_value_cache_try_get_or_set_with_missing() { + let mut c: ExpiringValueCache = ExpiringValueCache::with_size(3); + + assert_eq!( + c.cache_try_get_or_set_with(1, || Ok::<_, ()>(1)), + Ok(&mut 1) + ); + assert_eq!(c.cache_hits(), Some(0)); + assert_eq!(c.cache_misses(), Some(1)); + + assert_eq!(c.cache_try_get_or_set_with(1, || Err(())), Ok(&mut 1)); + assert_eq!(c.cache_hits(), Some(1)); + assert_eq!(c.cache_misses(), Some(1)); + + assert_eq!( + c.cache_try_get_or_set_with(2, || Ok::<_, ()>(2)), + Ok(&mut 2) + ); + assert_eq!(c.cache_hits(), Some(1)); + assert_eq!(c.cache_misses(), Some(2)); + } + #[test] fn flush_expired() { let mut c: ExpiringValueCache = ExpiringValueCache::with_size(3); diff --git a/src/stores/mod.rs b/src/stores/mod.rs index d030816..d0f7f3a 100644 --- a/src/stores/mod.rs +++ b/src/stores/mod.rs @@ -1,6 +1,5 @@ use crate::Cached; use std::cmp::Eq; -#[cfg(feature = "async")] use std::collections::hash_map::Entry; use std::collections::HashMap; use std::hash::Hash; @@ -73,6 +72,18 @@ where fn cache_get_or_set_with V>(&mut self, key: K, f: F) -> &mut V { self.entry(key).or_insert_with(f) } + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + let v = match self.entry(k) { + Entry::Occupied(occupied) => occupied.into_mut(), + Entry::Vacant(vacant) => vacant.insert(f()?), + }; + + Ok(v) + } fn cache_remove(&mut self, k: &Q) -> Option where K: std::borrow::Borrow, diff --git a/src/stores/sized.rs b/src/stores/sized.rs index d536bc7..804f829 100644 --- a/src/stores/sized.rs +++ b/src/stores/sized.rs @@ -295,8 +295,7 @@ impl SizedCache { } } - #[allow(dead_code)] - fn try_get_or_set_with_if Result, FC: FnOnce(&V) -> bool>( + pub(super) fn try_get_or_set_with_if Result, FC: FnOnce(&V) -> bool>( &mut self, key: K, f: F, @@ -456,6 +455,15 @@ impl Cached for SizedCache { v } + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + let (_, _, v) = self.try_get_or_set_with_if(k, f, |_| true)?; + Ok(v) + } + fn cache_remove(&mut self, k: &Q) -> Option where K: std::borrow::Borrow, @@ -759,6 +767,23 @@ mod tests { assert_eq!(c.cache_get_or_set_with(1, || 1), &1); assert_eq!(c.cache_misses(), Some(8)); + + c.cache_reset(); + fn _try_get(n: usize) -> Result { + if n < 10 { + Ok(n) + } else { + Err("dead".to_string()) + } + } + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10)); + assert!(res.is_err()); + assert!(c.key_order().next().is_none()); + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1)); + assert_eq!(res.unwrap(), &1); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &1); } #[cfg(feature = "async")] diff --git a/src/stores/timed.rs b/src/stores/timed.rs index 68b3e30..77d57a4 100644 --- a/src/stores/timed.rs +++ b/src/stores/timed.rs @@ -191,6 +191,33 @@ impl Cached for TimedCache { } } + fn cache_try_get_or_set_with Result, E>( + &mut self, + key: K, + f: F, + ) -> Result<&mut V, E> { + match self.store.entry(key) { + Entry::Occupied(mut occupied) => { + if occupied.get().0.elapsed().as_secs() < self.seconds { + if self.refresh { + occupied.get_mut().0 = Instant::now(); + } + self.hits += 1; + } else { + self.misses += 1; + let val = f()?; + occupied.insert((Instant::now(), val)); + } + Ok(&mut occupied.into_mut().1) + } + Entry::Vacant(vacant) => { + self.misses += 1; + let val = f()?; + Ok(&mut vacant.insert((Instant::now(), val)).1) + } + } + } + fn cache_set(&mut self, key: K, val: V) -> Option { let stamped = (Instant::now(), val); self.store.insert(key, stamped).and_then(|(instant, v)| { @@ -539,5 +566,25 @@ mod tests { assert_eq!(c.cache_get_or_set_with(1, || 42), &42); assert_eq!(c.cache_misses(), Some(7)); + + c.cache_reset(); + fn _try_get(n: usize) -> Result { + if n < 10 { + Ok(n) + } else { + Err("dead".to_string()) + } + } + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10)); + assert!(res.is_err()); + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1)); + assert_eq!(res.unwrap(), &1); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &1); + sleep(Duration::new(2, 0)); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &5); } } diff --git a/src/stores/timed_sized.rs b/src/stores/timed_sized.rs index 27b6506..1e3716b 100644 --- a/src/stores/timed_sized.rs +++ b/src/stores/timed_sized.rs @@ -212,6 +212,28 @@ impl Cached for TimedSizedCache { &mut stamped.1 } + fn cache_try_get_or_set_with Result, E>( + &mut self, + key: K, + f: F, + ) -> Result<&mut V, E> { + let setter = || Ok((Instant::now(), f()?)); + let max_seconds = self.seconds; + let (was_present, was_valid, stamped) = + self.store.try_get_or_set_with_if(key, setter, |stamped| { + stamped.0.elapsed().as_secs() < max_seconds + })?; + if was_present && was_valid { + if self.refresh { + stamped.0 = Instant::now(); + } + self.hits += 1; + } else { + self.misses += 1; + } + Ok(&mut stamped.1) + } + fn cache_set(&mut self, key: K, val: V) -> Option { let stamped = self.store.cache_set(key, (Instant::now(), val)); stamped.and_then(|(instant, v)| { @@ -634,6 +656,27 @@ mod tests { assert_eq!(c.cache_get_or_set_with(6, || 42), &6); assert_eq!(c.cache_misses(), Some(11)); + + c.cache_reset(); + fn _try_get(n: usize) -> Result { + if n < 10 { + Ok(n) + } else { + Err("dead".to_string()) + } + } + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10)); + assert!(res.is_err()); + assert!(c.key_order().next().is_none()); + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1)); + assert_eq!(res.unwrap(), &1); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &1); + sleep(Duration::new(2, 0)); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &5); } #[cfg(feature = "async")] diff --git a/src/stores/unbound.rs b/src/stores/unbound.rs index 766b1d8..c2be7d4 100644 --- a/src/stores/unbound.rs +++ b/src/stores/unbound.rs @@ -120,6 +120,23 @@ impl Cached for UnboundCache { } } } + fn cache_try_get_or_set_with Result, E>( + &mut self, + k: K, + f: F, + ) -> Result<&mut V, E> { + match self.store.entry(k) { + Entry::Occupied(occupied) => { + self.hits += 1; + Ok(occupied.into_mut()) + } + + Entry::Vacant(vacant) => { + self.misses += 1; + Ok(vacant.insert(f()?)) + } + } + } fn cache_remove(&mut self, k: &Q) -> Option where K: std::borrow::Borrow, @@ -345,5 +362,21 @@ mod tests { assert_eq!(c.cache_get_or_set_with(1, || 1), &1); assert_eq!(c.cache_misses(), Some(6)); + + c.cache_reset(); + fn _try_get(n: usize) -> Result { + if n < 10 { + Ok(n) + } else { + Err("dead".to_string()) + } + } + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10)); + assert!(res.is_err()); + + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1)); + assert_eq!(res.unwrap(), &1); + let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5)); + assert_eq!(res.unwrap(), &1); } }