Merge pull request databendlabs#1164 from drmingdrmer/ks

[store] feature: SledTree: add range_keys() to retrieve keys in a range

fusestore/store/src/meta_service/sled_tree.rs

@@ -50,6 +50,18 @@ impl<K: SledOrderedSerde + Display + Debug, V: SledSerde> SledTree<K, V> {
         Ok(rl)
     }
 
+    /// Return true if the tree contains the key.
+    pub fn contains_key(&self, key: &K) -> common_exception::Result<bool> {
+        let got = self
+            .tree
+            .contains_key(key.ser()?)
+            .map_err_to_code(ErrorCode::MetaStoreDamaged, || {
+                format!("contains_key: {}:{}", self.name, key)
+            })?;
+
+        Ok(got)
+    }
+
     /// Retrieve the value of key.
     pub fn get(&self, key: &K) -> common_exception::Result<Option<V>> {
         let got = self
@@ -123,6 +135,28 @@ impl<K: SledOrderedSerde + Display + Debug, V: SledSerde> SledTree<K, V> {
         Ok(())
     }
 
+    /// Get keys in `range`
+    pub fn range_keys<R>(&self, range: R) -> common_exception::Result<Vec<K>>
+    where R: RangeBounds<K> {
+        // TODO(xp): pre alloc vec space
+        let mut res = vec![];
+
+        let range_mes = self.range_message(&range);
+
+        // Convert K range into sled::IVec range
+        let range = range.ser()?;
+        for item in self.tree.range(range) {
+            let (k, _) = item.map_err_to_code(ErrorCode::MetaStoreDamaged, || {
+                format!("range_get: {}", range_mes,)
+            })?;
+
+            let key = K::de(&k)?;
+            res.push(key);
+        }
+
+        Ok(res)
+    }
+
     /// Get values of key in `range`
     pub fn range_get<R>(&self, range: R) -> common_exception::Result<Vec<V>>
     where R: RangeBounds<K> {

fusestore/store/src/meta_service/sled_tree_test.rs

@@ -152,6 +152,56 @@ async fn test_sled_tree_append_values_and_range_get() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_sled_tree_range_keys() -> anyhow::Result<()> {
+    let tc = new_sled_test_context();
+    let db = &tc.db;
+    let rl = SledTree::<LogIndex, Entry<LogEntry>>::open(db, "log").await?;
+
+    let logs: Vec<Entry<LogEntry>> = vec![
+        Entry {
+            log_id: LogId { term: 1, index: 2 },
+            payload: EntryPayload::Blank,
+        },
+        Entry {
+            log_id: LogId { term: 1, index: 9 },
+            payload: EntryPayload::Blank,
+        },
+        Entry {
+            log_id: LogId { term: 1, index: 10 },
+            payload: EntryPayload::Blank,
+        },
+    ];
+
+    rl.append_values(&logs).await?;
+
+    let got = rl.range_keys(0..)?;
+    assert_eq!(vec![2, 9, 10], got);
+
+    let got = rl.range_keys(0..=2)?;
+    assert_eq!(vec![2], got);
+
+    let got = rl.range_keys(0..3)?;
+    assert_eq!(vec![2], got);
+
+    let got = rl.range_keys(0..10)?;
+    assert_eq!(vec![2, 9], got);
+
+    let got = rl.range_keys(0..11)?;
+    assert_eq!(vec![2, 9, 10], got);
+
+    let got = rl.range_keys(9..11)?;
+    assert_eq!(vec![9, 10], got);
+
+    let got = rl.range_keys(10..256)?;
+    assert_eq!(vec![10], got);
+
+    let got = rl.range_keys(11..)?;
+    assert_eq!(Vec::<LogIndex>::new(), got);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn test_sled_tree_insert() -> anyhow::Result<()> {
     let tc = new_sled_test_context();
@@ -187,6 +237,43 @@ async fn test_sled_tree_insert() -> anyhow::Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_sled_tree_contains_key() -> anyhow::Result<()> {
+    let tc = new_sled_test_context();
+    let db = &tc.db;
+    let rl = SledTree::<LogIndex, Entry<LogEntry>>::open(db, "log").await?;
+
+    assert_eq!(None, rl.get(&5)?);
+
+    let logs: Vec<Entry<LogEntry>> = vec![
+        Entry {
+            log_id: LogId { term: 1, index: 2 },
+            payload: EntryPayload::Blank,
+        },
+        Entry {
+            log_id: LogId { term: 3, index: 4 },
+            payload: EntryPayload::Normal(EntryNormal {
+                data: LogEntry {
+                    txid: None,
+                    cmd: Cmd::IncrSeq {
+                        key: "foo".to_string(),
+                    },
+                },
+            }),
+        },
+    ];
+
+    rl.append_values(&logs).await?;
+
+    assert!(!rl.contains_key(&1)?);
+    assert!(rl.contains_key(&2)?);
+    assert!(!rl.contains_key(&3)?);
+    assert!(rl.contains_key(&4)?);
+    assert!(!rl.contains_key(&5)?);
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 async fn test_sled_tree_get() -> anyhow::Result<()> {
     let tc = new_sled_test_context();