perf: keep walk workers off the shared heap lock

packages/python/e2e-attestations/bot-1478-perdir-walk-1020.json

@@ -0,0 +1,7 @@
+{
+  "command": "uv run python -m pytest tests/e2e/glob_speed_depth_bound_test.py tests/e2e/glob_speed_one_level_test.py tests/e2e/glob_speed_recursive_test.py tests/e2e/glob_speed_matchers_test.py tests/e2e/glob_speed_dotfiles_test.py tests/e2e/glob_speed_gitignore_test.py tests/e2e/glob_speed_node_modules_test.py tests/e2e/glob_speed_symlink_test.py tests/e2e/query_contract_test.py tests/e2e/fanout_test.py tests/e2e/use_case_join_test.py tests/e2e/stat_speed_test.py tests/e2e/release_speed_test.py -q",
+  "ran_at": 1792499019,
+  "exit_code": 0,
+  "commit": "9647c7ae1dab28acb9c211eced8f82624376b3b2",
+  "branch": "bot/1478-perdir-walk-1020"
+}

packages/rust/src/tree_walk.rs

@@ -55,11 +55,54 @@ struct Open<D, L> {
 }
 
 /// A found directory, ordered so the heap's greatest is the earliest key.
-struct Pending<D, L>(Arc<Slot<D, L>>);
+/// The key's leading bytes ride along as `rank`, so most comparisons never
+/// leave the heap's own memory: the heap is ordered under the lock every
+/// worker takes, and a comparison reaching into a slot another core wrote
+/// held that lock long enough to serialize the workers.
+struct Pending<D, L> {
+    rank: u128,
+    slot: Arc<Slot<D, L>>,
+}
+
+impl<D, L> Pending<D, L> {
+    fn new(slot: Arc<Slot<D, L>>) -> Self {
+        Pending {
+            rank: rank(&slot.key),
+            slot,
+        }
+    }
+}
+
+/// The first sixteen bytes of `key` encoded so that bytes order as keys do:
+/// each index plus one, in one byte below 0x80, two below 0x4000, else a
+/// 0xC0 byte and four more. No encoded index starts with a zero byte, so a
+/// key ranks below every key it is a prefix of.
+fn rank(key: &[u32]) -> u128 {
+    let mut bytes = [0u8; 16];
+    let mut at = 0;
+    for &index in key {
+        let [b0, b1, b2, b3] = index.saturating_add(1).to_be_bytes();
+        let encoded: &[u8] = if index < 0x7F {
+            &[b3]
+        } else if index < 0x3FFF {
+            &[0x80 + b2, b3]
+        } else {
+            &[0xC0, b0, b1, b2, b3]
+        };
+        for &byte in encoded {
+            let Some(slot) = bytes.get_mut(at) else {
+                return u128::from_be_bytes(bytes);
+            };
+            *slot = byte;
+            at += 1;
+        }
+    }
+    u128::from_be_bytes(bytes)
+}
 
 impl<D, L> PartialEq for Pending<D, L> {
     fn eq(&self, other: &Self) -> bool {
-        self.0.key == other.0.key
+        self.cmp(other) == Order::Equal
     }
 }
 
@@ -73,10 +116,16 @@ impl<D, L> PartialOrd for Pending<D, L> {
 
 impl<D, L> Ord for Pending<D, L> {
     fn cmp(&self, other: &Self) -> Order {
-        other.0.key.cmp(&self.0.key)
+        other
+            .rank
+            .cmp(&self.rank)
+            .then_with(|| other.slot.key.cmp(&self.slot.key))
     }
 }
 
+/// A directory's items, and the directories among them.
+type Found<D, L> = (Vec<Item<D, L>>, Vec<Arc<Slot<D, L>>>);
+
 struct Shared<'e, D, L> {
     open: Mutex<Open<D, L>>,
     found_dir: Condvar,
@@ -156,28 +205,38 @@ impl<D, L> Shared<'_, D, L> {
     }
 
     fn work(&self) {
-        while let Some(slot) = self.claim() {
+        let mut next = self.claim();
+        while let Some(slot) = next {
             let unclaimed = lock(&slot.dir).take();
-            let Some(dir) = unclaimed else {
-                continue;
-            };
-            let items = self.explore(dir, &slot.key);
-            let mut explored = lock(&slot.explored);
-            explored.items = Some(items);
-            if explored.awaited {
-                slot.ready.notify_one();
-            }
+            next = unclaimed
+                .and_then(|dir| self.explore_taking_next(&slot, dir))
+                .or_else(|| self.claim());
         }
     }
 
+    /// Explore `dir`, `slot`'s directory, for whoever awaits it, taking back
+    /// the next directory to explore if one is open. Registering what it
+    /// found and taking the next under one hold halves how often a worker
+    /// takes the lock every worker shares.
+    fn explore_taking_next(&self, slot: &Slot<D, L>, dir: D) -> Option<Arc<Slot<D, L>>> {
+        let (items, dirs) = self.explore_only(dir, &slot.key);
+        let next = self.register(dirs, true);
+        let mut explored = lock(&slot.explored);
+        explored.items = Some(items);
+        if explored.awaited {
+            slot.ready.notify_one();
+        }
+        next
+    }
+
     /// The next directory found, or `None` once the walk is over.
     fn claim(&self) -> Option<Arc<Slot<D, L>>> {
         let mut open = self.open_lock();
         loop {
             if open.done {
                 return None;
             }
-            if let Some(Pending(slot)) = open.slots.pop() {
+            if let Some(Pending { slot, .. }) = open.slots.pop() {
                 return Some(slot);
             }
             open = self
@@ -190,6 +249,31 @@ impl<D, L> Shared<'_, D, L> {
     /// Explore `dir`, at `key`, registering the directories it holds for any
     /// thread to claim.
     fn explore(&self, dir: D, key: &[u32]) -> Vec<Item<D, L>> {
+        let (items, dirs) = self.explore_only(dir, key);
+        if !dirs.is_empty() {
+            self.register(dirs, false);
+        }
+        items
+    }
+
+    /// Open `dirs` to any thread, taking the earliest open directory back
+    /// when `take` is set.
+    fn register(&self, dirs: Vec<Arc<Slot<D, L>>>, take: bool) -> Option<Arc<Slot<D, L>>> {
+        let found = dirs.len();
+        let mut open = self.open_lock();
+        open.slots.extend(dirs.into_iter().map(Pending::new));
+        let taken = take.then(|| open.slots.pop()).flatten();
+        drop(open);
+        // The thread that found them takes one next, so only the rest need
+        // a waiting worker.
+        for _ in 1..found {
+            self.found_dir.notify_one();
+        }
+        taken.map(|pending| pending.slot)
+    }
+
+    /// Explore `dir`, at `key`.
+    fn explore_only(&self, dir: D, key: &[u32]) -> Found<D, L> {
         let steps = (self.explore)(dir);
         let mut dirs = Vec::new();
         let items: Vec<Item<D, L>> = steps
@@ -212,18 +296,10 @@ impl<D, L> Shared<'_, D, L> {
                 }
             })
             .collect();
-        if dirs.is_empty() {
-            return items;
-        }
-        self.found.fetch_add(dirs.len(), Ordering::Relaxed);
-        let found = dirs.len();
-        self.open_lock().slots.extend(dirs.into_iter().map(Pending));
-        // The thread that found them takes one next, so only the rest need
-        // a waiting worker.
-        for _ in 1..found {
-            self.found_dir.notify_one();
+        if !dirs.is_empty() {
+            self.found.fetch_add(dirs.len(), Ordering::Relaxed);
         }
-        items
+        (items, dirs)
     }
 }
 
@@ -290,16 +366,42 @@ mod tests {
         threads.into_inner().unwrap().len()
     }
 
-    fn pending(key: &[u32]) -> Pending<(), ()> {
-        Pending(Arc::new(Slot {
+    fn slot(key: &[u32]) -> Arc<Slot<(), ()>> {
+        Arc::new(Slot {
             key: key.to_vec(),
             dir: Mutex::new(None),
             explored: Mutex::new(Explored {
                 items: None,
                 awaited: false,
             }),
             ready: Condvar::new(),
-        }))
+        })
+    }
+
+    fn pending(key: &[u32]) -> Pending<(), ()> {
+        Pending::new(slot(key))
+    }
+
+    #[test]
+    fn registering_opens_the_directories_and_may_take_back_the_earliest() {
+        let explore = |(): ()| Vec::new();
+        let shared = Shared {
+            open: Mutex::new(Open {
+                slots: BinaryHeap::new(),
+                done: false,
+            }),
+            found_dir: Condvar::new(),
+            found: AtomicUsize::new(0),
+            explore: &explore,
+        };
+        assert!(
+            shared
+                .register(vec![slot(&[1]), slot(&[0])], false)
+                .is_none()
+        );
+        let taken = shared.register(vec![slot(&[2])], true);
+        assert_eq!(taken.map(|slot| slot.key.clone()), Some(vec![0]));
+        assert_eq!(shared.open_lock().slots.len(), 2);
     }
 
     #[test]
@@ -315,11 +417,36 @@ mod tests {
             open.push(pending(key));
         }
         let order: Vec<Vec<u32>> = std::iter::from_fn(|| open.pop())
-            .map(|Pending(slot)| slot.key.clone())
+            .map(|pending| pending.slot.key.clone())
             .collect();
         assert_eq!(order, vec![vec![0], vec![0, 1, 5], vec![0, 2], vec![1]]);
     }
 
+    #[test]
+    fn ranks_order_as_their_keys_do() {
+        let boundaries = [0, 1, 0x7E, 0x7F, 0x80, 0x3FFE, 0x3FFF, 0x4000, u32::MAX - 1];
+        let mut keys: Vec<Vec<u32>> = vec![vec![]];
+        for &a in &boundaries {
+            keys.push(vec![a]);
+            for &b in &boundaries {
+                keys.push(vec![a, b]);
+            }
+        }
+        keys.sort();
+        let ranks: Vec<u128> = keys.iter().map(|key| rank(key)).collect();
+        assert!(ranks.windows(2).all(|pair| pair[0] < pair[1]));
+    }
+
+    #[test]
+    fn keys_past_sixteen_encoded_bytes_are_ordered_by_the_whole_key() {
+        let deep = [0x4000; 4];
+        let mut open = BinaryHeap::new();
+        for last in [2, 1] {
+            open.push(pending(&[&deep[..], &[last]].concat()));
+        }
+        assert_eq!(open.pop().unwrap().slot.key.last(), Some(&1));
+    }
+
     #[test]
     fn a_walk_of_few_directories_stays_on_the_calling_thread() {
         assert_eq!(threads_exploring(3, 1), 1);

packages/ts/e2e-attestations/bot-1478-perdir-walk-1020.json

@@ -0,0 +1,7 @@
+{
+  "command": "pnpm test:e2e",
+  "ran_at": 1792499101,
+  "exit_code": 0,
+  "commit": "861a18d5fe6be842da12b40d85e8331a9ef849e5",
+  "branch": "bot/1478-perdir-walk-1020"
+}