feat: dump radix tree as router events (#2057)

Author: PeaBrane

lib/llm/src/kv_router/approx.rs

@@ -26,8 +26,8 @@ use tokio_util::sync::CancellationToken;
 use crate::tokens::TokenBlockSequence;
 
 use crate::kv_router::indexer::{
-    compute_block_hash_for_seq, KvIndexerInterface, KvRouterError, OverlapScores, RadixTree,
-    WorkerId,
+    compute_block_hash_for_seq, DumpRequest, KvIndexerInterface, KvRouterError, OverlapScores,
+    RadixTree, WorkerId,
 };
 use crate::kv_router::protocols::{
     ExternalSequenceBlockHash, KvCacheEvent, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData,
@@ -172,6 +172,8 @@ pub struct ApproxKvIndexer {
     route_tx: mpsc::Sender<RouterResult>,
     /// A sender for remove worker requests.
     remove_worker_tx: mpsc::Sender<WorkerId>,
+    /// A sender for dump requests.
+    dump_tx: mpsc::Sender<DumpRequest>,
     /// A handle to the background task managing the KV store.
     task: OnceLock<std::thread::JoinHandle<()>>,
     /// The size of the KV block this indexer can handle.
@@ -183,6 +185,7 @@ impl ApproxKvIndexer {
         let (match_tx, mut match_rx) = mpsc::channel::<MatchRequest>(2048);
         let (route_tx, mut route_rx) = mpsc::channel::<RouterResult>(2048);
         let (remove_worker_tx, mut remove_worker_rx) = mpsc::channel::<WorkerId>(16);
+        let (dump_tx, mut dump_rx) = mpsc::channel::<DumpRequest>(16);
         let cancel_clone = token.clone();
         let task = std::thread::spawn(move || {
             // create a new tokio runtime which will only perform work on a single thread
@@ -240,6 +243,10 @@ impl ApproxKvIndexer {
                         Some(worker) = remove_worker_rx.recv() => {
                             trie.remove_worker(worker);
                         }
+                        Some(dump_req) = dump_rx.recv() => {
+                            let events = trie.dump_tree_as_events();
+                            let _ = dump_req.resp.send(events);
+                        }
 
                         _ = expiry_fut => {
                             let expired = timer_manager.pop_expired();
@@ -278,6 +285,7 @@ impl ApproxKvIndexer {
             match_tx,
             route_tx,
             remove_worker_tx,
+            dump_tx,
             task: once,
             kv_block_size,
         }
@@ -355,6 +363,20 @@ impl KvIndexerInterface for ApproxKvIndexer {
         self.remove_worker_tx.send(worker).await.unwrap();
     }
 
+    async fn dump_events(&self) -> Result<Vec<RouterEvent>, KvRouterError> {
+        let (resp_tx, resp_rx) = oneshot::channel();
+        let dump_req = DumpRequest { resp: resp_tx };
+
+        if let Err(e) = self.dump_tx.send(dump_req).await {
+            tracing::error!("Failed to send dump request: {:?}", e);
+            return Err(KvRouterError::IndexerOffline);
+        }
+
+        resp_rx
+            .await
+            .map_err(|_| KvRouterError::IndexerDroppedRequest)
+    }
+
     fn shutdown(&mut self) {
         self.cancel.cancel();
         if let Some(task) = self.task.take() {