refactor: simplify MapApi

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

@@ -191,11 +191,22 @@ impl<'a> Applier<'a> {
         Change::new(prev, result).into()
     }
 
+    // TODO(1): when get an applier, pass in a now_ms to ensure all expired are cleaned.
+    /// Update or insert a kv entry.
+    ///
+    /// If the input entry has expired, it performs a delete operation.
     #[minitrace::trace]
-    async fn upsert_kv(&mut self, upsert_kv: &UpsertKV) -> (Option<SeqV>, Option<SeqV>) {
+    pub(crate) async fn upsert_kv(&mut self, upsert_kv: &UpsertKV) -> (Option<SeqV>, Option<SeqV>) {
         debug!(upsert_kv = as_debug!(upsert_kv); "upsert_kv");
 
-        let (prev, result) = self.sm.upsert_kv(upsert_kv.clone()).await;
+        let (prev, result) = self.sm.upsert_kv_primary_index(upsert_kv).await;
+
+        self.sm
+            .update_expire_index(&upsert_kv.key, &prev, &result)
+            .await;
+
+        let prev = Into::<Option<SeqV>>::into(prev);
+        let result = Into::<Option<SeqV>>::into(result);
 
         debug!(
             "applied UpsertKV: {:?}; prev: {:?}; result: {:?}",
@@ -209,7 +220,6 @@ impl<'a> Applier<'a> {
     }
 
     #[minitrace::trace]
-
     async fn apply_txn(&mut self, req: &TxnRequest) -> AppliedState {
         debug!(txn = as_display!(req); "apply txn cmd");
 
@@ -445,9 +455,7 @@ impl<'a> Applier<'a> {
                 assert_eq!(expire_key.seq, seq_v.seq);
                 info!("clean expired: {}, {}", key, expire_key);
 
-                self.sm.upsert_kv(UpsertKV::delete(key.clone())).await;
-                // dbg!("clean_expired", &key, &curr);
-                self.push_change(key, curr, None);
+                self.upsert_kv(&UpsertKV::delete(key.clone())).await;
             } else {
                 unreachable!(
                     "trying to remove un-cleanable: {}, {}, kv-entry: {:?}",

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 #![allow(clippy::uninlined_format_args)]
+#![feature(impl_trait_in_assoc_type)]
 // #![feature(type_alias_impl_trait)]
 
 // #![allow(incomplete_features)]

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

@@ -22,11 +22,19 @@ use futures_util::StreamExt;
 
 use crate::sm_v002::leveled_store::map_api::MapApi;
 use crate::sm_v002::leveled_store::map_api::MapApiRO;
+use crate::sm_v002::leveled_store::map_api::MapKey;
 use crate::sm_v002::leveled_store::sys_data::SysData;
 use crate::sm_v002::leveled_store::sys_data_api::SysDataApiRO;
 use crate::sm_v002::marked::Marked;
 use crate::state_machine::ExpireKey;
 
+impl MapKey for String {
+    type V = Vec<u8>;
+}
+impl MapKey for ExpireKey {
+    type V = String;
+}
+
 /// A single level of state machine data.
 ///
 /// State machine data is composed of multiple levels.
@@ -75,24 +83,25 @@ impl LevelData {
 
 #[async_trait::async_trait]
 impl MapApiRO<String> for LevelData {
-    type V = Vec<u8>;
-
-    async fn get<Q>(&self, key: &Q) -> Marked<Self::V>
+    async fn get<Q>(&self, key: &Q) -> Marked<<String as MapKey>::V>
     where
         String: Borrow<Q>,
         Q: Ord + Send + Sync + ?Sized,
     {
         self.kv.get(key).cloned().unwrap_or(Marked::empty())
     }
 
-    async fn range<'a, T, R>(&'a self, range: R) -> BoxStream<'a, (String, Marked)>
+    async fn range<'f, Q, R>(
+        &'f self,
+        range: R,
+    ) -> BoxStream<'f, (String, Marked<<String as MapKey>::V>)>
     where
-        String: 'a,
-        String: Borrow<T>,
-        T: Ord + ?Sized,
-        R: RangeBounds<T> + Send,
+        String: Borrow<Q>,
+        Q: Ord + Send + Sync + ?Sized,
+        R: RangeBounds<Q> + Clone + Send + Sync,
     {
-        futures::stream::iter(self.kv.range(range).map(|(k, v)| (k.clone(), v.clone()))).boxed()
+        let it = self.kv.range(range).map(|(k, v)| (k.clone(), v.clone()));
+        futures::stream::iter(it).boxed()
     }
 }
 
@@ -101,8 +110,8 @@ impl MapApi<String> for LevelData {
     async fn set(
         &mut self,
         key: String,
-        value: Option<(Self::V, Option<KVMeta>)>,
-    ) -> (Marked<Self::V>, Marked<Self::V>) {
+        value: Option<(<String as MapKey>::V, Option<KVMeta>)>,
+    ) -> (Marked<<String as MapKey>::V>, Marked<<String as MapKey>::V>) {
         // The chance it is the bottom level is very low in a loaded system.
         // Thus we always tombstone the key if it is None.
 
@@ -117,33 +126,30 @@ impl MapApi<String> for LevelData {
             Marked::new_tomb_stone(seq)
         };
 
-        let prev = MapApiRO::<String>::get(self, key.as_str()).await;
+        let prev = MapApiRO::<String>::get(&*self, key.as_str()).await;
         self.kv.insert(key, marked.clone());
         (prev, marked)
     }
 }
 
 #[async_trait::async_trait]
 impl MapApiRO<ExpireKey> for LevelData {
-    type V = String;
-
-    async fn get<Q>(&self, key: &Q) -> Marked<Self::V>
+    async fn get<Q>(&self, key: &Q) -> Marked<<ExpireKey as MapKey>::V>
     where
         ExpireKey: Borrow<Q>,
         Q: Ord + Send + Sync + ?Sized,
     {
         self.expire.get(key).cloned().unwrap_or(Marked::empty())
     }
 
-    async fn range<'a, T: ?Sized, R>(
-        &'a self,
+    async fn range<'f, Q, R>(
+        &'f self,
         range: R,
-    ) -> BoxStream<'a, (ExpireKey, Marked<String>)>
+    ) -> BoxStream<'f, (ExpireKey, Marked<<ExpireKey as MapKey>::V>)>
     where
-        ExpireKey: 'a,
-        ExpireKey: Borrow<T>,
-        T: Ord,
-        R: RangeBounds<T> + Send,
+        ExpireKey: Borrow<Q>,
+        Q: Ord + Send + Sync + ?Sized,
+        R: RangeBounds<Q> + Clone + Send + Sync,
     {
         let it = self
             .expire
@@ -159,8 +165,11 @@ impl MapApi<ExpireKey> for LevelData {
     async fn set(
         &mut self,
         key: ExpireKey,
-        value: Option<(Self::V, Option<KVMeta>)>,
-    ) -> (Marked<Self::V>, Marked<Self::V>) {
+        value: Option<(<ExpireKey as MapKey>::V, Option<KVMeta>)>,
+    ) -> (
+        Marked<<ExpireKey as MapKey>::V>,
+        Marked<<ExpireKey as MapKey>::V>,
+    ) {
         // dbg!("set expire", &key, &value);
 
         let seq = self.curr_seq();
@@ -171,7 +180,7 @@ impl MapApi<ExpireKey> for LevelData {
             Marked::TombStone { internal_seq: seq }
         };
 
-        let prev = MapApiRO::<ExpireKey>::get(self, &key).await;
+        let prev = MapApiRO::<ExpireKey>::get(&*self, &key).await;
         self.expire.insert(key, marked.clone());
         (prev, marked)
     }

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

@@ -19,14 +19,16 @@ use std::sync::Arc;
 
 use common_meta_types::KVMeta;
 use futures_util::stream::BoxStream;
-use stream_more::KMerge;
-use stream_more::StreamMore;
 
 use crate::sm_v002::leveled_store::level_data::LevelData;
+use crate::sm_v002::leveled_store::map_api::compacted_get;
+use crate::sm_v002::leveled_store::map_api::compacted_range;
 use crate::sm_v002::leveled_store::map_api::MapApi;
 use crate::sm_v002::leveled_store::map_api::MapApiRO;
+use crate::sm_v002::leveled_store::map_api::MapKey;
+use crate::sm_v002::leveled_store::ref_::Ref;
+use crate::sm_v002::leveled_store::ref_mut::RefMut;
 use crate::sm_v002::leveled_store::static_leveled_map::StaticLeveledMap;
-use crate::sm_v002::leveled_store::util;
 use crate::sm_v002::marked::Marked;
 
 /// State machine data organized in multiple levels.
@@ -89,76 +91,59 @@ impl LeveledMap {
     pub(crate) fn replace_frozen_levels(&mut self, b: StaticLeveledMap) {
         self.frozen = b;
     }
+
+    pub(crate) fn leveled_ref_mut(&mut self) -> RefMut {
+        RefMut::new(&mut self.writable, &self.frozen)
+    }
+
+    pub(crate) fn leveled_ref(&self) -> Ref {
+        Ref::new(Some(&self.writable), &self.frozen)
+    }
 }
 
 #[async_trait::async_trait]
 impl<K> MapApiRO<K> for LeveledMap
 where
-    K: Ord + fmt::Debug + Send + Sync + Unpin + 'static,
+    K: MapKey + fmt::Debug,
     LevelData: MapApiRO<K>,
 {
-    type V = <LevelData as MapApiRO<K>>::V;
-
-    async fn get<Q>(&self, key: &Q) -> Marked<Self::V>
+    async fn get<Q>(&self, key: &Q) -> Marked<K::V>
     where
         K: Borrow<Q>,
         Q: Ord + Send + Sync + ?Sized,
     {
-        for level_data in self.iter_levels() {
-            let got = level_data.get(key).await;
-            if !got.is_not_found() {
-                return got;
-            }
-        }
-        return Marked::empty();
+        let levels = self.iter_levels();
+        compacted_get(key, levels).await
     }
 
-    async fn range<'a, T: ?Sized, R>(&'a self, range: R) -> BoxStream<'a, (K, Marked<Self::V>)>
+    async fn range<'f, Q, R>(&'f self, range: R) -> BoxStream<'f, (K, Marked<K::V>)>
     where
-        K: 'a,
-        K: Borrow<T> + Clone,
-        Self::V: Unpin,
-        T: Ord,
-        R: RangeBounds<T> + Clone + Send + Sync,
+        K: Borrow<Q>,
+        Q: Ord + Send + Sync + ?Sized,
+        R: RangeBounds<Q> + Clone + Send + Sync,
     {
-        let mut km = KMerge::by(util::by_key_seq);
-
-        for api in self.iter_levels() {
-            let a = api.range(range.clone()).await;
-            km = km.merge(a);
-        }
-
-        // Merge entries with the same key, keep the one with larger internal-seq
-        let m = km.coalesce(util::choose_greater);
-
-        Box::pin(m)
+        let levels = self.iter_levels();
+        compacted_range(range, levels).await
     }
 }
 
 #[async_trait::async_trait]
 impl<K> MapApi<K> for LeveledMap
 where
-    K: Ord + fmt::Debug + Send + Sync + Unpin + 'static,
+    K: MapKey,
     LevelData: MapApi<K>,
 {
     async fn set(
         &mut self,
         key: K,
-        value: Option<(Self::V, Option<KVMeta>)>,
-    ) -> (Marked<Self::V>, Marked<Self::V>)
+        value: Option<(K::V, Option<KVMeta>)>,
+    ) -> (Marked<K::V>, Marked<K::V>)
     where
         K: Ord,
     {
-        // Get from this level or the base level.
-        let prev = self.get(&key).await.clone();
-
-        // No such entry at all, no need to create a tombstone for delete
-        if prev.is_not_found() && value.is_none() {
-            return (prev, Marked::new_tomb_stone(0));
-        }
+        let mut l = self.leveled_ref_mut();
+        MapApi::set(&mut l, key, value).await
 
-        // The data is a single level map and the returned `_prev` is only from that level.
-        let (_prev, inserted) = self.writable_mut().set(key, value).await;
-        (prev, inserted)
+        // (&mut l).set(key, value).await
     }
 }

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

@@ -17,6 +17,7 @@ use futures_util::StreamExt;
 
 use crate::sm_v002::leveled_store::leveled_map::LeveledMap;
 use crate::sm_v002::leveled_store::map_api::MapApi;
+use crate::sm_v002::leveled_store::map_api::MapApiExt;
 use crate::sm_v002::leveled_store::map_api::MapApiRO;
 use crate::sm_v002::marked::Marked;
 
@@ -204,9 +205,11 @@ async fn test_two_levels() -> anyhow::Result<()> {
 
 /// Create multi levels store:
 ///
+/// ```text
 /// l2 |         c(D) d
 /// l1 |    b(D) c        e
 /// l0 | a  b    c    d
+/// ```
 async fn build_3_levels() -> LeveledMap {
     let mut l = LeveledMap::default();
     // internal_seq: 0
@@ -356,8 +359,10 @@ async fn test_three_levels_delete() -> anyhow::Result<()> {
     Ok(())
 }
 
+/// ```text
 /// |      b(m) c
 /// | a(m) b    c(m)
+/// ```
 async fn build_2_level_with_meta() -> LeveledMap {
     let mut l = LeveledMap::default();
 
@@ -401,7 +406,7 @@ async fn test_two_level_update_value() -> anyhow::Result<()> {
     {
         let mut l = build_2_level_with_meta().await;
 
-        let (prev, result) = MapApi::<String>::upsert_value(&mut l, s("a"), b("a1")).await;
+        let (prev, result) = MapApiExt::upsert_value(&mut l, s("a"), b("a1")).await;
         assert_eq!(
             prev,
             Marked::new_normal(1, b("a0"), Some(KVMeta { expire_at: Some(1) }))
@@ -422,7 +427,7 @@ async fn test_two_level_update_value() -> anyhow::Result<()> {
     {
         let mut l = build_2_level_with_meta().await;
 
-        let (prev, result) = MapApi::<String>::upsert_value(&mut l, s("b"), b("x1")).await;
+        let (prev, result) = MapApiExt::upsert_value(&mut l, s("b"), b("x1")).await;
         assert_eq!(
             prev,
             Marked::new_normal(
@@ -461,7 +466,7 @@ async fn test_two_level_update_value() -> anyhow::Result<()> {
     {
         let mut l = build_2_level_with_meta().await;
 
-        let (prev, result) = MapApi::<String>::upsert_value(&mut l, s("d"), b("d1")).await;
+        let (prev, result) = MapApiExt::upsert_value(&mut l, s("d"), b("d1")).await;
         assert_eq!(prev, Marked::new_tomb_stone(0));
         assert_eq!(result, Marked::new_normal(6, b("d1"), None));
 
@@ -479,8 +484,7 @@ async fn test_two_level_update_meta() -> anyhow::Result<()> {
         let mut l = build_2_level_with_meta().await;
 
         let (prev, result) =
-            MapApi::<String>::update_meta(&mut l, s("a"), Some(KVMeta { expire_at: Some(2) }))
-                .await;
+            MapApiExt::update_meta(&mut l, s("a"), Some(KVMeta { expire_at: Some(2) })).await;
         assert_eq!(
             prev,
             Marked::new_normal(1, b("a0"), Some(KVMeta { expire_at: Some(1) }))
@@ -501,7 +505,7 @@ async fn test_two_level_update_meta() -> anyhow::Result<()> {
     {
         let mut l = build_2_level_with_meta().await;
 
-        let (prev, result) = MapApi::<String>::update_meta(&mut l, s("b"), None).await;
+        let (prev, result) = MapApiExt::update_meta(&mut l, s("b"), None).await;
         assert_eq!(
             prev,
             Marked::new_normal(
@@ -522,7 +526,7 @@ async fn test_two_level_update_meta() -> anyhow::Result<()> {
     {
         let mut l = build_2_level_with_meta().await;
 
-        let (prev, result) = MapApi::<String>::update_meta(
+        let (prev, result) = MapApiExt::update_meta(
             &mut l,
             s("c"),
             Some(KVMeta {
@@ -560,8 +564,7 @@ async fn test_two_level_update_meta() -> anyhow::Result<()> {
         let mut l = build_2_level_with_meta().await;
 
         let (prev, result) =
-            MapApi::<String>::update_meta(&mut l, s("d"), Some(KVMeta { expire_at: Some(2) }))
-                .await;
+            MapApiExt::update_meta(&mut l, s("d"), Some(KVMeta { expire_at: Some(2) })).await;
         assert_eq!(prev, Marked::new_tomb_stone(0));
         assert_eq!(result, Marked::new_tomb_stone(0));