Fix evicting map get lock and fast slow store EOF

Author: amankrx

nativelink-store/src/fast_slow_store.rs

@@ -396,7 +396,7 @@ impl StoreDriver for FastSlowStore {
             return self.slow_store.update(key, reader, size_info).await;
         }
 
-        let (mut fast_tx, fast_rx) = make_buf_channel_pair();
+        let (fast_tx, fast_rx) = make_buf_channel_pair();
         let (mut slow_tx, slow_rx) = make_buf_channel_pair();
 
         let key_debug = format!("{key:?}");
@@ -408,16 +408,17 @@ impl StoreDriver for FastSlowStore {
         let mut bytes_sent: u64 = 0;
 
         let data_stream_fut = async move {
+            let mut fast_tx = Some(fast_tx);
             loop {
                 let buffer = reader
                     .recv()
                     .await
                     .err_tip(|| "Failed to read buffer in fastslow store")?;
                 if buffer.is_empty() {
                     // EOF received.
-                    fast_tx.send_eof().err_tip(
-                        || "Failed to write eof to fast store in fast_slow store update",
-                    )?;
+                    if let Some(mut ftx) = fast_tx.take() {
+                        drop(ftx.send_eof());
+                    }
                     slow_tx
                         .send_eof()
                         .err_tip(|| "Failed to write eof to writer in fast_slow store update")?;
@@ -429,34 +430,43 @@ impl StoreDriver for FastSlowStore {
                 }
 
                 let chunk_len = buffer.len();
-                let send_start = std::time::Instant::now();
-                let (fast_result, slow_result) =
-                    join!(fast_tx.send(buffer.clone()), slow_tx.send(buffer));
-                let send_elapsed = send_start.elapsed();
-                if send_elapsed.as_secs() >= 5 {
-                    warn!(
-                        chunk_len,
-                        send_elapsed_ms = send_elapsed.as_millis(),
-                        total_bytes = bytes_sent,
-                        "FastSlowStore::update: channel send stalled (>5s). A downstream store may be hanging",
-                    );
-                }
-                bytes_sent += u64::try_from(chunk_len).unwrap_or(u64::MAX);
-                fast_result
-                    .map_err(|e| {
+                if let Some(ref mut ftx) = fast_tx {
+                    let send_start = std::time::Instant::now();
+                    let (fast_result, slow_result) =
+                        join!(ftx.send(buffer.clone()), slow_tx.send(buffer));
+                    let send_elapsed = send_start.elapsed();
+                    if send_elapsed.as_secs() >= 5 {
+                        warn!(
+                            chunk_len,
+                            send_elapsed_ms = send_elapsed.as_millis(),
+                            total_bytes = bytes_sent,
+                            "FastSlowStore::update: channel send stalled (>5s). A downstream store may be hanging",
+                        );
+                    }
+                    if fast_result.is_err() {
+                        warn!(
+                            total_bytes = bytes_sent,
+                            "FastSlowStore::update: fast store channel failed, continuing with slow store only",
+                        );
+                        fast_tx = None;
+                    }
+                    slow_result.map_err(|e| {
                         make_err!(
                             Code::Internal,
-                            "Failed to send message to fast_store in fast_slow_store {:?}",
+                            "Failed to send message to slow_store in fast_slow store {:?}",
                             e
                         )
-                    })
-                    .merge(slow_result.map_err(|e| {
+                    })?;
+                } else {
+                    slow_tx.send(buffer).await.map_err(|e| {
                         make_err!(
                             Code::Internal,
                             "Failed to send message to slow_store in fast_slow store {:?}",
                             e
                         )
-                    }))?;
+                    })?;
+                }
+                bytes_sent += u64::try_from(chunk_len).unwrap_or(u64::MAX);
             }
         };
 
@@ -483,7 +493,15 @@ impl StoreDriver for FastSlowStore {
                 "FastSlowStore::update: completed successfully",
             );
         }
-        data_stream_res.merge(fast_res).merge(slow_res)?;
+        // Slow store success is required; fast store failure is tolerated since it's a cache.
+        data_stream_res.merge(slow_res)?;
+        if let Err(err) = fast_res {
+            warn!(
+                ?err,
+                key = %key_debug,
+                "FastSlowStore::update: fast store failed during upload; data stored in slow store",
+            );
+        }
         Ok(())
     }
 
@@ -591,16 +609,23 @@ impl StoreDriver for FastSlowStore {
         // TODO(palfrey) Investigate if we should maybe ignore errors here instead of
         // forwarding them up.
         if self.fast_store.has(key.borrow()).await?.is_some() {
-            self.metrics
-                .fast_store_hit_count
-                .fetch_add(1, Ordering::Acquire);
-            self.fast_store
-                .get_part(key, writer.borrow_mut(), offset, length)
-                .await?;
-            self.metrics
-                .fast_store_downloaded_bytes
-                .fetch_add(writer.get_bytes_written(), Ordering::Acquire);
-            return Ok(());
+            match self
+                .fast_store
+                .get_part(key.borrow(), writer.borrow_mut(), offset, length)
+                .await
+            {
+                Ok(()) => {
+                    self.metrics
+                        .fast_store_hit_count
+                        .fetch_add(1, Ordering::Acquire);
+                    self.metrics
+                        .fast_store_downloaded_bytes
+                        .fetch_add(writer.get_bytes_written(), Ordering::Acquire);
+                    return Ok(());
+                }
+                Err(err) if err.code == Code::NotFound => {}
+                Err(err) => return Err(err),
+            }
         }
 
         // If the fast store is noop or read only or update only then bypass it.

nativelink-store/src/filesystem_store.rs

@@ -816,7 +816,8 @@ impl<Fe: FileEntry> FilesystemStore<Fe> {
             // it still exists in there. But first, get the lock...
             let mut encoded_file_path = entry.get_encoded_file_path().write().await;
             // Then check it's still in there...
-            if evicting_map.get(&key).await.is_none() {
+            // Use size_for_key instead of get() to avoid triggering bulk eviction of other entries.
+            if evicting_map.size_for_key(&key).await.is_none() {
                 info!(%key, "Got eviction while emplacing, dropping");
                 return Ok(());
             }

nativelink-util/src/evicting_map.rs

@@ -435,41 +435,50 @@ where
     }
 
     pub async fn get(&self, key: &Q) -> Option<T> {
-        // Fast path: Check if we need eviction before acquiring lock for eviction
-        let needs_eviction = {
-            let state = self.state.lock();
-            if let Some((_, peek_entry)) = state.lru.peek_lru() {
-                self.should_evict(
-                    state.lru.len(),
-                    peek_entry,
-                    state.sum_store_size,
-                    self.max_bytes,
-                )
+        let (result, items_to_unref, removal_futures) = {
+            let mut state = self.state.lock();
+            // Check if the requested item is expired before promoting it.
+            let result = if let Some(entry) = state.lru.peek(key.borrow()) {
+                if self.should_evict(state.lru.len(), entry, 0, u64::MAX) {
+                    // Item is expired, remove it.
+                    if let Some((k, eviction_item)) = state.lru.pop_entry(key.borrow()) {
+                        let (data, futures) = state.remove(k.borrow(), &eviction_item, false);
+                        let (mut items, mut removals) = self.evict_items(&mut *state);
+                        items.push(data);
+                        removals.extend(futures);
+                        return {
+                            Self::unref_items(items, removals).await;
+                            None
+                        };
+                    }
+                    None
+                } else {
+                    // Item is valid. Promote it in LRU so it's safe from eviction.
+                    state.lru.get_mut(key.borrow()).map(|entry| {
+                        entry.seconds_since_anchor =
+                            i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX);
+                        entry.data.clone()
+                    })
+                }
             } else {
-                false
-            }
-        };
-
-        // Perform eviction if needed
-        if needs_eviction {
-            let (items_to_unref, removal_futures) = {
-                let mut state = self.state.lock();
-                self.evict_items(&mut *state)
+                None
             };
-            // Unref items outside of lock
-            let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
-            while callbacks.next().await.is_some() {}
-            let mut callbacks: FuturesUnordered<_> =
-                items_to_unref.iter().map(LenEntry::unref).collect();
-            while callbacks.next().await.is_some() {}
-        }
+            // Evict other items if needed
+            let (items_to_unref, removal_futures) = self.evict_items(&mut *state);
+            (result, items_to_unref, removal_futures)
+        };
+        // Unref items outside of lock
+        Self::unref_items(items_to_unref, removal_futures).await;
+        result
+    }
 
-        // Now get the item
-        let mut state = self.state.lock();
-        let entry = state.lru.get_mut(key.borrow())?;
-        entry.seconds_since_anchor =
-            i32::try_from(self.anchor_time.elapsed().as_secs()).unwrap_or(i32::MAX);
-        Some(entry.data.clone())
+    /// Helper to unref evicted items outside of lock.
+    async fn unref_items(items_to_unref: Vec<T>, removal_futures: Vec<RemoveFuture>) {
+        let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
+        while callbacks.next().await.is_some() {}
+        let mut callbacks: FuturesUnordered<_> =
+            items_to_unref.iter().map(LenEntry::unref).collect();
+        while callbacks.next().await.is_some() {}
     }
 
     /// Returns the replaced item if any.