diff --git a/src/mito/src/table/test_util/mock_engine.rs b/src/mito/src/table/test_util/mock_engine.rs index 7725a955a92..d097f2a9be6 100644 --- a/src/mito/src/table/test_util/mock_engine.rs +++ b/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 { diff --git a/src/storage/src/compaction/strategy.rs b/src/storage/src/compaction/strategy.rs index 24186b19ef1..52181e58df3 100644 --- a/src/storage/src/compaction/strategy.rs +++ b/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, diff --git a/src/storage/src/compaction/task.rs b/src/storage/src/compaction/task.rs index 3cee9ea0497..89705a5800e 100644 --- a/src/storage/src/compaction/task.rs +++ b/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, }) } } diff --git a/src/storage/src/compaction/writer.rs b/src/storage/src/compaction/writer.rs index b8fa8cb8047..b09d2184e5f 100644 --- a/src/storage/src/compaction/writer.rs +++ b/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(), diff --git a/src/storage/src/file_purger.rs b/src/storage/src/file_purger.rs index 78334396703..f9a9840197b 100644 --- a/src/storage/src/file_purger.rs +++ b/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, diff --git a/src/storage/src/flush.rs b/src/storage/src/flush.rs index 39890110109..abeac00e952 100644 --- a/src/storage/src/flush.rs +++ b/src/storage/src/flush.rs @@ -203,7 +203,10 @@ impl FlushJob { 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 FlushJob { file_name, time_range, level: 0, + file_size, }) }); } diff --git a/src/storage/src/manifest/test_utils.rs b/src/storage/src/manifest/test_utils.rs index 1ff4a710284..d75f20f160f 100644 --- a/src/storage/src/manifest/test_utils.rs +++ b/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(), } diff --git a/src/storage/src/region.rs b/src/storage/src/region.rs index d91fad2de0c..edff614f6b2 100644 --- a/src/storage/src/region.rs +++ b/src/storage/src/region.rs @@ -125,6 +125,16 @@ impl Region for RegionImpl { 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::()) + .sum() + } } /// Storage related config for region. diff --git a/src/storage/src/region/tests.rs b/src/storage/src/region/tests.rs index f9921041475..adf6c9c8aca 100644 --- a/src/storage/src/region/tests.rs +++ b/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() diff --git a/src/storage/src/sst.rs b/src/storage/src/sst.rs index f7ebdc1af25..eb8cc7971c5 100644 --- a/src/storage/src/sst.rs +++ b/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, } } diff --git a/src/storage/src/sst/parquet.rs b/src/storage/src/sst/parquet.rs index eeae7396584..6aa346670ee 100644 --- a/src/storage/src/sst/parquet.rs +++ b/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()); diff --git a/src/store-api/src/storage/region.rs b/src/store-api/src/storage/region.rs index 032fbc5c2f5..3d267579bb2 100644 --- a/src/store-api/src/storage/region.rs +++ b/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.