From c6a565816fcfa6d7882073a3bc84abb313ac99b3 Mon Sep 17 00:00:00 2001
From: Bharath Kumarasubramanian <bkumarasubramanian@linkedin.com>
Date: Wed, 4 Mar 2026 18:39:45 -0800
Subject: [PATCH] Fix RangeIterator O(n) drain and query_next_batch routing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- RangeIterator: replace Vec::remove(0) with cursor index — O(1) per
  element instead of O(n) shift. Eliminates ~500K element moves when
  draining a 1000-entry batch.
- Btree::query_next_batch(): fix inverted is_sweep_query routing —
  sweep queries now continue with sweep_query_internal and traversal
  queries continue with traversal_query_internal (was swapped).

 ┌─────────────┬────────────┬───────────────────────┬─────────────────────┐
 │  Scan Size  │ Before Fix │       After Fix       │ vs RocksDB PT+noWAL │
 ├─────────────┼────────────┼───────────────────────┼─────────────────────┤
 │ 100 keys    │ 20.0 us    │ 12.2 us (1.6x faster) │ 1.9x slower         │
 ├─────────────┼────────────┼───────────────────────┼─────────────────────┤
 │ 1,000 keys  │ 276 us     │ 152 us (1.8x faster)  │ 2.5x slower         │
 ├─────────────┼────────────┼───────────────────────┼─────────────────────┤
 │ 5,000 keys  │ 1,380 us   │ 454 us (3.0x faster)  │ 1.5x slower         │
 ├─────────────┼────────────┼───────────────────────┼─────────────────────┤
 │ 10,000 keys │ 2,810 us   │ 454 us (6.2x faster)  │ HomeDb 1.3x faster! │
 └─────────────┴────────────┴───────────────────────┴─────────────────────┘
---
 src/homedb/src/iterator.rs         | 16 +++++++++++-----
 src/homestore/index/btree/btree.rs |  8 ++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/homedb/src/iterator.rs b/src/homedb/src/iterator.rs
index ee879274f..24fea37d3 100644
--- a/src/homedb/src/iterator.rs
+++ b/src/homedb/src/iterator.rs
@@ -18,6 +18,7 @@ pub struct RangeIterator {
     btree: Arc<Btree<DbKey, DbValue>>,
     handle: Option<QueryResultHandle<'static, DbKey, DbValue>>,
     current_batch: Vec<(DbKey, DbValue)>,
+    cursor: usize,
     start_key: Vec<u8>,
     end_key: Vec<u8>,
     batch_size: u32,
@@ -47,6 +48,7 @@ impl RangeIterator {
             btree,
             handle: if has_more { Some(handle) } else { None },
             current_batch: results,
+            cursor: 0,
             start_key,
             end_key,
             batch_size,
@@ -64,13 +66,16 @@ impl RangeIterator {
     #[maybe_async_cfg::maybe(keep_self, sync(feature = "sync_frontend"), async(feature = "async_frontend"))]
     pub async fn next(&mut self) -> Result<Option<(Vec<u8>, Vec<u8>)>> {
         loop {
-            // Try to get from current batch
-            if !self.current_batch.is_empty() {
-                let (key, value) = self.current_batch.remove(0);
-                return Ok(Some((key.into_vec(), value.into_vec())));
+            if self.cursor < self.current_batch.len() {
+                let (key, value) = &self.current_batch[self.cursor];
+                let result = (key.clone().into_vec(), value.clone().into_vec());
+                self.cursor += 1;
+                return Ok(Some(result));
             }
 
-            // Current batch exhausted, try to fetch next batch
+            // Batch exhausted — drop it and fetch the next one.
+            self.current_batch.clear();
+            self.cursor = 0;
             match self.handle.take() {
                 Some(handle) if handle.has_more() => {
                     let next_handle = self.btree
@@ -104,6 +109,7 @@ impl RangeIterator {
     pub async fn seek(&mut self, key: &[u8]) -> Result<bool> {
         self.handle = None;
         self.current_batch.clear();
+        self.cursor = 0;
 
         let (range_start, range_end) = if self.reverse {
             (
diff --git a/src/homestore/index/btree/btree.rs b/src/homestore/index/btree/btree.rs
index dd2e0b04b..6c7d0cc38 100644
--- a/src/homestore/index/btree/btree.rs
+++ b/src/homestore/index/btree/btree.rs
@@ -510,11 +510,11 @@ where
     ) -> Result<QueryResultHandle<K, V>, BtreeError> {
         let req = handle.request();
         if req.is_sweep_query() {
-            // Use traversal query for reverse iteration
-            self.traversal_query_internal(req).await
-        } else {
-            // Use sweep query for forward iteration
+            // Sweep query continues with sweep (forward, sibling-link based)
             self.sweep_query_internal(req).await
+        } else {
+            // Traversal query continues with traversal (supports reverse, no sibling links)
+            self.traversal_query_internal(req).await
         }
     }
 }