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.
1415//! A TTL cache which can be used to time out repeated operations that might
16//! experience intermittent failures.
1718use std::{borrow::Borrow, collections::HashMap, hash::Hash, time::Duration};
1920use ruma::time::Instant;
2122// One day is the default lifetime.
23const DEFAULT_LIFETIME: Duration = Duration::from_secs(24 * 60 * 60);
2425#[derive(Debug)]
26struct TtlItem<V: Clone> {
27 value: V,
28 insertion_time: Instant,
29 lifetime: Duration,
30}
3132impl<V: Clone> TtlItem<V> {
33fn expired(&self) -> bool {
34self.insertion_time.elapsed() >= self.lifetime
35 }
36}
3738/// 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}
4445impl<K, V> TtlCache<K, V>
46where
47K: Eq + Hash,
48 V: Clone,
49{
50/// Create a new, empty, [`TtlCache`].
51pub fn new() -> Self {
52Self { items: Default::default(), lifetime: DEFAULT_LIFETIME }
53 }
5455/// Does the cache contain an non-expired item with the matching key.
56pub fn contains<Q>(&self, key: &Q) -> bool
57where
58K: Borrow<Q>,
59 Q: Hash + Eq + ?Sized,
60 {
61let cache = &self.items;
62let contains = if let Some(item) = cache.get(key) { !item.expired() } else { false };
6364 contains
65 }
6667/// Add a single item to the cache.
68pub fn insert(&mut self, key: K, value: V) {
69self.extend([(key, value)]);
70 }
7172/// Extend the cache with the given iterator of items.
73pub fn extend(&mut self, iterator: impl IntoIterator<Item = (K, V)>) {
74let cache = &mut self.items;
7576let now = Instant::now();
7778for (key, value) in iterator {
79let item = TtlItem { value, insertion_time: now, lifetime: self.lifetime };
8081 cache.insert(key, item);
82 }
83 }
8485/// Remove the item that matches the given key.
86pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
87where
88K: Borrow<Q>,
89 Q: Hash + Eq + ?Sized,
90 {
91self.items.remove(key.borrow()).map(|item| item.value)
92 }
9394/// Get the item that matches the given key, if the item has expired `None`
95 /// will be returned and the item will be evicted from the cache.
96pub fn get<Q>(&mut self, key: &Q) -> Option<V>
97where
98K: Borrow<Q>,
99 Q: Hash + Eq + ?Sized,
100 {
101// Remove all expired items.
102self.items.retain(|_, value| !value.expired());
103// Now get the wanted item.
104self.items.get(key.borrow()).map(|item| item.value.clone())
105 }
106107/// Force the expiry of the given item, if it is present in the cache.
108 ///
109 /// This doesn't remove the item, it just marks it as expired.
110#[doc(hidden)]
111pub fn expire<Q>(&mut self, key: &Q)
112where
113K: Borrow<Q>,
114 Q: Hash + Eq + ?Sized,
115 {
116if let Some(item) = self.items.get_mut(key) {
117 item.lifetime = Duration::from_secs(0);
118 }
119 }
120}
121122impl<K: Eq + Hash, V: Clone> Default for TtlCache<K, V> {
123fn default() -> Self {
124Self::new()
125 }
126}
127128#[cfg(test)]
129mod tests {
130131use super::TtlCache;
132133#[test]
134fn test_ttl_cache_insertion() {
135let mut cache = TtlCache::new();
136assert!(!cache.contains("A"));
137138 cache.insert("A", 1);
139assert!(cache.contains("A"));
140141let value = cache.get("A").expect("The value should be in the cache");
142assert_eq!(value, 1);
143144 cache.expire("A");
145146assert!(!cache.contains("A"));
147assert!(cache.get("A").is_none(), "The item should have been removed from the cache");
148 }
149}