diff --git a/rust/src/router.rs b/rust/src/router.rs index 74e47209b..2b6b8d52c 100644 --- a/rust/src/router.rs +++ b/rust/src/router.rs @@ -168,7 +168,7 @@ impl Router { let locked_tree_clone = tree_clone.lock().unwrap(); // Run eviction - locked_tree_clone.evict_tenant_data(max_tree_size); + locked_tree_clone.evict_tenant_by_size(max_tree_size); // Print the process queue let locked_processed_queue = processed_queue_clone.lock().unwrap(); diff --git a/rust/src/tree.rs b/rust/src/tree.rs index 1e39f02da..e8dc8b7a0 100644 --- a/rust/src/tree.rs +++ b/rust/src/tree.rs @@ -5,6 +5,7 @@ use log::info; use std::cmp::Reverse; use std::collections::BinaryHeap; use std::collections::HashMap; +use std::collections::VecDeque; use std::sync::Arc; use std::sync::RwLock; @@ -404,7 +405,7 @@ impl Tree { .collect() } - pub fn evict_tenant_data(&self, max_size: usize) { + pub fn evict_tenant_by_size(&self, max_size: usize) { // Calculate used size and collect leaves let mut stack = vec![Arc::clone(&self.root)]; let mut pq = BinaryHeap::new(); @@ -483,6 +484,46 @@ impl Tree { } } + pub fn remove_tenant(&self, tenant: &str) { + // 1. Find all the leaves for the tenant + let mut stack = vec![Arc::clone(&self.root)]; + let mut queue = VecDeque::new(); + + while let Some(curr) = stack.pop() { + for child in curr.children.iter() { + stack.push(Arc::clone(child.value())); + } + + if Tree::leaf_of(&curr).contains(&tenant.to_string()) { + queue.push_back(Arc::clone(&curr)); + } + } + + // 2. Start from the leaves and traverse up to the root, removing the tenant from each node + while let Some(curr) = queue.pop_front() { + // remove tenant from node + curr.tenant_last_access_time.remove(&tenant.to_string()); + + // remove empty nodes + if curr.children.is_empty() && curr.tenant_last_access_time.is_empty() { + if let Some(parent) = curr.parent.read().unwrap().as_ref() { + let first_char = curr.text.read().unwrap().chars().next().unwrap(); + parent.children.remove(&first_char); + } + } + + // add parent to queue if it becomes a leaf + if let Some(parent) = curr.parent.read().unwrap().as_ref() { + if Tree::leaf_of(parent).contains(&tenant.to_string()) { + queue.push_back(Arc::clone(&parent)); + } + } + } + + // 3. Remove the tenant from the tenant_char_count map + self.tenant_char_count.remove(&tenant.to_string()); + } + pub fn get_tenant_char_count(&self) -> HashMap { self.tenant_char_count .iter() @@ -673,7 +714,7 @@ mod tests { ); // Test eviction - tree.evict_tenant_data(3); // This should evict tenants with more than 3 chars + tree.evict_tenant_by_size(3); // This should evict tenants with more than 3 chars let post_eviction_smallest = tree.get_smallest_tenant(); println!("Smallest tenant after eviction: {}", post_eviction_smallest); @@ -754,7 +795,7 @@ mod tests { ); // Phase 4: Eviction test - tree.evict_tenant_data(10); + tree.evict_tenant_by_size(10); let computed_sizes = tree.get_used_size_per_tenant(); let maintained_counts: HashMap = tree @@ -1132,7 +1173,7 @@ mod tests { assert_eq!(sizes_before.get("tenant2").unwrap(), &10); // "hello" + "world" = 10 // Evict - should remove "hello" from tenant2 as it's the oldest - tree.evict_tenant_data(max_size); + tree.evict_tenant_by_size(max_size); tree.pretty_print(); @@ -1168,7 +1209,7 @@ mod tests { } // Perform eviction - tree.evict_tenant_data(max_size); + tree.evict_tenant_by_size(max_size); // Check sizes after eviction let sizes_after = tree.get_used_size_per_tenant(); @@ -1200,7 +1241,7 @@ mod tests { let handle = thread::spawn(move || { while start_time.elapsed() < test_duration { // Run eviction - tree.evict_tenant_data(max_size); + tree.evict_tenant_by_size(max_size); // Sleep for 5 seconds thread::sleep(Duration::from_secs(5)); @@ -1245,7 +1286,7 @@ mod tests { } // final eviction - tree.evict_tenant_data(max_size); + tree.evict_tenant_by_size(max_size); // Final size check let final_sizes = tree.get_used_size_per_tenant(); @@ -1352,4 +1393,91 @@ mod tests { assert_eq!(tree.prefix_match_tenant("hello", "tenant3"), ""); // Non-existent tenant assert_eq!(tree.prefix_match_tenant("help", "tenant3"), ""); // Non-existent tenant } + + #[test] + fn test_simple_tenant_eviction() { + let tree = Tree::new(); + + // Insert data for multiple tenants + tree.insert("hello", "tenant1"); + tree.insert("world", "tenant1"); + tree.insert("hello", "tenant2"); + tree.insert("help", "tenant2"); + + // Verify initial state + let initial_sizes = tree.get_used_size_per_tenant(); + assert_eq!(initial_sizes.get("tenant1").unwrap(), &10); // "hello" + "world" + assert_eq!(initial_sizes.get("tenant2").unwrap(), &6); // "hello" + "p" + + // Evict tenant1 + tree.remove_tenant("tenant1"); + + // Verify after eviction + let final_sizes = tree.get_used_size_per_tenant(); + assert!( + !final_sizes.contains_key("tenant1"), + "tenant1 should be completely removed" + ); + assert_eq!( + final_sizes.get("tenant2").unwrap(), + &6, + "tenant2 should be unaffected" + ); + + // Verify tenant1's data is inaccessible + assert_eq!(tree.prefix_match_tenant("hello", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("world", "tenant1"), ""); + + // Verify tenant2's data is still accessible + assert_eq!(tree.prefix_match_tenant("hello", "tenant2"), "hello"); + assert_eq!(tree.prefix_match_tenant("help", "tenant2"), "help"); + } + + #[test] + fn test_complex_tenant_eviction() { + let tree = Tree::new(); + + // Create a more complex tree structure with shared prefixes + tree.insert("apple", "tenant1"); + tree.insert("application", "tenant1"); + tree.insert("apple", "tenant2"); + tree.insert("appetite", "tenant2"); + tree.insert("banana", "tenant1"); + tree.insert("banana", "tenant2"); + tree.insert("ball", "tenant2"); + + // Verify initial state + let initial_sizes = tree.get_used_size_per_tenant(); + println!("Initial sizes: {:?}", initial_sizes); + tree.pretty_print(); + + // Evict tenant1 + tree.remove_tenant("tenant1"); + + // Verify final state + let final_sizes = tree.get_used_size_per_tenant(); + println!("Final sizes: {:?}", final_sizes); + tree.pretty_print(); + + // Verify tenant1 is completely removed + assert!( + !final_sizes.contains_key("tenant1"), + "tenant1 should be completely removed" + ); + + // Verify all tenant1's data is inaccessible + assert_eq!(tree.prefix_match_tenant("apple", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("application", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("banana", "tenant1"), ""); + + // Verify tenant2's data is intact + assert_eq!(tree.prefix_match_tenant("apple", "tenant2"), "apple"); + assert_eq!(tree.prefix_match_tenant("appetite", "tenant2"), "appetite"); + assert_eq!(tree.prefix_match_tenant("banana", "tenant2"), "banana"); + assert_eq!(tree.prefix_match_tenant("ball", "tenant2"), "ball"); + + // Verify the tree structure is still valid for tenant2 + let tenant2_size = final_sizes.get("tenant2").unwrap(); + assert_eq!(tenant2_size, &(5 + 5 + 6 + 2)); // "apple" + "etite" + "banana" + "ll" + } }