Skip to content

Commit

Permalink
add Cached::cache_try_get_or_set_with
Browse files Browse the repository at this point in the history
  • Loading branch information
wbenny committed Sep 8, 2024
1 parent a953546 commit dd73b56
Show file tree
Hide file tree
Showing 9 changed files with 233 additions and 5 deletions.
14 changes: 13 additions & 1 deletion examples/kitchen_sink.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -83,6 +83,18 @@ impl<K: Hash + Eq, V> Cached<K, V> for MyCache<K, V> {
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V {
self.store.entry(k).or_insert_with(f)
}
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<V> {
self.store.insert(k, v)
}
Expand Down
14 changes: 13 additions & 1 deletion examples/kitchen_sink_proc_macro.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -100,6 +100,18 @@ impl<K: Hash + Eq, V> Cached<K, V> for MyCache<K, V> {
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V {
self.store.entry(k).or_insert_with(f)
}
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<V> {
self.store.insert(k, v)
}
Expand Down
7 changes: 7 additions & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,13 @@ pub trait Cached<K, V> {
/// Get or insert a key, value pair
fn cache_get_or_set_with<F: FnOnce() -> 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<F: FnOnce() -> Result<V, E>, E>(
&mut self,
k: K,
f: F,
) -> Result<&mut V, E>;

/// Remove a cached value
///
/// ```rust
Expand Down
38 changes: 38 additions & 0 deletions src/stores/expiring_value_cache.rs
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,21 @@ impl<K: Hash + Eq + Clone, V: CanExpire> Cached<K, V> for ExpiringValueCache<K,
}
v
}
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<V> {
self.store.cache_set(k, v)
}
Expand Down Expand Up @@ -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<u8, ExpiredU8> = 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<u8, ExpiredU8> = ExpiringValueCache::with_size(3);
Expand Down
13 changes: 12 additions & 1 deletion src/stores/mod.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -73,6 +72,18 @@ where
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, key: K, f: F) -> &mut V {
self.entry(key).or_insert_with(f)
}
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<Q>(&mut self, k: &Q) -> Option<V>
where
K: std::borrow::Borrow<Q>,
Expand Down
29 changes: 27 additions & 2 deletions src/stores/sized.rs
Original file line number Diff line number Diff line change
Expand Up @@ -295,8 +295,7 @@ impl<K: Hash + Eq + Clone, V> SizedCache<K, V> {
}
}

#[allow(dead_code)]
fn try_get_or_set_with_if<E, F: FnOnce() -> Result<V, E>, FC: FnOnce(&V) -> bool>(
pub(super) fn try_get_or_set_with_if<E, F: FnOnce() -> Result<V, E>, FC: FnOnce(&V) -> bool>(
&mut self,
key: K,
f: F,
Expand Down Expand Up @@ -456,6 +455,15 @@ impl<K: Hash + Eq + Clone, V> Cached<K, V> for SizedCache<K, V> {
v
}

fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<Q>(&mut self, k: &Q) -> Option<V>
where
K: std::borrow::Borrow<Q>,
Expand Down Expand Up @@ -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<usize, String> {
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")]
Expand Down
47 changes: 47 additions & 0 deletions src/stores/timed.rs
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,33 @@ impl<K: Hash + Eq, V> Cached<K, V> for TimedCache<K, V> {
}
}

fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<V> {
let stamped = (Instant::now(), val);
self.store.insert(key, stamped).and_then(|(instant, v)| {
Expand Down Expand Up @@ -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<usize, String> {
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);
}
}
43 changes: 43 additions & 0 deletions src/stores/timed_sized.rs
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,28 @@ impl<K: Hash + Eq + Clone, V> Cached<K, V> for TimedSizedCache<K, V> {
&mut stamped.1
}

fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<V> {
let stamped = self.store.cache_set(key, (Instant::now(), val));
stamped.and_then(|(instant, v)| {
Expand Down Expand Up @@ -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<usize, String> {
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")]
Expand Down
33 changes: 33 additions & 0 deletions src/stores/unbound.rs
Original file line number Diff line number Diff line change
Expand Up @@ -120,6 +120,23 @@ impl<K: Hash + Eq, V> Cached<K, V> for UnboundCache<K, V> {
}
}
}
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, 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<Q>(&mut self, k: &Q) -> Option<V>
where
K: std::borrow::Borrow<Q>,
Expand Down Expand Up @@ -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<usize, String> {
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);
}
}

0 comments on commit dd73b56

Please sign in to comment.