feat: track disk usage of regions (#1125)

* feat: track disk usage of regions

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

* calculate disk usage when call

* add default on file meta

---------

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

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/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

@@ -125,6 +125,16 @@ impl<S: LogStore> Region for RegionImpl<S> {
     async fn close(&self) -> Result<()> {
         self.inner.close().await
     }
+
+    fn disk_usage_bytes(&self) -> u64 {
+        let version = self.inner.version_control().current();
+        version
+            .ssts()
+            .levels()
+            .iter()
+            .map(|level_ssts| level_ssts.files().map(|sst| sst.file_size()).sum::<u64>())
+            .sum()
+    }
 }
 
 /// Storage related config for region.

src/storage/src/region/tests.rs

@@ -302,7 +302,7 @@ async fn test_recover_region_manifets() {
         &manifest,
         &memtable_builder,
         &sst_layer,
-        &file_purger
+        &file_purger,
     )
     .await
     .unwrap()

src/storage/src/sst.rs

@@ -232,6 +232,11 @@ impl FileHandle {
     pub fn meta(&self) -> FileMeta {
         self.inner.meta.clone()
     }
+
+    #[inline]
+    pub fn file_size(&self) -> u64 {
+        self.inner.meta.file_size
+    }
 }
 
 /// Actually data of [FileHandle].
@@ -286,7 +291,8 @@ impl FileHandleInner {
 }
 
 /// Immutable metadata of a sst file.
-#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
+#[serde(default)]
 pub struct FileMeta {
     /// Region of file.
     pub region_id: RegionId,
@@ -296,6 +302,8 @@ pub struct FileMeta {
     pub time_range: Option<(Timestamp, Timestamp)>,
     /// SST level of the file.
     pub level: Level,
+    /// Size of the file.
+    pub file_size: u64,
 }
 
 #[derive(Debug, Default)]
@@ -317,6 +325,7 @@ pub struct ReadOptions {
 #[derive(Debug, PartialEq)]
 pub struct SstInfo {
     pub time_range: Option<(Timestamp, Timestamp)>,
+    pub file_size: u64,
 }
 
 /// SST access layer.
@@ -437,6 +446,7 @@ mod tests {
             file_name: name.to_string(),
             time_range: None,
             level,
+            file_size: 0,
         }
     }
 

src/storage/src/sst/parquet.rs

@@ -130,7 +130,17 @@ impl<'a> ParquetWriter<'a> {
         object.write(buf).await.context(WriteObjectSnafu {
             path: object.path(),
         })?;
-        Ok(SstInfo { time_range })
+        let file_size = object
+            .metadata()
+            .await
+            .context(WriteObjectSnafu {
+                path: object.path(),
+            })?
+            .content_length();
+        Ok(SstInfo {
+            time_range,
+            file_size,
+        })
     }
 }
 
@@ -667,7 +677,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();
@@ -679,6 +692,7 @@ mod tests {
             )),
             time_range
         );
+        assert_ne!(file_size, 0);
         let operator = ObjectStore::new(
             Fs::default()
                 .root(dir.path().to_str().unwrap())
@@ -740,7 +754,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();
@@ -752,6 +769,7 @@ mod tests {
             )),
             time_range
         );
+        assert_ne!(file_size, 0);
         let operator = ObjectStore::new(
             Fs::default()
                 .root(dir.path().to_str().unwrap())
@@ -853,7 +871,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();
@@ -865,6 +886,7 @@ mod tests {
             )),
             time_range
         );
+        assert_ne!(file_size, 0);
 
         let projected_schema =
             Arc::new(ProjectedSchema::new(schema, Some(vec![1, 0, 3, 2])).unwrap());

src/store-api/src/storage/region.rs

@@ -74,6 +74,8 @@ pub trait Region: Send + Sync + Clone + std::fmt::Debug + 'static {
     async fn alter(&self, request: AlterRequest) -> Result<(), Self::Error>;
 
     async fn close(&self) -> Result<(), Self::Error>;
+
+    fn disk_usage_bytes(&self) -> u64;
 }
 
 /// Context for write operations.