feat: track disk usage of regions

Signed-off-by: Zheming Li <nkdudu@126.com>

Author: lizhemingi

src/mito/src/table/test_util/mock_engine.rs

@@ -196,6 +196,10 @@ impl Region for MockRegion {
     async fn close(&self) -> Result<()> {
         Ok(())
     }
+
+    fn disk_usage_bytes(&self) -> u64 {
+        0
+    }
 }
 
 impl MockRegionInner {

src/storage/src/compaction/strategy.rs

@@ -239,6 +239,7 @@ mod tests {
                     Timestamp::new_millisecond(end_ts_millis),
                 )),
                 level: 0,
+                file_size: 0,
             },
             layer,
             file_purger,

src/storage/src/compaction/task.rs

@@ -173,7 +173,10 @@ impl CompactionOutput {
         let output_file_name = format!("{}.parquet", Uuid::new_v4().hyphenated());
         let opts = WriteOptions {};
 
-        let SstInfo { time_range } = sst_layer
+        let SstInfo {
+            time_range,
+            file_size,
+        } = sst_layer
             .write_sst(&output_file_name, Source::Reader(reader), &opts)
             .await?;
 
@@ -182,6 +185,7 @@ impl CompactionOutput {
             file_name: output_file_name,
             time_range,
             level: self.output_level,
+            file_size,
         })
     }
 }

src/storage/src/compaction/writer.rs

@@ -221,7 +221,10 @@ mod tests {
         let iter = memtable.iter(&IterContext::default()).unwrap();
         let writer = ParquetWriter::new(sst_file_name, Source::Iter(iter), object_store.clone());
 
-        let SstInfo { time_range } = writer
+        let SstInfo {
+            time_range,
+            file_size,
+        } = writer
             .write_sst(&sst::WriteOptions::default())
             .await
             .unwrap();
@@ -231,6 +234,7 @@ mod tests {
                 file_name: sst_file_name.to_string(),
                 time_range,
                 level: 0,
+                file_size,
             },
             Arc::new(crate::test_util::access_layer_util::MockAccessLayer {}),
             new_noop_file_purger(),
@@ -409,13 +413,11 @@ mod tests {
         .await
         .unwrap();
         assert_eq!(
-            SstInfo {
-                time_range: Some((
-                    Timestamp::new_millisecond(2000),
-                    Timestamp::new_millisecond(2000)
-                )),
-            },
-            s1
+            Some((
+                Timestamp::new_millisecond(2000),
+                Timestamp::new_millisecond(2000)
+            )),
+            s1.time_range,
         );
 
         let s2 = ParquetWriter::new(
@@ -427,13 +429,11 @@ mod tests {
         .await
         .unwrap();
         assert_eq!(
-            SstInfo {
-                time_range: Some((
-                    Timestamp::new_millisecond(3000),
-                    Timestamp::new_millisecond(5002)
-                )),
-            },
-            s2
+            Some((
+                Timestamp::new_millisecond(3000),
+                Timestamp::new_millisecond(5002)
+            )),
+            s2.time_range,
         );
 
         let s3 = ParquetWriter::new(
@@ -446,13 +446,11 @@ mod tests {
         .unwrap();
 
         assert_eq!(
-            SstInfo {
-                time_range: Some((
-                    Timestamp::new_millisecond(6000),
-                    Timestamp::new_millisecond(8000)
-                )),
-            },
-            s3
+            Some((
+                Timestamp::new_millisecond(6000),
+                Timestamp::new_millisecond(8000)
+            )),
+            s3.time_range
         );
 
         let output_files = ["o1.parquet", "o2.parquet", "o3.parquet"]
@@ -464,6 +462,7 @@ mod tests {
                         file_name: f.to_string(),
                         level: 1,
                         time_range: None,
+                        file_size: 0,
                     },
                     Arc::new(crate::test_util::access_layer_util::MockAccessLayer {}),
                     new_noop_file_purger(),

src/storage/src/file_purger.rs

@@ -143,7 +143,7 @@ mod tests {
         let iter = memtable.iter(&IterContext::default()).unwrap();
         let sst_path = "table1";
         let layer = Arc::new(FsAccessLayer::new(sst_path, os.clone()));
-        let _sst_info = layer
+        let sst_info = layer
             .write_sst(sst_file_name, Source::Iter(iter), &WriteOptions {})
             .await
             .unwrap();
@@ -155,6 +155,7 @@ mod tests {
                     file_name: sst_file_name.to_string(),
                     time_range: None,
                     level: 0,
+                    file_size: sst_info.file_size,
                 },
                 layer.clone(),
                 file_purger,

src/storage/src/flush.rs

@@ -203,7 +203,10 @@ impl<S: LogStore> FlushJob<S> {
             let sst_layer = self.sst_layer.clone();
 
             futures.push(async move {
-                let SstInfo { time_range } = sst_layer
+                let SstInfo {
+                    time_range,
+                    file_size,
+                } = sst_layer
                     .write_sst(&file_name, Source::Iter(iter), &WriteOptions::default())
                     .await?;
 
@@ -212,6 +215,7 @@ impl<S: LogStore> FlushJob<S> {
                     file_name,
                     time_range,
                     level: 0,
+                    file_size,
                 })
             });
         }

src/storage/src/lib.rs

@@ -41,3 +41,4 @@ pub mod write_batch;
 
 pub use engine::EngineImpl;
 mod file_purger;
+mod statistics_collector;

src/storage/src/manifest/test_utils.rs

@@ -20,6 +20,8 @@ use crate::metadata::RegionMetadata;
 use crate::sst::FileMeta;
 use crate::test_util::descriptor_util::RegionDescBuilder;
 
+pub const DEFAULT_TEST_FILE_SIZE: u64 = 1024;
+
 pub fn build_region_meta() -> RegionMetadata {
     let region_name = "region-0";
     let desc = RegionDescBuilder::new(region_name)
@@ -45,6 +47,7 @@ pub fn build_region_edit(
                 file_name: f.to_string(),
                 time_range: None,
                 level: 0,
+                file_size: DEFAULT_TEST_FILE_SIZE,
             })
             .collect(),
         files_to_remove: files_to_remove
@@ -54,6 +57,7 @@ pub fn build_region_edit(
                 file_name: f.to_string(),
                 time_range: None,
                 level: 0,
+                file_size: DEFAULT_TEST_FILE_SIZE,
             })
             .collect(),
     }

src/storage/src/region.rs

@@ -46,6 +46,7 @@ pub use crate::region::writer::{AlterContext, RegionWriter, RegionWriterRef, Wri
 use crate::schema::compat::CompatWrite;
 use crate::snapshot::SnapshotImpl;
 use crate::sst::AccessLayerRef;
+use crate::statistics_collector::StatisticsCollectorRef;
 use crate::version::{
     Version, VersionControl, VersionControlRef, VersionEdit, INIT_COMMITTED_SEQUENCE,
 };
@@ -125,6 +126,10 @@ impl<S: LogStore> Region for RegionImpl<S> {
     async fn close(&self) -> Result<()> {
         self.inner.close().await
     }
+
+    fn disk_usage_bytes(&self) -> u64 {
+        self.inner.statistics_collector.disk_usage_bytes()
+    }
 }
 
 /// Storage related config for region.
@@ -190,6 +195,7 @@ impl<S: LogStore> RegionImpl<S> {
         let name = metadata.name().to_string();
         let version_control = VersionControl::with_version(version);
         let wal = Wal::new(id, store_config.log_store);
+        let statistics_collector = StatisticsCollectorRef::default();
 
         let inner = Arc::new(RegionInner {
             shared: Arc::new(SharedData {
@@ -201,13 +207,15 @@ impl<S: LogStore> RegionImpl<S> {
                 store_config.memtable_builder,
                 store_config.engine_config.clone(),
                 store_config.ttl,
+                statistics_collector.clone(),
             )),
             wal,
             flush_strategy: store_config.flush_strategy,
             flush_scheduler: store_config.flush_scheduler,
             compaction_scheduler: store_config.compaction_scheduler,
             sst_layer: store_config.sst_layer,
             manifest: store_config.manifest,
+            statistics_collector,
         });
 
         RegionImpl { inner }
@@ -221,12 +229,14 @@ impl<S: LogStore> RegionImpl<S> {
         store_config: StoreConfig<S>,
         _opts: &OpenOptions,
     ) -> Result<Option<RegionImpl<S>>> {
+        let statistics_collector = StatisticsCollectorRef::default();
         // Load version meta data from manifest.
         let (version, mut recovered_metadata) = match Self::recover_from_manifest(
             &store_config.manifest,
             &store_config.memtable_builder,
             &store_config.sst_layer,
             &store_config.file_purger,
+            &statistics_collector,
         )
         .await?
         {
@@ -281,6 +291,7 @@ impl<S: LogStore> RegionImpl<S> {
             store_config.memtable_builder,
             store_config.engine_config.clone(),
             store_config.ttl,
+            statistics_collector.clone(),
         ));
         let writer_ctx = WriterContext {
             shared: &shared,
@@ -306,6 +317,7 @@ impl<S: LogStore> RegionImpl<S> {
             compaction_scheduler: store_config.compaction_scheduler,
             sst_layer: store_config.sst_layer,
             manifest: store_config.manifest,
+            statistics_collector,
         });
 
         Ok(Some(RegionImpl { inner }))
@@ -321,6 +333,7 @@ impl<S: LogStore> RegionImpl<S> {
         memtable_builder: &MemtableBuilderRef,
         sst_layer: &AccessLayerRef,
         file_purger: &FilePurgerRef,
+        statistics_collector: &StatisticsCollectorRef,
     ) -> Result<(Option<Version>, RecoveredMetadataMap)> {
         let (start, end) = Self::manifest_scan_range();
         let mut iter = manifest.scan(start, end).await?;
@@ -352,7 +365,12 @@ impl<S: LogStore> RegionImpl<S> {
                             file_purger.clone(),
                         ));
                         for (manifest_version, action) in actions.drain(..) {
-                            version = Self::replay_edit(manifest_version, action, version);
+                            version = Self::replay_edit(
+                                manifest_version,
+                                action,
+                                version,
+                                statistics_collector,
+                            );
                         }
                     }
                     (RegionMetaAction::Change(c), Some(v)) => {
@@ -365,7 +383,12 @@ impl<S: LogStore> RegionImpl<S> {
                         version = None;
                     }
                     (action, Some(v)) => {
-                        version = Self::replay_edit(manifest_version, action, Some(v));
+                        version = Self::replay_edit(
+                            manifest_version,
+                            action,
+                            Some(v),
+                            statistics_collector,
+                        );
                     }
                 }
             }
@@ -391,6 +414,7 @@ impl<S: LogStore> RegionImpl<S> {
         manifest_version: ManifestVersion,
         action: RegionMetaAction,
         version: Option<Version>,
+        statistics_collector: &StatisticsCollectorRef,
     ) -> Option<Version> {
         if let RegionMetaAction::Edit(e) = action {
             let edit = VersionEdit {
@@ -400,6 +424,10 @@ impl<S: LogStore> RegionImpl<S> {
                 manifest_version,
                 max_memtable_id: None,
             };
+            statistics_collector
+                .increase_disk_usage_bytes(edit.files_to_add.iter().map(|f| f.file_size).sum());
+            statistics_collector
+                .descrease_disk_usage_bytes(edit.files_to_remove.iter().map(|f| f.file_size).sum());
             version.map(|mut v| {
                 v.apply_edit(edit);
                 v
@@ -483,6 +511,7 @@ struct RegionInner<S: LogStore> {
     compaction_scheduler: CompactionSchedulerRef<S>,
     sst_layer: AccessLayerRef,
     manifest: RegionManifest,
+    statistics_collector: StatisticsCollectorRef,
 }
 
 impl<S: LogStore> RegionInner<S> {

src/storage/src/region/tests.rs

@@ -297,17 +297,20 @@ async fn test_recover_region_manifets() {
         SchedulerConfig::default(),
         NoopFilePurgeHandler,
     ));
+    let statistics_collector = StatisticsCollectorRef::default();
     // Recover from empty
     assert!(RegionImpl::<NoopLogStore>::recover_from_manifest(
         &manifest,
         &memtable_builder,
         &sst_layer,
-        &file_purger
+        &file_purger,
+        &statistics_collector,
     )
     .await
     .unwrap()
     .0
     .is_none());
+    assert_eq!(statistics_collector.disk_usage_bytes(), 0);
 
     {
         // save some actions into region_meta
@@ -340,12 +343,14 @@ async fn test_recover_region_manifets() {
             .unwrap();
     }
 
+    let statistics_collector = StatisticsCollectorRef::default();
     // try to recover
     let (version, recovered_metadata) = RegionImpl::<NoopLogStore>::recover_from_manifest(
         &manifest,
         &memtable_builder,
         &sst_layer,
         &file_purger,
+        &statistics_collector,
     )
     .await
     .unwrap();
@@ -368,4 +373,7 @@ async fn test_recover_region_manifets() {
 
     // check manifest state
     assert_eq!(3, manifest.last_version());
+
+    // 3 files in total
+    assert_eq!(statistics_collector.disk_usage_bytes(), 3 * DEFAULT_TEST_FILE_SIZE);
 }