diff --git a/Cargo.lock b/Cargo.lock index 3aed170a76c..07e2f4c3f25 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -912,6 +912,12 @@ dependencies = [ "thiserror 2.0.17", ] +[[package]] +name = "asyncband" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2d85fd3d291fabcc40c7232c92c280ec1754fd7b5d7ea769f143222143e179a" + [[package]] name = "asynchronous-codec" version = "0.7.0" @@ -2274,7 +2280,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -7918,7 +7924,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "07033963ba89ebaf1584d767badaa2e8fcec21aedea6b8c0346d487d49c28667" dependencies = [ "cfg-if", - "windows-targets 0.48.5", + "windows-targets 0.53.5", ] [[package]] @@ -9588,8 +9594,7 @@ dependencies = [ [[package]] name = "object_store_opendal" version = "0.58.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88f165780495c17aa3ce86846600504198c3fffd99073521552751c2430fa6ac" +source = "git+https://github.com/GreptimeTeam/opendal.git?rev=c03e56802b47de3ccfc8fe028e59a459b278a0d2#c03e56802b47de3ccfc8fe028e59a459b278a0d2" dependencies = [ "async-trait", "bytes", @@ -9676,9 +9681,9 @@ checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" [[package]] name = "opendal" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f20562cc7447fcc915fc5c23df305a412ea80a733c9f2fd9e2d267e2815be6d" +checksum = "9fe43e16d96bed57937eb7c4a28559bf4c5fc4a3951656505d1fdd4f856349b4" dependencies = [ "ctor", "opendal-core", @@ -9700,11 +9705,12 @@ dependencies = [ [[package]] name = "opendal-core" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ec75551ff4cf3e57da98979f6a937aaa9ddb3915bf68cc17d03df733be6646ed" +checksum = "de4566a412776e3f65d53dd9fceced5b432a0a9c9a1e8d50ccf58823795e5c5b" dependencies = [ "anyhow", + "asyncband", "base64 0.23.1", "bytes", "futures", @@ -9712,7 +9718,6 @@ dependencies = [ "jiff", "log", "md-5 0.11.0", - "mea", "percent-encoding", "quick-xml 0.41.0", "reqsign-core", @@ -9726,9 +9731,9 @@ dependencies = [ [[package]] name = "opendal-http-transport-reqwest" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad4d4f19c3ce01126a30611f8e544eaa217104a278c889ac17c9374fe4f9e4ef" +checksum = "f772ce6c137ab43647726116d505a649a88d41928fcd813438eb30b02312bcbb" dependencies = [ "bytes", "futures", @@ -9740,21 +9745,21 @@ dependencies = [ [[package]] name = "opendal-layer-concurrent-limit" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "249ac5b0aa5a7a6c3737342d10456067937f9c9a6f3f02544271f7908ab91081" +checksum = "3734c648617e5ad724d49175dde8323550624f9614f50e9e81b7aafae30b1734" dependencies = [ + "asyncband", "futures", "http 1.5.0", - "mea", "opendal-core", ] [[package]] name = "opendal-layer-logging" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c75411ab00f77851ff086b686c1e9ca8175ac18c15afa2cb75b9036436cb06c" +checksum = "0a8daab5a91eac76f94516e7c71e677195c9221840ce8f5e71815e3fc8207742" dependencies = [ "log", "opendal-core", @@ -9762,9 +9767,9 @@ dependencies = [ [[package]] name = "opendal-layer-observe-metrics-common" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "12dbc31b27ee8e5f658ae6c8514b8276059e5235e2f85138a42e6ed172324876" +checksum = "0fbe664be117c293cfa3980a0d2fc238bbb38eae5f6c191f376e88ccb830d26b" dependencies = [ "futures", "http 1.5.0", @@ -9773,9 +9778,9 @@ dependencies = [ [[package]] name = "opendal-layer-prometheus" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c14e6cb0b3ea8d83bef751613b8bf4512678afe8b49c2c96a633515bc0b34d79" +checksum = "b32078d8248555ee6271332ec480b78306a44f0d902ea4721ccac00d9b866a21" dependencies = [ "opendal-core", "opendal-layer-observe-metrics-common", @@ -9784,9 +9789,9 @@ dependencies = [ [[package]] name = "opendal-layer-retry" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "80b7738bd5f233ad8da39af9b9316b9b7a4eaddd91e8e32a1e19b7030688121d" +checksum = "eb8e22cb1f86ef9af771a8947d3ccf380c8f5b76236f628314dd41c9ba73331a" dependencies = [ "backon", "log", @@ -9795,9 +9800,9 @@ dependencies = [ [[package]] name = "opendal-layer-timeout" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a704141924500f3803c05ed871b53305d2a2f11cb5ef20160c3ee688a1857f66" +checksum = "08ac6c77fe0d1b18e0064a02c19b98dbebcab78bf5e7753ad1f3abbc18850ff4" dependencies = [ "opendal-core", "tokio", @@ -9805,9 +9810,9 @@ dependencies = [ [[package]] name = "opendal-layer-tracing" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4827cf8b2ebd9123b1fb7286c421be079948e2b963a452f1f7e8aee2c12e7997" +checksum = "603807927a799ea981f62a3b0012e17181e39f468793b78ca8bedda66d19813c" dependencies = [ "futures", "http 1.5.0", @@ -9817,9 +9822,9 @@ dependencies = [ [[package]] name = "opendal-service-azblob" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3310fbbb48f111c6f590473c2cd15e1b7f8e384444b0d4e328f0464c864d767" +checksum = "c6cf31b02e7b44191d97c07254eea109b5da9c17210a29e478c8d0dc26dffb1f" dependencies = [ "base64 0.23.1", "bytes", @@ -9838,9 +9843,9 @@ dependencies = [ [[package]] name = "opendal-service-azure-common" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7348c88edf15af435b7be930077746b569fac5e738c1bf6a363b675e7317c9df" +checksum = "8107324c3f3cb9970fe20526c2c341b7c5769acbd4699bab73bd525f5fd63d62" dependencies = [ "http 1.5.0", "opendal-core", @@ -9848,9 +9853,9 @@ dependencies = [ [[package]] name = "opendal-service-fs" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "826c4e17a30643b888fe983897f9a4b23b07066e1d069727a923cc8fb419a702" +checksum = "a0e7d8974aebd33be899200c9d0d50b2141fd5fa5c462b29e20653e3ddbfe339" dependencies = [ "bytes", "log", @@ -9862,9 +9867,9 @@ dependencies = [ [[package]] name = "opendal-service-gcs" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "007f3fba63c21e516c956b891e96ff9892d8175662bfb781cdada9d3766a11e6" +checksum = "3d80eca6f9227bdf65d43e5b8d5fe75dca8cfab1da07571442a0420400847bea" dependencies = [ "async-trait", "bytes", @@ -9879,13 +9884,14 @@ dependencies = [ "serde", "serde_json", "tokio", + "uuid", ] [[package]] name = "opendal-service-http" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2fd1f0f28bd052e068d9bb4bc5d30a09de22a4f7cea24558dbeef2b91090cc3" +checksum = "9ac1553d6b57ca79e200e5c7ded9ab708f2016dadb28ede0f2da71e02fd0b960" dependencies = [ "http 1.5.0", "log", @@ -9895,11 +9901,11 @@ dependencies = [ [[package]] name = "opendal-service-mysql" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d4a25f582068ce30ba24eb3efd4ac16ca833c87b5282d8d79634fdbead74d84" +checksum = "e15592be32b105cd152ab398b3db847e22c8c5f230d9b34ff095ea3d0dc6fd97" dependencies = [ - "mea", + "asyncband", "opendal-core", "serde", "sqlx", @@ -9907,9 +9913,9 @@ dependencies = [ [[package]] name = "opendal-service-oss" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd528ec2d49c5ca69e674ffed7b3e0686fb9cfcfea0596870de381467fda4f1b" +checksum = "a0ebfa355b1d4edb0e5b2a5b86d8de9f811fae1d7baca6d65246b230b2afc0bd" dependencies = [ "bytes", "http 1.5.0", @@ -9924,9 +9930,9 @@ dependencies = [ [[package]] name = "opendal-service-s3" -version = "0.58.1" +version = "0.59.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "58e80cdf192d7eff05feed747894d64f81905ac4eaf132edf7ea270abdd2d663" +checksum = "2942f2d8d3d4953c0e8d07cd1244879628d084948e48ce6b6d808c5938f7858d" dependencies = [ "base64 0.23.1", "bytes", @@ -12409,9 +12415,9 @@ dependencies = [ [[package]] name = "reqsign-aws-core" -version = "3.0.3" +version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4af084e1f3cbf3e67e0c972765399bce54ecec804cceba46b39a8331f3c1bff" +checksum = "bac4749b7dfa7bfaccd01eb03e9dc795ed37e3f20d6f0f38e2c67ee85ad6bc86" dependencies = [ "bytes", "form_urlencoded", @@ -12430,9 +12436,9 @@ dependencies = [ [[package]] name = "reqsign-aws-v4" -version = "3.1.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ac5b3b7cefa28933792b439186459f77f19f9b6edbeab41b8b187150361a206" +checksum = "ff250f0fd0b913fbd565e405acc553da0f13bde30bfb5403178c9d0313cdc15f" dependencies = [ "bytes", "http 1.5.0", @@ -12466,9 +12472,9 @@ dependencies = [ [[package]] name = "reqsign-core" -version = "3.2.1" +version = "3.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c07dd510b1e1b9b241883e483358147fb2ed2d497a7b39b065ba61eb93deceb0" +checksum = "ff052daffb0599681c50f85c59e7236438976efe991ab864edd9f3b235501a0f" dependencies = [ "anyhow", "base64 0.23.1", @@ -12479,6 +12485,7 @@ dependencies = [ "http 1.5.0", "jiff", "log", + "mea", "percent-encoding", "rsa", "serde", @@ -12490,9 +12497,9 @@ dependencies = [ [[package]] name = "reqsign-file-read-tokio" -version = "3.0.4" +version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "663d9d55abd0df0830ef0ae43708297cc1371cf4e8ca91f3ac813c309cca8c98" +checksum = "b3235df90a6bca681aa47dd86f2393d122a6d77042aa8a7c81e218cd45c5bfc0" dependencies = [ "anyhow", "reqsign-core", @@ -16815,7 +16822,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf221c93e13a30d793f7645a0e7762c55d169dbb0a49671918a2319d289b10bb" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fade46261ce..babbf32751e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -181,7 +181,11 @@ nalgebra = "0.33" nix = { version = "0.30.1", default-features = false, features = ["event", "fs", "process"] } notify = "8.0" num_cpus = "1.16" -object_store_opendal = "0.58" +# Keep object_store 0.13 compatibility while using OpenDAL 0.59. +# TODO(#9240): Switch back to the official bridge after the DataFusion/Arrow/Parquet +# upgrade aligns object_store versions. Keep the bridge and OpenDAL releases compatible. +# https://github.com/GreptimeTeam/greptimedb/issues/9240 +object_store_opendal = { git = "https://github.com/GreptimeTeam/opendal.git", rev = "c03e56802b47de3ccfc8fe028e59a459b278a0d2" } once_cell = "1.18" opentelemetry-proto = { version = "0.32", features = [ "gen-tonic", diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index e8146c11330..58b24c3c350 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -16,6 +16,7 @@ use std::sync::Arc; use std::time::{Duration, Instant}; use async_stream::try_stream; +use common_base::readable_size::ReadableSize; use common_runtime::runtime::RuntimeTrait; use common_telemetry::warn; use common_time::Timestamp; @@ -49,7 +50,7 @@ use crate::sst::location::{self, region_dir_from_table_dir}; use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::sst::parquet::writer::ParquetWriter; use crate::sst::parquet::{SstInfo, WriteOptions}; -use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FormatType}; +use crate::sst::{DEFAULT_WRITE_CONCURRENCY, FormatType}; pub type AccessLayerRef = Arc; /// SST write results. @@ -244,6 +245,7 @@ impl AccessLayer { } let attempted_files = file_ids.to_vec(); + // Deleter does not normalize leading slashes like Operator::delete does. let paths: Vec<_> = file_ids .iter() .map(|file_id| { @@ -252,6 +254,8 @@ impl AccessLayer { RegionFileId::new(region_id, *file_id), self.path_type, ) + .trim_start_matches('/') + .to_string() }) .collect(); @@ -264,7 +268,7 @@ impl AccessLayer { file_ids: attempted_files.clone(), })?; deleter - .delete_iter(paths.iter().map(String::as_str)) + .delete_iter(paths) .await .with_context(|_| DeleteSstsSnafu { region_id, @@ -292,7 +296,11 @@ impl AccessLayer { .collect(); let paths: Vec<_> = index_ids .iter() - .map(|index_id| location::index_file_path(&self.table_dir, *index_id, self.path_type)) + .map(|index_id| { + location::index_file_path(&self.table_dir, *index_id, self.path_type) + .trim_start_matches('/') + .to_string() + }) .collect(); let mut deleter = self @@ -303,7 +311,7 @@ impl AccessLayer { file_ids: file_ids.clone(), })?; deleter - .delete_iter(paths.iter().map(String::as_str)) + .delete_iter(paths) .await .context(DeleteIndexesSnafu { file_ids: file_ids.clone(), @@ -473,9 +481,9 @@ impl AccessLayer { pub(crate) async fn put_sst( &self, data: &bytes::Bytes, - region_id: RegionId, - sst_info: &SstInfo, + region_file_id: RegionFileId, cache_manager: &CacheManagerRef, + write_buffer_size: ReadableSize, ) -> Result { if let Some(write_cache) = cache_manager.write_cache() { // Write to cache and upload to remote store @@ -487,27 +495,27 @@ impl AccessLayer { remote_store: self.object_store.clone(), }; write_cache - .put_and_upload_sst(data, region_id, sst_info, upload_request) + .put_and_upload_sst(data, region_file_id, upload_request, write_buffer_size) .await } else { let start = Instant::now(); - let cleaner = TempFileCleaner::new(region_id, self.object_store.clone()); + let cleaner = + TempFileCleaner::new(region_file_id.region_id(), self.object_store.clone()); let path_provider = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type); - let sst_file_path = - path_provider.build_sst_file_path(RegionFileId::new(region_id, sst_info.file_id)); + let sst_file_path = path_provider.build_sst_file_path(region_file_id); let mut writer = self .object_store .writer_with(&sst_file_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(write_buffer_size.as_bytes() as usize) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .context(OpenDalSnafu)?; if let Err(err) = writer.write(data.clone()).await.context(OpenDalSnafu) { - cleaner.clean_by_file_id(sst_info.file_id).await; + cleaner.clean_by_file_id(region_file_id.file_id()).await; return Err(err); } if let Err(err) = writer.close().await.context(OpenDalSnafu) { - cleaner.clean_by_file_id(sst_info.file_id).await; + cleaner.clean_by_file_id(region_file_id.file_id()).await; return Err(err); } let mut metrics = Metrics::new(WriteType::Flush); @@ -767,3 +775,111 @@ impl FilePathProvider for RegionFilePathFactory { location::sst_file_path(&self.table_dir, file_id, self.path_type) } } + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use bytes::Bytes; + use common_test_util::temp_dir::create_temp_dir; + + use super::*; + use crate::cache::CacheManager; + use crate::cache::test_util::new_fs_store; + use crate::test_util::TestEnv; + use crate::test_util::sst_util::WriteChunkRecorder; + + #[rstest::rstest] + #[tokio::test] + async fn test_put_sst_uses_configured_buffer_size( + #[values(false, true)] enable_write_cache: bool, + #[values(1024, 4096)] chunk_size: usize, + ) { + let mut env = TestEnv::new().await; + let remote_chunks = WriteChunkRecorder::default(); + let remote_store = env.init_object_store_manager().layer(remote_chunks.layer()); + let local_dir = create_temp_dir("encoded-sst-cache"); + let local_chunks = WriteChunkRecorder::default(); + let local_store = + new_fs_store(local_dir.path().to_str().unwrap()).layer(local_chunks.layer()); + let write_cache = if enable_write_cache { + Some( + env.create_write_cache(local_store.clone(), ReadableSize::mb(10)) + .await, + ) + } else { + None + }; + let cache_manager = Arc::new( + CacheManager::builder() + .write_cache(write_cache.clone()) + .build(), + ); + let access_layer = AccessLayer::new( + "test", + PathType::Bare, + remote_store.clone(), + env.get_puffin_manager(), + env.get_intermediate_manager(), + ); + let file_id = RegionFileId::new(RegionId::new(1024, 1), FileId::random()); + // Two full chunks and a partial tail distinguish the configured size from the default. + let encoded = Bytes::from( + (0..2 * chunk_size + 17) + .map(|i| (i % 251) as u8) + .collect::>(), + ); + access_layer + .put_sst( + &encoded, + file_id, + &cache_manager, + ReadableSize(chunk_size as u64), + ) + .await + .unwrap(); + + let remote_path = location::sst_file_path("test", file_id, PathType::Bare); + remote_chunks.assert_chunks(&remote_path, chunk_size, encoded.len()); + assert_eq!( + remote_store.read(&remote_path).await.unwrap().to_bytes(), + encoded + ); + if let Some(write_cache) = write_cache { + let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); + let cache_path = write_cache.file_cache().cache_file_path(key); + local_chunks.assert_chunks(&cache_path, chunk_size, encoded.len()); + assert_eq!( + local_store.read(&cache_path).await.unwrap().to_bytes(), + encoded + ); + } + } + + #[tokio::test] + async fn test_new_fs_cache_store() { + let root = common_test_util::temp_dir::create_temp_dir("fs-cache-store"); + let dirs = [ + root.path().join(ATOMIC_WRITE_DIR), + PathBuf::from(join_dir( + root.path().to_str().unwrap(), + OLD_ATOMIC_WRITE_DIR, + )), + ]; + for dir in &dirs { + tokio::fs::create_dir_all(dir).await.unwrap(); + tokio::fs::write(dir.join("stale"), b"stale").await.unwrap(); + } + let store = new_fs_cache_store(root.path().to_str().unwrap()) + .await + .unwrap(); + for dir in &dirs { + assert!(!dir.join("stale").exists()); + } + store.write("index", "contents").await.unwrap(); + assert_eq!( + tokio::fs::read(root.path().join("index")).await.unwrap(), + b"contents" + ); + } +} diff --git a/src/mito2/src/cache/write_cache.rs b/src/mito2/src/cache/write_cache.rs index 03c7bdddd73..0e96c2f16a2 100644 --- a/src/mito2/src/cache/write_cache.rs +++ b/src/mito2/src/cache/write_cache.rs @@ -38,8 +38,8 @@ use crate::sst::file::RegionFileId; use crate::sst::index::IndexerBuilderImpl; use crate::sst::index::intermediate::IntermediateManager; use crate::sst::index::puffin_manager::{PuffinManagerFactory, SstPuffinManager}; +use crate::sst::parquet::WriteOptions; use crate::sst::parquet::writer::ParquetWriter; -use crate::sst::parquet::{SstInfo, WriteOptions}; use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY}; /// Wraps the remote object store used by write cache uploads. @@ -176,11 +176,12 @@ impl WriteCache { pub(crate) async fn put_and_upload_sst( &self, data: &bytes::Bytes, - region_id: RegionId, - sst_info: &SstInfo, + region_file_id: RegionFileId, upload_request: SstUploadRequest, + write_buffer_size: ReadableSize, ) -> Result { - let file_id = sst_info.file_id; + let region_id = region_file_id.region_id(); + let file_id = region_file_id.file_id(); let mut metrics = Metrics::new(WriteType::Flush); // Create index key for the SST file @@ -192,7 +193,8 @@ impl WriteCache { let store = self.file_cache.local_store(); let cleaner = TempFileCleaner::new(region_id, store.clone()); let write_res = store - .write(&cache_path, data.clone()) + .write_with(&cache_path, data.clone()) + .chunk(write_buffer_size.as_bytes() as usize) .await .context(crate::error::OpenDalSnafu); if let Err(e) = write_res { @@ -204,7 +206,6 @@ impl WriteCache { // Upload to remote store let upload_start = Instant::now(); - let region_file_id = RegionFileId::new(region_id, file_id); let remote_path = upload_request .dest_path_provider .build_sst_file_path(region_file_id); @@ -214,8 +215,11 @@ impl WriteCache { parquet_key, &remote_path, &upload_request.remote_store, - // `put_and_upload_sst` is only used by the flush path. - OperationType::Flush, + UploadOptions { + // `put_and_upload_sst` is only used by the flush path. + op_type: OperationType::Flush, + write_buffer_size, + }, ) .await { @@ -312,7 +316,15 @@ impl WriteCache { .build_sst_file_path(RegionFileId::new(region_id, sst.file_id)); let start = Instant::now(); if let Err(e) = self - .upload(parquet_key, &parquet_path, remote_store, op_type) + .upload( + parquet_key, + &parquet_path, + remote_store, + UploadOptions { + op_type, + write_buffer_size: write_opts.write_buffer_size, + }, + ) .await { err = Some(e); @@ -328,7 +340,15 @@ impl WriteCache { .build_index_file_path(RegionFileId::new(region_id, sst.file_id)); let start = Instant::now(); if let Err(e) = self - .upload(puffin_key, &puffin_path, remote_store, op_type) + .upload( + puffin_key, + &puffin_path, + remote_store, + UploadOptions { + op_type, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) .await { err = Some(e); @@ -399,7 +419,7 @@ impl WriteCache { index_key: IndexKey, upload_path: &str, remote_store: &ObjectStore, - op_type: OperationType, + options: UploadOptions, ) -> Result<()> { let region_id = index_key.region_id; let file_id = index_key.file_id; @@ -425,11 +445,11 @@ impl WriteCache { let upload_store = self.upload_store_wrapper.as_ref().map_or_else( || remote_store.clone(), - |wrapper| wrapper.wrap(remote_store.clone(), op_type), + |wrapper| wrapper.wrap(remote_store.clone(), options.op_type), ); let mut writer = upload_store .writer_with(upload_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(options.write_buffer_size.as_bytes() as usize) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .context(error::OpenDalSnafu)? @@ -486,6 +506,14 @@ pub struct SstUploadRequest { pub remote_store: ObjectStore, } +/// Policies for uploading a cached SST or index file. +pub(crate) struct UploadOptions { + /// Origin of the upload, used by the remote store wrapper. + pub op_type: OperationType, + /// Buffer size for the remote writer. + pub write_buffer_size: ReadableSize, +} + /// A structs to track files to upload and clean them if upload failed. pub(crate) struct UploadTracker { /// Id of the region to track. @@ -568,7 +596,7 @@ mod tests { use crate::sst::parquet::reader::ParquetReaderBuilder; use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - new_flat_source_from_record_batches, new_record_batch_by_range, + WriteChunkRecorder, new_flat_source_from_record_batches, new_record_batch_by_range, sst_file_handle_with_file_id, sst_region_metadata, }; @@ -617,7 +645,15 @@ mod tests { local_store.write(&cache_path, data.clone()).await.unwrap(); write_cache - .upload(key, upload_path, &original_store, OperationType::Compact) + .upload( + key, + upload_path, + &original_store, + UploadOptions { + op_type: OperationType::Compact, + write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE, + }, + ) .await .unwrap(); @@ -633,16 +669,22 @@ mod tests { ); } + #[rstest::rstest] #[tokio::test] - async fn test_write_and_upload_sst() { + async fn test_write_and_upload_sst( + #[values(OperationType::Flush, OperationType::Compact)] op_type: OperationType, + ) { // TODO(QuenKar): maybe find a way to create some object server for testing, // and now just use local file system to mock. let mut env = TestEnv::new().await; - let mock_store = env.init_object_store_manager(); + let remote_chunks = WriteChunkRecorder::default(); + let mock_store = env.init_object_store_manager().layer(remote_chunks.layer()); let path_provider = RegionFilePathFactory::new("test".to_string(), PathType::Bare); let local_dir = create_temp_dir(""); - let local_store = new_fs_store(local_dir.path().to_str().unwrap()); + let local_chunks = WriteChunkRecorder::default(); + let local_store = + new_fs_store(local_dir.path().to_str().unwrap()).layer(local_chunks.layer()); let write_cache = env .create_write_cache(local_store.clone(), ReadableSize::mb(10)) @@ -658,7 +700,7 @@ mod tests { ]); let write_request = SstWriteRequest { - op_type: OperationType::Flush, + op_type, metadata, source, storage: None, @@ -681,6 +723,7 @@ mod tests { }; let write_opts = WriteOptions { + write_buffer_size: ReadableSize::kb(1), row_group_size: 512, ..Default::default() }; @@ -710,6 +753,14 @@ mod tests { .await .unwrap(); assert_eq!(remote_data.to_vec(), cache_data.to_vec()); + let chunk_size = write_opts.write_buffer_size.as_bytes() as usize; + assert!(remote_data.len() > chunk_size); + remote_chunks.assert_chunks(&sst_upload_path, chunk_size, remote_data.len()); + local_chunks.assert_chunks( + &write_cache.file_cache.cache_file_path(key), + chunk_size, + cache_data.len(), + ); // Check write cache contains the index key let index_key = IndexKey::new(region_id, file_id, FileType::Puffin(0)); @@ -721,6 +772,11 @@ mod tests { .await .unwrap(); assert_eq!(remote_index_data.to_vec(), cache_index_data.to_vec()); + remote_chunks.assert_chunks( + &index_upload_path, + DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize, + remote_index_data.len(), + ); // Removes the file from the cache. let sst_index_key = IndexKey::new(region_id, file_id, FileType::Parquet); diff --git a/src/mito2/src/config.rs b/src/mito2/src/config.rs index 6ec790c52b0..21ecb266d9f 100644 --- a/src/mito2/src/config.rs +++ b/src/mito2/src/config.rs @@ -159,6 +159,7 @@ pub struct MitoConfig { // Other configs: /// Buffer size for SST writing. + // TODO(#9240): Support per-object-store write buffer sizes for backend-specific tuning. pub sst_write_buffer_size: ReadableSize, /// Maximum number of SST files to scan concurrently (default 384). pub max_concurrent_scan_files: usize, diff --git a/src/mito2/src/flush.rs b/src/mito2/src/flush.rs index 66df07788fc..883a44845aa 100644 --- a/src/mito2/src/flush.rs +++ b/src/mito2/src/flush.rs @@ -63,7 +63,7 @@ use crate::request::{ }; use crate::schedule::CancellableTaskState; use crate::schedule::scheduler::{Job, SchedulerRef}; -use crate::sst::file::{FileMeta, UncommittedSsts}; +use crate::sst::file::{FileMeta, RegionFileId, UncommittedSsts}; use crate::sst::parquet::metadata::extract_primary_key_range; use crate::sst::parquet::{ DEFAULT_READ_BATCH_SIZE, DEFAULT_ROW_GROUP_SIZE, SstInfo, WriteOptions, flat_format, @@ -737,12 +737,18 @@ impl RegionFlushTask { let access_layer = self.access_layer.clone(); let cache_manager = self.cache_manager.clone(); let region_id = version.metadata.region_id; + let write_buffer_size = write_opts.write_buffer_size; let semaphore = self.flush_semaphore.clone(); let uncommitted = uncommitted.clone(); let task = common_runtime::spawn_global(async move { let _permit = semaphore.acquire().await.unwrap(); let metrics = access_layer - .put_sst(&encoded.data, region_id, &encoded.sst_info, &cache_manager) + .put_sst( + &encoded.data, + RegionFileId::new(region_id, encoded.sst_info.file_id), + &cache_manager, + write_buffer_size, + ) .await?; uncommitted.track(std::slice::from_ref(&encoded.sst_info)); FLUSH_FILE_TOTAL.inc(); diff --git a/src/mito2/src/sst/index.rs b/src/mito2/src/sst/index.rs index 8a1db35aaa8..b55edeeba59 100644 --- a/src/mito2/src/sst/index.rs +++ b/src/mito2/src/sst/index.rs @@ -48,7 +48,7 @@ use vector_index::creator::VectorIndexer; use crate::access_layer::{AccessLayerRef, FilePathProvider, OperationType, RegionFilePathFactory}; use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey}; -use crate::cache::write_cache::{UploadTracker, WriteCacheRef}; +use crate::cache::write_cache::{UploadOptions, UploadTracker, WriteCacheRef}; use crate::cache::{CacheManagerRef, CacheStrategy}; #[cfg(feature = "vector_index")] use crate::config::VectorIndexConfig; @@ -1035,7 +1035,10 @@ impl IndexBuildTask { puffin_key, &puffin_path, remote_store, - OperationType::Compact, + UploadOptions { + op_type: OperationType::Compact, + write_buffer_size: crate::sst::DEFAULT_WRITE_BUFFER_SIZE, + }, ) .await { diff --git a/src/mito2/src/sst/parquet.rs b/src/mito2/src/sst/parquet.rs index 3d90a360925..31ca70442b1 100644 --- a/src/mito2/src/sst/parquet.rs +++ b/src/mito2/src/sst/parquet.rs @@ -178,10 +178,10 @@ mod tests { }; use crate::test_util::TestEnv; use crate::test_util::sst_util::{ - build_test_binary_test_region_metadata, new_flat_source_from_record_batches, - new_primary_key, new_record_batch_by_range, new_record_batch_with_custom_sequence, - new_sparse_primary_key, sst_file_handle, sst_file_handle_with_file_id, sst_region_metadata, - sst_region_metadata_with_encoding, + WriteChunkRecorder, build_test_binary_test_region_metadata, + new_flat_source_from_record_batches, new_primary_key, new_record_batch_by_range, + new_record_batch_with_custom_sequence, new_sparse_primary_key, sst_file_handle, + sst_file_handle_with_file_id, sst_region_metadata, sst_region_metadata_with_encoding, }; const FILE_DIR: &str = "/"; @@ -697,12 +697,14 @@ mod tests { .await; } + #[rstest::rstest] #[tokio::test] - async fn test_write_multiple_files() { + async fn test_write_multiple_files(#[values(1024, 4096)] write_buffer_size: usize) { common_telemetry::init_default_ut_logging(); // create test env let mut env = TestEnv::new().await; - let object_store = env.init_object_store_manager(); + let chunks = WriteChunkRecorder::default(); + let object_store = env.init_object_store_manager().layer(chunks.layer()); let metadata = Arc::new(sst_region_metadata()); let batches = vec![ new_record_batch_by_range(&["a", "a"], 0, 1000), @@ -717,9 +719,9 @@ mod tests { let source = new_flat_source_from_record_batches(batches); let write_opts = WriteOptions { + write_buffer_size: ReadableSize(write_buffer_size as u64), row_group_size: 50, max_file_size: Some(1024 * 16), - ..Default::default() }; let path_provider = RegionFilePathFactory { @@ -732,7 +734,7 @@ mod tests { metadata.clone(), IndexConfig::default(), NoopIndexBuilder, - path_provider, + path_provider.clone(), &mut metrics, ) .await; @@ -743,8 +745,18 @@ mod tests { .unwrap(); assert_eq!(2, files.len()); + // The configured buffer size must reach every writer, including split files. + assert_eq!(files.len(), chunks.num_files()); + let mut rows_read = 0; for f in &files { + assert!(f.file_size > write_buffer_size as u64); + chunks.assert_chunks( + &path_provider + .build_sst_file_path(RegionFileId::new(metadata.region_id, f.file_id)), + write_buffer_size, + f.file_size as usize, + ); let file_handle = sst_file_handle_with_file_id( f.file_id, f.time_range.0.value(), diff --git a/src/mito2/src/sst/parquet/writer.rs b/src/mito2/src/sst/parquet/writer.rs index b9d5016173e..45c8bdcbbe6 100644 --- a/src/mito2/src/sst/parquet/writer.rs +++ b/src/mito2/src/sst/parquet/writer.rs @@ -63,8 +63,7 @@ use crate::sst::parquet::flat_format::{ use crate::sst::parquet::format::{PrimaryKeyArray, PrimaryKeyWriteFormat}; use crate::sst::parquet::{PARQUET_METADATA_KEY, SstInfo, WriteOptions}; use crate::sst::{ - DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, - maybe_wrap_schema, + DEFAULT_WRITE_CONCURRENCY, FlatSchemaOptions, SeriesEstimator, maybe_wrap_schema, }; /// Converts a flat RecordBatch for writing to parquet. @@ -179,7 +178,11 @@ pub struct ParquetWriter<'a, F: WriterFactory, I: IndexerBuilder, P: FilePathPro pub trait WriterFactory { type Writer: AsyncWrite + Send + Unpin; - fn create(&mut self, file_path: &str) -> impl Future>; + fn create( + &mut self, + file_path: &str, + write_buffer_size: usize, + ) -> impl Future>; } pub struct ObjectStoreWriterFactory { @@ -189,10 +192,10 @@ pub struct ObjectStoreWriterFactory { impl WriterFactory for ObjectStoreWriterFactory { type Writer = Compat; - async fn create(&mut self, file_path: &str) -> Result { + async fn create(&mut self, file_path: &str, write_buffer_size: usize) -> Result { self.object_store .writer_with(file_path) - .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize) + .chunk(write_buffer_size) .concurrent(DEFAULT_WRITE_CONCURRENCY) .await .map(|v| v.into_futures_async_write().compat_write()) @@ -573,7 +576,9 @@ where self.current_file, )); let writer = SizeAwareWriter::new( - self.writer_factory.create(&sst_file_path).await?, + self.writer_factory + .create(&sst_file_path, opts.write_buffer_size.as_bytes() as usize) + .await?, self.bytes_written.clone(), ); let arrow_writer = diff --git a/src/mito2/src/test_util/sst_util.rs b/src/mito2/src/test_util/sst_util.rs index 3f635c6e949..fd6321b0895 100644 --- a/src/mito2/src/test_util/sst_util.rs +++ b/src/mito2/src/test_util/sst_util.rs @@ -14,7 +14,8 @@ //! Utilities for testing SSTs. -use std::sync::Arc; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; use api::v1::{OpType, SemanticType}; use arrow_schema::Schema; @@ -28,6 +29,7 @@ use datatypes::prelude::ConcreteDataType; use datatypes::schema::{ColumnSchema, SkippingIndexOptions}; use datatypes::value::ValueRef; use mito_codec::row_converter::{DensePrimaryKeyCodec, PrimaryKeyCodecExt, SortField}; +use object_store::layers::mock::{self, MockLayer, MockLayerBuilder, oio}; use store_api::metadata::{ ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef, }; @@ -45,6 +47,73 @@ use crate::test_util::{new_batch_builder, new_noop_file_purger}; /// Test region id. const REGION_ID: RegionId = RegionId::new(0, 0); +/// Records write chunks at the object-store boundary without replacing filesystem I/O. +#[derive(Default)] +pub(crate) struct WriteChunkRecorder { + chunk_sizes: Arc>>>, +} + +impl WriteChunkRecorder { + pub(crate) fn layer(&self) -> MockLayer { + let chunk_sizes = self.chunk_sizes.clone(); + MockLayerBuilder::default() + .writer_factory(Arc::new(move |path, _, inner| { + Box::new(ChunkRecordingWriter { + inner, + path: path.to_string(), + chunk_sizes: chunk_sizes.clone(), + }) + })) + .build() + .unwrap() + } + + pub(crate) fn num_files(&self) -> usize { + self.chunk_sizes.lock().unwrap().len() + } + + pub(crate) fn assert_chunks(&self, path: &str, chunk_size: usize, file_size: usize) { + let mut expected = vec![chunk_size; file_size / chunk_size]; + let remainder = file_size % chunk_size; + if remainder > 0 { + expected.push(remainder); + } + assert_eq!( + self.chunk_sizes.lock().unwrap().get(path), + Some(&expected), + "unexpected write chunks for {path}" + ); + } +} + +struct ChunkRecordingWriter { + inner: oio::Writer, + path: String, + chunk_sizes: Arc>>>, +} + +impl oio::Write for ChunkRecordingWriter { + async fn write(&mut self, buffer: mock::Buffer) -> mock::Result<()> { + let size = buffer.len(); + self.inner.write(buffer).await?; + self.chunk_sizes + .lock() + .unwrap() + .entry(self.path.clone()) + .or_default() + .push(size); + Ok(()) + } + + async fn close(&mut self) -> mock::Result { + self.inner.close().await + } + + async fn abort(&mut self) -> mock::Result<()> { + self.inner.abort().await + } +} + /// Creates a new region metadata for testing SSTs with specified encoding. /// /// Dense schema: tag_0, tag_1, field_0, ts diff --git a/src/object-store/Cargo.toml b/src/object-store/Cargo.toml index f9b367c3dce..412f3563298 100644 --- a/src/object-store/Cargo.toml +++ b/src/object-store/Cargo.toml @@ -27,7 +27,7 @@ derive_builder = { workspace = true, optional = true } futures.workspace = true humantime-serde.workspace = true lazy_static.workspace = true -opendal = { version = "0.58", features = [ +opendal = { version = "0.59.2", features = [ "layers-tracing", "layers-prometheus", "services-azblob", @@ -37,7 +37,7 @@ opendal = { version = "0.58", features = [ "services-oss", "services-s3", ] } -opendal-http-transport-reqwest = "0.58" +opendal-http-transport-reqwest = "0.59.2" prometheus.workspace = true reqwest.workspace = true serde.workspace = true diff --git a/src/object-store/src/layers/mock.rs b/src/object-store/src/layers/mock.rs index 6b919aabb69..270f776f935 100644 --- a/src/object-store/src/layers/mock.rs +++ b/src/object-store/src/layers/mock.rs @@ -96,6 +96,7 @@ impl Service for MockService { type Lister = oio::Lister; type Deleter = oio::Deleter; type Copier = oio::Copier; + type Composer = oio::Composer; fn info(&self) -> opendal::raw::ServiceInfo { self.inner.info() @@ -165,7 +166,6 @@ impl Service for MockService { from: &str, to: &str, args: OpCopy, - opts: opendal::raw::OpCopier, ) -> Result { if let Some(result) = self .copy_interceptor @@ -176,7 +176,16 @@ impl Service for MockService { return Ok(Box::new(oio::OneShotCopier::completed()) as oio::Copier); } - self.inner.copy(ctx, from, to, args, opts) + self.inner.copy(ctx, from, to, args) + } + + fn compose( + &self, + ctx: &OperationContext, + to: &str, + args: opendal::raw::OpCompose, + ) -> Result { + self.inner.compose(ctx, to, args) } async fn rename( diff --git a/src/object-store/src/secure_fs.rs b/src/object-store/src/secure_fs.rs index 581cc2837c0..e96a6852e1e 100644 --- a/src/object-store/src/secure_fs.rs +++ b/src/object-store/src/secure_fs.rs @@ -24,8 +24,8 @@ use cap_std::fs::{Dir, DirEntry, OpenOptions, ReadDir}; use opendal::layers::SimulateLayer; use opendal::raw::*; use opendal::{ - Buffer, BytesRange, Capability, EntryMode, Error, ErrorKind, Metadata, OperationContext, - Operator, Result, + Buffer, BytesRange, Capability, EntryMode, Error, ErrorKind, Metadata, MetadataBuilder, + OperationContext, Operator, Result, }; use tokio::io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}; @@ -151,19 +151,18 @@ fn parse_write_error(error: io::Error, if_not_exists: bool) -> Error { } fn metadata_from_fs(metadata: cap_std::fs::Metadata) -> Result { - let mode = if metadata.is_dir() { - EntryMode::DIR + let mut builder = if metadata.is_dir() { + MetadataBuilder::dir() } else if metadata.is_file() { - EntryMode::FILE + MetadataBuilder::file(metadata.len()) } else { - EntryMode::Unknown + MetadataBuilder::unknown() }; - Ok(Metadata::new(mode) - .with_content_length(metadata.len()) - .with_last_modified(Timestamp::try_from( - metadata.modified().map_err(new_std_io_error)?.into_std(), - )?)) + builder.last_modified(Timestamp::try_from( + metadata.modified().map_err(new_std_io_error)?.into_std(), + )?); + Ok(builder.build()) } #[derive(Clone, Debug)] @@ -205,6 +204,7 @@ impl Service for SecureFsBackend { type Lister = SecureFsLister; type Deleter = oio::OneShotDeleter; type Copier = (); + type Composer = (); fn info(&self) -> ServiceInfo { self.info.clone() @@ -287,14 +287,7 @@ impl Service for SecureFsBackend { }) } - fn copy( - &self, - _: &OperationContext, - _: &str, - _: &str, - _: OpCopy, - _: OpCopier, - ) -> Result { + fn copy(&self, _: &OperationContext, _: &str, _: &str, _: OpCopy) -> Result { Err(Error::new( ErrorKind::Unsupported, "operation is not supported", @@ -435,11 +428,11 @@ impl oio::Write for SecureFsWriter { file.flush().await.map_err(new_std_io_error)?; file.sync_all().await.map_err(new_std_io_error)?; let metadata = file.metadata().await.map_err(new_std_io_error)?; - Ok(Metadata::new(EntryMode::FILE) - .with_content_length(metadata.len()) - .with_last_modified(Timestamp::try_from( - metadata.modified().map_err(new_std_io_error)?, - )?)) + let mut builder = MetadataBuilder::file(metadata.len()); + builder.last_modified(Timestamp::try_from( + metadata.modified().map_err(new_std_io_error)?, + )?); + Ok(builder.build()) } async fn abort(&mut self) -> Result<()> { @@ -505,7 +498,7 @@ impl oio::List for SecureFsLister { } else { &display_prefix }, - Metadata::new(EntryMode::DIR), + MetadataBuilder::dir().build(), ); self.entries = vec![current_path].into_iter(); } @@ -575,7 +568,7 @@ fn read_list_entry(entry: DirEntry, display_prefix: &str) -> io::Result match metadata_from_fs(metadata) {