memory_store: skip writes larger than max_bytes instead of buffering them (TraceMachina#2473)

corcillo · web-flow · commit 1a10e3fa90c4 · 2026-06-24T11:05:30.000Z
`MemoryStore::update` read the entire payload into memory (`consume(None)`,
plus a defensive full copy) before inserting into the evicting map and
running eviction. For a write whose size already meets the store's
`max_bytes`, that's wasted work: the eviction loop drops the entry the
instant it's inserted (a single item over budget), so it was never
retainable.

Worse, under concurrent large writes it's an OOM vector. A `memory` fast
tier inside a `fast_slow` store fronting CAS receives every blob regardless
of size, so a burst of large ByteStream uploads each materializes its whole
payload before eviction runs — dozens of multi-hundred-MB uploads in
parallel can blow the cgroup limit and OOMKill the CAS.

Fix: when the eviction policy caps size (`max_bytes != 0`) and the write's
`ExactSize` is &gt;= it, drain the stream and return without buffering. The
`&gt;=` matches the eviction comparator (`sum_store_size &gt;= max_bytes`), so a
blob of exactly `max_bytes` is skipped too. Only `ExactSize` is trusted — a
`MaxSize` upper bound could over-estimate and wrongly skip a blob that fits.

Because the rejected write never enters the map, the insert-then-evict
removal callbacks don't fire on their own. A wrapping `ExistenceCacheStore`
relies on those to drop a just-written-then-evicted key, so the skip path
fires them explicitly via a new `EvictingMap::fire_remove_callbacks` (a
no-op when no callbacks are registered), preserving the "copes with dropped
items" guarantee.

Scope: this closes the single-blob-exceeds-cap vector. It does NOT address N
concurrent writes each under `max_bytes` but collectively over the cgroup
limit — each still materializes fully before eviction. The general fix is
incremental/streaming admission with mid-stream size accounting (which would
also cover `MaxSize` and unknown-size writes); this is an interim step.

Tests (nativelink-store/tests/memory_store_test.rs):
  * skips_writes_larger_than_max_bytes — an oversized write is skipped, and a
    within-budget entry survives it (the old path would have evicted it).
  * oversized_skip_fires_remove_callbacks — the skip fires the remove callback
    so a wrapping existence cache can invalidate.
  * max_size_over_budget_with_small_actual_is_stored — a `MaxSize` over the
    budget with small actual content is still cached (only `ExactSize` skips).
diff --git a/nativelink-store/src/memory_store.rs b/nativelink-store/src/memory_store.rs
@@ -33,6 +33,7 @@ use nativelink_util::health_utils::{
 use nativelink_util::store_trait::{
     RemoveItemCallback, StoreDriver, StoreKey, StoreKeyBorrow, StoreOptimizations, UploadSizeInfo,
 };
+use tracing::warn;
 
 use crate::callback_utils::RemoveItemCallbackHolder;
 use crate::cas_utils::is_zero_digest;
@@ -68,13 +69,19 @@ pub struct MemoryStore {
         SystemTime,
         RemoveItemCallbackHolder,
     >,
+    /// The eviction policy's `max_bytes` (0 = unbounded). Cached here so `update`
+    /// can skip writes larger than the entire store budget without buffering
+    /// them — see the note in `update`.
+    #[metric(help = "Maximum bytes this store will hold before eviction (0 = unbounded)")]
+    max_bytes: u64,
 }
 
 impl MemoryStore {
     pub fn new(spec: &MemorySpec) -> Arc<Self> {
         let empty_policy = nativelink_config::stores::EvictionPolicy::default();
         let eviction_policy = spec.eviction_policy.as_ref().unwrap_or(&empty_policy);
         Arc::new(Self {
+            max_bytes: eviction_policy.max_bytes as u64,
             evicting_map: EvictingMap::new(eviction_policy, SystemTime::now()),
         })
     }
@@ -138,8 +145,51 @@ impl StoreDriver for MemoryStore {
         self: Pin<&Self>,
         key: StoreKey<'_>,
         mut reader: DropCloserReadHalf,
-        _size_info: UploadSizeInfo,
+        size_info: UploadSizeInfo,
     ) -> Result<u64, Error> {
+        // A write whose exact size is at least this store's `max_bytes` can never
+        // be usefully cached: the moment it's inserted, eviction drops it, since one
+        // entry alone meets the budget. Buffering it into memory first is therefore
+        // pure waste — and under concurrent large writes (e.g. a `memory` fast tier
+        // inside a `fast_slow` store fronting CAS) it is a real OOM vector, because
+        // each in-flight write materializes its whole payload before eviction ever
+        // runs. Drain the stream and skip instead.
+        //
+        // `>=` deliberately matches the eviction comparator: `EvictingMap` evicts
+        // while `sum_store_size >= max_bytes`, so a blob of exactly `max_bytes` is
+        // also unstorable and must be skipped rather than buffered-then-evicted.
+        // Only `ExactSize` is trusted; a `MaxSize` upper bound could over-estimate
+        // and wrongly skip a blob that would actually fit.
+        //
+        // For CAS digest keys the size is part of the key, so a given key is either
+        // always oversized or never — there is no "small write for this key" that the
+        // removal callbacks fired below could spuriously invalidate.
+        if self.max_bytes != 0
+            && let UploadSizeInfo::ExactSize(sz) = size_info
+            && sz >= self.max_bytes
+        {
+            let drained = reader
+                .drain()
+                .await
+                .err_tip(|| "Failed to drain oversized write in memory_store::update")?;
+            warn!(
+                ?key,
+                size = sz,
+                max_bytes = self.max_bytes,
+                "Write is larger than this memory store's max_bytes; skipping it \
+                 (it would be evicted immediately). If this store is a cache, large \
+                 blobs are served from the backing store — raise max_bytes to cache \
+                 them, or route large blobs around the memory tier.",
+            );
+            // The write never enters the map, so the insert-then-evict removal
+            // callbacks (which a wrapping `ExistenceCacheStore` relies on to drop
+            // a just-written-then-evicted key) don't fire on their own. Fire them
+            // explicitly so downstream listeners don't keep a stale "exists"
+            // entry for a blob we didn't store.
+            let owned_key = key.into_owned();
+            self.evicting_map.fire_remove_callbacks(&owned_key).await;
+            return Ok(drained);
+        }
         // Internally Bytes might hold a reference to more data than just our data. To prevent
         // this potential case, we make a full copy of our data for long-term storage.
         let final_buffer = {
diff --git a/nativelink-store/tests/memory_store_test.rs b/nativelink-store/tests/memory_store_test.rs
@@ -12,8 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use core::future::Future;
 use core::ops::RangeBounds;
 use core::pin::Pin;
+use core::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
 
 use bytes::{BufMut, Bytes, BytesMut};
 use memory_stats::memory_stats;
@@ -24,7 +27,9 @@ use nativelink_store::memory_store::MemoryStore;
 use nativelink_util::buf_channel::make_buf_channel_pair;
 use nativelink_util::common::DigestInfo;
 use nativelink_util::spawn;
-use nativelink_util::store_trait::{StoreKey, StoreLike};
+use nativelink_util::store_trait::{
+    RemoveItemCallback, Store, StoreKey, StoreLike, UploadSizeInfo,
+};
 use pretty_assertions::assert_eq;
 use sha2::{Digest, Sha256};
 
@@ -82,6 +87,143 @@ async fn insert_one_item_then_update() -> Result<(), Error> {
     Ok(())
 }
 
+// Streams `data` into the store via the streaming `update` path. Tests must use
+// this (not `update_oneshot`, which takes the bytes already in memory and hits
+// MemoryStore's optimized direct-insert path, bypassing the oversized-skip logic
+// these tests exercise).
+async fn update_streamed(
+    store: &Store,
+    digest: DigestInfo,
+    data: &'static [u8],
+    size_info: UploadSizeInfo,
+) -> Result<(), Error> {
+    let (mut tx, rx) = make_buf_channel_pair();
+    let send_fut = async move {
+        tx.send(Bytes::from_static(data)).await?;
+        tx.send_eof()
+    };
+    let pinned_store = Pin::new(store);
+    let (res1, res2) = futures::join!(send_fut, pinned_store.update(digest, rx, size_info));
+    // `merge` returns its argument's Ok value (the byte count); discard it.
+    res1.merge(res2).map(|_| ())
+}
+
+fn store_with_max_bytes(max_bytes: usize) -> Store {
+    Store::new(MemoryStore::new(&MemorySpec {
+        eviction_policy: Some(nativelink_config::stores::EvictionPolicy {
+            max_bytes,
+            ..Default::default()
+        }),
+    }))
+}
+
+// A write whose exact size is >= max_bytes is skipped (drained, never buffered)
+// rather than materialized-then-evicted, and — crucially — it leaves the rest of
+// the cache untouched (the old buffer-then-evict path would have evicted the
+// within-budget entry trying to make room for an unstorable blob).
+#[nativelink_test]
+async fn skips_writes_larger_than_max_bytes() -> Result<(), Error> {
+    const SMALL: &[u8] = b"ab"; // 2 bytes, within the 4-byte budget
+    const BIG: &[u8] = b"0123456789"; // 10 bytes, over budget
+
+    let store = store_with_max_bytes(4);
+    let small_digest = DigestInfo::try_new(VALID_HASH1, SMALL.len() as u64)?;
+    let big_digest = DigestInfo::try_new(VALID_HASH2, BIG.len() as u64)?;
+
+    // Within budget → stored.
+    update_streamed(
+        &store,
+        small_digest,
+        SMALL,
+        UploadSizeInfo::ExactSize(SMALL.len() as u64),
+    )
+    .await?;
+    assert_eq!(store.has(small_digest).await, Ok(Some(SMALL.len() as u64)));
+
+    // Over budget → skipped, not stored.
+    update_streamed(
+        &store,
+        big_digest,
+        BIG,
+        UploadSizeInfo::ExactSize(BIG.len() as u64),
+    )
+    .await?;
+    assert_eq!(
+        store.has(big_digest).await,
+        Ok(None),
+        "oversized write should be skipped, not stored",
+    );
+
+    // The skip must leave the within-budget entry intact.
+    assert_eq!(
+        store.has(small_digest).await,
+        Ok(Some(SMALL.len() as u64)),
+        "within-budget entry must survive an oversized write",
+    );
+
+    Ok(())
+}
+
+#[derive(Debug)]
+struct CountingRemoveCallback(Arc<AtomicUsize>);
+
+impl RemoveItemCallback for CountingRemoveCallback {
+    fn callback<'a>(
+        &'a self,
+        _store_key: StoreKey<'a>,
+    ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
+        let count = self.0.clone();
+        Box::pin(async move {
+            count.fetch_add(1, Ordering::SeqCst);
+        })
+    }
+}
+
+// Skipping an oversized write must still fire the remove callbacks, so a wrapping
+// store that tracks existence (e.g. `ExistenceCacheStore`) doesn't keep a stale
+// entry for a blob that was never stored.
+#[nativelink_test]
+async fn oversized_skip_fires_remove_callbacks() -> Result<(), Error> {
+    const BIG: &[u8] = b"0123456789"; // 10 bytes, over the 4-byte budget
+
+    let store = store_with_max_bytes(4);
+    let count = Arc::new(AtomicUsize::new(0));
+    store.register_remove_callback(Arc::new(CountingRemoveCallback(count.clone())))?;
+
+    update_streamed(
+        &store,
+        DigestInfo::try_new(VALID_HASH1, BIG.len() as u64)?,
+        BIG,
+        UploadSizeInfo::ExactSize(BIG.len() as u64),
+    )
+    .await?;
+
+    assert_eq!(
+        count.load(Ordering::SeqCst),
+        1,
+        "skipping an oversized write must fire the remove callback",
+    );
+    Ok(())
+}
+
+// `MaxSize` is an upper bound, not the real size, so it must NOT trigger the skip:
+// a MaxSize over max_bytes whose actual content fits is still cached.
+#[nativelink_test]
+async fn max_size_over_budget_with_small_actual_is_stored() -> Result<(), Error> {
+    const DATA: &[u8] = b"ab"; // 2 bytes, well within the 4-byte budget
+
+    let store = store_with_max_bytes(4);
+    let digest = DigestInfo::try_new(VALID_HASH1, DATA.len() as u64)?;
+    update_streamed(&store, digest, DATA, UploadSizeInfo::MaxSize(100)).await?;
+
+    assert_eq!(
+        store.has(digest).await,
+        Ok(Some(DATA.len() as u64)),
+        "MaxSize over budget but small actual content must still be cached",
+    );
+    Ok(())
+}
+
 // Regression test for: https://github.com/TraceMachina/nativelink/issues/289.
 #[nativelink_test]
 async fn ensure_full_copy_of_bytes_is_made_test() -> Result<(), Error> {
diff --git a/nativelink-util/src/evicting_map.rs b/nativelink-util/src/evicting_map.rs
@@ -490,6 +490,25 @@ where
         while callbacks.next().await.is_some() {}
     }
 
+    /// 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| callback.callback(key))
+                .collect()
+        };
+        while callbacks.next().await.is_some() {}
+    }
+
     /// 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.
