Fix CAS upload, eviction and cleanup races, and cut evicting map lock contention

nativelink-service/src/bytestream_server.rs

@@ -944,6 +944,17 @@ impl ByteStreamServer {
                     return Err(make_input_err!("Received more bytes than expected"));
                 }
                 if write_request.finish_write {
+                    // The stream wrapper only notices a short upload if it is
+                    // polled again after `finish_write`, which never happens,
+                    // so check here. Sending EOF would commit the short data
+                    // under the full digest.
+                    if tx.get_bytes_written() != expected_size {
+                        return Err(make_input_err!(
+                            "Write finished after {} bytes, expected {}",
+                            tx.get_bytes_written(),
+                            expected_size
+                        ));
+                    }
                     // Gracefully close our stream.
                     tx.send_eof()
                         .err_tip(|| "Failed to send EOF in ByteStream::write")?;
@@ -1058,6 +1069,14 @@ impl ByteStreamServer {
             }
 
             if write_request.finish_write {
+                // See `process_client_stream` in `inner_write`.
+                if bytes_received != expected_size {
+                    return Err(make_input_err!(
+                        "Write finished after {} bytes, expected {}",
+                        bytes_received,
+                        expected_size
+                    ));
+                }
                 break;
             }
         }

nativelink-service/tests/bytestream_server_test.rs

@@ -1878,7 +1878,10 @@ pub async fn max_decoding_message_size_test() -> Result<(), Box<dyn core::error:
         // Test to ensure if we send exactly our max message size, it will succeed.
         let data = Bytes::from(vec![0u8; MAX_MESSAGE_SIZE - WRITE_REQUEST_MSG_WRAPPER_SIZE]);
         let write_request = WriteRequest {
-            resource_name: make_resource_name(MAX_MESSAGE_SIZE),
+            // The digest size must match the data, or the write is rejected
+            // as short. It has as many digits as `MAX_MESSAGE_SIZE`, so the
+            // wrapper size is unchanged.
+            resource_name: make_resource_name(data.len()),
             write_offset: 0,
             finish_write: true,
             data,
@@ -1959,6 +1962,99 @@ async fn write_too_many_bytes_fails() -> Result<(), Box<dyn core::error::Error>>
     Ok(())
 }
 
+// Regression test for: https://github.com/TraceMachina/nativelink/issues/2242.
+// A single message with `finish_write` takes the oneshot path.
+#[nativelink_test]
+async fn oneshot_write_finished_short_of_digest_size_fails()
+-> Result<(), Box<dyn core::error::Error>> {
+    const WRITE_DATA: &str = "12345";
+    // The resource name promises one byte more than the client sends.
+    const DIGEST_SIZE: usize = WRITE_DATA.len() + 1;
+
+    let store_manager = make_store_manager().await?;
+    let bs_server = Arc::new(
+        make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"),
+    );
+    let store = store_manager.get_store("main_cas").unwrap();
+
+    let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None);
+    tx.send(Frame::data(encode_stream_proto(&WriteRequest {
+        resource_name: make_resource_name(DIGEST_SIZE),
+        write_offset: 0,
+        finish_write: true,
+        data: WRITE_DATA.into(),
+    })?))
+    .await?;
+
+    let status = join_handle
+        .await?
+        .expect_err("Expected an error for finishing a write short");
+    assert_eq!(status.code(), Code::InvalidArgument, "{status:?}");
+    assert!(
+        status
+            .message()
+            .contains("Write finished after 5 bytes, expected 6"),
+        "Got wrong error: {status:?}"
+    );
+    assert_eq!(
+        store.has(DigestInfo::try_new(HASH1, DIGEST_SIZE)?).await?,
+        None,
+        "A short write must not be stored under the full digest"
+    );
+    drop(tx);
+    Ok(())
+}
+
+// Regression test for: https://github.com/TraceMachina/nativelink/issues/2242.
+// A write spread over several messages takes the streaming path.
+#[nativelink_test]
+async fn streamed_write_finished_short_of_digest_size_fails()
+-> Result<(), Box<dyn core::error::Error>> {
+    const WRITE_DATA: &str = "12456789abcdefghijk";
+    const BYTE_SPLIT_OFFSET: usize = 8;
+    // The resource name promises one byte more than the client sends.
+    const DIGEST_SIZE: usize = WRITE_DATA.len() + 1;
+
+    let store_manager = make_store_manager().await?;
+    let bs_server = Arc::new(
+        make_bytestream_server(store_manager.as_ref(), None).expect("Failed to make server"),
+    );
+    let store = store_manager.get_store("main_cas").unwrap();
+
+    let (tx, join_handle) = make_stream_and_writer_spawn(bs_server, None);
+    let mut write_request = WriteRequest {
+        resource_name: make_resource_name(DIGEST_SIZE),
+        write_offset: 0,
+        finish_write: false,
+        data: WRITE_DATA[..BYTE_SPLIT_OFFSET].into(),
+    };
+    tx.send(Frame::data(encode_stream_proto(&write_request)?))
+        .await?;
+    write_request.write_offset = BYTE_SPLIT_OFFSET.try_into()?;
+    write_request.finish_write = true;
+    write_request.data = WRITE_DATA[BYTE_SPLIT_OFFSET..].into();
+    tx.send(Frame::data(encode_stream_proto(&write_request)?))
+        .await?;
+
+    let status = join_handle
+        .await?
+        .expect_err("Expected an error for finishing a write short");
+    assert_eq!(status.code(), Code::InvalidArgument, "{status:?}");
+    assert!(
+        status
+            .message()
+            .contains("Write finished after 19 bytes, expected 20"),
+        "Got wrong error: {status:?}"
+    );
+    assert_eq!(
+        store.has(DigestInfo::try_new(HASH1, DIGEST_SIZE)?).await?,
+        None,
+        "A short write must not be stored under the full digest"
+    );
+    drop(tx);
+    Ok(())
+}
+
 // NOTE: UUID collision fix has been verified manually.
 // When two uploads use the same UUID and one is active, the server generates
 // a unique UUID using nanosecond timestamp for the second upload.

nativelink-store/src/existence_cache_store.rs

@@ -14,12 +14,12 @@
 
 use core::pin::Pin;
 use std::borrow::Cow;
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
 use std::sync::{Arc, Weak};
 use std::time::SystemTime;
 
 use async_trait::async_trait;
-use futures::StreamExt;
-use futures::stream::FuturesUnordered;
 use nativelink_config::stores::{EvictionPolicy, ExistenceCacheSpec};
 use nativelink_error::{Error, ResultExt, error_if};
 use nativelink_metric::MetricsComponent;
@@ -32,7 +32,7 @@ use nativelink_util::store_trait::{
     RemoveCallback, RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
 };
 use parking_lot::Mutex;
-use tracing::{debug, info, trace};
+use tracing::{debug, trace};
 
 #[derive(Clone, Debug)]
 struct ExistenceItem(u64);
@@ -49,17 +49,67 @@ impl LenEntry for ExistenceItem {
     }
 }
 
+/// The updates of one key that are in flight, see `in_flight_updates`.
+#[derive(Debug, Default)]
+struct InFlightUpdates {
+    count: usize,
+    /// Whether the inner store evicted the key since the oldest of these
+    /// updates started.
+    evicted: bool,
+}
+
 #[derive(Debug, MetricsComponent)]
 pub struct ExistenceCacheStore<I: InstantWrapper> {
     #[metric(group = "inner_store")]
     inner_store: Store,
     existence_cache: EvictingMap<DigestInfo, DigestInfo, ExistenceItem, I>,
 
-    // We need to pause them temporarily when inserting into the inner store
-    // as if it immediately expires them, we should only apply the remove callbacks
-    // afterwards. If this is None, we're not pausing; if it's Some it's the location to
-    // store them in temporarily
-    pause_remove_callbacks: Mutex<Option<Vec<StoreKey<'static>>>>,
+    // The inner store can evict a key while an update of it is in flight,
+    // even as part of inserting it. The remove callback then runs before the
+    // update adds the key to the existence cache, so it removes nothing and
+    // the cache would report a blob that is gone. Updates register their key
+    // here and remove callbacks flag it, so the update knows not to cache it.
+    in_flight_updates: Mutex<HashMap<DigestInfo, InFlightUpdates>>,
+}
+
+/// Registers an update of `digest` in `in_flight_updates` until dropped.
+#[derive(Debug)]
+struct InFlightUpdateGuard<'a, I: InstantWrapper> {
+    store: &'a ExistenceCacheStore<I>,
+    digest: DigestInfo,
+}
+
+impl<'a, I: InstantWrapper> InFlightUpdateGuard<'a, I> {
+    fn new(store: &'a ExistenceCacheStore<I>, digest: DigestInfo) -> Self {
+        store
+            .in_flight_updates
+            .lock()
+            .entry(digest)
+            .or_default()
+            .count += 1;
+        Self { store, digest }
+    }
+
+    /// Whether the inner store evicted the key since the update started.
+    fn evicted(&self) -> bool {
+        self.store
+            .in_flight_updates
+            .lock()
+            .get(&self.digest)
+            .is_some_and(|in_flight| in_flight.evicted)
+    }
+}
+
+impl<I: InstantWrapper> Drop for InFlightUpdateGuard<'_, I> {
+    fn drop(&mut self) {
+        let mut in_flight_updates = self.store.in_flight_updates.lock();
+        if let Entry::Occupied(mut entry) = in_flight_updates.entry(self.digest) {
+            entry.get_mut().count -= 1;
+            if entry.get().count == 0 {
+                entry.remove();
+            }
+        }
+    }
 }
 
 impl ExistenceCacheStore<SystemTime> {
@@ -75,10 +125,15 @@ impl<I: InstantWrapper> RemoveItemCallback for ExistenceCacheStore<I> {
     ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
         debug!(?store_key, "Removing item from cache due to callback");
         let digest = store_key.borrow().into_digest();
+        if let Some(in_flight) = self.in_flight_updates.lock().get_mut(&digest) {
+            in_flight.evicted = true;
+        }
         Box::pin(async move {
             let deleted_key = self.existence_cache.remove(&digest).await;
             if !deleted_key {
-                info!(?store_key, "Failed to delete key from cache on callback");
+                // Expected for most keys: the inner store evicts plenty that
+                // were never queried through this cache.
+                trace!(?store_key, "Failed to delete key from cache on callback");
             }
         })
     }
@@ -96,17 +151,12 @@ impl<I: InstantWrapper> RemoveItemCallback for ExistenceCacheCallback<I> {
     ) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
         let cache = self.cache.upgrade();
         if let Some(local_cache) = cache {
-            if let Some(callbacks) = local_cache.pause_remove_callbacks.lock().as_mut() {
-                callbacks.push(store_key.into_owned());
-            } else {
-                let store_key = store_key.into_owned();
-                return Box::pin(async move {
-                    local_cache.callback(store_key).await;
-                });
-            }
-        } else {
-            debug!("Cache dropped, so not doing callback");
+            let store_key = store_key.into_owned();
+            return Box::pin(async move {
+                local_cache.callback(store_key).await;
+            });
         }
+        debug!("Cache dropped, so not doing callback");
         Box::pin(async {})
     }
 }
@@ -122,7 +172,7 @@ impl<I: InstantWrapper> ExistenceCacheStore<I> {
         let existence_cache_store = Arc::new(Self {
             inner_store,
             existence_cache: EvictingMap::new(eviction_policy, anchor_time),
-            pause_remove_callbacks: Mutex::new(None),
+            in_flight_updates: Mutex::new(HashMap::new()),
         });
         let other_ref = Arc::downgrade(&existence_cache_store);
         existence_cache_store
@@ -251,29 +301,21 @@ impl<I: InstantWrapper> StoreDriver for ExistenceCacheStore<I> {
                 .err_tip(|| "In ExistenceCacheStore::update")?;
             return Ok(size);
         }
-        {
-            let mut locked_callbacks = self.pause_remove_callbacks.lock();
-            if locked_callbacks.is_none() {
-                locked_callbacks.replace(vec![]);
-            }
-        }
+        let in_flight_update = InFlightUpdateGuard::new(self.get_ref(), digest);
         trace!(?digest, "Inserting into inner cache");
         let result = self.inner_store.update(digest, reader, size_info).await;
-        if let Ok(size) = &result {
+        if let Ok(size) = &result
+            && !in_flight_update.evicted()
+        {
             trace!(?digest, "Inserting into existence cache");
             let _ = self
                 .existence_cache
                 .insert(digest, ExistenceItem(*size))
                 .await;
-        }
-        {
-            let maybe_keys = self.pause_remove_callbacks.lock().take();
-            if let Some(keys) = maybe_keys {
-                let mut callbacks: FuturesUnordered<_> = keys
-                    .into_iter()
-                    .map(|store_key| self.callback(store_key))
-                    .collect();
-                while callbacks.next().await.is_some() {}
+            // A remove callback that raced the insert may have run first and
+            // removed nothing, so check again now that the key is cached.
+            if in_flight_update.evicted() {
+                self.existence_cache.remove(&digest).await;
             }
         }
         result