Line data Source code
1 : use std::{
2 : borrow::Borrow,
3 : hash::Hash,
4 : time::{Duration, Instant},
5 : };
6 : use tracing::debug;
7 :
8 : // This seems to make more sense than `lru` or `cached`:
9 : //
10 : // * `near/nearcore` ditched `cached` in favor of `lru`
11 : // (https://github.com/near/nearcore/issues?q=is%3Aissue+lru+is%3Aclosed).
12 : //
13 : // * `lru` methods use an obscure `KeyRef` type in their contraints (which is deliberately excluded from docs).
14 : // This severely hinders its usage both in terms of creating wrappers and supported key types.
15 : //
16 : // On the other hand, `hashlink` has good download stats and appears to be maintained.
17 : use hashlink::{linked_hash_map::RawEntryMut, LruCache};
18 :
19 : use super::{common::Cached, *};
20 :
21 : /// An implementation of timed LRU cache with fixed capacity.
22 : /// Key properties:
23 : ///
24 : /// * Whenever a new entry is inserted, the least recently accessed one is evicted.
25 : /// The cache also keeps track of entry's insertion time (`created_at`) and TTL (`expires_at`).
26 : ///
27 : /// * If `update_ttl_on_retrieval` is `true`. When the entry is about to be retrieved, we check its expiration timestamp.
28 : /// If the entry has expired, we remove it from the cache; Otherwise we bump the
29 : /// expiration timestamp (e.g. +5mins) and change its place in LRU list to prolong
30 : /// its existence.
31 : ///
32 : /// * There's an API for immediate invalidation (removal) of a cache entry;
33 : /// It's useful in case we know for sure that the entry is no longer correct.
34 : /// See [`timed_lru::LookupInfo`] & [`timed_lru::Cached`] for more information.
35 : ///
36 : /// * Expired entries are kept in the cache, until they are evicted by the LRU policy,
37 : /// or by a successful lookup (i.e. the entry hasn't expired yet).
38 : /// There is no background job to reap the expired records.
39 : ///
40 : /// * It's possible for an entry that has not yet expired entry to be evicted
41 : /// before expired items. That's a bit wasteful, but probably fine in practice.
42 : pub struct TimedLru<K, V> {
43 : /// Cache's name for tracing.
44 : name: &'static str,
45 :
46 : /// The underlying cache implementation.
47 : cache: parking_lot::Mutex<LruCache<K, Entry<V>>>,
48 :
49 : /// Default time-to-live of a single entry.
50 : ttl: Duration,
51 :
52 : update_ttl_on_retrieval: bool,
53 : }
54 :
55 : impl<K: Hash + Eq, V> Cache for TimedLru<K, V> {
56 : type Key = K;
57 : type Value = V;
58 : type LookupInfo<Key> = LookupInfo<Key>;
59 :
60 0 : fn invalidate(&self, info: &Self::LookupInfo<K>) {
61 0 : self.invalidate_raw(info)
62 0 : }
63 : }
64 :
65 : struct Entry<T> {
66 : created_at: Instant,
67 : expires_at: Instant,
68 : value: T,
69 : }
70 :
71 : impl<K: Hash + Eq, V> TimedLru<K, V> {
72 : /// Construct a new LRU cache with timed entries.
73 1 : pub fn new(
74 1 : name: &'static str,
75 1 : capacity: usize,
76 1 : ttl: Duration,
77 1 : update_ttl_on_retrieval: bool,
78 1 : ) -> Self {
79 1 : Self {
80 1 : name,
81 1 : cache: LruCache::new(capacity).into(),
82 1 : ttl,
83 1 : update_ttl_on_retrieval,
84 1 : }
85 1 : }
86 :
87 : /// Drop an entry from the cache if it's outdated.
88 0 : #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
89 : fn invalidate_raw(&self, info: &LookupInfo<K>) {
90 : let now = Instant::now();
91 :
92 : // Do costly things before taking the lock.
93 : let mut cache = self.cache.lock();
94 : let raw_entry = match cache.raw_entry_mut().from_key(&info.key) {
95 : RawEntryMut::Vacant(_) => return,
96 : RawEntryMut::Occupied(x) => x,
97 : };
98 :
99 : // Remove the entry if it was created prior to lookup timestamp.
100 : let entry = raw_entry.get();
101 : let (created_at, expires_at) = (entry.created_at, entry.expires_at);
102 : let should_remove = created_at <= info.created_at || expires_at <= now;
103 :
104 : if should_remove {
105 : raw_entry.remove();
106 : }
107 :
108 : drop(cache); // drop lock before logging
109 0 : debug!(
110 0 : created_at = format_args!("{created_at:?}"),
111 0 : expires_at = format_args!("{expires_at:?}"),
112 0 : entry_removed = should_remove,
113 0 : "processed a cache entry invalidation event"
114 0 : );
115 : }
116 :
117 : /// Try retrieving an entry by its key, then execute `extract` if it exists.
118 0 : #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
119 : fn get_raw<Q, R>(&self, key: &Q, extract: impl FnOnce(&K, &Entry<V>) -> R) -> Option<R>
120 : where
121 : K: Borrow<Q>,
122 : Q: Hash + Eq + ?Sized,
123 : {
124 : let now = Instant::now();
125 : let deadline = now.checked_add(self.ttl).expect("time overflow");
126 :
127 : // Do costly things before taking the lock.
128 : let mut cache = self.cache.lock();
129 : let mut raw_entry = match cache.raw_entry_mut().from_key(key) {
130 : RawEntryMut::Vacant(_) => return None,
131 : RawEntryMut::Occupied(x) => x,
132 : };
133 :
134 : // Immeditely drop the entry if it has expired.
135 : let entry = raw_entry.get();
136 : if entry.expires_at <= now {
137 : raw_entry.remove();
138 : return None;
139 : }
140 :
141 : let value = extract(raw_entry.key(), entry);
142 : let (created_at, expires_at) = (entry.created_at, entry.expires_at);
143 :
144 : // Update the deadline and the entry's position in the LRU list.
145 : if self.update_ttl_on_retrieval {
146 : raw_entry.get_mut().expires_at = deadline;
147 : }
148 : raw_entry.to_back();
149 :
150 : drop(cache); // drop lock before logging
151 0 : debug!(
152 0 : created_at = format_args!("{created_at:?}"),
153 0 : old_expires_at = format_args!("{expires_at:?}"),
154 0 : new_expires_at = format_args!("{deadline:?}"),
155 0 : "accessed a cache entry"
156 0 : );
157 :
158 : Some(value)
159 : }
160 :
161 : /// Insert an entry to the cache. If an entry with the same key already
162 : /// existed, return the previous value and its creation timestamp.
163 0 : #[tracing::instrument(level = "debug", fields(cache = self.name), skip_all)]
164 : fn insert_raw(&self, key: K, value: V) -> (Instant, Option<V>) {
165 : let created_at = Instant::now();
166 : let expires_at = created_at.checked_add(self.ttl).expect("time overflow");
167 :
168 : let entry = Entry {
169 : created_at,
170 : expires_at,
171 : value,
172 : };
173 :
174 : // Do costly things before taking the lock.
175 : let old = self
176 : .cache
177 : .lock()
178 : .insert(key, entry)
179 0 : .map(|entry| entry.value);
180 :
181 0 : debug!(
182 0 : created_at = format_args!("{created_at:?}"),
183 0 : expires_at = format_args!("{expires_at:?}"),
184 0 : replaced = old.is_some(),
185 0 : "created a cache entry"
186 0 : );
187 :
188 : (created_at, old)
189 : }
190 : }
191 :
192 : impl<K: Hash + Eq + Clone, V: Clone> TimedLru<K, V> {
193 0 : pub fn insert(&self, key: K, value: V) -> (Option<V>, Cached<&Self>) {
194 0 : let (created_at, old) = self.insert_raw(key.clone(), value.clone());
195 0 :
196 0 : let cached = Cached {
197 0 : token: Some((self, LookupInfo { created_at, key })),
198 0 : value,
199 0 : };
200 0 :
201 0 : (old, cached)
202 0 : }
203 : }
204 :
205 : impl<K: Hash + Eq, V: Clone> TimedLru<K, V> {
206 : /// Retrieve a cached entry in convenient wrapper.
207 0 : pub fn get<Q>(&self, key: &Q) -> Option<timed_lru::Cached<&Self>>
208 0 : where
209 0 : K: Borrow<Q> + Clone,
210 0 : Q: Hash + Eq + ?Sized,
211 0 : {
212 0 : self.get_raw(key, |key, entry| {
213 0 : let info = LookupInfo {
214 0 : created_at: entry.created_at,
215 0 : key: key.clone(),
216 0 : };
217 0 :
218 0 : Cached {
219 0 : token: Some((self, info)),
220 0 : value: entry.value.clone(),
221 0 : }
222 0 : })
223 0 : }
224 :
225 : /// Retrieve a cached entry in convenient wrapper, ignoring its TTL.
226 0 : pub fn get_ignoring_ttl<Q>(&self, key: &Q) -> Option<timed_lru::Cached<&Self>>
227 0 : where
228 0 : K: Borrow<Q>,
229 0 : Q: Hash + Eq + ?Sized,
230 0 : {
231 0 : let mut cache = self.cache.lock();
232 0 : cache
233 0 : .get(key)
234 0 : .map(|entry| Cached::new_uncached(entry.value.clone()))
235 0 : }
236 :
237 : /// Remove an entry from the cache.
238 0 : pub fn remove<Q>(&self, key: &Q) -> Option<V>
239 0 : where
240 0 : K: Borrow<Q> + Clone,
241 0 : Q: Hash + Eq + ?Sized,
242 0 : {
243 0 : let mut cache = self.cache.lock();
244 0 : cache.remove(key).map(|entry| entry.value)
245 0 : }
246 : }
247 :
248 : /// Lookup information for key invalidation.
249 : pub struct LookupInfo<K> {
250 : /// Time of creation of a cache [`Entry`].
251 : /// We use this during invalidation lookups to prevent eviction of a newer
252 : /// entry sharing the same key (it might've been inserted by a different
253 : /// task after we got the entry we're trying to invalidate now).
254 : created_at: Instant,
255 :
256 : /// Search by this key.
257 : key: K,
258 : }
|