Coverage Report

Created: 2026-07-21 15:28

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/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
}