/build/source/nativelink-util/src/evicting_map.rs
Line | Count | Source |
1 | | // Copyright 2024 The NativeLink Authors. All rights reserved. |
2 | | // |
3 | | // Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (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 | | // See LICENSE file for details |
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 | | use core::borrow::Borrow; |
16 | | use core::cmp::Eq; |
17 | | use core::fmt::{Debug, Display}; |
18 | | use core::future::Future; |
19 | | use core::hash::Hash; |
20 | | use core::marker::PhantomData; |
21 | | use core::ops::RangeBounds; |
22 | | use core::pin::Pin; |
23 | | use std::collections::BTreeSet; |
24 | | use std::sync::Arc; |
25 | | |
26 | | use futures::StreamExt; |
27 | | use futures::stream::FuturesUnordered; |
28 | | use lru::LruCache; |
29 | | use nativelink_config::stores::EvictionPolicy; |
30 | | use nativelink_metric::MetricsComponent; |
31 | | use parking_lot::Mutex; |
32 | | use serde::{Deserialize, Serialize}; |
33 | | use tracing::{debug, info}; |
34 | | |
35 | | use crate::instant_wrapper::InstantWrapper; |
36 | | use crate::metrics_utils::{Counter, CounterWithTime}; |
37 | | |
38 | | #[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)] |
39 | | pub struct SerializedLRU<K> { |
40 | | pub data: Vec<(K, i32)>, |
41 | | pub anchor_time: u64, |
42 | | } |
43 | | |
44 | | #[derive(Debug)] |
45 | | struct EvictionItem<T: LenEntry + Debug> { |
46 | | seconds_since_anchor: i32, |
47 | | data: T, |
48 | | } |
49 | | |
50 | | pub trait LenEntry: 'static { |
51 | | /// Length of referenced data. |
52 | | fn len(&self) -> u64; |
53 | | |
54 | | /// Returns `true` if `self` has zero length. |
55 | | fn is_empty(&self) -> bool; |
56 | | |
57 | | /// This will be called when object is removed from map. |
58 | | /// Note: There may still be a reference to it held somewhere else, which |
59 | | /// is why it can't be mutable. This is a good place to mark the item |
60 | | /// to be deleted and then in the Drop call actually do the deleting. |
61 | | /// This will ensure nowhere else in the program still holds a reference |
62 | | /// to this object. |
63 | | /// You should not rely only on the Drop trait. Doing so might result in the |
64 | | /// program safely shutting down and calling the Drop method on each object, |
65 | | /// which if you are deleting items you may not want to do. |
66 | | /// It is undefined behavior to have `unref()` called more than once. |
67 | | /// During the execution of `unref()` no items can be added or removed to/from |
68 | | /// the `EvictionMap` globally (including inside `unref()`). |
69 | | #[inline] |
70 | 381 | fn unref(&self) -> impl Future<Output = ()> + Send { |
71 | 381 | core::future::ready(()) |
72 | 381 | } |
73 | | } |
74 | | |
75 | | impl<T: LenEntry + Send + Sync> LenEntry for Arc<T> { |
76 | | #[inline] |
77 | 1.88k | fn len(&self) -> u64 { |
78 | 1.88k | T::len(self.as_ref()) |
79 | 1.88k | } |
80 | | |
81 | | #[inline] |
82 | 0 | fn is_empty(&self) -> bool { |
83 | 0 | T::is_empty(self.as_ref()) |
84 | 0 | } |
85 | | |
86 | | #[inline] |
87 | 17 | async fn unref(&self) { |
88 | 17 | self.as_ref().unref().await; |
89 | 17 | } |
90 | | } |
91 | | |
92 | | // Callback to be called when the EvictingMap removes an item |
93 | | // either via eviction or direct deletion. This will be called with |
94 | | // whatever key type the EvictingMap uses. |
95 | | pub trait RemoveItemCallback<Q>: Debug + Send + Sync { |
96 | | fn callback(&self, store_key: &Q) -> Pin<Box<dyn Future<Output = ()> + Send>>; |
97 | | } |
98 | | |
99 | | #[derive(Debug, MetricsComponent)] |
100 | | struct State< |
101 | | K: Ord + Hash + Eq + Clone + Debug + Send + Borrow<Q>, |
102 | | Q: Ord + Hash + Eq + Debug, |
103 | | T: LenEntry + Debug + Send, |
104 | | C: RemoveItemCallback<Q>, |
105 | | > { |
106 | | lru: LruCache<K, EvictionItem<T>>, |
107 | | btree: Option<BTreeSet<K>>, |
108 | | #[metric(help = "Total size of all items in the store")] |
109 | | sum_store_size: u64, |
110 | | |
111 | | #[metric(help = "Number of bytes evicted from the store")] |
112 | | evicted_bytes: Counter, |
113 | | #[metric(help = "Number of items evicted from the store")] |
114 | | evicted_items: CounterWithTime, |
115 | | #[metric(help = "Number of bytes replaced in the store")] |
116 | | replaced_bytes: Counter, |
117 | | #[metric(help = "Number of items replaced in the store")] |
118 | | replaced_items: CounterWithTime, |
119 | | #[metric(help = "Number of bytes inserted into the store since it was created")] |
120 | | lifetime_inserted_bytes: Counter, |
121 | | |
122 | | _key_type: PhantomData<Q>, |
123 | | remove_callbacks: Vec<C>, |
124 | | } |
125 | | |
126 | | type RemoveFuture = Pin<Box<dyn Future<Output = ()> + Send>>; |
127 | | |
128 | | impl< |
129 | | K: Ord + Hash + Eq + Clone + Debug + Send + Sync + Borrow<Q>, |
130 | | Q: Ord + Hash + Eq + Debug + Sync, |
131 | | T: LenEntry + Debug + Sync + Send, |
132 | | C: RemoveItemCallback<Q>, |
133 | | > State<K, Q, T, C> |
134 | | { |
135 | | /// Removes an item from the cache and returns the data for deferred cleanup. |
136 | | /// The caller is responsible for calling `unref()` on the returned data outside of the lock. |
137 | | #[must_use] |
138 | 395 | fn remove( |
139 | 395 | &mut self, |
140 | 395 | key: &Q, |
141 | 395 | eviction_item: &EvictionItem<T>, |
142 | 395 | replaced: bool, |
143 | 395 | ) -> (T, Vec<RemoveFuture>) |
144 | 395 | where |
145 | 395 | T: Clone, |
146 | | { |
147 | 395 | if let Some(btree0 ) = &mut self.btree { |
148 | 0 | btree.remove(key); |
149 | 395 | } |
150 | 395 | self.sum_store_size -= eviction_item.data.len(); |
151 | 395 | if replaced { |
152 | 357 | self.replaced_items.inc(); |
153 | 357 | self.replaced_bytes.add(eviction_item.data.len()); |
154 | 357 | } else { |
155 | 38 | self.evicted_items.inc(); |
156 | 38 | self.evicted_bytes.add(eviction_item.data.len()); |
157 | 38 | } |
158 | | |
159 | 395 | let callbacks = self |
160 | 395 | .remove_callbacks |
161 | 395 | .iter() |
162 | 395 | .map(|callback| callback13 .callback13 (key13 )) |
163 | 395 | .collect(); |
164 | | |
165 | | // Return the data for deferred unref outside of lock |
166 | 395 | (eviction_item.data.clone(), callbacks) |
167 | 395 | } |
168 | | |
169 | | /// Inserts a new item into the cache. If the key already exists, the old item is returned |
170 | | /// for deferred cleanup. |
171 | | #[must_use] |
172 | 9.13k | fn put(&mut self, key: &K, eviction_item: EvictionItem<T>) -> Option<(T, Vec<RemoveFuture>)> |
173 | 9.13k | where |
174 | 9.13k | K: Clone, |
175 | 9.13k | T: Clone, |
176 | | { |
177 | | // If we are maintaining a btree index, we need to update it. |
178 | 9.13k | if let Some(btree0 ) = &mut self.btree { |
179 | 0 | btree.insert(key.clone()); |
180 | 9.13k | } |
181 | 9.13k | self.lru |
182 | 9.13k | .put(key.clone(), eviction_item) |
183 | 9.13k | .map(|old_item| self357 .remove357 (key.borrow()357 , &old_item357 , true)) |
184 | 9.13k | } |
185 | | |
186 | 103 | fn add_remove_callback(&mut self, callback: C) { |
187 | 103 | self.remove_callbacks.push(callback); |
188 | 103 | } |
189 | | } |
190 | | |
191 | | #[derive(Debug, Clone, Copy)] |
192 | | pub struct NoopRemove; |
193 | | |
194 | | impl<Q> RemoveItemCallback<Q> for NoopRemove { |
195 | 0 | fn callback(&self, _store_key: &Q) -> Pin<Box<dyn Future<Output = ()> + Send>> { |
196 | 0 | Box::pin(async {}) |
197 | 0 | } |
198 | | } |
199 | | |
200 | | #[derive(Debug, MetricsComponent)] |
201 | | pub struct EvictingMap< |
202 | | K: Ord + Hash + Eq + Clone + Debug + Send + Borrow<Q>, |
203 | | Q: Ord + Hash + Eq + Debug, |
204 | | T: LenEntry + Debug + Send, |
205 | | I: InstantWrapper, |
206 | | C: RemoveItemCallback<Q> = NoopRemove, |
207 | | > { |
208 | | #[metric] |
209 | | state: Mutex<State<K, Q, T, C>>, |
210 | | anchor_time: I, |
211 | | #[metric(help = "Maximum size of the store in bytes")] |
212 | | max_bytes: u64, |
213 | | #[metric(help = "Number of bytes to evict when the store is full")] |
214 | | evict_bytes: u64, |
215 | | #[metric(help = "Maximum number of seconds to keep an item in the store")] |
216 | | max_seconds: i32, |
217 | | #[metric(help = "Maximum number of items to keep in the store")] |
218 | | max_count: u64, |
219 | | } |
220 | | |
221 | | // debugging helper used mostly to get a snapshot of what eviction threshold might be causing issues |
222 | | #[derive(Debug, Copy, Clone)] |
223 | | pub struct EvictionSnapshot { |
224 | | max_bytes: u64, |
225 | | current_bytes: u64, |
226 | | max_items: u64, |
227 | | current_items: usize, |
228 | | max_seconds: i32, |
229 | | } |
230 | | |
231 | | impl Display for EvictionSnapshot { |
232 | 4 | fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { |
233 | 4 | if self.max_bytes != 0 { |
234 | 2 | write!( |
235 | 2 | f, |
236 | | "Bytes: {} of {} ({:.3}%); ", |
237 | | self.current_bytes, |
238 | | self.max_bytes, |
239 | 2 | (self.current_bytes as f64 * 100.0) / (self.max_bytes as f64) |
240 | 0 | )?; |
241 | | } else { |
242 | 2 | write!(f, "Bytes: {} of unlimited; ", self.current_bytes)?0 ; |
243 | | } |
244 | 4 | if self.max_items != 0 { |
245 | 2 | write!( |
246 | 2 | f, |
247 | | "Items: {} of {} ({:.3}%); ", |
248 | | self.current_items, |
249 | | self.max_items, |
250 | 2 | (self.current_items as f64 * 100.0) / (self.max_items as f64) |
251 | 0 | )?; |
252 | | } else { |
253 | 2 | write!(f, "Items: {} of unlimited; ", self.current_items)?0 ; |
254 | | } |
255 | 4 | if self.max_seconds > 0 { |
256 | 1 | write!(f, "Timeout: {}s", self.max_seconds)?0 ; |
257 | | } else { |
258 | 3 | write!(f, "Timeout: unlimited")?0 ; |
259 | | } |
260 | 4 | Ok(()) |
261 | 4 | } |
262 | | } |
263 | | |
264 | | impl<K, Q, T, I, C> EvictingMap<K, Q, T, I, C> |
265 | | where |
266 | | K: Ord + Hash + Eq + Clone + Debug + Send + Sync + Borrow<Q>, |
267 | | Q: Ord + Hash + Eq + Debug + Sync, |
268 | | T: LenEntry + Debug + Clone + Send + Sync, |
269 | | I: InstantWrapper, |
270 | | C: RemoveItemCallback<Q>, |
271 | | { |
272 | 606 | pub fn new(config: &EvictionPolicy, anchor_time: I) -> Self { |
273 | 606 | Self { |
274 | 606 | // We use unbounded because if we use the bounded version we can't call the delete |
275 | 606 | // function on the LenEntry properly. |
276 | 606 | state: Mutex::new(State { |
277 | 606 | lru: LruCache::unbounded(), |
278 | 606 | btree: None, |
279 | 606 | sum_store_size: 0, |
280 | 606 | evicted_bytes: Counter::default(), |
281 | 606 | evicted_items: CounterWithTime::default(), |
282 | 606 | replaced_bytes: Counter::default(), |
283 | 606 | replaced_items: CounterWithTime::default(), |
284 | 606 | lifetime_inserted_bytes: Counter::default(), |
285 | 606 | _key_type: PhantomData, |
286 | 606 | remove_callbacks: Vec::new(), |
287 | 606 | }), |
288 | 606 | anchor_time, |
289 | 606 | max_bytes: config.max_bytes as u64, |
290 | 606 | evict_bytes: config.evict_bytes as u64, |
291 | 606 | max_seconds: config.max_seconds as i32, |
292 | 606 | max_count: config.max_count, |
293 | 606 | } |
294 | 606 | } |
295 | | |
296 | 0 | pub async fn enable_filtering(&self) { |
297 | 0 | let mut state = self.state.lock(); |
298 | 0 | if state.btree.is_none() { |
299 | 0 | Self::rebuild_btree_index(&mut state); |
300 | 0 | } |
301 | 0 | } |
302 | | |
303 | 2 | fn rebuild_btree_index(state: &mut State<K, Q, T, C>) { |
304 | 2 | state.btree = Some(state.lru.iter().map(|(k, _)| k).cloned().collect()); |
305 | 2 | } |
306 | | |
307 | | /// Run the `handler` function on each key-value pair that matches the `prefix_range` |
308 | | /// and return the number of items that were processed. |
309 | | /// The `handler` function should return `true` to continue processing the next item |
310 | | /// or `false` to stop processing. |
311 | 11 | pub fn range<F>(&self, prefix_range: impl RangeBounds<Q> + Send, mut handler: F) -> u64 |
312 | 11 | where |
313 | 11 | F: FnMut(&K, &T) -> bool + Send, |
314 | 11 | K: Ord, |
315 | | { |
316 | 11 | let mut state = self.state.lock(); |
317 | 11 | let btree = if let Some(ref btree9 ) = state.btree { |
318 | 9 | btree |
319 | | } else { |
320 | 2 | Self::rebuild_btree_index(&mut state); |
321 | 2 | state.btree.as_ref().unwrap() |
322 | | }; |
323 | 11 | let mut continue_count = 0; |
324 | 22 | for key in btree11 .range11 (prefix_range11 ) { |
325 | 22 | let value = &state.lru.peek(key.borrow()).unwrap().data; |
326 | 22 | let should_continue = handler(key, value); |
327 | 22 | if !should_continue { |
328 | 0 | break; |
329 | 22 | } |
330 | 22 | continue_count += 1; |
331 | | } |
332 | 11 | continue_count |
333 | 11 | } |
334 | | |
335 | | /// Returns the number of key-value pairs that are currently in the the cache. |
336 | | /// Function is not for production code paths. |
337 | 30 | pub fn len_for_test(&self) -> usize { |
338 | 30 | self.state.lock().lru.len() |
339 | 30 | } |
340 | | |
341 | 19.8k | fn should_evict( |
342 | 19.8k | &self, |
343 | 19.8k | lru_len: usize, |
344 | 19.8k | peek_entry: &EvictionItem<T>, |
345 | 19.8k | sum_store_size: u64, |
346 | 19.8k | max_bytes: u64, |
347 | 19.8k | ) -> bool { |
348 | 19.8k | let is_over_size = max_bytes != 0 && sum_store_size >= max_bytes10.7k ; |
349 | | |
350 | 19.8k | let elapsed_seconds = |
351 | 19.8k | i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX); |
352 | 19.8k | let evict_older_than_seconds = elapsed_seconds.saturating_sub(self.max_seconds); |
353 | 19.8k | let old_item_exists = |
354 | 19.8k | self.max_seconds != 0 && peek_entry.seconds_since_anchor < evict_older_than_seconds131 ; |
355 | | |
356 | 19.8k | let is_over_count = |
357 | 19.8k | self.max_count != 0 && u64::try_from36 (lru_len36 ).unwrap_or(u64::MAX) > self.max_count; |
358 | | |
359 | 19.8k | is_over_size || old_item_exists19.8k || is_over_count19.8k |
360 | 19.8k | } |
361 | | |
362 | | // Gets a debugging snapshot of the state of the map. It's inevitably out of date by the time it gets sent to the user |
363 | | // but does provide a momentary glimpse into possible issues e.g. if current_bytes is close to max_bytes |
364 | 4 | pub fn get_snapshot(&self) -> EvictionSnapshot { |
365 | 4 | let state = self.state.lock(); |
366 | 4 | EvictionSnapshot { |
367 | 4 | max_bytes: self.max_bytes, |
368 | 4 | current_bytes: state.sum_store_size, |
369 | 4 | max_items: self.max_count, |
370 | 4 | current_items: state.lru.len(), |
371 | 4 | max_seconds: self.max_seconds, |
372 | 4 | } |
373 | 4 | } |
374 | | |
375 | | #[must_use] |
376 | 9.15k | fn evict_items(&self, state: &mut State<K, Q, T, C>) -> (Vec<T>, Vec<RemoveFuture>) { |
377 | 9.15k | let Some((_, mut peek_entry)) = state.lru.peek_lru() else { |
378 | 0 | return (Vec::new(), Vec::new()); |
379 | | }; |
380 | | |
381 | 9.15k | let max_bytes = if self.max_bytes != 0 |
382 | 21 | && self.evict_bytes != 0 |
383 | 4 | && self.should_evict( |
384 | 4 | state.lru.len(), |
385 | 4 | peek_entry, |
386 | 4 | state.sum_store_size, |
387 | 4 | self.max_bytes, |
388 | | ) { |
389 | 1 | self.max_bytes.saturating_sub(self.evict_bytes) |
390 | | } else { |
391 | 9.15k | self.max_bytes |
392 | | }; |
393 | | |
394 | 9.15k | let mut items_to_unref = Vec::new(); |
395 | 9.15k | let mut removal_futures = Vec::new(); |
396 | | |
397 | 9.16k | while self.should_evict(state.lru.len(), peek_entry, state.sum_store_size, max_bytes) { |
398 | 16 | let (key, eviction_item) = state |
399 | 16 | .lru |
400 | 16 | .pop_lru() |
401 | 16 | .expect("Tried to peek() then pop() but failed"); |
402 | 16 | debug!(?key, "Evicting"); |
403 | 16 | let (data, futures) = state.remove(key.borrow(), &eviction_item, false); |
404 | 16 | items_to_unref.push(data); |
405 | 16 | removal_futures.extend(futures); |
406 | | |
407 | 16 | peek_entry = if let Some((_, entry13 )) = state.lru.peek_lru() { |
408 | 13 | entry |
409 | | } else { |
410 | 3 | break; |
411 | | }; |
412 | | } |
413 | | |
414 | 9.15k | (items_to_unref, removal_futures) |
415 | 9.15k | } |
416 | | |
417 | | /// Return the size of a `key`, if not found `None` is returned. |
418 | 25 | pub async fn size_for_key(&self, key: &Q) -> Option<u64> { |
419 | 25 | let mut results = [None]; |
420 | 25 | self.sizes_for_keys([key], &mut results[..], false).await; |
421 | 25 | results[0] |
422 | 25 | } |
423 | | |
424 | | /// Return the sizes of a collection of `keys`. Expects `results` collection |
425 | | /// to be provided for storing the resulting key sizes. Each index value in |
426 | | /// `keys` maps directly to the size value for the key in `results`. |
427 | | /// If no key is found in the internal map, `None` is filled in its place. |
428 | | /// If `peek` is set to `true`, the items are not promoted to the front of the |
429 | | /// LRU cache. Note: peek may still evict, but won't promote. |
430 | 3.51k | pub async fn sizes_for_keys<It, R>(&self, keys: It, results: &mut [Option<u64>], peek: bool) |
431 | 3.51k | where |
432 | 3.51k | It: IntoIterator<Item = R> + Send, |
433 | 3.51k | // Note: It's not enough to have the inserts themselves be Send. The |
434 | 3.51k | // returned iterator should be Send as well. |
435 | 3.51k | <It as IntoIterator>::IntoIter: Send, |
436 | 3.51k | // This may look strange, but what we are doing is saying: |
437 | 3.51k | // * `K` must be able to borrow `Q` |
438 | 3.51k | // * `R` (the input stream item type) must also be able to borrow `Q` |
439 | 3.51k | // Note: That K and R do not need to be the same type, they just both need |
440 | 3.51k | // to be able to borrow a `Q`. |
441 | 3.51k | R: Borrow<Q> + Send, |
442 | 3.51k | { |
443 | 3.51k | let (removal_futures, data_to_unref) = { |
444 | 3.51k | let mut state = self.state.lock(); |
445 | | |
446 | 3.51k | let lru_len = state.lru.len(); |
447 | 3.51k | let mut data_to_unref = Vec::new(); |
448 | 3.51k | let mut removal_futures = Vec::new(); |
449 | 3.57k | for (key, result) in keys3.51k .into_iter3.51k ().zip3.51k (results3.51k .iter_mut3.51k ()) { |
450 | 3.57k | let maybe_entry = if peek { |
451 | 13 | state.lru.peek_mut(key.borrow()) |
452 | | } else { |
453 | 3.55k | state.lru.get_mut(key.borrow()) |
454 | | }; |
455 | 3.57k | match maybe_entry { |
456 | 1.87k | Some(entry) => { |
457 | | // Note: We need to check eviction because the item might be expired |
458 | | // based on the current time. In such case, we remove the item while |
459 | | // we are here. |
460 | 1.87k | if self.should_evict(lru_len, entry, 0, u64::MAX) { |
461 | 1 | *result = None; |
462 | 1 | if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) { |
463 | 1 | info!(?key, "Item expired, evicting"); |
464 | 1 | let (data, futures) = |
465 | 1 | state.remove(key.borrow(), &eviction_item, false); |
466 | | // Store data for later unref - we can't drop state here as we're still iterating |
467 | 1 | data_to_unref.push(data); |
468 | 1 | removal_futures.extend(futures); |
469 | 0 | } |
470 | | } else { |
471 | 1.87k | if !peek { |
472 | 1.87k | entry.seconds_since_anchor = |
473 | 1.87k | i32::try_from(self.anchor_time.elapsed().as_secs()) |
474 | 1.87k | .unwrap_or(i32::MAX); |
475 | 1.87k | }3 |
476 | 1.87k | *result = Some(entry.data.len()); |
477 | | } |
478 | | } |
479 | 1.69k | None => *result = None, |
480 | | } |
481 | | } |
482 | 3.51k | (removal_futures, data_to_unref) |
483 | | }; |
484 | | |
485 | | // Perform the async callbacks outside of the lock |
486 | 3.51k | let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
487 | 3.51k | while callbacks.next().await.is_some() {}0 |
488 | 3.51k | let mut callbacks: FuturesUnordered<_> = |
489 | 3.51k | data_to_unref.iter().map(LenEntry::unref).collect(); |
490 | 3.51k | while callbacks.next().await.is_some() {}1 |
491 | 3.51k | } |
492 | | |
493 | | /// Fires the registered remove callbacks for `key` without touching the map. |
494 | | /// A store uses this to invalidate downstream listeners (e.g. an |
495 | | /// `ExistenceCacheStore`) when a write is rejected outright and never |
496 | | /// inserted — mirroring the callbacks an insert-then-immediate-evict would |
497 | | /// otherwise have fired. It only *notifies*; it does not remove anything from |
498 | | /// the map (there is nothing to remove). A no-op when no callbacks are |
499 | | /// registered. |
500 | 3 | pub async fn fire_remove_callbacks(&self, key: &Q) { |
501 | 3 | let mut callbacks: FuturesUnordered<_> = { |
502 | 3 | let state = self.state.lock(); |
503 | 3 | state |
504 | 3 | .remove_callbacks |
505 | 3 | .iter() |
506 | 3 | .map(|callback| callback2 .callback2 (key2 )) |
507 | 3 | .collect() |
508 | | }; |
509 | 5 | while callbacks.next().await.is_some() {}2 |
510 | 3 | } |
511 | | |
512 | | /// Returns the value for `key` if present and not expired, refreshing |
513 | | /// its LRU/atime position. If the entry is present but TTL- or |
514 | | /// count-expired, it is reaped and `None` is returned. |
515 | | /// |
516 | | /// A read never cascades into other entries — only the queried key is |
517 | | /// ever touched. Global eviction (size/count overflow trim) runs on |
518 | | /// inserts; it is not driven by reads, since `sum_store_size` cannot |
519 | | /// grow without an insert. |
520 | 10.3k | pub async fn get(&self, key: &Q) -> Option<T> { |
521 | | // Lazily reap *only* the requested entry if it is itself expired; |
522 | | // leave the rest for inserts (which already run the global eviction |
523 | | // loop). |
524 | 8.81k | let (data, expired_data, removal_futures) = { |
525 | 10.3k | let mut state = self.state.lock(); |
526 | 10.3k | let lru_len = state.lru.len(); |
527 | 10.3k | let entry8.81k = state.lru.get_mut(key.borrow())?1.49k ; |
528 | | // Pass `sum_store_size=0` and `max_bytes=u64::MAX` so we only |
529 | | // consult TTL / count predicates — never the global byte budget. |
530 | | // Mirrors the per-key reap path in `sizes_for_keys`. |
531 | 8.81k | if self.should_evict(lru_len, entry, 0, u64::MAX) { |
532 | 3 | let (popped_key, eviction_item) = state |
533 | 3 | .lru |
534 | 3 | .pop_entry(key.borrow()) |
535 | 3 | .expect("entry was just observed via get_mut"); |
536 | 3 | info!(?popped_key, "Item expired, evicting"); |
537 | 3 | let (data, futures) = state.remove(popped_key.borrow(), &eviction_item, false); |
538 | 3 | (None, Some(data), futures) |
539 | | } else { |
540 | 8.80k | entry.seconds_since_anchor = |
541 | 8.80k | i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX); |
542 | 8.80k | (Some(entry.data.clone()), None, Vec::new()) |
543 | | } |
544 | | }; |
545 | | |
546 | | // Drain remove_callbacks and unref the reaped entry outside the lock. |
547 | 8.81k | let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
548 | 8.81k | while callbacks.next().await.is_some() {}0 |
549 | 8.81k | if let Some(d3 ) = expired_data { |
550 | 3 | d.unref().await; |
551 | 8.80k | } |
552 | | |
553 | 8.81k | data |
554 | 10.3k | } |
555 | | |
556 | | /// Returns the replaced item if any. |
557 | 9.12k | pub async fn insert(&self, key: K, data: T) -> Option<T> |
558 | 9.12k | where |
559 | 9.12k | K: 'static, |
560 | 9.12k | { |
561 | 9.12k | self.insert_with_time( |
562 | 9.12k | key, |
563 | 9.12k | data, |
564 | 9.12k | i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX), |
565 | 9.12k | ) |
566 | 9.12k | .await |
567 | 9.12k | } |
568 | | |
569 | | /// Returns the replaced item if any. |
570 | 9.12k | pub async fn insert_with_time(&self, key: K, data: T, seconds_since_anchor: i32) -> Option<T> { |
571 | 9.12k | let (items_to_unref, removal_futures) = { |
572 | 9.12k | let mut state = self.state.lock(); |
573 | 9.12k | self.inner_insert_many(&mut state, [(key, data)], seconds_since_anchor) |
574 | 9.12k | }; |
575 | | |
576 | 9.12k | let mut futures: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
577 | 9.13k | while futures.next().await.is_some() {}8 |
578 | | |
579 | | // Unref items outside of lock |
580 | 9.12k | let futures: FuturesUnordered<_> = items_to_unref |
581 | 9.12k | .into_iter() |
582 | 9.12k | .map(|item| async move {372 |
583 | 372 | item.unref().await; |
584 | 372 | item |
585 | 744 | }) |
586 | 9.12k | .collect(); |
587 | 9.12k | futures.collect::<Vec<_>>().await.into_iter().next() |
588 | 9.12k | } |
589 | | |
590 | | /// Same as `insert()`, but optimized for multiple inserts. |
591 | | /// Returns the replaced items if any. |
592 | 9 | pub async fn insert_many<It>(&self, inserts: It) -> Vec<T> |
593 | 9 | where |
594 | 9 | It: IntoIterator<Item = (K, T)> + Send, |
595 | 9 | // Note: It's not enough to have the inserts themselves be Send. The |
596 | 9 | // returned iterator should be Send as well. |
597 | 9 | <It as IntoIterator>::IntoIter: Send, |
598 | 9 | K: 'static, |
599 | 9 | { |
600 | 9 | let mut inserts = inserts.into_iter().peekable(); |
601 | | // Shortcut for cases where there are no inserts, so we don't need to lock. |
602 | 9 | if inserts.peek().is_none() { |
603 | 5 | return Vec::new(); |
604 | 4 | } |
605 | | |
606 | 4 | let (items_to_unref, removal_futures) = { |
607 | 4 | let mut state = self.state.lock(); |
608 | 4 | self.inner_insert_many( |
609 | 4 | &mut state, |
610 | 4 | inserts, |
611 | 4 | i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX), |
612 | 4 | ) |
613 | 4 | }; |
614 | | |
615 | 4 | let mut futures: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
616 | 4 | while futures.next().await.is_some() {}0 |
617 | | |
618 | | // Unref items outside of lock |
619 | 4 | items_to_unref |
620 | 4 | .into_iter() |
621 | 4 | .map(|item| async move {0 |
622 | 0 | item.unref().await; |
623 | 0 | item |
624 | 0 | }) |
625 | 4 | .collect::<FuturesUnordered<_>>() |
626 | 4 | .collect::<Vec<_>>() |
627 | 4 | .await |
628 | 9 | } |
629 | | |
630 | 9.13k | fn inner_insert_many<It>( |
631 | 9.13k | &self, |
632 | 9.13k | state: &mut State<K, Q, T, C>, |
633 | 9.13k | inserts: It, |
634 | 9.13k | seconds_since_anchor: i32, |
635 | 9.13k | ) -> (Vec<T>, Vec<RemoveFuture>) |
636 | 9.13k | where |
637 | 9.13k | It: IntoIterator<Item = (K, T)> + Send, |
638 | 9.13k | // Note: It's not enough to have the inserts themselves be Send. The |
639 | 9.13k | // returned iterator should be Send as well. |
640 | 9.13k | <It as IntoIterator>::IntoIter: Send, |
641 | | { |
642 | 9.13k | let mut replaced_items = Vec::new(); |
643 | 9.13k | let mut removal_futures = Vec::new(); |
644 | 9.13k | for (key, data) in inserts { |
645 | 9.13k | let new_item_size = data.len(); |
646 | 9.13k | let eviction_item = EvictionItem { |
647 | 9.13k | seconds_since_anchor, |
648 | 9.13k | data, |
649 | 9.13k | }; |
650 | | |
651 | 9.13k | if let Some((old_item357 , futures357 )) = state.put(&key, eviction_item) { |
652 | 357 | removal_futures.extend(futures); |
653 | 357 | debug!(?key, "Evicting old item"); |
654 | 357 | replaced_items.push(old_item); |
655 | 8.77k | } |
656 | 9.13k | state.sum_store_size += new_item_size; |
657 | 9.13k | state.lifetime_inserted_bytes.add(new_item_size); |
658 | | } |
659 | | |
660 | | // Perform eviction after all insertions |
661 | 9.13k | let (items_to_unref, futures) = self.evict_items(state); |
662 | 9.13k | removal_futures.extend(futures); |
663 | | |
664 | | // Note: We cannot drop the state lock here since we're borrowing it, |
665 | | // but the caller will handle unreffing these items after releasing the lock |
666 | 9.13k | replaced_items.extend(items_to_unref); |
667 | | |
668 | 9.13k | (replaced_items, removal_futures) |
669 | 9.13k | } |
670 | | |
671 | 18 | pub async fn remove(&self, key: &Q) -> bool { |
672 | 18 | let (items_to_unref, removed_item, removal_futures) = { |
673 | 18 | let mut state = self.state.lock(); |
674 | | |
675 | | // First perform eviction |
676 | 18 | let (evicted_items, mut removal_futures) = self.evict_items(&mut *state); |
677 | | |
678 | | // Then try to remove the requested item |
679 | 18 | let removed = if let Some(entry17 ) = state.lru.pop(key.borrow()) { |
680 | 17 | let (removed_item, more_removal_futures) = state.remove(key, &entry, false); |
681 | 17 | removal_futures.extend(more_removal_futures); |
682 | 17 | Some(removed_item) |
683 | | } else { |
684 | 1 | None |
685 | | }; |
686 | | |
687 | 18 | (evicted_items, removed, removal_futures) |
688 | | }; |
689 | | |
690 | 18 | let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
691 | 22 | while callbacks.next().await.is_some() {}4 |
692 | | |
693 | | // Unref evicted items outside of lock |
694 | 18 | let mut callbacks: FuturesUnordered<_> = |
695 | 18 | items_to_unref.iter().map(LenEntry::unref).collect(); |
696 | 19 | while callbacks.next().await.is_some() {}1 |
697 | | |
698 | | // Unref removed item if any |
699 | 18 | if let Some(item17 ) = removed_item { |
700 | 17 | debug!(?key, "Evicting (direct remove)"); |
701 | 17 | item.unref().await; |
702 | 17 | return true; |
703 | 1 | } |
704 | | |
705 | 1 | false |
706 | 18 | } |
707 | | |
708 | | /// Same as `remove()`, but allows for a conditional to be applied to the |
709 | | /// entry before removal in an atomic fashion. |
710 | 1 | pub async fn remove_if<F>(&self, key: &Q, cond: F) -> bool |
711 | 1 | where |
712 | 1 | F: FnOnce(&T) -> bool + Send, |
713 | 1 | { |
714 | 1 | let (evicted_items, removal_futures, removed_item) = { |
715 | 1 | let mut state = self.state.lock(); |
716 | 1 | if let Some(entry) = state.lru.get(key.borrow()) { |
717 | 1 | if !cond(&entry.data) { |
718 | 0 | return false; |
719 | 1 | } |
720 | | // First perform eviction |
721 | 1 | let (evicted_items, mut removal_futures) = self.evict_items(&mut state); |
722 | | |
723 | | // Then try to remove the requested item |
724 | 1 | let removed_item = if let Some(entry) = state.lru.pop(key.borrow()) { |
725 | 1 | let (item, more_removal_futures) = state.remove(key, &entry, false); |
726 | 1 | removal_futures.extend(more_removal_futures); |
727 | 1 | Some(item) |
728 | | } else { |
729 | 0 | None |
730 | | }; |
731 | | |
732 | 1 | (evicted_items, removal_futures, removed_item) |
733 | | } else { |
734 | 0 | (vec![], vec![].into_iter().collect(), None) |
735 | | } |
736 | | }; |
737 | | |
738 | | // Perform the async callbacks outside of the lock |
739 | 1 | let mut removal_futures: FuturesUnordered<_> = removal_futures.into_iter().collect(); |
740 | 2 | while removal_futures.next().await.is_some() {}1 |
741 | | |
742 | | // Unref evicted items |
743 | 1 | let mut callbacks: FuturesUnordered<_> = |
744 | 1 | evicted_items.iter().map(LenEntry::unref).collect(); |
745 | 1 | while callbacks.next().await.is_some() {}0 |
746 | | |
747 | | // Unref removed item if any |
748 | 1 | if let Some(item) = removed_item { |
749 | 1 | debug!(?key, "Evicting (conditional remove)"); |
750 | 1 | item.unref().await; |
751 | 1 | true |
752 | | } else { |
753 | 0 | false |
754 | | } |
755 | 1 | } |
756 | | |
757 | 103 | pub fn add_remove_callback(&self, callback: C) { |
758 | 103 | self.state.lock().add_remove_callback(callback); |
759 | 103 | } |
760 | | } |