refactor(meta-service): replace Vec with BTreeMap in ImmutableLevels (#18608)

* refactor(meta-service): replace Vec with BTreeMap in ImmutableLevels

Replace `Vec<Immutable>` with `BTreeMap<LevelIndex, Immutable>` in
`ImmutableLevels` to improve level management and ordering guarantees.

Key changes:
- Use `LevelIndex` as map key for O(1) lookups and automatic ordering
- Add level index validation in insert operations
- Update compaction logic to work with indexed levels
- Rename methods for clarity: `new` -> `new_form_iter`, `push` -> `insert`
- Add `remove_levels_upto` method for efficient level cleanup
- Update all tests and dependent code

* chore: fix export bug

Author: drmingdrmer

src/meta/raft-store/src/leveled_store/immutable_levels.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::collections::BTreeMap;
 use std::io;
 use std::ops::RangeBounds;
 
@@ -22,52 +23,82 @@ use seq_marked::SeqMarked;
 
 use crate::leveled_store::immutable::Immutable;
 use crate::leveled_store::level::Level;
+use crate::leveled_store::level_index::LevelIndex;
 use crate::leveled_store::map_api::KVResultStream;
 use crate::leveled_store::map_api::MapKey;
 
 /// A readonly leveled map that owns the data.
 #[derive(Debug, Default, Clone)]
 pub struct ImmutableLevels {
-    /// From oldest to newest, i.e., levels[0] is the oldest
-    levels: Vec<Immutable>,
+    /// From oldest to newest, i.e., first is the oldest
+    immutables: BTreeMap<LevelIndex, Immutable>,
 }
 
 impl ImmutableLevels {
-    pub(crate) fn new(levels: impl IntoIterator<Item = Immutable>) -> Self {
+    pub(crate) fn new_form_iter(immutables: impl IntoIterator<Item = Immutable>) -> Self {
         Self {
-            levels: levels.into_iter().collect(),
+            immutables: immutables
+                .into_iter()
+                .map(|immu| (*immu.level_index(), immu))
+                .collect(),
         }
     }
 
+    pub(crate) fn indexes(&self) -> impl Iterator<Item = LevelIndex> + use<'_> {
+        self.immutables.keys().cloned()
+    }
+
     /// Return an iterator of all Arc of levels from newest to oldest.
     pub(crate) fn iter_immutable_levels(&self) -> impl Iterator<Item = &Immutable> {
-        self.levels.iter().rev()
+        self.immutables.values().rev()
     }
 
     /// Return an iterator of all levels from newest to oldest.
     pub(crate) fn iter_levels(&self) -> impl Iterator<Item = &Level> {
-        self.levels.iter().map(|x| x.as_ref()).rev()
+        self.immutables.values().map(|x| x.as_ref()).rev()
     }
 
     pub(crate) fn newest(&self) -> Option<&Immutable> {
-        self.levels.last()
+        self.immutables.values().next_back()
     }
 
-    pub(crate) fn push(&mut self, level: Immutable) {
-        self.levels.push(level);
+    pub(crate) fn newest_level_index(&self) -> Option<LevelIndex> {
+        self.newest().map(|x| *x.level_index())
     }
 
-    pub(crate) fn len(&self) -> usize {
-        self.levels.len()
+    pub(crate) fn insert(&mut self, level: Immutable) {
+        let key = *level.level_index();
+
+        let last_key = self.newest_level_index().unwrap_or_default();
+
+        // The newly added must have greater data index or there are no data added.
+        assert!(
+            key > last_key || (key == LevelIndex::default() && key == last_key),
+            "new level to insert {:?} must have greater index than the newest level {:?}",
+            key,
+            last_key
+        );
+
+        self.immutables.insert(key, level);
     }
 
-    pub(crate) fn levels(&self) -> &Vec<Immutable> {
-        &self.levels
+    /// Remove all levels up to and including the given level index.
+    pub(crate) fn remove_levels_upto(&mut self, level_index: LevelIndex) {
+        assert!(
+            self.immutables.contains_key(&level_index),
+            "level_index to remove {:?} must exist",
+            level_index
+        );
+
+        let mut left = self.immutables.split_off(&level_index);
+        // split_off() also returns the given key, remove it.
+        left.pop_first();
+        self.immutables = left;
     }
 
     #[allow(dead_code)]
-    pub(crate) fn levels_mut(&mut self) -> &mut Vec<Immutable> {
-        &mut self.levels
+    pub(crate) fn levels_mut(&mut self) -> &mut BTreeMap<LevelIndex, Immutable> {
+        &mut self.immutables
     }
 }
 
@@ -89,3 +120,28 @@ where
         compacted_range::<_, _, Level, _, Level>(range, None, levels, []).await
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use crate::leveled_store::immutable::Immutable;
+    use crate::leveled_store::immutable_levels::ImmutableLevels;
+    use crate::leveled_store::level::Level;
+
+    #[test]
+    fn test_remove_levels_upto() {
+        let mut immutables = ImmutableLevels::new_form_iter(vec![
+            Immutable::new(Arc::new(Level::default())),
+            Immutable::new(Arc::new(Level::default())),
+            Immutable::new(Arc::new(Level::default())),
+        ]);
+
+        let index = immutables.indexes().collect::<Vec<_>>()[1];
+        let left_index = immutables.indexes().collect::<Vec<_>>()[2];
+
+        immutables.remove_levels_upto(index);
+
+        assert_eq!(immutables.indexes().collect::<Vec<_>>(), vec![left_index]);
+    }
+}

src/meta/raft-store/src/leveled_store/level_index.rs

@@ -15,7 +15,7 @@
 /// Represents a unique identifier for a level.
 ///
 /// The magnitude of the index indicates the relative freshness of the level.
-#[derive(Clone, Debug, Copy, PartialEq, Eq, PartialOrd, Ord)]
+#[derive(Clone, Debug, Copy, Default, PartialEq, Eq, PartialOrd, Ord)]
 pub struct LevelIndex {
     /// The internal sequence number of the level.
     internal_seq: u64,

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

@@ -41,6 +41,7 @@ use crate::utils::add_cooperative_yielding;
 /// The data to compact.
 ///
 /// Including several in-memory immutable levels and an optional persisted db.
+#[derive(Debug)]
 pub(crate) struct CompactingData {
     pub(crate) immutable_levels: Arc<ImmutableLevels>,
     pub(crate) persisted: Option<Arc<DB>>,
@@ -81,7 +82,10 @@ impl CompactingData {
         let table = mvcc::Table::from_stream(strm).await?;
         data.replace_kv(table);
 
-        self.immutable_levels = Arc::new(ImmutableLevels::new([Immutable::new_from_level(data)]));
+        self.immutable_levels =
+            Arc::new(ImmutableLevels::new_form_iter([Immutable::new_from_level(
+                data,
+            )]));
         Ok(())
     }
 

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

@@ -28,17 +28,16 @@ use crate::leveled_store::leveled_map::compactor_acquirer::CompactorPermit;
 /// Compactor is responsible for compacting the immutable levels and db.
 ///
 /// Only one Compactor can be running at a time.
+#[derive(Debug)]
 pub struct Compactor {
     /// Acquired permit for this compactor.
     ///
     /// This is used to ensure that only one compactor can run at a time.
     pub(crate) _permit: CompactorPermit,
 
     pub(crate) compacting_data: CompactingData,
-    // pub(crate) immutable_levels: Arc<ImmutableLevels>,
-    // pub(crate) persisted: Option<Arc<DB>>,
     /// Remember the newest level included in this compactor.
-    pub(super) since: Option<LevelIndex>,
+    pub(super) upto: Option<LevelIndex>,
 }
 
 impl Compactor {

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

@@ -35,6 +35,7 @@ impl CompactorAcquirer {
     }
 }
 
+#[derive(Debug)]
 pub struct CompactorPermit {
     _permit: OwnedSemaphorePermit,
 }

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

@@ -38,6 +38,7 @@ where
     K: MapKeyDecode,
     SeqMarked<K::V>: ValueConvert<SeqMarked>,
     Level: GetTable<K, K::V>,
+    <K as MapKey>::V: fmt::Debug,
     for<'a> MapView<'a>: MapApiRO<K>,
 {
     async fn get(&self, key: &K) -> Result<SeqMarked<K::V>, io::Error> {
@@ -58,6 +59,7 @@ where
     K: MapKeyDecode,
     SeqMarked<K::V>: ValueConvert<SeqMarked>,
     Level: GetTable<K, K::V>,
+    <K as MapKey>::V: fmt::Debug,
     for<'a> MapView<'a>: MapApiRO<K>,
 {
     async fn get(&self, key: &K) -> Result<SeqMarked<K::V>, io::Error> {

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use std::collections::BTreeMap;
+use std::fmt;
 use std::io;
 use std::ops::RangeBounds;
 use std::sync::Arc;
@@ -185,6 +186,7 @@ impl LeveledMapData {
         K: MapKeyDecode,
         SeqMarked<K::V>: ValueConvert<SeqMarked>,
         Level: GetTable<K, K::V>,
+        <K as MapKey>::V: fmt::Debug,
     {
         // TODO: test it
 
@@ -207,18 +209,32 @@ impl LeveledMapData {
             );
         }
 
-        let strm = futures::stream::iter(vec).map(Ok).boxed();
+        let strm = futures::stream::iter(vec)
+            // .map(|x| {
+            //     debug!("range-item from writable: {:?}", x);
+            //     x
+            // })
+            .map(Ok)
+            .boxed();
         kmerge = kmerge.merge(strm);
 
         // Immutable levels
 
         let immutable_levels = self.immutable_levels();
 
         for level in immutable_levels.iter_levels() {
+            let index = level.with_sys_data(|s| s.curr_seq());
+            let _ = index;
             let it = level.get_table().range(range.clone(), upto_seq);
 
             let vec = it.map(|(k, v)| (k.clone(), v.cloned())).collect::<Vec<_>>();
-            let strm = futures::stream::iter(vec).map(Ok).boxed();
+            let strm = futures::stream::iter(vec)
+                // .map(move |x| {
+                //     debug!("range-item from immutable seq={}: {:?}", index, x);
+                //     x
+                // })
+                .map(Ok)
+                .boxed();
             kmerge = kmerge.merge(strm);
         }
 
@@ -227,6 +243,10 @@ impl LeveledMapData {
         if let Some(db) = self.persisted() {
             let map_view = MapView(&db);
             let strm = map_view.range(range.clone()).await?;
+            // let strm = strm.map(|x| {
+            //     debug!("range-item from db: {:?}", x);
+            //     x
+            // });
             kmerge = kmerge.merge(strm);
         };
 
@@ -309,8 +329,7 @@ impl LeveledMap {
     /// Return the [`LevelIndex`] of the newest **immutable** data
     pub(crate) fn immutable_level_index(&self) -> Option<LevelIndex> {
         let immutables = self.data.immutable_levels();
-        let newest = immutables.newest()?;
-        Some(*newest.level_index())
+        immutables.newest_level_index()
     }
 
     /// Freeze the current writable level and create a new empty writable level.
@@ -326,23 +345,24 @@ impl LeveledMap {
     }
 
     pub fn do_freeze_writable(&self) -> Arc<ImmutableLevels> {
-        let new_immutable = {
-            let mut writable = self.data.writable();
-            let new_writable = writable.new_level();
+        let mut writable = self.data.writable();
+        let mut immutables = self.data.immutable_levels.lock().unwrap();
 
-            std::mem::replace(&mut *writable, new_writable)
-        };
+        let new_writable = writable.new_level();
+        let new_immutable = std::mem::replace(&mut *writable, new_writable);
 
-        let mut immutable_levels = self.data.immutable_levels().as_ref().clone();
+        let mut immutable_levels = immutables.as_ref().clone();
+        immutable_levels.insert(Immutable::new_from_level(new_immutable));
 
-        immutable_levels.push(Immutable::new_from_level(new_immutable));
+        *immutables = Arc::new(immutable_levels);
 
-        {
-            let mut immutables = self.data.immutable_levels.lock().unwrap();
-            *immutables = Arc::new(immutable_levels)
-        }
+        info!(
+            "do_freeze_writable: after writable: {:?}, immutables: {:?}",
+            writable,
+            immutables.indexes().collect::<Vec<_>>()
+        );
 
-        self.data.immutable_levels()
+        immutables.clone()
     }
 
     pub fn persisted(&self) -> Option<Arc<DB>> {
@@ -385,38 +405,33 @@ impl LeveledMap {
     /// **Important**: Do not drop the compactor within this function when called
     /// under a state machine lock, as dropping may take ~250ms.
     pub fn replace_with_compacted(&self, compactor: &mut Compactor, db: DB) {
-        let len = compactor.immutable_levels().len();
-
-        // Get the level index of the newest level.
-        let corresponding_index = compactor
-            .immutable_levels()
-            .levels()
-            .get(len - 1)
-            .map(|l| l.level_index())
-            .copied();
-
-        assert_eq!(
-            compactor.since, corresponding_index,
-            "unexpected change to sm during compaction"
+        info!(
+            "replace_with_compacted: compacted upto {:?} immutable levels; my levels: {:?}; compacted levels: {:?}",
+            compactor.upto,
+            self.data.immutable_levels().indexes().collect::<Vec<_>>(),
+            compactor.compacting_data.immutable_levels.indexes().collect::<Vec<_>>(),
         );
 
-        let my_immutables = self.data.immutable_levels();
-        let mut levels = my_immutables.levels().clone();
+        let mut immutables = self.data.immutable_levels().as_ref().clone();
 
-        // Remove the levels already included in `persisted`
-        let uncompacted = levels.split_off(len);
+        // If there is immutable levels compacted, remove them.
+        if let Some(upto) = compactor.upto {
+            immutables.remove_levels_upto(upto);
+        }
 
-        let uncompacted_immutables = ImmutableLevels::new(uncompacted);
-
-        // replace the immutables
-        self.data
-            .with_immutable_levels(|x| *x = Arc::new(uncompacted_immutables));
+        // NOTE: Replace data from bottom to top.
+        // replace the db first, db contains more data.
+        // Otherwise, there is a chance some data is removed from immutables and the new db containing this data is not inserted.
 
         // replace the persisted
         self.data.with_persisted(|p| {
             *p = Some(Arc::new(db));
         });
 
+        // replace the immutables
+        self.data
+            .with_immutable_levels(|x| *x = Arc::new(immutables));
+
         info!("replace_with_compacted: finished replacing the db");
     }
 
@@ -437,10 +452,12 @@ impl LeveledMap {
     }
 
     pub fn new_compactor(&self, permit: CompactorPermit) -> Compactor {
+        let level_index = self.immutable_level_index();
+
         Compactor {
             _permit: permit,
             compacting_data: CompactingData::new(self.immutable_levels(), self.persisted()),
-            since: self.immutable_level_index(),
+            upto: level_index,
         }
     }
 }