feat: Add a new method fork to create a mutable memtable

src/mito2/src/flush.rs

@@ -31,7 +31,6 @@ use crate::config::MitoConfig;
 use crate::error::{
     Error, FlushRegionSnafu, RegionClosedSnafu, RegionDroppedSnafu, RegionTruncatedSnafu, Result,
 };
-use crate::memtable::MemtableBuilderRef;
 use crate::metrics::{FLUSH_BYTES_TOTAL, FLUSH_ELAPSED, FLUSH_ERRORS_TOTAL, FLUSH_REQUESTS_TOTAL};
 use crate::read::Source;
 use crate::region::version::{VersionControlData, VersionControlRef, VersionRef};
@@ -197,7 +196,6 @@ pub(crate) struct RegionFlushTask {
     pub(crate) request_sender: mpsc::Sender<WorkerRequest>,
 
     pub(crate) access_layer: AccessLayerRef,
-    pub(crate) memtable_builder: MemtableBuilderRef,
     pub(crate) file_purger: FilePurgerRef,
     pub(crate) listener: WorkerListener,
     pub(crate) engine_config: Arc<MitoConfig>,
@@ -461,7 +459,9 @@ impl FlushScheduler {
         }
 
         // Now we can flush the region directly.
-        version_control.freeze_mutable(&task.memtable_builder);
+        version_control.freeze_mutable().inspect(|e| {
+            error!(e; "Failed to freeze the mutable memtable for region {}", region_id);
+        })?;
         // Submit a flush job.
         let job = task.into_flush_job(version_control);
         if let Err(e) = self.scheduler.schedule(job) {
@@ -760,7 +760,6 @@ mod tests {
             senders: Vec::new(),
             request_sender: tx,
             access_layer: env.access_layer.clone(),
-            memtable_builder: builder.memtable_builder(),
             file_purger: builder.file_purger(),
             listener: WorkerListener::default(),
             engine_config: Arc::new(MitoConfig::default()),

src/mito2/src/memtable.rs

@@ -81,11 +81,14 @@ pub trait Memtable: Send + Sync + fmt::Debug {
     /// Returns true if the memtable is empty.
     fn is_empty(&self) -> bool;
 
-    /// Mark the memtable as immutable.
-    fn mark_immutable(&self);
+    /// Marks the memtable as immutable.
+    fn freeze(&self) -> Result<()>;
 
     /// Returns the [MemtableStats] info of Memtable.
     fn stats(&self) -> MemtableStats;
+
+    /// Forks this memtable and returns a new mutable memtable with specific memtable `id`.
+    fn fork(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef;
 }
 
 pub type MemtableRef = Arc<dyn Memtable>;
@@ -157,6 +160,11 @@ impl AllocTracker {
     pub(crate) fn bytes_allocated(&self) -> usize {
         self.bytes_allocated.load(Ordering::Relaxed)
     }
+
+    /// Returns the write buffer manager.
+    pub(crate) fn write_buffer_manager(&self) -> Option<WriteBufferManagerRef> {
+        self.write_buffer_manager.clone()
+    }
 }
 
 impl Drop for AllocTracker {

src/mito2/src/memtable/merge_tree.rs

@@ -100,8 +100,11 @@ impl Memtable for MergeTreeMemtable {
         self.tree.is_empty()
     }
 
-    fn mark_immutable(&self) {
+    fn freeze(&self) -> Result<()> {
         self.alloc_tracker.done_allocating();
+
+        // TODO(yingwen): Freeze the tree.
+        Ok(())
     }
 
     fn stats(&self) -> MemtableStats {
@@ -131,6 +134,10 @@ impl Memtable for MergeTreeMemtable {
             time_range: Some((min_timestamp, max_timestamp)),
         }
     }
+
+    fn fork(&self, _id: MemtableId, _metadata: &RegionMetadataRef) -> MemtableRef {
+        unimplemented!()
+    }
 }
 
 impl MergeTreeMemtable {

src/mito2/src/memtable/time_series.rs

@@ -234,8 +234,10 @@ impl Memtable for TimeSeriesMemtable {
         self.series_set.series.read().unwrap().is_empty()
     }
 
-    fn mark_immutable(&self) {
+    fn freeze(&self) -> Result<()> {
         self.alloc_tracker.done_allocating();
+
+        Ok(())
     }
 
     fn stats(&self) -> MemtableStats {
@@ -263,6 +265,14 @@ impl Memtable for TimeSeriesMemtable {
             time_range: Some((min_timestamp, max_timestamp)),
         }
     }
+
+    fn fork(&self, id: MemtableId, metadata: &RegionMetadataRef) -> MemtableRef {
+        Arc::new(TimeSeriesMemtable::new(
+            metadata.clone(),
+            id,
+            self.alloc_tracker.write_buffer_manager(),
+        ))
+    }
 }
 
 type SeriesRwLockMap = RwLock<BTreeMap<Vec<u8>, Arc<RwLock<Series>>>>;

src/mito2/src/memtable/version.rs

@@ -17,7 +17,9 @@
 use std::sync::Arc;
 
 use smallvec::SmallVec;
+use store_api::metadata::RegionMetadataRef;
 
+use crate::error::Result;
 use crate::memtable::{MemtableId, MemtableRef};
 
 /// A version of current memtables in a region.
@@ -61,28 +63,32 @@ impl MemtableVersion {
     /// memtable.
     ///
     /// Returns `None` if the mutable memtable is empty.
-    #[must_use]
-    pub(crate) fn freeze_mutable(&self, mutable: MemtableRef) -> Option<MemtableVersion> {
-        debug_assert!(mutable.is_empty());
+    pub(crate) fn freeze_mutable(
+        &self,
+        metadata: &RegionMetadataRef,
+    ) -> Result<Option<MemtableVersion>> {
         if self.mutable.is_empty() {
             // No need to freeze the mutable memtable.
-            return None;
+            return Ok(None);
         }
 
         // Marks the mutable memtable as immutable so it can free the memory usage from our
         // soft limit.
-        self.mutable.mark_immutable();
+        self.mutable.freeze()?;
+        // Fork the memtable.
+        let mutable = self.mutable.fork(self.mutable.id() + 1, metadata);
+
         // Pushes the mutable memtable to immutable list.
         let immutables = self
             .immutables
             .iter()
             .cloned()
             .chain([self.mutable.clone()])
             .collect();
-        Some(MemtableVersion {
+        Ok(Some(MemtableVersion {
             mutable,
             immutables,
-        })
+        }))
     }
 
     /// Removes memtables by ids from immutable memtables.

src/mito2/src/region/version.rs

@@ -29,6 +29,7 @@ use std::time::Duration;
 use store_api::metadata::RegionMetadataRef;
 use store_api::storage::SequenceNumber;
 
+use crate::error::Result;
 use crate::manifest::action::RegionEdit;
 use crate::memtable::version::{MemtableVersion, MemtableVersionRef};
 use crate::memtable::{MemtableBuilderRef, MemtableId, MemtableRef};
@@ -76,14 +77,16 @@ impl VersionControl {
     }
 
     /// Freezes the mutable memtable if it is not empty.
-    pub(crate) fn freeze_mutable(&self, builder: &MemtableBuilderRef) {
+    pub(crate) fn freeze_mutable(&self) -> Result<()> {
         let version = self.current().version;
         if version.memtables.mutable.is_empty() {
-            return;
+            return Ok(());
         }
-        let new_mutable = builder.build(&version.metadata);
         // Safety: Immutable memtable is None.
-        let new_memtables = version.memtables.freeze_mutable(new_mutable).unwrap();
+        let new_memtables = version
+            .memtables
+            .freeze_mutable(&version.metadata)?
+            .unwrap();
         // Create a new version with memtable switched.
         let new_version = Arc::new(
             VersionBuilder::from_version(version)
@@ -93,6 +96,8 @@ impl VersionControl {
 
         let mut version_data = self.data.write().unwrap();
         version_data.version = new_version;
+
+        Ok(())
     }
 
     /// Apply edit to current version.

src/mito2/src/test_util/memtable_util.rs

@@ -67,11 +67,17 @@ impl Memtable for EmptyMemtable {
         true
     }
 
-    fn mark_immutable(&self) {}
+    fn freeze(&self) -> Result<()> {
+        Ok(())
+    }
 
     fn stats(&self) -> MemtableStats {
         MemtableStats::default()
     }
+
+    fn fork(&self, id: MemtableId, _metadata: &RegionMetadataRef) -> MemtableRef {
+        Arc::new(EmptyMemtable::new(id))
+    }
 }
 
 /// Empty memtable builder.

src/mito2/src/worker/handle_flush.rs

@@ -143,7 +143,6 @@ impl<S> RegionWorkerLoop<S> {
             senders: Vec::new(),
             request_sender: self.sender.clone(),
             access_layer: region.access_layer.clone(),
-            memtable_builder: self.memtable_builder.clone(),
             file_purger: region.file_purger.clone(),
             listener: self.listener.clone(),
             engine_config,