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