Explicitly separate state locks and awaits

Author: palfrey

.bazelrc

@@ -102,13 +102,16 @@ build --@rules_rust//:clippy_flag=-Wclippy::nursery
 build --@rules_rust//:clippy_flag=-Wclippy::pedantic
 build --@rules_rust//:clippy_flag=-Dclippy::alloc_instead_of_core
 build --@rules_rust//:clippy_flag=-Dclippy::as_underscore
+build --@rules_rust//:clippy_flag=-Dclippy::await_holding_lock
 build --@rules_rust//:clippy_flag=-Wclippy::dbg_macro
 build --@rules_rust//:clippy_flag=-Wclippy::decimal_literal_representation
 build --@rules_rust//:clippy_flag=-Dclippy::elidable_lifetime_names
+build --@rules_rust//:clippy_flag=-Dclippy::explicit_into_iter_loop
 build --@rules_rust//:clippy_flag=-Aclippy::get_unwrap
 build --@rules_rust//:clippy_flag=-Dclippy::missing_const_for_fn
 build --@rules_rust//:clippy_flag=-Aclippy::missing_docs_in_private_items
 build --@rules_rust//:clippy_flag=-Wclippy::print_stdout
+build --@rules_rust//:clippy_flag=-Dclippy::redundant_closure_for_method_calls
 build --@rules_rust//:clippy_flag=-Dclippy::semicolon_if_nothing_returned
 build --@rules_rust//:clippy_flag=-Dclippy::std_instead_of_core
 build --@rules_rust//:clippy_flag=-Dclippy::todo

Cargo.toml

@@ -139,7 +139,10 @@ pedantic = { level = "warn", priority = -1 }
 # Restriction Denies with default priority
 alloc-instead-of-core = "deny"
 as-underscore = "deny"
+await-holding-lock = "deny"
 elidable-lifetime-names = "deny"
+explicit-into-iter-loop = "deny"
+redundant-closure-for-method-calls = "deny"
 semicolon-if-nothing-returned = "deny"
 std-instead-of-core = "deny"
 todo = "deny"

nativelink-store/src/ref_store.rs

@@ -88,7 +88,7 @@ impl RefStore {
             .err_tip(|| "Store manager is gone")?;
         if let Some(store) = store_manager.get_store(&self.name) {
             let remove_callbacks = self.remove_callbacks.lock().clone();
-            for callback in remove_callbacks.into_iter() {
+            for callback in remove_callbacks {
                 store.register_remove_callback(callback)?;
             }
             unsafe {

nativelink-util/src/evicting_map.rs

@@ -383,49 +383,51 @@ where
         // to be able to borrow a `Q`.
         R: Borrow<Q> + Send,
     {
-        let mut state = self.state.lock();
+        let (mut callbacks, data_to_unref) = {
+            let mut state = self.state.lock();
 
-        let lru_len = state.lru.len();
-        let mut data_to_unref = Vec::new();
-        let mut removal_futures = Vec::new();
-        for (key, result) in keys.into_iter().zip(results.iter_mut()) {
-            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 self.should_evict(lru_len, entry, 0, u64::MAX) {
-                        *result = None;
-                        if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) {
-                            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.into_iter());
+            let lru_len = state.lru.len();
+            let mut data_to_unref = Vec::new();
+            let mut removal_futures = Vec::new();
+            for (key, result) in keys.into_iter().zip(results.iter_mut()) {
+                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 self.should_evict(lru_len, entry, 0, u64::MAX) {
+                            *result = None;
+                            if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) {
+                                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.into_iter());
+                            }
+                        } else {
+                            if !peek {
+                                entry.seconds_since_anchor =
+                                    self.anchor_time.elapsed().as_secs() as i32;
+                            }
+                            *result = Some(entry.data.len());
                         }
-                    } else {
-                        if !peek {
-                            entry.seconds_since_anchor =
-                                self.anchor_time.elapsed().as_secs() as i32;
-                        }
-                        *result = Some(entry.data.len());
                     }
+                    None => *result = None,
                 }
-                None => *result = None,
             }
-        }
+            let callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
+            (callbacks, data_to_unref)
+        };
 
-        // Drop the state and perform the async callbacks.
-        drop(state);
-        let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
         while callbacks.next().await.is_some() {}
         let mut callbacks: FuturesUnordered<_> =
-            data_to_unref.iter().map(|item| item.unref()).collect();
+            data_to_unref.iter().map(LenEntry::unref).collect();
         while callbacks.next().await.is_some() {}
     }
 
@@ -455,7 +457,7 @@ where
             let mut callbacks: FuturesUnordered<_> = removal_futures.into_iter().collect();
             while callbacks.next().await.is_some() {}
             let mut callbacks: FuturesUnordered<_> =
-                items_to_unref.iter().map(|item| item.unref()).collect();
+                items_to_unref.iter().map(LenEntry::unref).collect();
             while callbacks.next().await.is_some() {}
         }
 
@@ -567,11 +569,11 @@ where
 
         // Perform eviction after all insertions
         let (items_to_unref, futures) = self.evict_items(state);
-        removal_futures.extend(futures.into_iter());
+        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.into_iter());
+        replaced_items.extend(items_to_unref);
 
         (replaced_items, removal_futures)
     }
@@ -600,7 +602,7 @@ where
 
         // Unref evicted items outside of lock
         let mut callbacks: FuturesUnordered<_> =
-            items_to_unref.iter().map(|item| item.unref()).collect();
+            items_to_unref.iter().map(LenEntry::unref).collect();
         while callbacks.next().await.is_some() {}
 
         // Unref removed item if any
@@ -618,41 +620,44 @@ where
     where
         F: FnOnce(&T) -> bool + Send,
     {
-        let mut state = self.state.lock();
-        if let Some(entry) = state.lru.get(key.borrow()) {
-            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.into_iter());
-                Some(item)
+        let (evicted_items, mut removal_futures, removed_item) = {
+            let mut state = self.state.lock();
+            if let Some(entry) = state.lru.get(key.borrow()) {
+                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.into_iter());
+                    Some(item)
+                } else {
+                    None
+                };
+
+                let removal_futures: FuturesUnordered<Pin<Box<dyn Future<Output = ()> + Send>>> =
+                    removal_futures.into_iter().collect();
+
+                (evicted_items, removal_futures, removed_item)
             } else {
-                None
-            };
-
-            // Drop the lock before unref operations
-            drop(state);
+                (vec![], vec![].into_iter().collect(), None)
+            }
+        };
 
-            let mut removal_futures: FuturesUnordered<_> = removal_futures.into_iter().collect();
-            while removal_futures.next().await.is_some() {}
+        while removal_futures.next().await.is_some() {}
 
-            // Unref evicted items
-            let mut callbacks: FuturesUnordered<_> =
-                evicted_items.iter().map(|item| item.unref()).collect();
-            while callbacks.next().await.is_some() {}
+        // Unref evicted items
+        let mut callbacks: FuturesUnordered<_> =
+            evicted_items.iter().map(LenEntry::unref).collect();
+        while callbacks.next().await.is_some() {}
 
-            // Unref removed item if any
-            if let Some(item) = removed_item {
-                item.unref().await;
-                true
-            } else {
-                false
-            }
+        // Unref removed item if any
+        if let Some(item) = removed_item {
+            item.unref().await;
+            true
         } else {
             false
         }