fix: correct the range behavior in MemoryKvBackend & RaftEngineBackend

src/common/meta/src/key.rs

@@ -258,7 +258,6 @@ impl<T: Serialize + DeserializeOwned> DeserializedValueWithBytes<T> {
         self.bytes.to_vec()
     }
 
-    #[cfg(feature = "testing")]
     /// Notes: used for test purpose.
     pub fn from_inner(inner: T) -> Self {
         let bytes = serde_json::to_vec(&inner).unwrap();

src/common/meta/src/kv_backend/memory.rs

@@ -17,7 +17,6 @@ use std::collections::btree_map::Entry;
 use std::collections::BTreeMap;
 use std::fmt::{Display, Formatter};
 use std::marker::PhantomData;
-use std::ops::Range;
 use std::sync::RwLock;
 
 use async_trait::async_trait;
@@ -85,20 +84,14 @@ impl<T: ErrorExt + Send + Sync + 'static> KvBackend for MemoryKvBackend<T> {
     }
 
     async fn range(&self, req: RangeRequest) -> Result<RangeResponse, Self::Error> {
+        let range = req.range();
         let RangeRequest {
-            key,
-            range_end,
-            limit,
-            keys_only,
+            limit, keys_only, ..
         } = req;
 
         let kvs = self.kvs.read().unwrap();
+        let iter = kvs.range(range);
 
-        let iter: Box<dyn Iterator<Item = (&Vec<u8>, &Vec<u8>)>> = if range_end.is_empty() {
-            Box::new(kvs.get_key_value(&key).into_iter())
-        } else {
-            Box::new(kvs.range(key..range_end))
-        };
         let mut kvs = iter
             .map(|(k, v)| {
                 let key = k.clone();
@@ -215,31 +208,23 @@ impl<T: ErrorExt + Send + Sync + 'static> KvBackend for MemoryKvBackend<T> {
         &self,
         req: DeleteRangeRequest,
     ) -> Result<DeleteRangeResponse, Self::Error> {
-        let DeleteRangeRequest {
-            key,
-            range_end,
-            prev_kv,
-        } = req;
+        let range = req.range();
+        let DeleteRangeRequest { prev_kv, .. } = req;
 
         let mut kvs = self.kvs.write().unwrap();
 
-        let prev_kvs = if range_end.is_empty() {
-            kvs.remove(&key)
-                .into_iter()
-                .map(|value| KeyValue {
-                    key: key.clone(),
-                    value,
-                })
-                .collect::<Vec<_>>()
-        } else {
-            let range = Range {
-                start: key,
-                end: range_end,
-            };
-            kvs.extract_if(|key, _| range.contains(key))
-                .map(Into::into)
-                .collect::<Vec<_>>()
-        };
+        let keys = kvs
+            .range(range)
+            .map(|(key, _)| key.clone())
+            .collect::<Vec<_>>();
+
+        let mut prev_kvs = Vec::with_capacity(keys.len());
+
+        for key in keys {
+            if let Some(value) = kvs.remove(&key) {
+                prev_kvs.push((key.clone(), value).into())
+            }
+        }
 
         Ok(DeleteRangeResponse {
             deleted: prev_kvs.len() as i64,
@@ -504,6 +489,43 @@ mod tests {
         assert_eq!(b"val1", resp.kvs[0].value());
     }
 
+    #[tokio::test]
+    async fn test_range_2() {
+        let kv = MemoryKvBackend::<Error>::new();
+
+        kv.put(PutRequest::new().with_key("atest").with_value("value"))
+            .await
+            .unwrap();
+
+        kv.put(PutRequest::new().with_key("test").with_value("value"))
+            .await
+            .unwrap();
+
+        // If both key and range_end are ‘\0’, then range represents all keys.
+        let result = kv
+            .range(RangeRequest::new().with_range(b"\0".to_vec(), b"\0".to_vec()))
+            .await
+            .unwrap();
+
+        assert_eq!(result.kvs.len(), 2);
+
+        // If range_end is ‘\0’, the range is all keys greater than or equal to the key argument.
+        let result = kv
+            .range(RangeRequest::new().with_range(b"a".to_vec(), b"\0".to_vec()))
+            .await
+            .unwrap();
+
+        assert_eq!(result.kvs.len(), 2);
+
+        let result = kv
+            .range(RangeRequest::new().with_range(b"b".to_vec(), b"\0".to_vec()))
+            .await
+            .unwrap();
+
+        assert_eq!(result.kvs.len(), 1);
+        assert_eq!(result.kvs[0].key, b"test");
+    }
+
     #[tokio::test]
     async fn test_batch_get() {
         let kv_store = mock_mem_store_with_data().await;

src/common/meta/src/rpc/store.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::fmt::{Display, Formatter};
+use std::ops::Bound;
 
 use api::v1::meta::{
     BatchDeleteRequest as PbBatchDeleteRequest, BatchDeleteResponse as PbBatchDeleteResponse,
@@ -30,6 +31,22 @@ use crate::error;
 use crate::error::Result;
 use crate::rpc::{util, KeyValue};
 
+pub fn to_range(key: Vec<u8>, range_end: Vec<u8>) -> (Bound<Vec<u8>>, Bound<Vec<u8>>) {
+    if range_end.is_empty() {
+        (Bound::Included(key.clone()), Bound::Included(key))
+    } else if range_end.len() == 1 && range_end[0] == 0 {
+        // If both key and range_end are ‘\0’, then range represents all keys.
+        if key.len() == 1 && key[0] == 0 {
+            (Bound::Unbounded, Bound::Unbounded)
+        } else {
+            // If range_end is ‘\0’, the range is all keys greater than or equal to the key argument.
+            (Bound::Included(key), Bound::Unbounded)
+        }
+    } else {
+        (Bound::Included(key), Bound::Excluded(range_end))
+    }
+}
+
 #[derive(Debug, Clone, Default)]
 pub struct RangeRequest {
     /// key is the first key for the range, If range_end is not given, the
@@ -96,6 +113,11 @@ impl RangeRequest {
         }
     }
 
+    /// Returns the `RangeBounds`.
+    pub fn range(&self) -> (Bound<Vec<u8>>, Bound<Vec<u8>>) {
+        to_range(self.key.clone(), self.range_end.clone())
+    }
+
     /// key is the first key for the range, If range_end is not given, the
     /// request only looks up key.
     #[inline]
@@ -690,6 +712,11 @@ impl DeleteRangeRequest {
         }
     }
 
+    /// Returns the `RangeBounds`.
+    pub fn range(&self) -> (Bound<Vec<u8>>, Bound<Vec<u8>>) {
+        to_range(self.key.clone(), self.range_end.clone())
+    }
+
     /// key is the first key to delete in the range. If range_end is not given,
     /// the range is defined to contain only the key argument.
     #[inline]

src/log-store/src/raft_engine/backend.rs

@@ -137,29 +137,48 @@ impl KvBackend for RaftEngineBackend {
 
     async fn range(&self, req: RangeRequest) -> Result<RangeResponse, Self::Error> {
         let mut res = vec![];
+        let (start, end) = req.range();
+        let RangeRequest { limit, .. } = req;
+
+        let start_key = match start {
+            std::ops::Bound::Included(key) => Some(key),
+            std::ops::Bound::Excluded(_) => unreachable!(),
+            std::ops::Bound::Unbounded => None,
+        };
+
+        let end_key = match end {
+            std::ops::Bound::Included(_) => unreachable!(),
+            std::ops::Bound::Excluded(key) => Some(key),
+            std::ops::Bound::Unbounded => None,
+        };
+        let mut more = false;
+
         self.engine
             .read()
             .unwrap()
             .scan_raw_messages(
                 SYSTEM_NAMESPACE,
-                Some(&req.key),
-                Some(&req.range_end),
+                start_key.as_deref(),
+                end_key.as_deref(),
                 false,
                 |key, value| {
                     res.push(KeyValue {
                         key: key.to_vec(),
                         value: value.to_vec(),
                     });
-                    true
+                    if limit > 0 && limit as usize == res.len() {
+                        more = true;
+                        false
+                    } else {
+                        // continue
+                        true
+                    }
                 },
             )
             .context(RaftEngineSnafu)
             .map_err(BoxedError::new)
             .context(meta_error::ExternalSnafu)?;
-        Ok(RangeResponse {
-            kvs: res,
-            more: false,
-        })
+        Ok(RangeResponse { kvs: res, more })
     }
 
     async fn put(&self, req: PutRequest) -> Result<PutResponse, Self::Error> {