matrix_sdk_common/
ttl_cache.rs

1// Copyright 2025 The Matrix.org Foundation C.I.C.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! A TTL cache which can be used to time out repeated operations that might
16//! experience intermittent failures.
17
18use std::{borrow::Borrow, collections::HashMap, hash::Hash, time::Duration};
19
20use ruma::time::Instant;
21
22// One day is the default lifetime.
23const DEFAULT_LIFETIME: Duration = Duration::from_secs(24 * 60 * 60);
24
25#[derive(Debug)]
26struct TtlItem<V: Clone> {
27    value: V,
28    insertion_time: Instant,
29    lifetime: Duration,
30}
31
32impl<V: Clone> TtlItem<V> {
33    fn expired(&self) -> bool {
34        self.insertion_time.elapsed() >= self.lifetime
35    }
36}
37
38/// A TTL cache where items get removed deterministically in the `get()` call.
39#[derive(Debug)]
40pub struct TtlCache<K: Eq + Hash, V: Clone> {
41    lifetime: Duration,
42    items: HashMap<K, TtlItem<V>>,
43}
44
45impl<K, V> TtlCache<K, V>
46where
47    K: Eq + Hash,
48    V: Clone,
49{
50    /// Create a new, empty, [`TtlCache`].
51    pub fn new() -> Self {
52        Self { items: Default::default(), lifetime: DEFAULT_LIFETIME }
53    }
54
55    /// Does the cache contain an non-expired item with the matching key.
56    pub fn contains<Q>(&self, key: &Q) -> bool
57    where
58        K: Borrow<Q>,
59        Q: Hash + Eq + ?Sized,
60    {
61        let cache = &self.items;
62        if let Some(item) = cache.get(key) { !item.expired() } else { false }
63    }
64
65    /// Add a single item to the cache.
66    pub fn insert(&mut self, key: K, value: V) {
67        self.extend([(key, value)]);
68    }
69
70    /// Extend the cache with the given iterator of items.
71    pub fn extend(&mut self, iterator: impl IntoIterator<Item = (K, V)>) {
72        let cache = &mut self.items;
73
74        let now = Instant::now();
75
76        for (key, value) in iterator {
77            let item = TtlItem { value, insertion_time: now, lifetime: self.lifetime };
78
79            cache.insert(key, item);
80        }
81    }
82
83    /// Remove the item that matches the given key.
84    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
85    where
86        K: Borrow<Q>,
87        Q: Hash + Eq + ?Sized,
88    {
89        self.items.remove(key.borrow()).map(|item| item.value)
90    }
91
92    /// Get the item that matches the given key, if the item has expired `None`
93    /// will be returned and the item will be evicted from the cache.
94    pub fn get<Q>(&mut self, key: &Q) -> Option<V>
95    where
96        K: Borrow<Q>,
97        Q: Hash + Eq + ?Sized,
98    {
99        // Remove all expired items.
100        self.items.retain(|_, value| !value.expired());
101        // Now get the wanted item.
102        self.items.get(key.borrow()).map(|item| item.value.clone())
103    }
104
105    /// Force the expiry of the given item, if it is present in the cache.
106    ///
107    /// This doesn't remove the item, it just marks it as expired.
108    #[doc(hidden)]
109    pub fn expire<Q>(&mut self, key: &Q)
110    where
111        K: Borrow<Q>,
112        Q: Hash + Eq + ?Sized,
113    {
114        if let Some(item) = self.items.get_mut(key) {
115            item.lifetime = Duration::from_secs(0);
116        }
117    }
118}
119
120impl<K: Eq + Hash, V: Clone> Default for TtlCache<K, V> {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126#[cfg(test)]
127mod tests {
128
129    use super::TtlCache;
130
131    #[test]
132    fn test_ttl_cache_insertion() {
133        let mut cache = TtlCache::new();
134        assert!(!cache.contains("A"));
135
136        cache.insert("A", 1);
137        assert!(cache.contains("A"));
138
139        let value = cache.get("A").expect("The value should be in the cache");
140        assert_eq!(value, 1);
141
142        cache.expire("A");
143
144        assert!(!cache.contains("A"));
145        assert!(cache.get("A").is_none(), "The item should have been removed from the cache");
146    }
147}