[Router] Shard the cache-aware KV tree by chain root (#39167)

Co-authored-by: Kangyan Zhou <kangyan.zhou@radixark.ai>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>

experimental/sgl-router/tests/component/policies/kv_events_tree_concurrent.rs

@@ -3,7 +3,7 @@
 
 //! Concurrent-mutation stress test for `HashTree`.
 //!
-//! The 19 inline tests in `policies::kv_events::tree` are all
+//! The inline tests in `policies::kv_events::tree` are all
 //! single-threaded.  Under production load, multiple worker subscribers
 //! drive `insert` / `remove` / `clear_worker` against the same tree from
 //! tokio worker threads while the chat handler simultaneously calls
@@ -155,7 +155,7 @@ fn match_prefix_is_consistent_with_concurrent_clear() {
         // consistent.
         if m.matched_blocks == chain.len() {
             assert!(
-                m.workers().contains(&w),
+                m.holds(&w),
                 "full match must include worker; got {:?}",
                 m.workers(),
             );
@@ -165,3 +165,102 @@ fn match_prefix_is_consistent_with_concurrent_clear() {
     stop.store(true, std::sync::atomic::Ordering::Relaxed);
     mutator.join().unwrap();
 }
+
+/// Reader storm concurrent with a writer hammering insert/remove on
+/// DISTINCT chain roots — the pattern sharding targets. Asserts
+/// CORRECTNESS under that contention: warm chains are pre-inserted and
+/// never removed, so a reader must always get a full match with the warm
+/// worker present, and the writer's scratch chains are fully removed each
+/// round, so after join only the warm chains remain.
+///
+/// No sleeps, no wall-clock — the readers run a fixed number of bounded
+/// iterations and the writer churns until they are done, so nothing here
+/// can flake on timing.
+#[test]
+fn readers_unaffected_by_writer_on_distinct_roots() {
+    let tree = Arc::new(HashTree::new());
+
+    // Fixed chains the readers query and the writer never touches, on
+    // distinct roots (1_000 apart) so they spread across shards.
+    let warm = KvWorkerId {
+        url: "http://warm:30000".into(),
+        dp_rank: 0,
+    };
+    const WARM_CHAINS: i64 = 16;
+    let warm_chain = |c: i64| -> Vec<i64> { vec![c * 1_000, c * 1_000 + 1, c * 1_000 + 2] };
+    for c in 0..WARM_CHAINS {
+        tree.insert(&warm, None, &warm_chain(c));
+    }
+
+    let stop = Arc::new(std::sync::atomic::AtomicBool::new(false));
+
+    // Scratch chains on distinct roots far from the warm space, insert then
+    // immediate remove. Runs until the readers are done rather than for a
+    // fixed count the writer could burn through early, leaving the test with
+    // no contention at all. Every round is a complete pair, so stopping at
+    // any point leaves no residue.
+    let writer = {
+        let tree = tree.clone();
+        let stop = stop.clone();
+        thread::spawn(move || {
+            let scratch = KvWorkerId {
+                url: "http://scratch:30000".into(),
+                dp_rank: 0,
+            };
+            let mut round = 0i64;
+            while !stop.load(std::sync::atomic::Ordering::Relaxed) {
+                // Cycled so the arithmetic cannot drift into the warm space.
+                let base = 1_000_000 + (round % 100_000) * 7;
+                let chain = [base, base + 1, base + 2, base + 3];
+                tree.insert(&scratch, None, &chain);
+                tree.remove(&scratch, &chain);
+                round = round.wrapping_add(1);
+            }
+        })
+    };
+
+    // Each reader asserts the warm worker is present at full depth on every
+    // warm chain, regardless of writer churn.
+    let mut readers = Vec::new();
+    for _ in 0..4 {
+        let tree = tree.clone();
+        let warm = warm.clone();
+        readers.push(thread::spawn(move || {
+            for _ in 0..2_000 {
+                for c in 0..WARM_CHAINS {
+                    let chain = warm_chain(c);
+                    let m = tree.match_prefix(None, &chain);
+                    assert_eq!(
+                        m.matched_blocks,
+                        chain.len(),
+                        "warm chain {c} must always fully match despite writer churn",
+                    );
+                    assert!(
+                        m.holds(&warm),
+                        "warm worker must always hold its own untouched chain",
+                    );
+                }
+            }
+        }));
+    }
+
+    for r in readers {
+        r.join().expect("reader thread panicked under contention");
+    }
+    stop.store(true, std::sync::atomic::Ordering::Relaxed);
+    writer
+        .join()
+        .expect("writer thread panicked under contention");
+
+    // Only the warm chains remain: 16 chains x 3 nodes = 48 non-root nodes.
+    assert_eq!(
+        tree.node_count(),
+        (WARM_CHAINS * 3) as usize,
+        "writer's scratch churn must leave no residual nodes",
+    );
+    for c in 0..WARM_CHAINS {
+        let m = tree.match_prefix(None, &warm_chain(c));
+        assert_eq!(m.matched_blocks, 3);
+        assert!(m.holds(&warm));
+    }
+}

experimental/sgl-router/tests/component/policies/kv_events_two_subscribers.rs

@@ -84,8 +84,8 @@ async fn two_independent_subscribers_converge_to_same_tree_state() {
         let mb = router_b.tree().match_prefix(None, &hashes);
         let converged = ma.matched_blocks == target
             && mb.matched_blocks == target
-            && ma.workers().contains(&key)
-            && mb.workers().contains(&key);
+            && ma.holds(&key)
+            && mb.holds(&key);
         if converged {
             // Both trees agree on count AND on the worker that holds the
             // prefix. This is what the Radix Tree provider reads to
@@ -253,32 +253,32 @@ async fn two_subscribers_merge_events_from_two_publishers() {
             && ay.matched_blocks == target_y
             && bx.matched_blocks == target_x
             && by.matched_blocks == target_y
-            && ax.workers().contains(&key_x)
-            && ay.workers().contains(&key_y)
-            && bx.workers().contains(&key_x)
-            && by.workers().contains(&key_y);
+            && ax.holds(&key_x)
+            && ay.holds(&key_y)
+            && bx.holds(&key_x)
+            && by.holds(&key_y);
         if converged {
             // Negative attribution: prefix X must not be attributed to
             // worker_y in either tree, and vice versa. A regression that
             // keyed events by arriving socket rather than announced
             // worker URL would set BOTH worker keys on each prefix.
             assert!(
-                !ax.workers().contains(&key_y),
+                !ax.holds(&key_y),
                 "router_a cross-attributed worker_y to prefix X: {:?}",
                 ax.workers(),
             );
             assert!(
-                !ay.workers().contains(&key_x),
+                !ay.holds(&key_x),
                 "router_a cross-attributed worker_x to prefix Y: {:?}",
                 ay.workers(),
             );
             assert!(
-                !bx.workers().contains(&key_y),
+                !bx.holds(&key_y),
                 "router_b cross-attributed worker_y to prefix X: {:?}",
                 bx.workers(),
             );
             assert!(
-                !by.workers().contains(&key_x),
+                !by.holds(&key_x),
                 "router_b cross-attributed worker_x to prefix Y: {:?}",
                 by.workers(),
             );