Skip to content

Commit 227846f

Browse files
PeaBranePeaBrane
authored
chore: use Rc to do reference counting in Router slot manager (#3545)
Signed-off-by: PeaBrane <[email protected]>
1 parent b94ecd1 commit 227846f

File tree

2 files changed

+182
-198
lines changed

2 files changed

+182
-198
lines changed

‎lib/llm/src/kv_router/indexer.rs‎

Lines changed: 63 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -801,57 +801,50 @@ impl KvIndexer {
801801
let cancel_clone = token.clone();
802802

803803
let task = std::thread::spawn(move || {
804-
// create a new tokio runtime which will only perform work on a single thread
805-
let runtime = tokio::runtime::Builder::new_multi_thread()
806-
.worker_threads(1) // Single-threaded environment
804+
// Create a single-threaded tokio runtime
805+
let runtime = tokio::runtime::Builder::new_current_thread()
807806
.enable_all()
808807
.build()
809808
.unwrap();
810809

811-
let local_set = tokio::task::LocalSet::new();
812-
813-
runtime.block_on(local_set.run_until(async move {
814-
tokio::task::spawn_local(async move {
815-
let cancel = cancel_clone;
816-
let mut match_rx = match_rx;
817-
let mut event_rx = event_rx;
818-
let mut remove_worker_rx = remove_worker_rx;
819-
let mut dump_rx = dump_rx;
820-
let mut trie = RadixTree::new_with_frequency(expiration_duration);
821-
loop {
822-
tokio::select! {
823-
biased;
824-
825-
_ = cancel.cancelled() => {
826-
tracing::debug!("KvCacheIndexer progress loop shutting down");
827-
return;
828-
}
810+
runtime.block_on(async move {
811+
let cancel = cancel_clone;
812+
let mut match_rx = match_rx;
813+
let mut event_rx = event_rx;
814+
let mut remove_worker_rx = remove_worker_rx;
815+
let mut dump_rx = dump_rx;
816+
let mut trie = RadixTree::new_with_frequency(expiration_duration);
817+
loop {
818+
tokio::select! {
819+
biased;
820+
821+
_ = cancel.cancelled() => {
822+
tracing::debug!("KvCacheIndexer progress loop shutting down");
823+
return;
824+
}
829825

830-
Some(worker) = remove_worker_rx.recv() => {
831-
trie.remove_worker(worker);
832-
}
826+
Some(worker) = remove_worker_rx.recv() => {
827+
trie.remove_worker(worker);
828+
}
833829

834-
Some(event) = event_rx.recv() => {
835-
let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
836-
let result = trie.apply_event(event);
837-
metrics.increment_event_applied(event_type, result);
838-
}
830+
Some(event) = event_rx.recv() => {
831+
let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
832+
let result = trie.apply_event(event);
833+
metrics.increment_event_applied(event_type, result);
834+
}
839835

840-
Some(dump_req) = dump_rx.recv() => {
841-
let events = trie.dump_tree_as_events();
842-
let _ = dump_req.resp.send(events);
843-
}
836+
Some(dump_req) = dump_rx.recv() => {
837+
let events = trie.dump_tree_as_events();
838+
let _ = dump_req.resp.send(events);
839+
}
844840

845-
Some(req) = match_rx.recv() => {
846-
let matches = trie.find_matches(req.sequence, req.early_exit);
847-
let _ = req.resp.send(matches);
848-
}
841+
Some(req) = match_rx.recv() => {
842+
let matches = trie.find_matches(req.sequence, req.early_exit);
843+
let _ = req.resp.send(matches);
849844
}
850845
}
851-
})
852-
.await
853-
.unwrap()
854-
}));
846+
}
847+
});
855848

856849
tracing::debug!("KvCacheIndexer task completed");
857850
});
@@ -1058,54 +1051,47 @@ impl KvIndexerSharded {
10581051
remove_worker_tx.push(shard_remove_worker_tx);
10591052
dump_tx.push(shard_dump_tx); // Store dump sender
10601053

1061-
let runtime = tokio::runtime::Builder::new_multi_thread()
1062-
.worker_threads(1)
1054+
let runtime = tokio::runtime::Builder::new_current_thread()
10631055
.enable_all()
10641056
.build()
10651057
.unwrap();
10661058

10671059
tasks.push(std::thread::spawn(move || {
1068-
let local_set = tokio::task::LocalSet::new();
1069-
1070-
runtime.block_on(local_set.run_until(async move {
1071-
tokio::task::spawn_local(async move {
1072-
let mut trie = RadixTree::new_with_frequency(expiration_duration);
1073-
loop {
1074-
tokio::select! {
1075-
biased;
1076-
1077-
_ = cancel.cancelled() => {
1078-
tracing::trace!("KvCacheIndexer progress loop shutting down");
1079-
return;
1080-
}
1060+
runtime.block_on(async move {
1061+
let mut trie = RadixTree::new_with_frequency(expiration_duration);
1062+
loop {
1063+
tokio::select! {
1064+
biased;
10811065

1082-
Some(worker) = shard_remove_worker_rx.recv() => {
1083-
trie.remove_worker(worker);
1084-
}
1066+
_ = cancel.cancelled() => {
1067+
tracing::trace!("KvCacheIndexer progress loop shutting down");
1068+
return;
1069+
}
10851070

1086-
Some(event) = shard_event_rx.recv() => {
1087-
let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
1088-
let result = trie.apply_event(event);
1089-
metrics.increment_event_applied(event_type, result);
1090-
}
1071+
Some(worker) = shard_remove_worker_rx.recv() => {
1072+
trie.remove_worker(worker);
1073+
}
10911074

1092-
Some(dump_req) = shard_dump_rx.recv() => {
1093-
let events = trie.dump_tree_as_events();
1094-
let _ = dump_req.resp.send(events);
1095-
}
1075+
Some(event) = shard_event_rx.recv() => {
1076+
let event_type = KvIndexerMetrics::get_event_type(&event.event.data);
1077+
let result = trie.apply_event(event);
1078+
metrics.increment_event_applied(event_type, result);
1079+
}
10961080

1097-
Ok(req) = shard_broadcast_rx.recv() => {
1098-
let matches = trie.find_matches(req.sequence, req.early_exit);
1099-
if let Err(e) = req.resp.send(matches).await {
1100-
tracing::trace!("Failed to send match response: {:?}", e);
1101-
}
1081+
Some(dump_req) = shard_dump_rx.recv() => {
1082+
let events = trie.dump_tree_as_events();
1083+
let _ = dump_req.resp.send(events);
1084+
}
1085+
1086+
Ok(req) = shard_broadcast_rx.recv() => {
1087+
let matches = trie.find_matches(req.sequence, req.early_exit);
1088+
if let Err(e) = req.resp.send(matches).await {
1089+
tracing::trace!("Failed to send match response: {:?}", e);
11021090
}
11031091
}
11041092
}
1105-
})
1106-
.await
1107-
.unwrap()
1108-
}));
1093+
}
1094+
});
11091095

11101096
tracing::debug!("KvCacheIndexer task completed");
11111097
}));

0 commit comments

Comments
 (0)