Skip to content

Commit 6e8ad31

Browse files
jaemkjaemk
authored
Merge pull request #224 from wbenny/master
add Cached::cache_try_get_or_set_with
2 parents a953546 + dd73b56 commit 6e8ad31

9 files changed

+233
-5
lines changed

‎examples/kitchen_sink.rs

+13-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
extern crate cached;
33

44
use std::cmp::Eq;
5-
use std::collections::HashMap;
5+
use std::collections::{hash_map::Entry, HashMap};
66
use std::hash::Hash;
77

88
use std::thread::sleep;
@@ -83,6 +83,18 @@ impl<K: Hash + Eq, V> Cached<K, V> for MyCache<K, V> {
8383
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V {
8484
self.store.entry(k).or_insert_with(f)
8585
}
86+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
87+
&mut self,
88+
k: K,
89+
f: F,
90+
) -> Result<&mut V, E> {
91+
let v = match self.store.entry(k) {
92+
Entry::Occupied(occupied) => occupied.into_mut(),
93+
Entry::Vacant(vacant) => vacant.insert(f()?),
94+
};
95+
96+
Ok(v)
97+
}
8698
fn cache_set(&mut self, k: K, v: V) -> Option<V> {
8799
self.store.insert(k, v)
88100
}

‎examples/kitchen_sink_proc_macro.rs

+13-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ use cached::proc_macro::cached;
22
use cached::Return;
33
use cached::{Cached, SizedCache, UnboundCache};
44
use std::cmp::Eq;
5-
use std::collections::HashMap;
5+
use std::collections::{hash_map::Entry, HashMap};
66
use std::hash::Hash;
77
use std::thread::{sleep, spawn};
88
use std::time::Duration;
@@ -100,6 +100,18 @@ impl<K: Hash + Eq, V> Cached<K, V> for MyCache<K, V> {
100100
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V {
101101
self.store.entry(k).or_insert_with(f)
102102
}
103+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
104+
&mut self,
105+
k: K,
106+
f: F,
107+
) -> Result<&mut V, E> {
108+
let v = match self.store.entry(k) {
109+
Entry::Occupied(occupied) => occupied.into_mut(),
110+
Entry::Vacant(vacant) => vacant.insert(f()?),
111+
};
112+
113+
Ok(v)
114+
}
103115
fn cache_set(&mut self, k: K, v: V) -> Option<V> {
104116
self.store.insert(k, v)
105117
}

‎src/lib.rs

+7
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,13 @@ pub trait Cached<K, V> {
318318
/// Get or insert a key, value pair
319319
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, k: K, f: F) -> &mut V;
320320

321+
/// Get or insert a key, value pair with error handling
322+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
323+
&mut self,
324+
k: K,
325+
f: F,
326+
) -> Result<&mut V, E>;
327+
321328
/// Remove a cached value
322329
///
323330
/// ```rust

‎src/stores/expiring_value_cache.rs

+38
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,21 @@ impl<K: Hash + Eq + Clone, V: CanExpire> Cached<K, V> for ExpiringValueCache<K,
116116
}
117117
v
118118
}
119+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
120+
&mut self,
121+
k: K,
122+
f: F,
123+
) -> Result<&mut V, E> {
124+
let (was_present, was_valid, v) = self
125+
.store
126+
.try_get_or_set_with_if(k, f, |v| !v.is_expired())?;
127+
if was_present && was_valid {
128+
self.hits += 1;
129+
} else {
130+
self.misses += 1;
131+
}
132+
Ok(v)
133+
}
119134
fn cache_set(&mut self, k: K, v: V) -> Option<V> {
120135
self.store.cache_set(k, v)
121136
}
@@ -278,6 +293,29 @@ mod tests {
278293
assert_eq!(c.cache_misses(), Some(1));
279294
}
280295

296+
#[test]
297+
fn expiring_value_cache_try_get_or_set_with_missing() {
298+
let mut c: ExpiringValueCache<u8, ExpiredU8> = ExpiringValueCache::with_size(3);
299+
300+
assert_eq!(
301+
c.cache_try_get_or_set_with(1, || Ok::<_, ()>(1)),
302+
Ok(&mut 1)
303+
);
304+
assert_eq!(c.cache_hits(), Some(0));
305+
assert_eq!(c.cache_misses(), Some(1));
306+
307+
assert_eq!(c.cache_try_get_or_set_with(1, || Err(())), Ok(&mut 1));
308+
assert_eq!(c.cache_hits(), Some(1));
309+
assert_eq!(c.cache_misses(), Some(1));
310+
311+
assert_eq!(
312+
c.cache_try_get_or_set_with(2, || Ok::<_, ()>(2)),
313+
Ok(&mut 2)
314+
);
315+
assert_eq!(c.cache_hits(), Some(1));
316+
assert_eq!(c.cache_misses(), Some(2));
317+
}
318+
281319
#[test]
282320
fn flush_expired() {
283321
let mut c: ExpiringValueCache<u8, ExpiredU8> = ExpiringValueCache::with_size(3);

‎src/stores/mod.rs

+12-1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
use crate::Cached;
22
use std::cmp::Eq;
3-
#[cfg(feature = "async")]
43
use std::collections::hash_map::Entry;
54
use std::collections::HashMap;
65
use std::hash::Hash;
@@ -73,6 +72,18 @@ where
7372
fn cache_get_or_set_with<F: FnOnce() -> V>(&mut self, key: K, f: F) -> &mut V {
7473
self.entry(key).or_insert_with(f)
7574
}
75+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
76+
&mut self,
77+
k: K,
78+
f: F,
79+
) -> Result<&mut V, E> {
80+
let v = match self.entry(k) {
81+
Entry::Occupied(occupied) => occupied.into_mut(),
82+
Entry::Vacant(vacant) => vacant.insert(f()?),
83+
};
84+
85+
Ok(v)
86+
}
7687
fn cache_remove<Q>(&mut self, k: &Q) -> Option<V>
7788
where
7889
K: std::borrow::Borrow<Q>,

‎src/stores/sized.rs

+27-2
Original file line numberDiff line numberDiff line change
@@ -295,8 +295,7 @@ impl<K: Hash + Eq + Clone, V> SizedCache<K, V> {
295295
}
296296
}
297297

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

458+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
459+
&mut self,
460+
k: K,
461+
f: F,
462+
) -> Result<&mut V, E> {
463+
let (_, _, v) = self.try_get_or_set_with_if(k, f, |_| true)?;
464+
Ok(v)
465+
}
466+
459467
fn cache_remove<Q>(&mut self, k: &Q) -> Option<V>
460468
where
461469
K: std::borrow::Borrow<Q>,
@@ -759,6 +767,23 @@ mod tests {
759767
assert_eq!(c.cache_get_or_set_with(1, || 1), &1);
760768

761769
assert_eq!(c.cache_misses(), Some(8));
770+
771+
c.cache_reset();
772+
fn _try_get(n: usize) -> Result<usize, String> {
773+
if n < 10 {
774+
Ok(n)
775+
} else {
776+
Err("dead".to_string())
777+
}
778+
}
779+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10));
780+
assert!(res.is_err());
781+
assert!(c.key_order().next().is_none());
782+
783+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1));
784+
assert_eq!(res.unwrap(), &1);
785+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
786+
assert_eq!(res.unwrap(), &1);
762787
}
763788

764789
#[cfg(feature = "async")]

‎src/stores/timed.rs

+47
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,33 @@ impl<K: Hash + Eq, V> Cached<K, V> for TimedCache<K, V> {
191191
}
192192
}
193193

194+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
195+
&mut self,
196+
key: K,
197+
f: F,
198+
) -> Result<&mut V, E> {
199+
match self.store.entry(key) {
200+
Entry::Occupied(mut occupied) => {
201+
if occupied.get().0.elapsed().as_secs() < self.seconds {
202+
if self.refresh {
203+
occupied.get_mut().0 = Instant::now();
204+
}
205+
self.hits += 1;
206+
} else {
207+
self.misses += 1;
208+
let val = f()?;
209+
occupied.insert((Instant::now(), val));
210+
}
211+
Ok(&mut occupied.into_mut().1)
212+
}
213+
Entry::Vacant(vacant) => {
214+
self.misses += 1;
215+
let val = f()?;
216+
Ok(&mut vacant.insert((Instant::now(), val)).1)
217+
}
218+
}
219+
}
220+
194221
fn cache_set(&mut self, key: K, val: V) -> Option<V> {
195222
let stamped = (Instant::now(), val);
196223
self.store.insert(key, stamped).and_then(|(instant, v)| {
@@ -539,5 +566,25 @@ mod tests {
539566
assert_eq!(c.cache_get_or_set_with(1, || 42), &42);
540567

541568
assert_eq!(c.cache_misses(), Some(7));
569+
570+
c.cache_reset();
571+
fn _try_get(n: usize) -> Result<usize, String> {
572+
if n < 10 {
573+
Ok(n)
574+
} else {
575+
Err("dead".to_string())
576+
}
577+
}
578+
579+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10));
580+
assert!(res.is_err());
581+
582+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1));
583+
assert_eq!(res.unwrap(), &1);
584+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
585+
assert_eq!(res.unwrap(), &1);
586+
sleep(Duration::new(2, 0));
587+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
588+
assert_eq!(res.unwrap(), &5);
542589
}
543590
}

‎src/stores/timed_sized.rs

+43
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,28 @@ impl<K: Hash + Eq + Clone, V> Cached<K, V> for TimedSizedCache<K, V> {
212212
&mut stamped.1
213213
}
214214

215+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
216+
&mut self,
217+
key: K,
218+
f: F,
219+
) -> Result<&mut V, E> {
220+
let setter = || Ok((Instant::now(), f()?));
221+
let max_seconds = self.seconds;
222+
let (was_present, was_valid, stamped) =
223+
self.store.try_get_or_set_with_if(key, setter, |stamped| {
224+
stamped.0.elapsed().as_secs() < max_seconds
225+
})?;
226+
if was_present && was_valid {
227+
if self.refresh {
228+
stamped.0 = Instant::now();
229+
}
230+
self.hits += 1;
231+
} else {
232+
self.misses += 1;
233+
}
234+
Ok(&mut stamped.1)
235+
}
236+
215237
fn cache_set(&mut self, key: K, val: V) -> Option<V> {
216238
let stamped = self.store.cache_set(key, (Instant::now(), val));
217239
stamped.and_then(|(instant, v)| {
@@ -634,6 +656,27 @@ mod tests {
634656
assert_eq!(c.cache_get_or_set_with(6, || 42), &6);
635657

636658
assert_eq!(c.cache_misses(), Some(11));
659+
660+
c.cache_reset();
661+
fn _try_get(n: usize) -> Result<usize, String> {
662+
if n < 10 {
663+
Ok(n)
664+
} else {
665+
Err("dead".to_string())
666+
}
667+
}
668+
669+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10));
670+
assert!(res.is_err());
671+
assert!(c.key_order().next().is_none());
672+
673+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1));
674+
assert_eq!(res.unwrap(), &1);
675+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
676+
assert_eq!(res.unwrap(), &1);
677+
sleep(Duration::new(2, 0));
678+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
679+
assert_eq!(res.unwrap(), &5);
637680
}
638681

639682
#[cfg(feature = "async")]

‎src/stores/unbound.rs

+33
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,23 @@ impl<K: Hash + Eq, V> Cached<K, V> for UnboundCache<K, V> {
120120
}
121121
}
122122
}
123+
fn cache_try_get_or_set_with<F: FnOnce() -> Result<V, E>, E>(
124+
&mut self,
125+
k: K,
126+
f: F,
127+
) -> Result<&mut V, E> {
128+
match self.store.entry(k) {
129+
Entry::Occupied(occupied) => {
130+
self.hits += 1;
131+
Ok(occupied.into_mut())
132+
}
133+
134+
Entry::Vacant(vacant) => {
135+
self.misses += 1;
136+
Ok(vacant.insert(f()?))
137+
}
138+
}
139+
}
123140
fn cache_remove<Q>(&mut self, k: &Q) -> Option<V>
124141
where
125142
K: std::borrow::Borrow<Q>,
@@ -345,5 +362,21 @@ mod tests {
345362
assert_eq!(c.cache_get_or_set_with(1, || 1), &1);
346363

347364
assert_eq!(c.cache_misses(), Some(6));
365+
366+
c.cache_reset();
367+
fn _try_get(n: usize) -> Result<usize, String> {
368+
if n < 10 {
369+
Ok(n)
370+
} else {
371+
Err("dead".to_string())
372+
}
373+
}
374+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(10));
375+
assert!(res.is_err());
376+
377+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(1));
378+
assert_eq!(res.unwrap(), &1);
379+
let res: Result<&mut usize, String> = c.cache_try_get_or_set_with(0, || _try_get(5));
380+
assert_eq!(res.unwrap(), &1);
348381
}
349382
}

0 commit comments

Comments
 (0)