/build/source/nativelink-util/src/evicting_map.rs

Source
// Copyright 2024 Trace Machina, Inc. All rights reserved.
//
// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    See LICENSE file for details
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use core::borrow::Borrow;
use core::cmp::Eq;
use core::fmt::{Debug, Display};
use core::future::Future;
use core::hash::Hash;
use core::marker::PhantomData;
use core::ops::RangeBounds;
use core::pin::Pin;
use std::collections::{BTreeSet, HashMap};
use std::sync::{Arc, OnceLock};

use futures::StreamExt;
use futures::stream::FuturesUnordered;
use lru::LruCache;
use nativelink_config::stores::EvictionPolicy;
use nativelink_metric::MetricsComponent;
use opentelemetry::KeyValue;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use tracing::{debug, info};

use crate::instant_wrapper::InstantWrapper;
use crate::metrics::{record_cache_entries_delta, saturating_i64};
use crate::metrics_utils::{Counter, CounterWithTime};

#[derive(Serialize, Deserialize, PartialEq, Eq, Debug, Clone)]
pub struct SerializedLRU<K> {
    pub data: Vec<(K, i32)>,
    pub anchor_time: u64,
}

#[derive(Debug)]
struct EvictionItem<T: LenEntry + Debug> {
    seconds_since_anchor: i32,
    data: T,
}

pub trait LenEntry: 'static {
    /// Length of referenced data.
    fn len(&self) -> u64;

    /// Length of the content a read of the entry returns. This, not `len()`,
    /// is the size `sizes_for_keys` reports, since callers size reads by it.
    #[inline]
    fn content_len(&self) -> u64 {
        self.len()
    }

    /// Returns `true` if `self` has zero length.
    fn is_empty(&self) -> bool;

    /// This will be called when object is removed from map.
    /// Note: There may still be a reference to it held somewhere else, which
    /// is why it can't be mutable. This is a good place to mark the item
    /// to be deleted and then in the Drop call actually do the deleting.
    /// This will ensure nowhere else in the program still holds a reference
    /// to this object.
    /// You should not rely only on the Drop trait. Doing so might result in the
    /// program safely shutting down and calling the Drop method on each object,
    /// which if you are deleting items you may not want to do.
    /// It is undefined behavior to have `unref()` called more than once.
    /// Runs outside the map lock. Another task may insert a replacement for
    /// the same key before this call finishes; cleanup must only affect the
    /// removed entry's resources.
    #[inline]
    fn unref(&self) -> impl Future<Output = ()> + Send {
        core::future::ready(())
    }
}

impl<T: LenEntry + Send + Sync> LenEntry for Arc<T> {
    #[inline]
    fn len(&self) -> u64 {
        T::len(self.as_ref())
    }

    #[inline]
    fn content_len(&self) -> u64 {
        T::content_len(self.as_ref())
    }

    #[inline]
    fn is_empty(&self) -> bool {
        T::is_empty(self.as_ref())
    }

    #[inline]
    async fn unref(&self) {
        self.as_ref().unref().await;
    }
}

// Callback to be called when the EvictingMap removes an item
// either via eviction or direct deletion. This will be called with
// whatever key type the EvictingMap uses.
pub trait RemoveItemCallback<Q>: Debug + Send + Sync {
    fn callback(&self, store_key: &Q) -> Pin<Box<dyn Future<Output = ()> + Send>>;
}

#[derive(Debug, MetricsComponent)]
struct State<
    K: Ord + Hash + Eq + Clone + Debug + Send + Borrow<Q>,
    Q: Ord + Hash + Eq + Debug,
    T: LenEntry + Debug + Send,
    C: RemoveItemCallback<Q>,
> {
    lru: LruCache<K, EvictionItem<T>>,
    btree: Option<BTreeSet<K>>,
    /// A mirror of `lru` that contains only currently unleased keys.
    /// Keeping this index in LRU order lets eviction visit candidates without
    /// scanning entries that an active operation has pinned.
    evictable_lru: LruCache<K, ()>,
    /// Reference counts for keys that must not be evicted while an active
    /// operation is using them. A key can be leased before it is inserted,
    /// which closes the admission-to-insert race for cache populations.
    leases: HashMap<K, u32>,
    #[metric(help = "Total size of all items in the store")]
    sum_store_size: u64,
    #[metric(help = "Number of items currently leased and not evictable")]
    leased_items: u64,
    #[metric(help = "Number of bytes currently leased and not evictable")]
    leased_bytes: u64,

    #[metric(help = "Number of bytes evicted from the store")]
    evicted_bytes: Counter,
    #[metric(help = "Number of items evicted from the store")]
    evicted_items: CounterWithTime,
    #[metric(help = "Number of bytes replaced in the store")]
    replaced_bytes: Counter,
    #[metric(help = "Number of items replaced in the store")]
    replaced_items: CounterWithTime,
    #[metric(help = "Number of bytes inserted into the store since it was created")]
    lifetime_inserted_bytes: Counter,

    _key_type: PhantomData<Q>,
    remove_callbacks: Vec<C>,
    /// Size and entry-count changes waiting to be reported as `cache.size` and
    /// `cache.entries`. Recording to `OpenTelemetry` costs far more than the
    /// plain atomic counters beside it, so the mutation paths only do integer
    /// arithmetic here and the totals are flushed once the lock is released.
    /// A flush that is missed only delays the next update, it never loses a
    /// change.
    pending_size_delta: i64,
    pending_entries_delta: i64,
}

type RemoveFuture = Pin<Box<dyn Future<Output = ()> + Send>>;

impl<
    K: Ord + Hash + Eq + Clone + Debug + Send + Sync + Borrow<Q>,
    Q: Ord + Hash + Eq + Debug + Sync,
    T: LenEntry + Debug + Sync + Send,
    C: RemoveItemCallback<Q>,
> State<K, Q, T, C>
{
    fn is_leased(&self, key: &Q) -> bool {
        self.leases.contains_key(key)
    }

    /// Keep candidate recency aligned with a normal read from `lru`.
    /// Returns whether the key is leased so callers do not repeat the lookup.
    fn touch_evictable(&mut self, key: &Q) -> bool {
        if self.is_leased(key) {
            return true;
        }
        self.evictable_lru.get(key);
        false
    }

    /// Add a newly written key to the candidate index unless it is reserved
    /// by a lease that was acquired before insertion.
    fn insert_evictable(&mut self, key: K) {
        if !self.is_leased(key.borrow()) {
            self.evictable_lru.put(key, ());
        }11
    }

    fn lease(&mut self, key: K) {
        if let Some(lease_count2) = self.leases.get_mut(key.borrow()) {
            *lease_count = lease_count.checked_add(1).expect("lease count overflow");
            return;
        }

        self.evictable_lru.pop(key.borrow());
        if let Some(entry28) = self.lru.peek(key.borrow()) {
            self.leased_bytes += entry.data.len();
        }19
        self.leased_items += 1;
        self.leases.insert(key, 1);
    }

    /// Release one reference and return the owned key when the lease ends.
    fn release_lease(&mut self, key: &Q) -> Option<K> {
        let lease_count = self.leases.get_mut(key)?0;
        if *lease_count > 1 {
            *lease_count -= 1;
            return None;
        }

        let leased_key = self
            .leases
            .remove_entry(key)
            .expect("Lease must exist when its final reference is released")
            .0;
        self.leased_items -= 1;
        if let Some(entry35) = self.lru.peek(leased_key.borrow()) {
            self.leased_bytes -= entry.data.len();
        }9
        Some(leased_key)
    }

    /// Re-enter a resident key as the newest eviction candidate. Its age is
    /// intentionally preserved, so a long lease does not extend the TTL.
    fn reinsert_evictable(&mut self, key: K) {
        if self.lru.peek(key.borrow()).is_some() {
            self.evictable_lru.put(key, ());
        }9
    }

    fn peek_evictable(&self) -> Option<&EvictionItem<T>> {
        let (key141k, ()) = self.evictable_lru.peek_lru()?29;
        Some(
            self.lru
                .peek(key.borrow())
                .expect("Evictable LRU key must be resident in the main LRU"),
        )
    }

    fn pop_evictable(&mut self) -> Option<(K, EvictionItem<T>)> {
        let (candidate_key, ()) = self.evictable_lru.pop_lru()?0;
        let (key, entry) = self
            .lru
            .pop_entry(candidate_key.borrow())
            .expect("Evictable LRU key must be resident in the main LRU");
        Some((key, entry))
    }

    /// Removes an item from the cache and returns the data for deferred cleanup.
    /// The caller is responsible for calling `unref()` on the returned data outside of the lock.
    #[must_use]
    fn remove(
        &mut self,
        key: &Q,
        eviction_item: &EvictionItem<T>,
        replaced: bool,
    ) -> (T, Vec<RemoveFuture>)
    where
        T: Clone,
    {
        if let Some(btree1) = &mut self.btree {
            btree.remove(key);
        }
        // Keep every auxiliary resident index in sync with the main LRU.
        self.evictable_lru.pop(key);
        if self.is_leased(key) {
            self.leased_bytes -= eviction_item.data.len();
        }
        self.sum_store_size -= eviction_item.data.len();
        self.pending_size_delta -= saturating_i64(eviction_item.data.len());
        self.pending_entries_delta -= 1;
        if replaced {
            self.replaced_items.inc();
            self.replaced_bytes.add(eviction_item.data.len());
        } else {
            self.evicted_items.inc();
            self.evicted_bytes.add(eviction_item.data.len());
        }

        let callbacks = self
            .remove_callbacks
            .iter()
            .map(|callback| callback51.callback51(key51))
            .collect();

        // Return the data for deferred unref outside of lock
        (eviction_item.data.clone(), callbacks)
    }

    /// Inserts a new item into the cache. If the key already exists, the old item is returned
    /// for deferred cleanup.
    #[must_use]
    fn put(&mut self, key: &K, eviction_item: EvictionItem<T>) -> Option<(T, Vec<RemoveFuture>)>
    where
        K: Clone,
        T: Clone,
    {
        let is_leased = self.is_leased(key.borrow());
        let new_item_size = eviction_item.data.len();
        let replaced_item = self
            .lru
            .put(key.clone(), eviction_item)
            .map(|old_item| self16.2k.remove16.2k(key.borrow()16.2k, &old_item16.2k, true));

        if is_leased {
            self.leased_bytes += new_item_size;
        }
        // `remove()` drops the old key from both indexes. Reinsert it after
        // replacement so an unleased write is MRU in both LRU indexes.
        self.insert_evictable(key.clone());
        if let Some(btree6) = &mut self.btree {
            btree.insert(key.clone());
        }

        replaced_item
    }

    fn add_remove_callback(&mut self, callback: C) {
        self.remove_callbacks.push(callback);
    }
}

#[derive(Debug, Clone, Copy)]
pub struct NoopRemove;

impl<Q> RemoveItemCallback<Q> for NoopRemove {
    fn callback(&self, _store_key: &Q) -> Pin<Box<dyn Future<Output = ()> + Send>> {
        Box::pin(async {})
    }
}

#[derive(Debug, MetricsComponent)]
pub struct EvictingMap<
    K: Ord + Hash + Eq + Clone + Debug + Send + Borrow<Q>,
    Q: Ord + Hash + Eq + Debug,
    T: LenEntry + Debug + Send,
    I: InstantWrapper,
    C: RemoveItemCallback<Q> = NoopRemove,
> {
    #[metric]
    state: Mutex<State<K, Q, T, C>>,
    anchor_time: I,
    #[metric(help = "Maximum size of the store in bytes")]
    max_bytes: u64,
    #[metric(help = "Number of bytes to evict when the store is full")]
    evict_bytes: u64,
    #[metric(help = "Maximum number of seconds to keep an item in the store")]
    max_seconds: i32,
    #[metric(help = "Maximum number of items to keep in the store")]
    max_count: u64,
    /// Attributes to report `cache.size` and `cache.entries` under. Unset until
    /// a `cache_metrics` wrapper enables it, which is what keeps those two
    /// instruments off by default.
    cache_size_attrs: OnceLock<Vec<KeyValue>>,
}

// debugging helper used mostly to get a snapshot of what eviction threshold might be causing issues
#[derive(Debug, Copy, Clone)]
pub struct EvictionSnapshot {
    max_bytes: u64,
    current_bytes: u64,
    max_items: u64,
    current_items: usize,
    max_seconds: i32,
}

impl Display for EvictionSnapshot {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        if self.max_bytes != 0 {
            write!(
                f,
                "Bytes: {} of {} ({:.3}%); ",
                self.current_bytes,
                self.max_bytes,
                (self.current_bytes as f64 * 100.0) / (self.max_bytes as f64)
            )?;
        } else {
            write!(f, "Bytes: {} of unlimited; ", self.current_bytes)?0;
        }
        if self.max_items != 0 {
            write!(
                f,
                "Items: {} of {} ({:.3}%); ",
                self.current_items,
                self.max_items,
                (self.current_items as f64 * 100.0) / (self.max_items as f64)
            )?;
        } else {
            write!(f, "Items: {} of unlimited; ", self.current_items)?0;
        }
        if self.max_seconds > 0 {
            write!(f, "Timeout: {}s", self.max_seconds)?0;
        } else {
            write!(f, "Timeout: unlimited")?0;
        }
        Ok(())
    }
}

impl<K, Q, T, I, C> EvictingMap<K, Q, T, I, C>
where
    K: Ord + Hash + Eq + Clone + Debug + Send + Sync + Borrow<Q>,
    Q: Ord + Hash + Eq + Debug + Sync,
    T: LenEntry + Debug + Clone + Send + Sync,
    I: InstantWrapper,
    C: RemoveItemCallback<Q>,
{
    pub fn new(config: &EvictionPolicy, anchor_time: I) -> Self {
        Self {
            // We use unbounded because if we use the bounded version we can't call the delete
            // function on the LenEntry properly.
            state: Mutex::new(State {
                lru: LruCache::unbounded(),
                btree: None,
                evictable_lru: LruCache::unbounded(),
                leases: HashMap::new(),
                sum_store_size: 0,
                leased_items: 0,
                leased_bytes: 0,
                evicted_bytes: Counter::default(),
                evicted_items: CounterWithTime::default(),
                replaced_bytes: Counter::default(),
                replaced_items: CounterWithTime::default(),
                lifetime_inserted_bytes: Counter::default(),
                _key_type: PhantomData,
                remove_callbacks: Vec::new(),
                pending_size_delta: 0,
                pending_entries_delta: 0,
            }),
            anchor_time,
            max_bytes: config.max_bytes as u64,
            evict_bytes: config.evict_bytes as u64,
            max_seconds: config.max_seconds.try_into().unwrap_or(i32::MAX),
            max_count: config.max_count,
            cache_size_attrs: OnceLock::new(),
        }
    }

    /// Reports this map's size and entry count as `cache.size` and
    /// `cache.entries` under `attrs`, starting from what it already holds.
    /// Only the first call takes effect, so a map is never counted twice.
    pub fn enable_cache_size_metrics(&self, attrs: Vec<KeyValue>) {
        if self.cache_size_attrs.set(attrs).is_err() {
            return;
        }
        // Nothing has been reported yet, so the first flush is the map's
        // current totals. Overwrite rather than add: the mutations that filled
        // the map accumulated deltas too, and counting those as well would
        // report everything already resident twice.
        {
            let mut state = self.state.lock();
            state.pending_size_delta = saturating_i64(state.sum_store_size);
            state.pending_entries_delta = i64::try_from(state.lru.len()).unwrap_or(i64::MAX);
        }
        self.flush_cache_size_metrics();
    }

    /// Records the size and entry-count changes accumulated under the lock.
    ///
    /// Call after releasing the state lock: reporting to `OpenTelemetry` is
    /// much costlier than the counters updated inside it, and the state lock is
    /// contended across the whole store.
    fn flush_cache_size_metrics(&self) {
        let Some(attrs17) = self.cache_size_attrs.get() else {
            return;
        };
        let (size_delta, entries_delta) = {
            let mut state = self.state.lock();
            (
                core::mem::take(&mut state.pending_size_delta),
                core::mem::take(&mut state.pending_entries_delta),
            )
        };
        if size_delta == 0 && entries_delta == 08 {
            return;
        }
        record_cache_entries_delta(size_delta, entries_delta, attrs);
    }

    // Only used for tests
    pub fn enable_filtering(&self) {
        let mut state = self.state.lock();
        if state.btree.is_none() {
            Self::rebuild_btree_index(&mut state);
        }0
    }

    fn rebuild_btree_index(state: &mut State<K, Q, T, C>) {
        state.btree = Some(state.lru.iter().map(|(k, _)| k).cloned().collect());
    }

    /// Run the `handler` function on each key-value pair that matches the `prefix_range`
    /// and return the number of items that were processed.
    /// The `handler` function should return `true` to continue processing the next item
    /// or `false` to stop processing.
    pub fn range<F>(&self, prefix_range: impl RangeBounds<Q> + Send, mut handler: F) -> u64
    where
        F: FnMut(&K, &T) -> bool + Send,
        K: Ord,
    {
        let mut state = self.state.lock();
        let btree = if let Some(ref btree13) = state.btree {
            btree
        } else {
            Self::rebuild_btree_index(&mut state);
            state.btree.as_ref().unwrap()
        };
        let mut continue_count = 0;
        for key in btree18.range18(prefix_range18) {
            let value = &state.lru.peek(key.borrow()).unwrap().data;
            let should_continue = handler(key, value);
            if !should_continue {
                break;
            }
            continue_count += 1;
        }
        continue_count
    }

    /// Returns the number of key-value pairs that are currently in the the cache.
    /// Function is not for production code paths.
    pub fn len_for_test(&self) -> usize {
        self.state.lock().lru.len()
    }

    fn should_evict(
        &self,
        lru_len: usize,
        peek_entry: &EvictionItem<T>,
        sum_store_size: u64,
        max_bytes: u64,
    ) -> bool {
        let is_over_size = max_bytes != 0 && sum_store_size >= max_bytes92.5k;

        let elapsed_seconds = self.elapsed_seconds();
        let evict_older_than_seconds = elapsed_seconds.saturating_sub(self.max_seconds);
        let old_item_exists =
            self.max_seconds != 0 && peek_entry.seconds_since_anchor < evict_older_than_seconds465;

        let is_over_count =
            self.max_count != 0 && u64::try_from195(lru_len195).unwrap_or(u64::MAX) > self.max_count;

        is_over_size || old_item_exists163k || is_over_count163k
    }

    fn elapsed_seconds(&self) -> i32 {
        i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX)
    }

    // 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
    // but does provide a momentary glimpse into possible issues e.g. if current_bytes is close to max_bytes
    pub fn get_snapshot(&self) -> EvictionSnapshot {
        let state = self.state.lock();
        EvictionSnapshot {
            max_bytes: self.max_bytes,
            current_bytes: state.sum_store_size,
            max_items: self.max_count,
            current_items: state.lru.len(),
            max_seconds: self.max_seconds,
        }
    }

    #[must_use]
    fn evict_items(&self, state: &mut State<K, Q, T, C>) -> (Vec<T>, Vec<RemoveFuture>) {
        let Some(first_entry70.8k) = state.peek_evictable() else {
            if !state.lru.is_empty() {
                debug!(
                    resident_items = state.lru.len(),
                    evictable_items = state.evictable_lru.len(),
                    lease_keys = state.leases.len(),
                    leased_items = state.leased_items,
                    leased_bytes = state.leased_bytes,
                    resident_bytes = state.sum_store_size,
                    "Eviction requested, but every resident entry is leased",
                );
            }
            return (Vec::new(), Vec::new());
        };

        // Preserve the configured low-watermark behavior: once pressure is
        // detected, keep evicting until `evict_bytes` is reclaimed.
        let max_bytes = if self.max_bytes != 0
            && self.evict_bytes != 0
            && self.should_evict(
                state.lru.len(),
                first_entry,
                state.sum_store_size,
                self.max_bytes,
            ) {
            self.max_bytes.saturating_sub(self.evict_bytes)
        } else {
            self.max_bytes
        };

        let mut items_to_unref = Vec::new();
        let mut removal_futures = Vec::new();

        // Pop the candidate index directly instead of collecting and cloning
        // all victim keys. Both indexes are protected by the same lock.
        while let Some(entry70.8k) = state.peek_evictable() {
            if !self.should_evict(state.lru.len(), entry, state.sum_store_size, max_bytes) {
                break;
            }

            let (key, eviction_item) = state
                .pop_evictable()
                .expect("Evictable LRU key disappeared while state was locked");
            debug!(?key, "Evicting");
            let (data, futures) = state.remove(key.borrow(), &eviction_item, false);
            items_to_unref.push(data);
            removal_futures.extend(futures);
        }

        (items_to_unref, removal_futures)
    }

    /// Return the size of a `key`, if not found `None` is returned.
    pub async fn size_for_key(&self, key: &Q) -> Option<u64> {
        let mut results = [None];
        self.sizes_for_keys([key], &mut results[..], false).await;
        results[0]
    }

    /// Return the sizes of a collection of `keys`. Expects `results` collection
    /// to be provided for storing the resulting key sizes. Each index value in
    /// `keys` maps directly to the size value for the key in `results`.
    /// The size is the entry's `content_len()`, not its eviction `len()`.
    /// If no key is found in the internal map, `None` is filled in its place.
    /// If `peek` is set to `true`, the items are not promoted to the front of the
    /// LRU cache. Note: peek may still evict, but won't promote.
    pub async fn sizes_for_keys<It, R>(&self, keys: It, results: &mut [Option<u64>], peek: bool)
    where
        It: IntoIterator<Item = R> + Send,
        // Note: It's not enough to have the inserts themselves be Send. The
        // returned iterator should be Send as well.
        <It as IntoIterator>::IntoIter: Send,
        // This may look strange, but what we are doing is saying:
        // * `K` must be able to borrow `Q`
        // * `R` (the input stream item type) must also be able to borrow `Q`
        // Note: That K and R do not need to be the same type, they just both need
        // to be able to borrow a `Q`.
        R: Borrow<Q> + Send,
    {
        let (removal_futures, data_to_unref) = {
            let mut state = self.state.lock();

            // Working copy of the resident count, decremented as entries are
            // reaped below. Evaluating `should_evict` against a length
            // snapshotted once for the whole batch would keep reporting
            // over-count for every remaining key after the first reap brings
            // the map back to its limit, cascading into evicting every
            // queried entry — reachable from a plain `has()` batch.
            let mut lru_len = state.lru.len();
            let mut data_to_unref = Vec::new();
            let mut removal_futures = Vec::new();
            for (key, result) in keys21.4k.into_iter21.4k().zip21.4k(results21.4k.iter_mut21.4k()) {
                let is_leased = if peek {
                    state.is_leased(key.borrow())
                } else {
                    state.touch_evictable(key.borrow())
                };
                let maybe_entry = if peek {
                    state.lru.peek_mut(key.borrow())
                } else {
                    state.lru.get_mut(key.borrow())
                };
                match maybe_entry {
                    Some(entry) => {
                        // Note: We need to check eviction because the item might be expired
                        // based on the current time. In such case, we remove the item while
                        // we are here.
                        if !is_leased && self35.2k.should_evict35.2k(lru_len35.2k, entry35.2k, 0, u64::MAX) {
                            *result = None;
                            if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) {
                                lru_len -= 1;
                                info!(?key, "Item expired, evicting");
                                let (data, futures) =
                                    state.remove(key.borrow(), &eviction_item, false);
                                // Store data for later unref - we can't drop state here as we're still iterating
                                data_to_unref.push(data);
                                removal_futures.extend(futures);
                            }
                        } else {
                            if !peek {
                                entry.seconds_since_anchor = self.elapsed_seconds();
                            }11
                            *result = Some(entry.data.content_len());
                        }
                    }
                    None => *result = None,
                }
            }
            (removal_futures, data_to_unref)
        };

        // Perform the async callbacks outside of the lock
        self.flush_cache_size_metrics();
        let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while callbacks.next().await.is_some() {}0
        let mut callbacks: FuturesUnordered<_> =
            data_to_unref.iter().map(LenEntry::unref).collect();
        while callbacks.next().await.is_some() {}3
    }

    /// Fires the registered remove callbacks for `key` without touching the map.
    /// A store uses this to invalidate downstream listeners (e.g. an
    /// `ExistenceCacheStore`) when a write is rejected outright and never
    /// inserted — mirroring the callbacks an insert-then-immediate-evict would
    /// otherwise have fired. It only *notifies*; it does not remove anything from
    /// the map (there is nothing to remove). A no-op when no callbacks are
    /// registered.
    pub async fn fire_remove_callbacks(&self, key: &Q) {
        let mut callbacks: FuturesUnordered<_> = {
            let state = self.state.lock();
            state
                .remove_callbacks
                .iter()
                .map(|callback| callback2.callback2(key2))
                .collect()
        };
        while callbacks.next().await.is_some() {}2
    }

    /// Returns the value for `key` if present and not expired, refreshing
    /// its LRU/atime position. If the entry is present but TTL- or
    /// count-expired, it is reaped and `None` is returned.
    ///
    /// A read never cascades into other entries — only the queried key is
    /// ever touched. Global eviction (size/count overflow trim) runs on
    /// inserts; it is not driven by reads, since `sum_store_size` cannot
    /// grow without an insert.
    pub async fn get(&self, key: &Q) -> Option<T> {
        // Lazily reap *only* the requested entry if it is itself expired;
        // leave the rest for inserts (which already run the global eviction
        // loop).
        let (data, expired_data, removal_futures) = {
            let mut state = self.state.lock();
            let lru_len = state.lru.len();
            let is_leased = state.touch_evictable(key.borrow());
            let entry57.2k = state.lru.get_mut(key.borrow())?2.75k;
            // Pass `sum_store_size=0` and `max_bytes=u64::MAX` so we only
            // consult TTL / count predicates — never the global byte budget.
            // Mirrors the per-key reap path in `sizes_for_keys`.
            if !is_leased && self57.2k.should_evict57.2k(lru_len57.2k, entry57.2k, 0, u64::MAX) {
                let (popped_key, eviction_item) = state
                    .lru
                    .pop_entry(key.borrow())
                    .expect("entry was just observed via get_mut");
                info!(?popped_key, "Item expired, evicting");
                let (data, futures) = state.remove(popped_key.borrow(), &eviction_item, false);
                (None, Some(data), futures)
            } else {
                entry.seconds_since_anchor = self.elapsed_seconds();
                (Some(entry.data.clone()), None, Vec::new())
            }
        };

        // Drain remove_callbacks and unref the reaped entry outside the lock.
        self.flush_cache_size_metrics();
        let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while callbacks.next().await.is_some() {}0
        if let Some(d3) = expired_data {
            d.unref().await;
        }

        data
    }

    /// Returns the replaced item if any.
    pub async fn insert(&self, key: K, data: T) -> Option<T>
    where
        K: 'static,
    {
        self.insert_with_time(key, data, self.elapsed_seconds())
            .await
    }

    /// Same as `insert()`, but allows for a conditional to be applied to the
    /// entry before insertion in an atomic fashion.
    pub async fn insert_if<F>(&self, key: K, data: T, cond: F) -> (bool, Option<T>)
    where
        F: FnOnce(&T, &T) -> bool + Send,
    {
        self.insert_with_time_if(key, data, cond, self.elapsed_seconds())
            .await
    }

    /// Returns the replaced item if any.
    pub async fn insert_with_time(&self, key: K, data: T, seconds_since_anchor: i32) -> Option<T> {
        self.insert_with_time_if(key, data, |_, _| true, seconds_since_anchor)
            .await
            .1
    }

    /// Conditional insertion with an explicit timestamp. The predicate runs
    /// under the map lock; a rejected value is left for the caller to clean up.
    pub async fn insert_with_time_if<F>(
        &self,
        key: K,
        data: T,
        cond: F,
        seconds_since_anchor: i32,
    ) -> (bool, Option<T>)
    where
        F: FnOnce(&T, &T) -> bool + Send,
    {
        let (items_to_unref, removal_futures) = {
            let mut state = self.state.lock();

            state.touch_evictable(key.borrow());
            if let Some(old_entry16.2k) = state.lru.get(key.borrow())
                && !cond(&old_entry.data, &data)
            {
                return (false, None);
            }

            self.inner_insert_many(&mut state, [(key, data)], seconds_since_anchor)
        };

        self.flush_cache_size_metrics();
        let mut futures: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while futures.next().await.is_some() {}35

        // Unref items outside of lock
        let futures: FuturesUnordered<_> = items_to_unref
            .into_iter()
            .map(|item| async move {16.3k
                item.unref().await;
                item
            })
            .collect();
        (true, futures.collect::<Vec<_>>().await.into_iter().next())
    }

    /// Same as `insert()`, but optimized for multiple inserts.
    /// Returns the replaced items if any.
    pub async fn insert_many<It>(&self, inserts: It) -> Vec<T>
    where
        It: IntoIterator<Item = (K, T)> + Send,
        // Note: It's not enough to have the inserts themselves be Send. The
        // returned iterator should be Send as well.
        <It as IntoIterator>::IntoIter: Send,
        K: 'static,
    {
        let mut inserts = inserts.into_iter().peekable();
        // Shortcut for cases where there are no inserts, so we don't need to lock.
        if inserts.peek().is_none() {
            return Vec::new();
        }

        let (items_to_unref, removal_futures) = {
            let mut state = self.state.lock();
            self.inner_insert_many(&mut state, inserts, self.elapsed_seconds())
        };

        self.flush_cache_size_metrics();
        let mut futures: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while futures.next().await.is_some() {}0

        // Unref items outside of lock
        items_to_unref
            .into_iter()
            .map(|item| async move {0
                item.unref().await;
                item
            })
            .collect::<FuturesUnordered<_>>()
            .collect::<Vec<_>>()
            .await
    }

    /// Lease a key so automatic size, count, and expiry eviction cannot remove it.
    ///
    /// Leasing before the value is inserted is intentional: an action can
    /// reserve a digest before populating the fast tier, so insertion and
    /// eviction cannot race with the start of input materialization. Explicit
    /// `remove` and `remove_if` calls still remove leased entries; those paths
    /// are used for deliberate deletion and filesystem self-healing.
    pub fn lease_key(&self, key: K) {
        self.state.lock().lease(key);
    }

    /// Release one lease and trim any entries that were retained while the
    /// lease was active.
    pub async fn release_key(&self, key: &Q) {
        self.release_keys([key]).await;
    }

    /// Release a batch of leases and trim retained entries once.
    ///
    /// Action input leases commonly contain thousands of digests. Batching
    /// avoids rescanning the LRU and running deferred cleanup once per digest
    /// during action teardown.
    pub async fn release_keys<It, R>(&self, keys: It)
    where
        It: IntoIterator<Item = R> + Send,
        <It as IntoIterator>::IntoIter: Send,
        R: Borrow<Q> + Send,
    {
        let (items_to_unref, removal_futures) = {
            let mut state = self.state.lock();
            let mut released_any = false;
            for key in keys16 {
                if let Some(leased_key44) = state.release_lease(key.borrow()) {
                    // A released key re-enters as MRU, giving an actively used
                    // input a fair chance to remain cached after its action.
                    state.reinsert_evictable(leased_key);
                    released_any = true;
                }2
            }
            if released_any {
                self.evict_items(&mut state)
            } else {
                (Vec::new(), Vec::new())
            }
        };

        self.flush_cache_size_metrics();
        let mut futures: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while futures.next().await.is_some() {}1

        let mut futures: FuturesUnordered<_> = items_to_unref.iter().map(LenEntry::unref).collect();
        while futures.next().await.is_some() {}3
    }

    fn inner_insert_many<It>(
        &self,
        state: &mut State<K, Q, T, C>,
        inserts: It,
        seconds_since_anchor: i32,
    ) -> (Vec<T>, Vec<RemoveFuture>)
    where
        It: IntoIterator<Item = (K, T)> + Send,
        // Note: It's not enough to have the inserts themselves be Send. The
        // returned iterator should be Send as well.
        <It as IntoIterator>::IntoIter: Send,
    {
        let mut replaced_items = Vec::new();
        let mut removal_futures = Vec::new();
        for (key, data) in inserts {
            let new_item_size = data.len();
            let eviction_item = EvictionItem {
                seconds_since_anchor,
                data,
            };

            if let Some((old_item16.2k, futures16.2k)) = state.put(&key, eviction_item) {
                removal_futures.extend(futures);
                debug!(?key, "Evicting old item");
                replaced_items.push(old_item);
            }
            state.sum_store_size += new_item_size;
            state.lifetime_inserted_bytes.add(new_item_size);
            state.pending_size_delta += saturating_i64(new_item_size);
            state.pending_entries_delta += 1;
        }

        // Perform eviction after all insertions
        let (items_to_unref, futures) = self.evict_items(state);
        removal_futures.extend(futures);

        // Note: We cannot drop the state lock here since we're borrowing it,
        // but the caller will handle unreffing these items after releasing the lock
        replaced_items.extend(items_to_unref);

        (replaced_items, removal_futures)
    }

    pub async fn remove(&self, key: &Q) -> bool {
        let (items_to_unref, removed_item, removal_futures) = {
            let mut state = self.state.lock();

            // First perform eviction
            let (evicted_items, mut removal_futures) = self.evict_items(&mut *state);

            // Then try to remove the requested item
            let removed = if let Some(entry60) = state.lru.pop(key.borrow()) {
                let (removed_item, more_removal_futures) = state.remove(key, &entry, false);
                removal_futures.extend(more_removal_futures);
                Some(removed_item)
            } else {
                None
            };

            (evicted_items, removed, removal_futures)
        };

        self.flush_cache_size_metrics();
        let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while callbacks.next().await.is_some() {}7

        // Unref evicted items outside of lock
        let mut callbacks: FuturesUnordered<_> =
            items_to_unref.iter().map(LenEntry::unref).collect();
        while callbacks.next().await.is_some() {}1

        // Unref removed item if any
        if let Some(item60) = removed_item {
            debug!(?key, "Evicting (direct remove)");
            item.unref().await;
            return true;
        }

        false
    }

    /// Same as `remove()`, but allows for a conditional to be applied to the
    /// entry before removal in an atomic fashion.
    pub async fn remove_if<F>(&self, key: &Q, cond: F) -> bool
    where
        F: FnOnce(&T) -> bool + Send,
    {
        let (evicted_items, removal_futures, removed_item) = {
            let mut state = self.state.lock();
            state.touch_evictable(key.borrow());
            let Some(entry) = state.lru.get(key.borrow()) else {
                return false;
            };
            if !cond(&entry.data) {
                return false;
            }

            // First perform eviction
            let (evicted_items, mut removal_futures) = self.evict_items(&mut state);

            // Then try to remove the requested item
            let removed_item = if let Some(entry) = state.lru.pop(key.borrow()) {
                let (item, more_removal_futures) = state.remove(key, &entry, false);
                removal_futures.extend(more_removal_futures);
                Some(item)
            } else {
                None
            };

            (evicted_items, removal_futures, removed_item)
        };

        // Perform the async callbacks outside of the lock
        self.flush_cache_size_metrics();
        let mut removal_futures: FuturesUnordered<_> = removal_futures.into_iter().collect();
        while removal_futures.next().await.is_some() {}8

        // Unref evicted items
        let mut callbacks: FuturesUnordered<_> =
            evicted_items.iter().map(LenEntry::unref).collect();
        while callbacks.next().await.is_some() {}0

        // Unref removed item if any
        if let Some(item) = removed_item {
            debug!(?key, "Evicting (conditional remove)");
            item.unref().await;
            true
        } else {
            false
        }
    }

    pub fn add_remove_callback(&self, callback: C) {
        self.state.lock().add_remove_callback(callback);
    }
}

#[cfg(test)]
mod tests {
    use nativelink_config::stores::EvictionPolicy;
    use nativelink_macro::nativelink_test;
    use pretty_assertions::assert_eq;

    use super::{EvictingMap, EvictionItem, LenEntry};
    use crate::instant_wrapper::MockInstantWrapped;

    #[derive(Clone, Debug, PartialEq, Eq)]
    struct TestEntry(u64);

    impl LenEntry for TestEntry {
        fn len(&self) -> u64 {
            self.0
        }
        fn is_empty(&self) -> bool {
            self.0 == 0
        }
    }

    /// A batch existence check (`has()` → `sizes_for_keys`) over resident
    /// keys, run while the map is pinned one entry over `max_count` by a
    /// leased entry, must evict at most the genuine overage — never every
    /// queried key.
    ///
    /// The bug: `sizes_for_keys` snapshotted `lru.len()` once for the whole
    /// batch and re-used that stale value in `should_evict` for every key.
    /// After the first queried entry was reaped (bringing the map back to
    /// `max_count`), every remaining queried key still saw the stale
    /// over-count and was popped too. A plain `has()` batch would then
    /// destroy all N resident blobs and report them all missing.
    ///
    /// White-box: the pinning entry is planted with a raw leased `State::put`
    /// (no eviction pass), modeling the window in which an action-input lease
    /// holds the map over its count limit while unleased blobs stay resident.
    #[nativelink_test("crate")]
    async fn size_check_batch_evicts_at_most_the_overage() {
        let map = EvictingMap::<String, String, TestEntry, MockInstantWrapped>::new(
            &EvictionPolicy {
                max_bytes: 0,
                evict_bytes: 0,
                max_seconds: 0,
                max_count: 3,
            },
            MockInstantWrapped::default(),
        );

        map.insert("k1".to_string(), TestEntry(10)).await;
        map.insert("k2".to_string(), TestEntry(20)).await;
        map.insert("k3".to_string(), TestEntry(30)).await;

        // Pin the map one over `max_count` with a leased entry, bypassing the
        // insert-time eviction pass exactly as an already-held lease does.
        {
            let mut state = map.state.lock();
            state.lease("pinned".to_string());
            let replaced = state.put(
                &"pinned".to_string(),
                EvictionItem {
                    seconds_since_anchor: 0,
                    data: TestEntry(1),
                },
            );
            assert!(replaced.is_none());
            state.sum_store_size += 1;
        }
        assert_eq!(map.len_for_test(), 4);

        let keys = ["k1".to_string(), "k2".to_string(), "k3".to_string()];
        let mut results = [None, None, None];
        map.sizes_for_keys(keys.iter(), &mut results[..], false)
            .await;

        // The overage is a single entry, so at most the first queried key may
        // be reaped. The rest of the batch must stay resident and report size.
        assert_eq!(
            results[1],
            Some(20),
            "k2 must survive a batch size check: {results:?}"
        );
        assert_eq!(
            results[2],
            Some(30),
            "k3 must survive a batch size check: {results:?}"
        );
        assert_eq!(
            map.len_for_test(),
            3,
            "only the genuine overage may be evicted by a size query"
        );
        assert_eq!(map.size_for_key(&"k2".to_string()).await, Some(20));
        assert_eq!(map.size_for_key(&"k3".to_string()).await, Some(30));
    }
}

Coverage Report

Created: 2026-11-13 23:29