[model-gateway] optimize radix tree memory and reduce allocations (#15933)

sgl-model-gateway/src/policies/cache_aware.rs

@@ -327,15 +327,16 @@ impl LoadBalancingPolicy for CacheAwarePolicy {
 
         if let Some(tree) = tree {
             // Now we work with the tree without holding the HashMap lock
-            let (matched_text, matched_worker) = tree.prefix_match(text);
+            // Use prefix_match_with_counts to avoid redundant chars().count() calls
-            let match_rate = if text.is_empty() {
+            let result = tree.prefix_match_with_counts(text);
+            let match_rate = if result.input_char_count == 0 {
                 0.0
             } else {
-                matched_text.chars().count() as f32 / text.chars().count() as f32
+                result.matched_char_count as f32 / result.input_char_count as f32
             };
 
             let selected_url = if match_rate > self.config.cache_threshold {
-                matched_worker.to_string()
+                result.tenant
             } else {
                 let min_load_idx = *healthy_indices
                     .iter()

sgl-model-gateway/src/policies/mod.rs

@@ -17,7 +17,6 @@ mod random;
 mod registry;
 mod round_robin;
 pub mod tree;
-
 pub use bucket::BucketPolicy;
 pub use cache_aware::CacheAwarePolicy;
 pub use consistent_hashing::ConsistentHashingPolicy;
@@ -27,6 +26,7 @@ pub use power_of_two::PowerOfTwoPolicy;
 pub use random::RandomPolicy;
 pub use registry::PolicyRegistry;
 pub use round_robin::RoundRobinPolicy;
+pub use tree::PrefixMatchResult;
 
 /// Core trait for load balancing policies
 ///

sgl-model-gateway/src/policies/tree.rs

@@ -18,6 +18,19 @@ type NodeRef = Arc<Node>;
 /// Using Arc<str> allows cheap cloning and comparison.
 pub type TenantId = Arc<str>;
 
+/// Result of a prefix match operation, including char counts to avoid recomputation.
+#[derive(Debug, Clone)]
+pub struct PrefixMatchResult {
+    /// The matched prefix text
+    pub matched_text: String,
+    /// The tenant that owns the matched prefix
+    pub tenant: String,
+    /// Number of characters matched (avoids chars().count())
+    pub matched_char_count: usize,
+    /// Total number of characters in the input text
+    pub input_char_count: usize,
+}
+
 /// A fast identity hasher for single-character keys (used in children DashMap).
 /// Since chars have good distribution already, we use identity hashing with mixing.
 #[derive(Default)]
@@ -53,34 +66,30 @@ impl Hasher for CharHasher {
 
 type CharHasherBuilder = BuildHasherDefault<CharHasher>;
 
-/// Pre-indexed text for efficient character access.
+/// Advance a string slice by N characters, returning the remaining slice.
-/// Converts UTF-8 string to Vec<char> once to enable O(1) indexing.
+/// Returns empty string if n >= char count.
-struct CharIndexedText {
+#[inline]
-    chars: Vec<char>,
+fn advance_by_chars(s: &str, n: usize) -> &str {
+    if n == 0 {
+        return s;
+    }
+    s.char_indices()
+        .nth(n)
+        .map(|(idx, _)| &s[idx..])
+        .unwrap_or("")
 }
 
-impl CharIndexedText {
+/// Get the first N characters of a string as a new String.
-    #[inline]
+/// More efficient than chars().take(n).collect() for known bounds.
-    fn new(text: &str) -> Self {
+#[inline]
-        Self {
+fn take_chars(s: &str, n: usize) -> String {
-            chars: text.chars().collect(),
+    if n == 0 {
-        }
+        return String::new();
-    }
-
-    #[inline]
-    fn len(&self) -> usize {
-        self.chars.len()
-    }
-
-    #[inline]
-    fn get(&self, idx: usize) -> Option<char> {
-        self.chars.get(idx).copied()
-    }
-
-    #[inline]
-    fn slice_to_string(&self, start: usize, end: usize) -> String {
-        self.chars[start..end].iter().collect()
     }
+    s.char_indices()
+        .nth(n)
+        .map(|(idx, _)| s[..idx].to_string())
+        .unwrap_or_else(|| s.to_string())
 }
 
 /// Node text with cached character count to avoid repeated O(n) chars().count() calls.
@@ -251,23 +260,15 @@ impl PartialEq for EvictionEntry {
 // Note that in rust, `.len()` or slice is operated on the "byte" level. It causes issues for UTF-8 characters because one character might use multiple bytes.
 // https://en.wikipedia.org/wiki/UTF-8
 
-/// Efficient shared prefix count using pre-indexed chars for O(1) access.
+/// Count matching prefix characters between two strings.
-/// Returns the number of characters that match between `a` (starting at `a_start`) and `b`.
+/// Returns the number of characters that match from the start.
+/// Uses iterator-based comparison - no allocation required.
 #[inline]
-fn shared_prefix_count_indexed(a: &CharIndexedText, a_start: usize, b: &str) -> usize {
+fn shared_prefix_count(a: &str, b: &str) -> usize {
-    let mut i = 0;
+    a.chars()
-    let mut b_iter = b.chars();
+        .zip(b.chars())
-
+        .take_while(|(a_char, b_char)| a_char == b_char)
-    while a_start + i < a.len() {
+        .count()
-        match (a.get(a_start + i), b_iter.next()) {
-            (Some(a_char), Some(b_char)) if a_char == b_char => {
-                i += 1;
-            }
-            _ => break,
-        }
-    }
-
-    i
 }
 
 /// Intern a tenant string into an Arc<str> for efficient storage and comparison.
@@ -311,12 +312,7 @@ impl Tree {
 
     pub fn insert(&self, text: &str, tenant: &str) {
         // Insert text into tree with given tenant
-        // Pre-index text once for O(1) character access (avoids O(n²) chars().nth() calls)
+        // Use slice-based traversal to avoid Vec<char> allocation
-        let indexed_text = CharIndexedText::new(text);
-        let text_count = indexed_text.len();
-
-        let mut curr = Arc::clone(&self.root);
-        let mut curr_idx = 0;
 
         // Use cached timestamp to reduce syscalls
         let timestamp_ms = get_timestamp_ms();
@@ -324,84 +320,72 @@ impl Tree {
         // Intern the tenant ID once for reuse
         let tenant_id = intern_tenant(tenant);
 
-        curr.tenant_last_access_time
+        // Pre-compute total char count once (needed for size tracking)
+        let total_char_count = text.chars().count();
+
+        self.root
+            .tenant_last_access_time
             .insert(Arc::clone(&tenant_id), timestamp_ms);
 
         self.tenant_char_count
             .entry(Arc::clone(&tenant_id))
             .or_insert(0);
 
+        // Track remaining text as a slice - no allocation needed
+        let mut remaining = text;
+        let mut matched_chars = 0;
         let mut prev = Arc::clone(&self.root);
 
-        while curr_idx < text_count {
+        // Result type to carry state out of the match block
-            // O(1) character access instead of O(n) chars().nth()
+        // This allows the entry guard to be dropped before we update prev
-            let first_char = indexed_text.get(curr_idx).unwrap();
+        enum InsertStep {
+            Done,
+            Continue {
+                next_prev: NodeRef,
+                advance_chars: usize,
+            },
+        }
 
-            curr = prev;
+        while !remaining.is_empty() {
+            let first_char = remaining.chars().next().unwrap();
 
-            // dashmap.entry locks the entry until the op is done
+            // Use entry API for atomic check-and-insert semantics (required for thread safety)
-            // if using contains_key + insert, there will be an issue that
+            let step = match prev.children.entry(first_char) {
-            // 1. "apple" and "app" entered at the same time
-            // 2. and get inserted to the dashmap concurrently, so only one is inserted
-
-            match curr.children.entry(first_char) {
                 Entry::Vacant(entry) => {
-                    /*
+                    // No match - create new node with remaining text
-                       no matched
+                    let remaining_char_count = total_char_count - matched_chars;
-                       [curr]
-                       becomes
-                       [curr] => [new node]
-                    */
-
-                    // Use indexed slice for efficient string extraction
-                    let curr_text = indexed_text.slice_to_string(curr_idx, text_count);
-                    let curr_text_count = text_count - curr_idx;
                     let new_node = Arc::new(Node {
                         children: DashMap::with_hasher(CharHasherBuilder::default()),
-                        text: RwLock::new(NodeText::new(curr_text)),
+                        text: RwLock::new(NodeText::new(remaining.to_string())),
                         tenant_last_access_time: DashMap::new(),
-                        parent: RwLock::new(Some(Arc::clone(&curr))),
+                        parent: RwLock::new(Some(Arc::clone(&prev))),
                     });
 
-                    // Attach tenant to the new node (map is empty here) and increment count once
+                    // Attach tenant to the new node and increment count
                     self.tenant_char_count
                         .entry(Arc::clone(&tenant_id))
-                        .and_modify(|count| *count += curr_text_count)
+                        .and_modify(|count| *count += remaining_char_count)
-                        .or_insert(curr_text_count);
+                        .or_insert(remaining_char_count);
                     new_node
                         .tenant_last_access_time
                         .insert(Arc::clone(&tenant_id), timestamp_ms);
 
-                    entry.insert(Arc::clone(&new_node));
+                    entry.insert(new_node);
-
+                    InsertStep::Done
-                    prev = Arc::clone(&new_node);
-                    curr_idx = text_count;
                 }
 
                 Entry::Occupied(mut entry) => {
-                    // matched
                     let matched_node = entry.get().clone();
 
                     let matched_node_text = matched_node.text.read().unwrap();
-                    // Use cached char count instead of chars().count()
                     let matched_node_text_count = matched_node_text.char_count();
+                    let matched_node_text_str = matched_node_text.as_str();
 
-                    // Use indexed comparison to avoid creating intermediate string
+                    // Use slice-based comparison - no allocation
-                    let shared_count = shared_prefix_count_indexed(
+                    let shared_count = shared_prefix_count(remaining, matched_node_text_str);
-                        &indexed_text,
-                        curr_idx,
-                        matched_node_text.as_str(),
-                    );
 
                     if shared_count < matched_node_text_count {
-                        /*
+                        // Split the matched node
-                           split the matched node
-                           [curr] -> [matched_node] =>
-                           becomes
-                           [curr] -> [new_node] -> [contracted_matched_node]
-                        */
-
-                        // Use split_at_char for efficient splitting with cached counts
                         let (matched_text, contracted_text) =
                             matched_node_text.split_at_char(shared_count);
                         let matched_text_count = shared_count;
@@ -412,7 +396,7 @@ impl Tree {
                         let new_node = Arc::new(Node {
                             text: RwLock::new(matched_text),
                             children: DashMap::with_hasher(CharHasherBuilder::default()),
-                            parent: RwLock::new(Some(Arc::clone(&curr))),
+                            parent: RwLock::new(Some(Arc::clone(&prev))),
                             tenant_last_access_time: matched_node.tenant_last_access_time.clone(),
                         });
 
@@ -426,10 +410,11 @@ impl Tree {
                         *matched_node.text.write().unwrap() = contracted_text;
                         *matched_node.parent.write().unwrap() = Some(Arc::clone(&new_node));
 
-                        prev = Arc::clone(&new_node);
+                        // Atomically attach tenant to the new split node
-
+                        match new_node
-                        // Atomically attach tenant to the new split node and increment count once
+                            .tenant_last_access_time
-                        match prev.tenant_last_access_time.entry(Arc::clone(&tenant_id)) {
+                            .entry(Arc::clone(&tenant_id))
+                        {
                             Entry::Vacant(v) => {
                                 self.tenant_char_count
                                     .entry(Arc::clone(&tenant_id))
@@ -442,16 +427,19 @@ impl Tree {
                             }
                         }
 
-                        curr_idx += shared_count;
+                        InsertStep::Continue {
+                            next_prev: new_node,
+                            advance_chars: shared_count,
+                        }
                     } else {
-                        // move to next node
+                        // Full match - move to next node
-                        // Drop read lock before continuing
                         drop(matched_node_text);
 
-                        prev = Arc::clone(&matched_node);
+                        // Atomically attach tenant to existing node
-
+                        match matched_node
-                        // Atomically attach tenant to existing node and increment count once
+                            .tenant_last_access_time
-                        match prev.tenant_last_access_time.entry(Arc::clone(&tenant_id)) {
+                            .entry(Arc::clone(&tenant_id))
+                        {
                             Entry::Vacant(v) => {
                                 self.tenant_char_count
                                     .entry(Arc::clone(&tenant_id))
@@ -463,51 +451,62 @@ impl Tree {
                                 o.insert(timestamp_ms);
                             }
                         }
-                        curr_idx += shared_count;
+
+                        InsertStep::Continue {
+                            next_prev: matched_node,
+                            advance_chars: shared_count,
                         }
                     }
                 }
+            };
+
+            // Entry guard is now dropped - safe to update prev
+            match step {
+                InsertStep::Done => break,
+                InsertStep::Continue {
+                    next_prev,
+                    advance_chars,
+                } => {
+                    prev = next_prev;
+                    remaining = advance_by_chars(remaining, advance_chars);
+                    matched_chars += advance_chars;
+                }
+            }
         }
     }
 
-    #[allow(unused_assignments)]
+    /// Performs prefix matching and returns detailed result with char counts.
-    pub fn prefix_match(&self, text: &str) -> (String, String) {
+    /// This is the optimized version that avoids redundant chars().count() calls.
-        // Pre-index text once for O(1) character access
+    pub fn prefix_match_with_counts(&self, text: &str) -> PrefixMatchResult {
-        let indexed_text = CharIndexedText::new(text);
+        // Use slice-based traversal - no Vec<char> allocation
-        let text_count = indexed_text.len();
+        let input_char_count = text.chars().count();
-
-        let mut curr = Arc::clone(&self.root);
-        let mut curr_idx = 0;
 
+        let mut remaining = text;
+        let mut matched_chars = 0;
         let mut prev = Arc::clone(&self.root);
 
-        while curr_idx < text_count {
+        while !remaining.is_empty() {
-            // O(1) character access instead of O(n) chars().nth()
+            let first_char = remaining.chars().next().unwrap();
-            let first_char = indexed_text.get(curr_idx).unwrap();
 
-            curr = prev.clone();
+            let child_node = prev.children.get(&first_char).map(|e| e.value().clone());
 
-            if let Some(entry) = curr.children.get(&first_char) {
+            if let Some(matched_node) = child_node {
-                let matched_node = entry.value().clone();
                 let matched_text_guard = matched_node.text.read().unwrap();
-                // Use indexed comparison to avoid creating intermediate string
-                let shared_count = shared_prefix_count_indexed(
-                    &indexed_text,
-                    curr_idx,
-                    matched_text_guard.as_str(),
-                );
-                // Use cached char count instead of chars().count()
                 let matched_node_text_count = matched_text_guard.char_count();
+
+                // Use slice-based comparison - no allocation
+                let shared_count = shared_prefix_count(remaining, matched_text_guard.as_str());
                 drop(matched_text_guard);
 
                 if shared_count == matched_node_text_count {
                     // Full match with current node's text, continue to next node
-                    curr_idx += shared_count;
+                    matched_chars += shared_count;
-                    prev = Arc::clone(&matched_node);
+                    remaining = advance_by_chars(remaining, shared_count);
+                    prev = matched_node;
                 } else {
-                    // Partial match, stop here
+                    // Partial match - still use this node for tenant selection
-                    curr_idx += shared_count;
+                    matched_chars += shared_count;
-                    prev = Arc::clone(&matched_node);
+                    prev = matched_node;
                     break;
                 }
             } else {
@@ -516,9 +515,9 @@ impl Tree {
             }
         }
 
-        curr = prev.clone();
+        let curr = prev;
 
-        // Select the first tenant (key in the map) - use Arc<str> directly
+        // Select the first tenant (key in the map)
         let tenant: Option<TenantId> = curr
             .tenant_last_access_time
             .iter()
@@ -530,7 +529,7 @@ impl Tree {
 
         // Traverse from the curr node to the root and update the timestamp
         if let Some(ref tenant_id) = tenant {
-            let mut current_node = Some(curr);
+            let mut current_node = Some(Arc::clone(&curr));
             while let Some(node) = current_node {
                 node.tenant_last_access_time
                     .insert(Arc::clone(tenant_id), timestamp_ms);
@@ -538,39 +537,45 @@ impl Tree {
             }
         }
 
-        // Use indexed slice for result
+        // Build matched text from original input using char count
-        let ret_text = indexed_text.slice_to_string(0, curr_idx);
+        let matched_text = take_chars(text, matched_chars);
         let tenant_str = tenant
             .map(|t| t.to_string())
             .unwrap_or_else(|| "empty".to_string());
-        (ret_text, tenant_str)
+
+        PrefixMatchResult {
+            matched_text,
+            tenant: tenant_str,
+            matched_char_count: matched_chars,
+            input_char_count,
+        }
     }
 
-    #[allow(unused_assignments, dead_code)]
+    /// Legacy prefix_match API for backward compatibility.
+    /// Prefer prefix_match_with_counts() for better performance.
+    pub fn prefix_match(&self, text: &str) -> (String, String) {
+        let result = self.prefix_match_with_counts(text);
+        (result.matched_text, result.tenant)
+    }
+
+    #[allow(dead_code)]
     pub fn prefix_match_tenant(&self, text: &str, tenant: &str) -> String {
-        // Pre-index text once for O(1) character access
+        // Use slice-based traversal - no Vec<char> allocation
-        let indexed_text = CharIndexedText::new(text);
-        let text_count = indexed_text.len();
 
         // Intern tenant ID once for efficient lookups
         let tenant_id = intern_tenant(tenant);
 
-        let mut curr = Arc::clone(&self.root);
+        let mut remaining = text;
-        let mut curr_idx = 0;
+        let mut matched_chars = 0;
-
         let mut prev = Arc::clone(&self.root);
 
-        while curr_idx < text_count {
+        while !remaining.is_empty() {
-            // O(1) character access instead of O(n) chars().nth()
+            let first_char = remaining.chars().next().unwrap();
-            let first_char = indexed_text.get(curr_idx).unwrap();
 
-            curr = prev.clone();
+            let child_node = prev.children.get(&first_char).map(|e| e.value().clone());
-
-            if let Some(entry) = curr.children.get(&first_char) {
-                let matched_node = entry.value().clone();
 
+            if let Some(matched_node) = child_node {
                 // Only continue matching if this node belongs to the specified tenant
-                // Note: contains_key with &str works because Arc<str> implements Borrow<str>
                 if !matched_node
                     .tenant_last_access_time
                     .contains_key(tenant_id.as_ref())
@@ -579,24 +584,21 @@ impl Tree {
                 }
 
                 let matched_text_guard = matched_node.text.read().unwrap();
-                // Use indexed comparison to avoid creating intermediate string
-                let shared_count = shared_prefix_count_indexed(
-                    &indexed_text,
-                    curr_idx,
-                    matched_text_guard.as_str(),
-                );
-                // Use cached char count instead of chars().count()
                 let matched_node_text_count = matched_text_guard.char_count();
+
+                // Use slice-based comparison - no allocation
+                let shared_count = shared_prefix_count(remaining, matched_text_guard.as_str());
                 drop(matched_text_guard);
 
                 if shared_count == matched_node_text_count {
                     // Full match with current node's text, continue to next node
-                    curr_idx += shared_count;
+                    matched_chars += shared_count;
-                    prev = Arc::clone(&matched_node);
+                    remaining = advance_by_chars(remaining, shared_count);
+                    prev = matched_node;
                 } else {
-                    // Partial match, stop here
+                    // Partial match - still use this node for timestamp update
-                    curr_idx += shared_count;
+                    matched_chars += shared_count;
-                    prev = Arc::clone(&matched_node);
+                    prev = matched_node;
                     break;
                 }
             } else {
@@ -605,7 +607,7 @@ impl Tree {
             }
         }
 
-        curr = prev.clone();
+        let curr = prev;
 
         // Only update timestamp if we found a match for the specified tenant
         if curr
@@ -623,15 +625,13 @@ impl Tree {
             }
         }
 
-        // Use indexed slice for result
+        // Build result from original input using char count
-        indexed_text.slice_to_string(0, curr_idx)
+        take_chars(text, matched_chars)
     }
 
+    /// Return the list of tenants for which this node is a leaf.
+    /// A tenant is a leaf at this node if no children have that tenant.
     fn leaf_of(node: &NodeRef) -> Vec<TenantId> {
-        /*
-        Return the list of tenants if it's a leaf for the tenant.
-        A tenant is a "leaf" at this node if this node has the tenant but none of its children do.
-         */
         let mut candidates: HashMap<TenantId, bool> = node
             .tenant_last_access_time
             .iter()
@@ -689,35 +689,57 @@ impl Tree {
                 }
             }
 
+            // Verify this node is still a leaf for this tenant (may have changed)
+            // A node is a leaf for a tenant if no children have that tenant
+            let is_still_leaf = node.tenant_last_access_time.contains_key(tenant.as_ref())
+                && !node.children.iter().any(|child| {
+                    child
+                        .value()
+                        .tenant_last_access_time
+                        .contains_key(tenant.as_ref())
+                });
+            if !is_still_leaf {
+                continue;
+            }
+
             // Decrement when removing tenant from node
-            if node.tenant_last_access_time.contains_key(tenant.as_ref()) {
-                // Use cached char count instead of chars().count()
             let node_len = node.text.read().unwrap().char_count();
             self.tenant_char_count
                 .entry(Arc::clone(&tenant))
                 .and_modify(|count| {
                     *count = count.saturating_sub(node_len);
                 });
-            }
 
             // Remove tenant from node
             node.tenant_last_access_time.remove(tenant.as_ref());
 
+            // Get parent reference outside of the borrow scope
+            let parent_opt = node.parent.read().unwrap().clone();
+
             // Remove empty nodes
             if node.children.is_empty() && node.tenant_last_access_time.is_empty() {
-                if let Some(parent) = node.parent.read().unwrap().as_ref() {
+                if let Some(ref parent) = parent_opt {
-                    let text_guard = node.text.read().unwrap();
+                    if let Some(fc) = node.text.read().unwrap().first_char() {
-                    if let Some(first_char) = text_guard.first_char() {
+                        parent.children.remove(&fc);
-                        parent.children.remove(&first_char);
                     }
                 }
             }
 
-            // Add parent to queue if it becomes a leaf
+            // If parent has this tenant and no other children have it,
-            if let Some(parent) = node.parent.read().unwrap().as_ref() {
+            // parent becomes a new leaf - add to priority queue
-                let parent_leaves = Tree::leaf_of(parent);
+            if let Some(ref parent) = parent_opt {
-                if parent_leaves.iter().any(|t| t.as_ref() == tenant.as_ref()) {
+                if parent.tenant_last_access_time.contains_key(tenant.as_ref()) {
-                    if let Some(timestamp) = parent.tenant_last_access_time.get(tenant.as_ref()) {
+                    let has_child_with_tenant = parent.children.iter().any(|child| {
+                        child
+                            .value()
+                            .tenant_last_access_time
+                            .contains_key(tenant.as_ref())
+                    });
+
+                    if !has_child_with_tenant {
+                        // Add parent to priority queue as new leaf
+                        if let Some(timestamp) = parent.tenant_last_access_time.get(tenant.as_ref())
+                        {
                             pq.push(Reverse(EvictionEntry {
                                 timestamp: *timestamp,
                                 tenant: Arc::clone(&tenant),
@@ -725,7 +747,8 @@ impl Tree {
                             }));
                         }
                     }
-            };
+                }
+            }
         }
 
         debug!("After eviction - Used size per tenant:");
@@ -739,6 +762,7 @@ impl Tree {
         let tenant_id = intern_tenant(tenant);
 
         // 1. Find all the leaves for the tenant
+        // A leaf is a node that has this tenant but no children have it
         let mut stack = vec![Arc::clone(&self.root)];
         let mut queue = VecDeque::new();
 
@@ -747,38 +771,60 @@ impl Tree {
                 stack.push(Arc::clone(child.value()));
             }
 
-            let leaves = Tree::leaf_of(&curr);
+            // Check if this node is a leaf for the tenant
-            if leaves.iter().any(|t| t.as_ref() == tenant_id.as_ref()) {
+            if curr
+                .tenant_last_access_time
+                .contains_key(tenant_id.as_ref())
+            {
+                let has_child_with_tenant = curr.children.iter().any(|child| {
+                    child
+                        .value()
+                        .tenant_last_access_time
+                        .contains_key(tenant_id.as_ref())
+                });
+                if !has_child_with_tenant {
                     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
+            // Remove tenant from node
             curr.tenant_last_access_time.remove(tenant_id.as_ref());
 
-            // remove empty nodes
+            // Get parent reference outside of the borrow scope
+            let parent_opt = curr.parent.read().unwrap().clone();
+
+            // 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() {
+                if let Some(ref parent) = parent_opt {
-                    let text_guard = curr.text.read().unwrap();
+                    if let Some(fc) = curr.text.read().unwrap().first_char() {
-                    if let Some(first_char) = text_guard.first_char() {
+                        parent.children.remove(&fc);
-                        parent.children.remove(&first_char);
                     }
                 }
             }
 
-            // add parent to queue if it becomes a leaf
+            // If parent has this tenant and no other children have it,
-            if let Some(parent) = curr.parent.read().unwrap().as_ref() {
+            // parent becomes a new leaf - add to queue
-                let parent_leaves = Tree::leaf_of(parent);
+            if let Some(ref parent) = parent_opt {
-                if parent_leaves
+                if parent
-                    .iter()
+                    .tenant_last_access_time
-                    .any(|t| t.as_ref() == tenant_id.as_ref())
+                    .contains_key(tenant_id.as_ref())
                 {
+                    let has_child_with_tenant = parent.children.iter().any(|child| {
+                        child
+                            .value()
+                            .tenant_last_access_time
+                            .contains_key(tenant_id.as_ref())
+                    });
+
+                    if !has_child_with_tenant {
                         queue.push_back(Arc::clone(parent));
                     }
                 }
             }
+        }
 
         // 3. Remove the tenant from the tenant_char_count map
         self.tenant_char_count.remove(tenant_id.as_ref());