Individual ttl for entries

stusmall · Nov 22, 2016 · ebf2186 · ebf2186
1 parent 682f4ae
commit ebf2186
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 16 deletions.
diff --git a/src/lib.rs b/src/lib.rs
@@ -117,6 +117,31 @@ impl<K: Eq + Hash, V, S: BuildHasher> TtlCache<K, V, S> {
         old_val.and_then(|x| if x.is_expired() { None } else { Some(x.value) })
     }
 
+    /// Inserts a key-value pair into the cache with an individual ttl for the key. If the key already existed and hasn't expired,
+    /// the old value is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::time::Duration;
+    /// use ttl_cache::TtlCache;
+    ///
+    /// let mut cache = TtlCache::new(Duration::from_secs(30), 2);
+    ///
+    /// cache.insert_with_ttl(1, "a", Duration::from_secs(20));
+    /// cache.insert_with_ttl(2, "b", Duration::from_secs(60));
+    /// assert_eq!(cache.get_mut(&1), Some(&mut "a"));
+    /// assert_eq!(cache.get_mut(&2), Some(&mut "b"));
+    /// ```
+    pub fn insert_with_ttl(&mut self, k: K, v: V, ttl: Duration) -> Option<V> {
+        let to_insert = Entry::new(v, ttl);
+        let old_val = self.map.insert(k, to_insert);
+        if self.len() > self.capacity() {
+            self.remove_oldest();
+        }
+        old_val.and_then(|x| if x.is_expired() { None } else { Some(x.value) })
+    }
+
     /// Returns a mutable reference to the value corresponding to the given key in the cache, if
     /// it contains an unexpired entry.
     ///

diff --git a/tests/test.rs b/tests/test.rs
@@ -5,8 +5,8 @@ use std::time::Duration;
 use std::thread::sleep;
 
 #[test]
-fn test_put_and_get(){
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),2);
+fn test_put_and_get() {
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 2);
     cache.insert(1, 10);
     cache.insert(2, 20);
     assert_eq!(cache.get_mut(&1), Some(&mut 10));
@@ -15,30 +15,46 @@ fn test_put_and_get(){
 
 #[test]
 fn test_put_update() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),1);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 1);
     cache.insert("1", 10);
     cache.insert("1", 19);
     assert_eq!(cache.get_mut("1"), Some(&mut 19));
 }
 
 #[test]
 fn test_contains_key() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),1);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 1);
     cache.insert("1", 10);
     assert_eq!(cache.contains_key("1"), true);
 }
 
 #[test]
-fn test_expire_value(){
-    let mut cache = TtlCache::new(Duration::from_millis(1),1);
+fn test_expire_value() {
+    let mut cache = TtlCache::new(Duration::from_millis(1), 1);
     cache.insert("1", 10);
     sleep(Duration::from_millis(10));
     assert_eq!(cache.contains_key("1"), false);
 }
 
+#[test]
+fn test_individual_ttl_value() {
+    let mut cache = TtlCache::new(Duration::from_millis(1), 1);
+    cache.insert_with_ttl("1", 99, Duration::from_millis(15));
+    sleep(Duration::from_millis(10));
+    assert_eq!(cache.contains_key("1"), true);
+}
+
+#[test]
+fn test_individual_ttl_expired() {
+    let mut cache = TtlCache::new(Duration::from_millis(20), 1);
+    cache.insert_with_ttl("1", 99, Duration::from_millis(5));
+    sleep(Duration::from_millis(10));
+    assert_eq!(cache.contains_key("1"), false);
+}
+
 #[test]
 fn test_pop() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),2);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 2);
     cache.insert(1, 10);
     cache.insert(2, 20);
     let opt1 = cache.remove(&1);
@@ -49,7 +65,7 @@ fn test_pop() {
 
 #[test]
 fn test_change_capacity() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),2);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 2);
     assert_eq!(cache.capacity(), 2);
     cache.insert(1, 10);
     cache.insert(2, 20);
@@ -60,7 +76,7 @@ fn test_change_capacity() {
 
 #[test]
 fn test_remove() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),3);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 3);
     cache.insert(1, 10);
     cache.insert(2, 20);
     cache.insert(3, 30);
@@ -81,7 +97,7 @@ fn test_remove() {
 
 #[test]
 fn test_clear() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),2);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 2);
     cache.insert(1, 10);
     cache.insert(2, 20);
     cache.clear();
@@ -91,7 +107,7 @@ fn test_clear() {
 
 #[test]
 fn test_iter() {
-    let mut cache = TtlCache::new(Duration::from_secs(60*60),3);
+    let mut cache = TtlCache::new(Duration::from_secs(60 * 60), 3);
     cache.insert(1, 10);
     cache.insert(2, 20);
     cache.insert(3, 30);
@@ -110,17 +126,16 @@ fn test_iter() {
 
 #[test]
 fn test_iter_w_expired() {
-    let mut cache = TtlCache::new(Duration::from_millis(100),3);
+    let mut cache = TtlCache::new(Duration::from_millis(100), 3);
     cache.insert(1, 10);
     sleep(Duration::from_millis(200));
     cache.insert(2, 20);
     cache.insert(3, 30);
-    assert_eq!(cache.iter().collect::<Vec<_>>(),
-               [(&2, &20), (&3, &30)]);
+    assert_eq!(cache.iter().collect::<Vec<_>>(), [(&2, &20), (&3, &30)]);
     assert_eq!(cache.iter_mut().collect::<Vec<_>>(),
-                [(&2, &mut 20), (&3, &mut 30)]);
+               [(&2, &mut 20), (&3, &mut 30)]);
     assert_eq!(cache.iter().rev().collect::<Vec<_>>(),
-                [(&3, &30), (&2, &20)]);
+               [(&3, &30), (&2, &20)]);
     assert_eq!(cache.iter_mut().rev().collect::<Vec<_>>(),
                [(&3, &mut 30), (&2, &mut 20)]);
 }