fix: fix range bugs

GreptimeTeam · Oct 18, 2023 · f012bde · f012bde
1 parent 80fae1e
commit f012bde
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 16 deletions.
diff --git a/src/common/meta/src/kv_backend/memory.rs b/src/common/meta/src/kv_backend/memory.rs
@@ -90,16 +90,16 @@ impl<T: ErrorExt + Send + Sync + 'static> KvBackend for MemoryKvBackend<T> {
         } = req;
 
         let kvs = self.kvs.read().unwrap();
-        let iter = kvs.range(range);
+        let values = kvs.range(range);
 
         let mut more = false;
-        let mut took = 0;
+        let mut iter: i64 = 0;
 
-        let kvs = iter
+        let kvs = values
             .take_while(|_| {
-                let take = limit == 0 || took != limit;
-                took += 1;
-                more = limit > 0 && took > limit;
+                let take = limit == 0 || iter != limit;
+                iter += 1;
+                more = limit > 0 && iter > limit;
 
                 take
             })

diff --git a/src/common/meta/src/kv_backend/test.rs b/src/common/meta/src/kv_backend/test.rs
@@ -195,6 +195,18 @@ pub async fn test_kv_range_2(kv_store: impl KvBackend) {
     assert_eq!(result.kvs.len(), 1);
     assert!(result.more);
 
+    // Fetches the keys >= "a", set limit to 2, the `more` should be false.
+    let result = kv_store
+        .range(
+            RangeRequest::new()
+                .with_range(b"a".to_vec(), b"\0".to_vec())
+                .with_limit(2),
+        )
+        .await
+        .unwrap();
+    assert_eq!(result.kvs.len(), 2);
+    assert!(!result.more);
+
     // Fetches the keys >= "a", set limit to 3, the `more` should be false.
     let result = kv_store
         .range(

diff --git a/src/log-store/src/raft_engine/backend.rs b/src/log-store/src/raft_engine/backend.rs
@@ -152,6 +152,7 @@ impl KvBackend for RaftEngineBackend {
             _ => unreachable!(),
         };
         let mut more = false;
+        let mut iter = 0;
 
         self.engine
             .read()
@@ -162,17 +163,18 @@ impl KvBackend for RaftEngineBackend {
                 end_key.as_deref(),
                 false,
                 |key, value| {
-                    res.push(KeyValue {
-                        key: key.to_vec(),
-                        value: if keys_only { vec![] } else { value.to_vec() },
-                    });
-                    if limit > 0 && limit as usize == res.len() {
-                        more = true;
-                        false
-                    } else {
-                        // continue
-                        true
+                    let take = limit == 0 || iter != limit;
+                    iter += 1;
+                    more = limit > 0 && iter > limit;
+
+                    if take {
+                        res.push(KeyValue {
+                            key: key.to_vec(),
+                            value: if keys_only { vec![] } else { value.to_vec() },
+                        });
                     }
+
+                    take
                 },
             )
             .context(RaftEngineSnafu)