fix: remove memtables after flush

src/mito2/src/engine/flush_test.rs

@@ -53,6 +53,7 @@ async fn test_manual_flush() {
 
     let request = ScanRequest::default();
     let scanner = engine.scan(region_id, request).unwrap();
+    assert_eq!(0, scanner.num_memtables());
     assert_eq!(1, scanner.num_files());
     let stream = scanner.scan().await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
@@ -98,7 +99,7 @@ async fn test_flush_engine() {
 
     write_buffer_manager.set_should_flush(true);
 
-    // Writes and triggers flush.
+    // Writes to the mutable memtable and triggers flush.
     let rows = Rows {
         schema: column_schemas.clone(),
         rows: build_rows_for_key("b", 0, 2, 0),
@@ -110,6 +111,7 @@ async fn test_flush_engine() {
 
     let request = ScanRequest::default();
     let scanner = engine.scan(region_id, request).unwrap();
+    assert_eq!(1, scanner.num_memtables());
     assert_eq!(1, scanner.num_files());
     let stream = scanner.scan().await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
@@ -174,6 +176,7 @@ async fn test_write_stall() {
 
     let request = ScanRequest::default();
     let scanner = engine.scan(region_id, request).unwrap();
+    assert_eq!(1, scanner.num_memtables());
     assert_eq!(1, scanner.num_files());
     let stream = scanner.scan().await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();
@@ -209,6 +212,7 @@ async fn test_flush_empty() {
 
     let request = ScanRequest::default();
     let scanner = engine.scan(region_id, request).unwrap();
+    assert_eq!(0, scanner.num_memtables());
     assert_eq!(0, scanner.num_files());
     let stream = scanner.scan().await.unwrap();
     let batches = RecordBatches::try_collect(stream).await.unwrap();

src/mito2/src/flush.rs

@@ -479,8 +479,11 @@ impl FlushScheduler {
     }
 
     /// Returns true if the region has pending DDLs.
-    pub(crate) fn has_pending_ddls(&self, _region_id: RegionId) -> bool {
-        unimplemented!()
+    pub(crate) fn has_pending_ddls(&self, region_id: RegionId) -> bool {
+        self.region_status
+            .get(&region_id)
+            .map(|status| !status.pending_ddls.is_empty())
+            .unwrap_or(false)
     }
 
     /// Schedules a new flush task when the scheduler can submit next task.

src/mito2/src/memtable/version.rs

@@ -18,10 +18,10 @@ use std::sync::Arc;
 
 use smallvec::SmallVec;
 
-use crate::memtable::MemtableRef;
+use crate::memtable::{MemtableId, MemtableRef};
 
 /// A version of current memtables in a region.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(crate) struct MemtableVersion {
     /// Mutable memtable.
     pub(crate) mutable: MemtableRef,
@@ -85,7 +85,15 @@ impl MemtableVersion {
         })
     }
 
-    // TODO(yingwen): Remove memtables on flush done.
+    /// Removes memtables by ids from immutable memtables.
+    pub(crate) fn remove_memtables(&mut self, ids: &[MemtableId]) {
+        self.immutables = self
+            .immutables
+            .iter()
+            .filter(|mem| !ids.contains(&mem.id()))
+            .cloned()
+            .collect();
+    }
 
     /// Returns the memory usage of the mutable memtable.
     pub(crate) fn mutable_usage(&self) -> usize {

src/mito2/src/read/scan_region.rs

@@ -52,6 +52,13 @@ impl Scanner {
             Scanner::Seq(seq_scan) => seq_scan.num_files(),
         }
     }
+
+    /// Returns number of memtables to scan.
+    pub(crate) fn num_memtables(&self) -> usize {
+        match self {
+            Scanner::Seq(seq_scan) => seq_scan.num_memtables(),
+        }
+    }
 }
 
 #[cfg_attr(doc, aquamarine::aquamarine)]
@@ -139,6 +146,11 @@ impl ScanRegion {
         }
 
         let memtables = self.version.memtables.list_memtables();
+        // Skip empty memtables.
+        let memtables: Vec<_> = memtables
+            .into_iter()
+            .filter(|mem| !mem.is_empty())
+            .collect();
 
         debug!(
             "Seq scan region {}, memtables: {}, ssts_to_read: {}, total_ssts: {}",

src/mito2/src/read/seq_scan.rs

@@ -140,6 +140,11 @@ impl SeqScan {
 
 #[cfg(test)]
 impl SeqScan {
+    /// Returns number of memtables to scan.
+    pub(crate) fn num_memtables(&self) -> usize {
+        self.memtables.len()
+    }
+
     /// Returns number of SST files to scan.
     pub(crate) fn num_files(&self) -> usize {
         self.files.len()

src/mito2/src/region/version.rs

@@ -30,7 +30,7 @@ use store_api::storage::SequenceNumber;
 
 use crate::manifest::action::RegionEdit;
 use crate::memtable::version::{MemtableVersion, MemtableVersionRef};
-use crate::memtable::{MemtableBuilderRef, MemtableRef};
+use crate::memtable::{MemtableBuilderRef, MemtableId, MemtableRef};
 use crate::sst::file_purger::FilePurgerRef;
 use crate::sst::version::{SstVersion, SstVersionRef};
 use crate::wal::EntryId;
@@ -90,11 +90,17 @@ impl VersionControl {
     }
 
     /// Apply edit to current version.
-    pub(crate) fn apply_edit(&self, edit: RegionEdit, purger: FilePurgerRef) {
+    pub(crate) fn apply_edit(
+        &self,
+        edit: RegionEdit,
+        memtables_to_remove: &[MemtableId],
+        purger: FilePurgerRef,
+    ) {
         let version = self.current().version;
         let new_version = Arc::new(
             VersionBuilder::from_version(version)
                 .apply_edit(edit, purger)
+                .remove_memtables(memtables_to_remove)
                 .build(),
         );
 
@@ -208,6 +214,16 @@ impl VersionBuilder {
         self
     }
 
+    /// Remove memtables from the builder.
+    pub(crate) fn remove_memtables(mut self, ids: &[MemtableId]) -> VersionBuilder {
+        if !ids.is_empty() {
+            let mut memtables = (*self.memtables).clone();
+            memtables.remove_memtables(ids);
+            self.memtables = Arc::new(memtables);
+        }
+        self
+    }
+
     /// Builds a new [Version] from the builder.
     pub(crate) fn build(self) -> Version {
         Version {

src/mito2/src/worker/handle_alter.rs

@@ -63,7 +63,7 @@ impl<S> RegionWorkerLoop<S> {
                 return;
             }
 
-            // TODO(yingwen): Maybe assert in add_ddl_request_to_pending instead returning result.
+            // Safety: We have requested flush.
             self.flush_scheduler
                 .add_ddl_request_to_pending(SenderDdlRequest {
                     region_id,

src/mito2/src/worker/handle_flush.rs

@@ -56,9 +56,11 @@ impl<S: LogStore> RegionWorkerLoop<S> {
         }
 
         // Apply edit to region's version.
-        region
-            .version_control
-            .apply_edit(edit, region.file_purger.clone());
+        region.version_control.apply_edit(
+            edit,
+            &request.memtables_to_remove,
+            region.file_purger.clone(),
+        );
         region.update_flush_millis();
 
         // Delete wal.

src/mito2/src/worker/handle_write.rs

@@ -94,6 +94,7 @@ impl<S> RegionWorkerLoop<S> {
             // If region is waiting for alteration, add requests to pending writes.
             if self.flush_scheduler.has_pending_ddls(region_id) {
                 // TODO(yingwen): consider adding some metrics for this.
+                // Safety: The region has pending ddls.
                 self.flush_scheduler
                     .add_write_request_to_pending(sender_req);
                 continue;