From 499c85f1318d5ad914a599050bd3f616a28007e0 Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Fri, 6 Dec 2024 11:26:07 -0800 Subject: [PATCH] [Router] remove duplicate char count (#2378) --- rust/py_test/test_launch_server.py | 2 + rust/src/server.rs | 4 +- rust/src/tree.rs | 90 +++++++++++++----------------- 3 files changed, 41 insertions(+), 55 deletions(-) diff --git a/rust/py_test/test_launch_server.py b/rust/py_test/test_launch_server.py index dcfe42346..0dacc2c9f 100644 --- a/rust/py_test/test_launch_server.py +++ b/rust/py_test/test_launch_server.py @@ -45,6 +45,8 @@ def popen_launch_router( port, "--dp", str(dp_size), # Convert dp_size to string + "--router-eviction-interval", + "5", # frequent eviction for testing ] # Use current environment diff --git a/rust/src/server.rs b/rust/src/server.rs index 269214acf..7197b9a27 100644 --- a/rust/src/server.rs +++ b/rust/src/server.rs @@ -1,8 +1,6 @@ use crate::router::PolicyConfig; use crate::router::Router; -use actix_web::{ - delete, get, post, put, web, App, HttpRequest, HttpResponse, HttpServer, Responder, -}; +use actix_web::{get, post, web, App, HttpRequest, HttpResponse, HttpServer, Responder}; use bytes::Bytes; use env_logger::Builder; use log::{info, LevelFilter}; diff --git a/rust/src/tree.rs b/rust/src/tree.rs index 3c403676f..1e39f02da 100644 --- a/rust/src/tree.rs +++ b/rust/src/tree.rs @@ -24,7 +24,6 @@ struct Node { #[derive(Debug)] pub struct Tree { root: NodeRef, - // TODO: Char Count per tenant pub tenant_char_count: DashMap, } @@ -408,17 +407,9 @@ impl Tree { pub fn evict_tenant_data(&self, max_size: usize) { // Calculate used size and collect leaves let mut stack = vec![Arc::clone(&self.root)]; - let mut used_size_per_tenant: HashMap = HashMap::new(); let mut pq = BinaryHeap::new(); while let Some(curr) = stack.pop() { - for tenant in curr.tenant_last_access_time.iter() { - let size = used_size_per_tenant - .entry(tenant.key().clone()) - .or_insert(0); - *size += curr.text.read().unwrap().chars().count(); - } - for child in curr.children.iter() { stack.push(Arc::clone(child.value())); } @@ -436,64 +427,59 @@ impl Tree { } info!("Before eviction - Used size per tenant:"); - for (tenant, size) in &used_size_per_tenant { - info!("Tenant: {}, Size: {}", tenant, size); + for entry in self.tenant_char_count.iter() { + info!("Tenant: {}, Size: {}", entry.key(), entry.value()); } // Process eviction while let Some(Reverse(entry)) = pq.pop() { let EvictionEntry { tenant, node, .. } = entry; - if let Some(&used_size) = used_size_per_tenant.get(&tenant) { - if used_size <= max_size { + if let Some(used_size) = self.tenant_char_count.get(&tenant) { + if *used_size <= max_size { continue; } + } - // Update used size - if let Some(size) = used_size_per_tenant.get_mut(&tenant) { - *size -= node.text.read().unwrap().chars().count(); - } - - // Decrement when removing tenant from node - if node.tenant_last_access_time.contains_key(&tenant) { - self.tenant_char_count - .entry(tenant.clone()) - .and_modify(|count| { - if *count > 0 { - *count -= node.text.read().unwrap().chars().count(); - } - }); - } - - // Remove tenant from node - node.tenant_last_access_time.remove(&tenant); - - // Remove empty nodes - if node.children.is_empty() && node.tenant_last_access_time.is_empty() { - if let Some(parent) = node.parent.write().unwrap().as_ref() { - let first_char = node.text.read().unwrap().chars().next().unwrap(); - parent.children.remove(&first_char); - } - } - - // Add parent to queue if it becomes a leaf - if let Some(parent) = node.parent.read().unwrap().as_ref() { - if Tree::leaf_of(parent).contains(&tenant) { - if let Some(timestamp) = parent.tenant_last_access_time.get(&tenant) { - pq.push(Reverse(EvictionEntry { - timestamp: *timestamp, - tenant: tenant.clone(), - node: Arc::clone(parent), - })); + // Decrement when removing tenant from node + if node.tenant_last_access_time.contains_key(&tenant) { + self.tenant_char_count + .entry(tenant.clone()) + .and_modify(|count| { + if *count > 0 { + *count -= node.text.read().unwrap().chars().count(); } - } + }); + } + + // Remove tenant from node + node.tenant_last_access_time.remove(&tenant); + + // Remove empty nodes + if node.children.is_empty() && node.tenant_last_access_time.is_empty() { + if let Some(parent) = node.parent.write().unwrap().as_ref() { + let first_char = node.text.read().unwrap().chars().next().unwrap(); + parent.children.remove(&first_char); } } + + // Add parent to queue if it becomes a leaf + if let Some(parent) = node.parent.read().unwrap().as_ref() { + if Tree::leaf_of(parent).contains(&tenant) { + if let Some(timestamp) = parent.tenant_last_access_time.get(&tenant) { + pq.push(Reverse(EvictionEntry { + timestamp: *timestamp, + tenant: tenant.clone(), + node: Arc::clone(parent), + })); + } + } + }; } info!("After eviction - Used size per tenant:"); - for (tenant, size) in &used_size_per_tenant { - info!("Tenant: {}, Size: {}", tenant, size); + for entry in self.tenant_char_count.iter() { + info!("Tenant: {}, Size: {}", entry.key(), entry.value()); } }