fix: when listing-kv, limit the right bound if possible (#15372)

Author: drmingdrmer

src/meta/raft-store/src/lib.rs

@@ -28,3 +28,4 @@ pub mod ondisk;
 pub mod sm_v002;
 pub mod state;
 pub mod state_machine;
+pub mod utils;

src/meta/raft-store/src/sm_v002/leveled_store/level.rs

@@ -112,7 +112,7 @@ impl MapApiRO<String> for Level {
 
         if vec.len() > 1000 {
             warn!(
-                "Level::<ExpireKey>::range() returns big range of len={}",
+                "Level::<String>::range() returns big range of len={}",
                 vec.len()
             );
         }

src/meta/raft-store/src/sm_v002/sm_v002.rs

@@ -62,6 +62,7 @@ use crate::sm_v002::SnapshotViewV002;
 use crate::state_machine::sm::BlockingConfig;
 use crate::state_machine::ExpireKey;
 use crate::state_machine::StateMachineSubscriber;
+use crate::utils::prefix_right_bound;
 
 /// A wrapper that implements KVApi **readonly** methods for the state machine.
 pub struct SMV002KVApi<'a> {
@@ -290,7 +291,11 @@ impl SMV002 {
     pub async fn list_kv(&self, prefix: &str) -> Result<ResultStream<(String, SeqV)>, io::Error> {
         let p = prefix.to_string();
 
-        let strm = self.levels.str_map().range(p.clone()..).await?;
+        let strm = if let Some(right) = prefix_right_bound(&p) {
+            self.levels.str_map().range(p.clone()..right).await?
+        } else {
+            self.levels.str_map().range(p.clone()..).await?
+        };
 
         let strm = strm
             // Return only keys with the expected prefix

src/meta/raft-store/src/utils.rs

@@ -0,0 +1,40 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Return the right bound of the prefix, so that `p..right` will cover all strings with prefix `p`.
+pub fn prefix_right_bound(p: &str) -> Option<String> {
+    let last = p.chars().last()?;
+    let mut next_str = p[..p.len() - last.len_utf8()].to_owned();
+    let next_char = char::from_u32(last as u32 + 1)?;
+    next_str.push(next_char);
+    Some(next_str)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::prefix_right_bound;
+
+    #[test]
+    fn test_next_string() {
+        assert_eq!(None, prefix_right_bound(""));
+        assert_eq!(Some(s("b")), prefix_right_bound("a"));
+        assert_eq!(Some(s("{")), prefix_right_bound("z"));
+        assert_eq!(Some(s("foo0")), prefix_right_bound("foo/"));
+        assert_eq!(Some(s("foo💰")), prefix_right_bound("foo💯"));
+    }
+
+    fn s(s: impl ToString) -> String {
+        s.to_string()
+    }
+}