diff --git a/Cargo.lock b/Cargo.lock index d40ca8b5baa..59c8421e955 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8798,6 +8798,7 @@ dependencies = [ "bytes", "chrono", "common-base", + "common-cache", "common-datasource", "common-error", "common-function", diff --git a/Cargo.toml b/Cargo.toml index a5e3db37960..1f84eb7dbcb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -301,6 +301,7 @@ client = { path = "src/client" } cmd = { path = "src/cmd", default-features = false } common-base = { path = "src/common/base" } common-batcher = { path = "src/common/batcher" } +common-cache = { path = "src/common/cache" } common-catalog = { path = "src/common/catalog" } common-config = { path = "src/common/config" } common-datasource = { path = "src/common/datasource" } diff --git a/src/mito2/Cargo.toml b/src/mito2/Cargo.toml index a5e26fcd60f..411cf643145 100644 --- a/src/mito2/Cargo.toml +++ b/src/mito2/Cargo.toml @@ -28,6 +28,7 @@ bytemuck.workspace = true bytes.workspace = true chrono.workspace = true common-base.workspace = true +common-cache.workspace = true common-datasource.workspace = true common-error.workspace = true common-function.workspace = true diff --git a/src/mito2/src/access_layer.rs b/src/mito2/src/access_layer.rs index 3067c6098a1..7a80687958b 100644 --- a/src/mito2/src/access_layer.rs +++ b/src/mito2/src/access_layer.rs @@ -446,26 +446,29 @@ impl AccessLayer { let parquet_metadata = parquet_metadata.clone(); let region_metadata = region_metadata.clone(); let cache_manager = cache_manager.clone(); - // Compact cache preparation is best-effort. Run the entire operation in one - // detached blocking task so it neither blocks an async worker nor delays the - // SST write. - runtime.spawn_blocking(move || { - match prepare_sst_meta_sync( - &file_path, - Arc::unwrap_or_clone(parquet_metadata), - Some(region_metadata), - page_index_policy, - ) { + // Prepare and cache metadata in the background without delaying the SST write. + let task_runtime = runtime.clone(); + runtime.spawn(async move { + let preparation = task_runtime.spawn_blocking(move || { + prepare_sst_meta_sync( + &file_path, + Arc::unwrap_or_clone(parquet_metadata), + Some(region_metadata), + page_index_policy, + ) + }).await.context(crate::error::JoinSnafu).and_then(|result| result); + match preparation { Ok(SstMetaPreparation::Prepared(metadata)) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, true); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); + cache_manager.put_sst_meta_data(file_id, metadata.decoded).await; } Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( encoding_error; "Failed to encode parquet metadata for cache, file: {}", - file_path + file_id.file_id() ), Err(err) => { - warn!(err; "Failed to cache parquet metadata for {}", file_path); + warn!(err; "Failed to cache parquet metadata for {}", file_id.file_id()); } } }); diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index c4b8ffa0d73..4cff42043ce 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -241,8 +241,8 @@ pub(crate) struct CompactSstMeta { /// Both forms produced by decoding metadata after a cache miss. #[derive(Debug)] pub(crate) struct PreparedSstMeta { - compact: Arc, - decoded: Arc, + pub(crate) compact: Arc, + pub(crate) decoded: Arc, } impl PreparedSstMeta { @@ -731,17 +731,96 @@ impl CacheStrategy { } } - /// Gets fused SST metadata with cache metrics tracking. - pub(crate) async fn get_sst_meta_data( + #[cfg(test)] + pub(crate) async fn get_or_load_sst_meta( &self, file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + load: F, + ) -> Result<(Arc, bool)> + where + F: std::future::Future, bool)>>, + { + self.get_or_load_observed_sst_meta(file_id, page_index_policy, async { + (load.await, SstMetaResolution::default()) + }) + .await + .0 + } + + /// Shares the resolved source on success and failure, without sharing I/O accounting. + pub(crate) async fn get_or_load_observed_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + load: F, + ) -> (Result<(Arc, bool)>, SstMetaResolution) + where + F: std::future::Future, bool)>, SstMetaResolution)>, + { + let cache = match self { + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager.sst_decoded_meta_cache.as_ref() + } + Self::Disabled => None, + }; + let Some(cache) = cache else { + return load.await; + }; + let key = SstDecodedMetaKey::new(file_id, page_index_policy); + let result = cache + .try_get_with(&key, async { + let (result, resolution) = load.await; + let (metadata, cache_miss) = result.map_err(|source| SstMetaLoadError { + source: Arc::new(source), + resolution, + })?; + let value = SstMetaLoadResult::new(metadata, cache_miss, resolution); + Ok::<_, SstMetaLoadError>(value) + }) + .await; + match result { + Ok(value) => ( + Ok((value.metadata.clone(), value.cache_miss)), + value.resolution, + ), + Err(error) => ( + Err(error.source.clone()).context(crate::error::LoadSstMetaSnafu), + error.resolution, + ), + } + } + + /// Observes memory residency for this caller before sharing decode or I/O work. + pub(crate) async fn observe_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + metrics: &mut MetadataCacheMetrics, + ) -> SstMetaObservation { + match self { + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager + .observe_sst_meta(file_id, page_index_policy, metrics) + .await + } + Self::Disabled => SstMetaObservation::Miss, + } + } + + /// Resolves an observed metadata entry and updates this caller's cache metrics. + pub(crate) async fn resolve_sst_meta( + &self, + observation: SstMetaObservation, + file_id: RegionFileId, metrics: &mut MetadataCacheMetrics, page_index_policy: PageIndexPolicy, ) -> Option> { match self { - CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager - .get_sst_meta_data( + Self::EnableAll(manager) | Self::Compaction(manager) => { + manager + .resolve_sst_meta( + observation, file_id, metrics, page_index_policy, @@ -749,7 +828,7 @@ impl CacheStrategy { ) .await } - CacheStrategy::Disabled => { + Self::Disabled => { metrics.cache_miss += 1; None } @@ -757,45 +836,47 @@ impl CacheStrategy { } /// Calls [CacheManager::get_sst_meta_data_from_mem_cache()]. - pub(crate) fn get_sst_meta_data_from_mem_cache( + pub(crate) async fn get_sst_meta_data_from_mem_cache( &self, file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.get_sst_meta_data_from_mem_cache(file_id, page_index_policy) + cache_manager + .get_sst_meta_data_from_mem_cache(file_id, page_index_policy) + .await } CacheStrategy::Disabled => None, } } /// Calls [CacheManager::get_parquet_meta_data_from_mem_cache()]. - pub fn get_parquet_meta_data_from_mem_cache( + pub async fn get_parquet_meta_data_from_mem_cache( &self, file_id: RegionFileId, ) -> Option> { self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip) + .await .map(|metadata| metadata.parquet_metadata()) } - /// Puts compact and decoded forms of SST metadata into the shared cache capacity. - pub(crate) fn put_prepared_sst_meta( + /// Calls [CacheManager::put_compact_sst_meta()]. + pub(crate) fn put_compact_sst_meta( &self, file_id: RegionFileId, - metadata: PreparedSstMeta, - retain_decoded: bool, + metadata: Arc, ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, retain_decoded); + cache_manager.put_compact_sst_meta(file_id, metadata); } CacheStrategy::Disabled => {} } } /// Calls [CacheManager::put_parquet_meta_data()]. - pub fn put_parquet_meta_data( + pub async fn put_parquet_meta_data( &self, file_id: RegionFileId, metadata: Arc, @@ -803,7 +884,9 @@ impl CacheStrategy { ) { match self { CacheStrategy::EnableAll(cache_manager) | CacheStrategy::Compaction(cache_manager) => { - cache_manager.put_parquet_meta_data(file_id, metadata, region_metadata); + cache_manager + .put_parquet_meta_data(file_id, metadata, region_metadata) + .await; } CacheStrategy::Disabled => {} } @@ -827,13 +910,13 @@ impl CacheStrategy { } /// Calls [CacheManager::remove_parquet_meta_data()]. - pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + pub async fn remove_parquet_meta_data(&self, file_id: RegionFileId) { match self { CacheStrategy::EnableAll(cache_manager) => { - cache_manager.remove_parquet_meta_data(file_id); + cache_manager.remove_parquet_meta_data(file_id).await; } CacheStrategy::Compaction(cache_manager) => { - cache_manager.remove_parquet_meta_data(file_id); + cache_manager.remove_parquet_meta_data(file_id).await; } CacheStrategy::Disabled => {} } @@ -1100,8 +1183,8 @@ impl CacheManager { CacheManagerBuilder::default() } - /// Gets fused SST metadata with metrics tracking. - /// Tries in-memory cache first, then file cache, updating metrics accordingly. + /// Gets fused SST metadata with metrics tracking for direct cache tests. + #[cfg(test)] pub(crate) async fn get_sst_meta_data( &self, file_id: RegionFileId, @@ -1109,42 +1192,54 @@ impl CacheManager { page_index_policy: PageIndexPolicy, runtime: &Runtime, ) -> Option> { + let observation = self + .observe_sst_meta(file_id, page_index_policy, metrics) + .await; + self.resolve_sst_meta(observation, file_id, metrics, page_index_policy, runtime) + .await + } + + async fn observe_sst_meta( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + metrics: &mut MetadataCacheMetrics, + ) -> SstMetaObservation { let cache_key = SstMetaKey(file_id.region_id(), file_id.file_id()); let compact = self .get_compact_sst_meta(&cache_key) .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); let decoded = self - .get_decoded_sst_meta(&cache_key) + .get_decoded_sst_meta(file_id, page_index_policy) + .await .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + let observation = match (compact, decoded) { + (_, Some(decoded)) => SstMetaObservation::Decoded(decoded), + (Some(compact), None) => SstMetaObservation::Compact(compact), + (None, None) => SstMetaObservation::Miss, + }; + observation.record(metrics); + observation + } - if let Some(compact) = compact { - CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); - metrics.mem_cache_hit += 1; - if let Some(decoded) = decoded { - CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - return Some(decoded); - } - - CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - match compact.decode(runtime).await { - Ok(decoded) => { - self.put_sst_meta_data(file_id, decoded.clone()); - return Some(decoded); - } + async fn resolve_sst_meta( + &self, + observation: SstMetaObservation, + file_id: RegionFileId, + metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + runtime: &Runtime, + ) -> Option> { + match observation { + SstMetaObservation::Decoded(decoded) => return Some(decoded), + SstMetaObservation::Compact(compact) => match compact.decode(runtime).await { + Ok(decoded) => return Some(decoded), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); - self.remove_parquet_meta_data(file_id); + self.remove_parquet_meta_data(file_id).await; } - } - } else if let Some(decoded) = decoded { - // Metadata produced by a new SST writer can enter the decoded tier before its compact - // representation is prepared. - CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); - CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(); - metrics.mem_cache_hit += 1; - return Some(decoded); - } else { - CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(); + }, + SstMetaObservation::Miss => {} } let key = IndexKey::new(file_id.region_id(), file_id.file_id(), FileType::Parquet); @@ -1159,7 +1254,7 @@ impl CacheManager { let decoded = metadata.decoded(); match metadata { SstMetaPreparation::Prepared(metadata) => { - self.put_prepared_sst_meta(file_id, metadata, true); + self.put_compact_sst_meta(file_id, metadata.compact); } SstMetaPreparation::DecodedOnly { encoding_error, .. } => { warn!( @@ -1187,55 +1282,61 @@ impl CacheManager { /// Gets cached fused SST metadata from in-memory cache. /// This method does not perform I/O. - pub(crate) fn get_sst_meta_data_from_mem_cache( + pub(crate) async fn get_sst_meta_data_from_mem_cache( &self, file_id: RegionFileId, page_index_policy: PageIndexPolicy, ) -> Option> { - let key = SstMetaKey(file_id.region_id(), file_id.file_id()); let value = self - .get_decoded_sst_meta(&key) - .filter(|metadata| metadata.satisfies_page_index_policy(page_index_policy)); + .get_sst_meta_data_from_mem_cache_inner(file_id, page_index_policy) + .await; update_hit_miss(value, SST_META_DECODED_TYPE) } + async fn get_sst_meta_data_from_mem_cache_inner( + &self, + file_id: RegionFileId, + page_index_policy: PageIndexPolicy, + ) -> Option> { + self.get_decoded_sst_meta(file_id, page_index_policy).await + } + /// Gets cached [ParquetMetaData] from in-memory cache. /// This method does not perform I/O. - pub fn get_parquet_meta_data_from_mem_cache( + pub async fn get_parquet_meta_data_from_mem_cache( &self, file_id: RegionFileId, ) -> Option> { self.get_sst_meta_data_from_mem_cache(file_id, PageIndexPolicy::Skip) + .await .map(|metadata| metadata.parquet_metadata()) } /// Puts fused SST metadata into the cache. - pub(crate) fn put_sst_meta_data(&self, file_id: RegionFileId, metadata: Arc) { + pub(crate) async fn put_sst_meta_data( + &self, + file_id: RegionFileId, + metadata: Arc, + ) { if let Some(cache) = &self.sst_decoded_meta_cache { - let key = SstMetaKey(file_id.region_id(), file_id.file_id()); - CACHE_BYTES - .with_label_values(&[SST_META_DECODED_TYPE]) - .add(decoded_meta_cache_weight(&key, &metadata).into()); - cache.insert(key, metadata); + let key = SstDecodedMetaKey::new(file_id, metadata.page_index_policy); + let value = SstMetaLoadResult::new(metadata, false, SstMetaResolution::default()); + cache.insert(key, value).await; } } - /// Puts a compact metadata entry and optionally retains its decoded acceleration entry. - pub(crate) fn put_prepared_sst_meta( + /// Puts a compact SST metadata entry into the cache. + pub(crate) fn put_compact_sst_meta( &self, file_id: RegionFileId, - metadata: PreparedSstMeta, - retain_decoded: bool, + metadata: Arc, ) { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { CACHE_BYTES .with_label_values(&[SST_META_TYPE]) - .add(meta_cache_weight(&key, &metadata.compact).into()); - cache.insert(key.clone(), metadata.compact); - } - if retain_decoded { - self.put_sst_meta_data(file_id, metadata.decoded); + .add(meta_cache_weight(&key, &metadata).into()); + cache.insert(key, metadata); } } @@ -1243,8 +1344,26 @@ impl CacheManager { self.sst_meta_cache.as_ref()?.get(key) } - fn get_decoded_sst_meta(&self, key: &SstMetaKey) -> Option> { - self.sst_decoded_meta_cache.as_ref()?.get(key) + async fn get_decoded_sst_meta( + &self, + file_id: RegionFileId, + requested: PageIndexPolicy, + ) -> Option> { + let cache = self.sst_decoded_meta_cache.as_ref()?; + for policy in [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ] { + // The key describes the loading request, while compact metadata may + // already contain stronger page-index guarantees. + if let Some(value) = cache.get(&SstDecodedMetaKey::new(file_id, policy)).await + && value.metadata.satisfies_page_index_policy(requested) + { + return Some(value.metadata.clone()); + } + } + None } /// Gets and decodes only the compact tier without promoting into the decoded cache. @@ -1264,14 +1383,14 @@ impl CacheManager { Ok(metadata) => Some(metadata), Err(err) => { warn!(err; "Failed to decode compact SST metadata, region_id: {}, file_id: {}", file_id.region_id(), file_id.file_id()); - self.remove_parquet_meta_data(file_id); + self.remove_parquet_meta_data(file_id).await; None } } } /// Puts [ParquetMetaData] into the cache. - pub fn put_parquet_meta_data( + pub async fn put_parquet_meta_data( &self, file_id: RegionFileId, metadata: Arc, @@ -1288,7 +1407,7 @@ impl CacheManager { Arc::unwrap_or_clone(metadata), region_metadata, ) { - Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)), + Ok(metadata) => self.put_sst_meta_data(file_id, Arc::new(metadata)).await, Err(err) => warn!( err; "Failed to decode region metadata while caching parquet metadata, region_id: {}, file_id: {}", file_id.region_id(), @@ -1299,13 +1418,21 @@ impl CacheManager { } /// Removes [ParquetMetaData] from the cache. - pub fn remove_parquet_meta_data(&self, file_id: RegionFileId) { + pub async fn remove_parquet_meta_data(&self, file_id: RegionFileId) { let key = SstMetaKey(file_id.region_id(), file_id.file_id()); if let Some(cache) = &self.sst_meta_cache { cache.remove(&key); } if let Some(cache) = &self.sst_decoded_meta_cache { - cache.remove(&key); + for policy in [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ] { + cache + .invalidate(&SstDecodedMetaKey::new(file_id, policy)) + .await; + } } } @@ -1636,19 +1763,18 @@ impl CacheManagerBuilder { .build() }); let sst_decoded_meta_cache = (decoded_meta_capacity != 0).then(|| { - Cache::builder() - .max_capacity(decoded_meta_capacity) - .weigher(decoded_meta_cache_weight) - .eviction_listener(|k, v, cause| { - let size = decoded_meta_cache_weight(&k, &v); - CACHE_BYTES - .with_label_values(&[SST_META_DECODED_TYPE]) - .sub(size.into()); + common_cache::Cache::new( + "sst_decoded_metadata", + decoded_meta_capacity, + decoded_meta_cache_weight, + |_k, _v, cause| { CACHE_EVICTION - .with_label_values(&[SST_META_DECODED_TYPE, removal_cause_str(cause)]) + .with_label_values(&[SST_META_DECODED_TYPE, cause.as_str()]) .inc(); - }) - .build() + }, + ) + .with_size_metric(CACHE_BYTES.bind(&[SST_META_DECODED_TYPE])) + .with_metrics(crate::metrics::cache_loading_metrics(SST_META_DECODED_TYPE)) }); let vector_cache = (self.vector_cache_size != 0).then(|| { Cache::builder() @@ -1683,7 +1809,7 @@ impl CacheManagerBuilder { let prefilter_result_cache = (self.prefilter_result_cache_size != 0) .then(|| new_prefilter_result_cache(self.prefilter_result_cache_size)); let puffin_metadata_cache = - PuffinMetadataCache::new(self.puffin_metadata_size, &CACHE_BYTES); + PuffinMetadataCache::new(self.puffin_metadata_size, CACHE_BYTES.gauges()); let selector_result_cache = (self.selector_result_cache_size != 0).then(|| { Cache::builder() .max_capacity(self.selector_result_cache_size) @@ -1730,9 +1856,8 @@ fn meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { u32::try_from(size).unwrap_or(u32::MAX) } -fn decoded_meta_cache_weight(k: &SstMetaKey, v: &Arc) -> u32 { - let size = k.estimated_size() + v.parquet_metadata_size + v.region_metadata_weight; - u32::try_from(size).unwrap_or(u32::MAX) +fn decoded_meta_cache_weight(k: &SstDecodedMetaKey, v: &Arc) -> u32 { + SstMetaLoadResult::weight(k, &v.metadata) } fn vector_cache_weight(_k: &(ConcreteDataType, Value), v: &VectorRef) -> u32 { @@ -2119,10 +2244,96 @@ impl SelectorResultValue { } } -/// Maps (region id, file id) to fused SST metadata. +/// A caller's memory observation. Holding the entry avoids a second lookup after coordination. +pub(crate) enum SstMetaObservation { + Decoded(Arc), + Compact(Arc), + Miss, +} + +impl SstMetaObservation { + fn record(&self, metrics: &mut MetadataCacheMetrics) { + match self { + Self::Decoded(_) | Self::Compact(_) => { + metrics.mem_cache_hit += 1; + CACHE_HIT.with_label_values(&[SST_META_TYPE]).inc(); + match self { + Self::Decoded(_) => CACHE_HIT.with_label_values(&[SST_META_DECODED_TYPE]).inc(), + _ => CACHE_MISS.with_label_values(&[SST_META_DECODED_TYPE]).inc(), + } + } + Self::Miss => CACHE_MISS.with_label_values(&[SST_META_TYPE]).inc(), + } + } +} + +/// Resolved cache source, not physical I/O. Also returned when the subsequent load fails. +/// A caller that observed a memory miss can find a completed value at the native recheck; +/// it reuses that value's source without executing its initializer. This does not prove +/// that the caller waited for an in-flight load or performed a file-cache lookup. +/// Ordinary memory hits return before this resolution is applied. +/// A default value means no file-cache resolution was needed (or known). +#[derive(Clone, Copy, Default)] +pub(crate) struct SstMetaResolution { + pub(crate) file_cache_hit: usize, + pub(crate) cache_miss: usize, +} + +struct SstMetaLoadError { + source: Arc, + resolution: SstMetaResolution, +} + +/// Metadata returned by a shared load. Cache hits ignore the original load's miss flag. +struct SstMetaLoadResult { + metadata: Arc, + /// Preserves background-download eligibility for all requests sharing a storage load. + cache_miss: bool, + resolution: SstMetaResolution, +} + +impl SstMetaLoadResult { + fn weight(key: &SstDecodedMetaKey, metadata: &CachedSstMeta) -> u32 { + let size = size_of_val(key) + + size_of::() + + metadata.parquet_metadata_size + + metadata.region_metadata_weight; + u32::try_from(size).unwrap_or(u32::MAX) + } + + fn new( + metadata: Arc, + cache_miss: bool, + resolution: SstMetaResolution, + ) -> Arc { + Arc::new(Self { + metadata, + cache_miss, + resolution, + }) + } +} + +/// Separates loads that require different page-index guarantees. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +struct SstDecodedMetaKey { + file: SstMetaKey, + policy: mem::Discriminant, +} + +impl SstDecodedMetaKey { + fn new(file: RegionFileId, policy: PageIndexPolicy) -> Self { + Self { + file: SstMetaKey(file.region_id(), file.file_id()), + policy: mem::discriminant(&policy), + } + } +} + +/// Maps (region id, file id) to compact SST metadata. type SstMetaCache = Cache>; /// Maps (region id, file id) to decoded SST metadata retained as an acceleration tier. -type SstDecodedMetaCache = Cache>; +type SstDecodedMetaCache = common_cache::Cache>; /// Maps [Value] to a vector that holds this value repeatedly. /// /// e.g. `"hello" => ["hello", "hello", "hello"]` @@ -2169,7 +2380,7 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); let metadata = parquet_meta(); let mut metrics = MetadataCacheMetrics::default(); - cache.put_parquet_meta_data(file_id, metadata, None); + cache.put_parquet_meta_data(file_id, metadata, None).await; assert!( cache .get_sst_meta_data( @@ -2191,21 +2402,365 @@ mod tests { .is_none() ); - cache.put_page_ranges( - file_id.file_id(), - 1, - &[Range { start: 0, end: 5 }], - &[Bytes::from_static(b"abcde")], - ); - assert!( - cache - .get_page_ranges(file_id.file_id(), 1, &[Range { start: 0, end: 5 }]) - .is_none() - ); + assert!(cache.page_cache.is_none()); assert!(cache.write_cache().is_none()); } + #[tokio::test] + async fn test_sst_meta_load_coalescing() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let manager = Arc::new(CacheManager::builder().sst_meta_cache_size(2).build()); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let calls = AtomicUsize::new(0); + let gate = tokio::sync::Notify::new(); + let load = || async { + calls.fetch_add(1, Ordering::Relaxed); + gate.notified().await; + Ok((metadata.clone(), true)) + }; + let mut first = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + let mut second = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut second).is_pending()); + assert_eq!(1, calls.load(Ordering::Relaxed)); + gate.notify_waiters(); + let (first, second) = tokio::join!(first, second); + let first = first.unwrap(); + let second = second.unwrap(); + assert!(Arc::ptr_eq(&first.0, &second.0)); + assert!( + first.1 && second.1, + "all waiters preserve the remote miss flag" + ); + + // Moka applies capacity limits during maintenance, not synchronously on insertion. + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .run_pending_tasks() + .await; + // An entry larger than the decoded reservation must not be retained. + let mut next = + Box::pin(strategy.get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, load())); + assert!(futures::poll!(&mut next).is_pending()); + assert_eq!(2, calls.load(Ordering::Relaxed)); + drop(next); + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .run_pending_tasks() + .await; + assert_eq!( + 0, + manager + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .entry_count() + ); + } + + #[tokio::test] + async fn test_compact_sst_meta_stronger_policy_reused_from_skip_key() { + for policy in [PageIndexPolicy::Optional, PageIndexPolicy::Required] { + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (metadata, region) = sst_parquet_meta(); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let prepared = prepare_sst_meta( + "test.parquet", + with_offset_indexes(Arc::unwrap_or_clone(metadata)), + None, + policy, + &common_runtime::global_runtime(), + ) + .await + .unwrap(); + let SstMetaPreparation::Prepared(prepared) = prepared else { + panic!("valid metadata should produce a compact entry"); + }; + manager.put_compact_sst_meta(file_id, prepared.compact); + let mut metrics = MetadataCacheMetrics::default(); + let observation = strategy + .observe_sst_meta(file_id, PageIndexPolicy::Skip, &mut metrics) + .await; + assert!(matches!(&observation, SstMetaObservation::Compact(_))); + let (decoded, _) = strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + let value = strategy + .resolve_sst_meta( + observation, + file_id, + &mut MetadataCacheMetrics::default(), + PageIndexPolicy::Skip, + ) + .await + .expect("compact entry must decode without storage I/O"); + Ok((value, false)) + }) + .await + .unwrap(); + assert!(decoded.satisfies_page_index_policy(policy)); + let cache = manager.sst_decoded_meta_cache.as_ref().unwrap(); + cache.run_pending_tasks().await; + let weight = cache.weighted_size(); + assert_eq!(cache.entry_count(), 1); + + // Losing the compact copy must not force another decode or storage load. + manager + .sst_meta_cache + .as_ref() + .unwrap() + .invalidate(&SstMetaKey(file_id.region_id(), file_id.file_id())); + for requested in [PageIndexPolicy::Optional, policy] { + let mut metrics = MetadataCacheMetrics::default(); + let observation = strategy + .observe_sst_meta(file_id, requested, &mut metrics) + .await; + let SstMetaObservation::Decoded(cached) = observation else { + panic!("the completed Skip-key entry already satisfies this request"); + }; + assert!(Arc::ptr_eq(&decoded, &cached)); + assert_eq!(metrics.mem_cache_hit, 1); + assert_eq!(metrics.cache_miss, 0); + } + cache.run_pending_tasks().await; + assert_eq!(cache.entry_count(), 1); + assert_eq!(cache.weighted_size(), weight); + assert!( + cache + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Optional)) + .await + .is_none() + ); + assert!( + cache + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Required)) + .await + .is_none() + ); + } + } + + #[tokio::test] + async fn test_async_sst_meta_policy_reuse_and_invalidation() { + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (parquet_metadata, region) = sst_parquet_meta(); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let policies = [ + PageIndexPolicy::Skip, + PageIndexPolicy::Optional, + PageIndexPolicy::Required, + ]; + for loaded in policies { + let metadata = Arc::new( + CachedSstMeta::try_new_with_page_index_policy( + "test.parquet", + (*parquet_metadata).clone(), + None, + loaded, + ) + .unwrap(), + ); + let (result, miss) = strategy + .get_or_load_sst_meta(file_id, loaded, async { Ok((metadata.clone(), true)) }) + .await + .unwrap(); + assert!(Arc::ptr_eq(&metadata, &result)); + assert!(miss); + for requested in policies { + let mut metrics = MetadataCacheMetrics::default(); + let cached = strategy + .observe_sst_meta(file_id, requested, &mut metrics) + .await; + let cached = match cached { + SstMetaObservation::Decoded(value) => Some(value), + _ => None, + }; + assert_eq!( + satisfies_page_index_policy(loaded, requested), + cached.is_some() + ); + if let Some(cached) = cached { + assert!(Arc::ptr_eq(&metadata, &cached)); + assert_eq!(1, metrics.mem_cache_hit); + } + } + manager.remove_parquet_meta_data(file_id).await; + } + + // Populate all policy keys before removing the file, not just the last inserted key. + for loaded in policies { + let metadata = Arc::new( + CachedSstMeta::try_new_with_page_index_policy( + "test.parquet", + (*parquet_metadata).clone(), + None, + loaded, + ) + .unwrap(), + ); + manager.put_sst_meta_data(file_id, metadata).await; + } + manager.remove_parquet_meta_data(file_id).await; + for requested in policies { + assert!( + manager + .get_sst_meta_data_from_mem_cache(file_id, requested) + .await + .is_none() + ); + } + } + + #[tokio::test] + async fn test_sst_meta_load_cancellation_and_policy_isolation() { + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().sst_meta_cache_size(2000).build(), + )); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let mut first = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + std::future::pending(), + )); + assert!(futures::poll!(&mut first).is_pending()); + // A different policy or file does not wait for this load. + for (id, policy) in [ + (file_id, PageIndexPolicy::Optional), + (file_id, PageIndexPolicy::Required), + ( + RegionFileId::new(region.region_id, FileId::random()), + PageIndexPolicy::Skip, + ), + ] { + let mut independent = Box::pin( + strategy.get_or_load_sst_meta(id, policy, async { Ok((metadata.clone(), false)) }), + ); + assert!(futures::poll!(&mut independent).is_ready()); + } + let mut waiter = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { Ok((metadata.clone(), false)) }, + )); + assert!(futures::poll!(&mut waiter).is_pending()); + drop(first); + let result = tokio::time::timeout(std::time::Duration::from_secs(5), waiter) + .await + .unwrap() + .unwrap(); + assert!(Arc::ptr_eq(&metadata, &result.0)); + } + + #[tokio::test] + async fn test_sst_meta_load_failure_retry_and_disabled_cache() { + use common_error::ext::ErrorExt; + + let manager = Arc::new(CacheManager::builder().sst_meta_cache_size(2000).build()); + let strategy = CacheStrategy::EnableAll(manager); + let (metadata, region) = sst_parquet_meta(); + let metadata = Arc::new( + CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), + ); + let file_id = RegionFileId::new(region.region_id, FileId::random()); + let gate = tokio::sync::Notify::new(); + let mut first = Box::pin(strategy.get_or_load_observed_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { + gate.notified().await; + ( + UnexpectedSnafu { + reason: "load failed", + } + .fail(), + SstMetaResolution { + file_cache_hit: 0, + cache_miss: 1, + }, + ) + }, + )); + let mut waiter = Box::pin(strategy.get_or_load_observed_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { panic!("waiter must share the failure") }, + )); + assert!(futures::poll!(&mut first).is_pending()); + assert!(futures::poll!(&mut waiter).is_pending()); + gate.notify_waiters(); + let (first, first_resolution) = first.await; + let (second, second_resolution) = waiter.await; + for resolution in [first_resolution, second_resolution] { + assert_eq!(1, resolution.cache_miss); + assert_eq!(0, resolution.file_cache_hit); + } + let first = first.unwrap_err(); + let second = second.unwrap_err(); + assert_eq!(first.status_code(), second.status_code()); + match (first, second) { + ( + crate::error::Error::LoadSstMeta { source: a, .. }, + crate::error::Error::LoadSstMeta { source: b, .. }, + ) => assert!(Arc::ptr_eq(&a, &b)), + errors => panic!("unexpected errors: {errors:?}"), + } + assert!( + strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + Ok((metadata.clone(), true)) + }) + .await + .is_ok() + ); + + for strategy in [ + CacheStrategy::Disabled, + CacheStrategy::EnableAll(Arc::new(CacheManager::default())), + CacheStrategy::EnableAll(Arc::new( + CacheManager::builder().sst_meta_cache_size(1).build(), + )), + ] { + let mut blocked = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + std::future::pending(), + )); + assert!(futures::poll!(&mut blocked).is_pending()); + let mut independent = Box::pin(strategy.get_or_load_sst_meta( + file_id, + PageIndexPolicy::Skip, + async { Ok((metadata.clone(), true)) }, + )); + assert!(futures::poll!(&mut independent).is_ready()); + } + } + #[test] fn test_sst_meta_cache_splits_capacity() { let cache = CacheManager::builder().sst_meta_cache_size(101).build(); @@ -2220,13 +2775,8 @@ mod tests { .max_capacity() ); assert_eq!( - Some(50), - cache - .sst_decoded_meta_cache - .as_ref() - .unwrap() - .policy() - .max_capacity() + 50, + cache.sst_decoded_meta_cache.as_ref().unwrap().capacity() ); } @@ -2248,7 +2798,7 @@ mod tests { .is_none() ); let (metadata, region_metadata) = sst_parquet_meta(); - cache.put_parquet_meta_data(file_id, metadata, None); + cache.put_parquet_meta_data(file_id, metadata, None).await; let cached = cache .get_sst_meta_data( file_id, @@ -2270,7 +2820,7 @@ mod tests { .all(|key_value| key_value.key != PARQUET_METADATA_KEY) }) ); - cache.remove_parquet_meta_data(file_id); + cache.remove_parquet_meta_data(file_id).await; assert!( cache .get_sst_meta_data( @@ -2292,7 +2842,9 @@ mod tests { let file_id = RegionFileId::new(region_id, FileId::random()); let (metadata, region_metadata) = sst_parquet_meta(); - cache.put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())); + cache + .put_parquet_meta_data(file_id, metadata, Some(region_metadata.clone())) + .await; let cached = cache .get_sst_meta_data( @@ -2306,15 +2858,7 @@ mod tests { assert!(Arc::ptr_eq(®ion_metadata, &cached.region_metadata())); } - #[tokio::test] - async fn test_compact_sst_meta_round_trip() { - let cache = CacheManager::builder() - .sst_meta_cache_size(1024 * 1024) - .build(); - let region_id = RegionId::new(1, 1); - let file_id = RegionFileId::new(region_id, FileId::random()); - let (metadata, expected_region_metadata) = sst_parquet_meta(); - let metadata = Arc::unwrap_or_clone(metadata); + fn with_offset_indexes(metadata: ParquetMetaData) -> ParquetMetaData { let offset_indexes = metadata .row_groups() .iter() @@ -2333,10 +2877,22 @@ mod tests { .collect() }) .collect(); - let metadata = metadata + metadata .into_builder() .set_offset_index(Some(offset_indexes)) + .build() + } + + #[tokio::test] + async fn test_compact_sst_meta_round_trip() { + let cache = CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) .build(); + let region_id = RegionId::new(1, 1); + let file_id = RegionFileId::new(region_id, FileId::random()); + let (metadata, expected_region_metadata) = sst_parquet_meta(); + let metadata = Arc::unwrap_or_clone(metadata); + let metadata = with_offset_indexes(metadata); let expected_rows = metadata.file_metadata().num_rows(); let prepared = prepare_sst_meta( "test.parquet", @@ -2352,7 +2908,7 @@ mod tests { }; // Retain only the compact form so the lookup must exercise decompression and decoding. - cache.put_prepared_sst_meta(file_id, prepared, false); + cache.put_compact_sst_meta(file_id, prepared.compact); let mut metrics = MetadataCacheMetrics::default(); let cached = cache .get_sst_meta_data( @@ -2371,6 +2927,16 @@ mod tests { cached.parquet_metadata.file_metadata().num_rows() ); assert!(cached.parquet_metadata.offset_index().is_some()); + assert!( + cache + .sst_decoded_meta_cache + .as_ref() + .unwrap() + .get(&SstDecodedMetaKey::new(file_id, PageIndexPolicy::Required)) + .await + .is_none(), + "compact inspection must not fill the decoded reservation" + ); assert_eq!(expected_region_metadata, cached.region_metadata()); assert!( cached @@ -2383,8 +2949,8 @@ mod tests { ); } - #[test] - fn test_cache_encoding_failure_preserves_decoded_metadata() { + #[tokio::test] + async fn test_cache_encoding_failure_preserves_decoded_metadata() { let (metadata, expected_region_metadata) = sst_parquet_meta(); let expected_rows = metadata.file_metadata().num_rows(); let cache_encoding: Result<(Bytes, usize)> = Err(UnexpectedSnafu { @@ -2418,6 +2984,36 @@ mod tests { .to_string() .contains("injected compact metadata encoding failure") ); + let file_id = RegionFileId::new(expected_region_metadata.region_id, FileId::random()); + let manager = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let strategy = CacheStrategy::EnableAll(manager.clone()); + let (shared, miss) = strategy + .get_or_load_sst_meta(file_id, PageIndexPolicy::Skip, async { + Ok((decoded.clone(), true)) + }) + .await + .unwrap(); + assert!(miss); + assert!(Arc::ptr_eq(&decoded, &shared)); + assert!( + manager + .get_compact_sst_meta(&SstMetaKey(file_id.region_id(), file_id.file_id())) + .is_none() + ); + let mut metrics = MetadataCacheMetrics::default(); + let SstMetaObservation::Decoded(cached) = strategy + .observe_sst_meta(file_id, PageIndexPolicy::Skip, &mut metrics) + .await + else { + panic!("decoded fallback must be retained without a compact entry"); + }; + assert!(Arc::ptr_eq(&decoded, &cached)); + assert_eq!(1, metrics.mem_cache_hit); + assert_eq!(0, metrics.cache_miss); } #[tokio::test] @@ -2436,7 +3032,7 @@ mod tests { ) .unwrap(), ); - cache.put_sst_meta_data(file_id, skip_metadata); + cache.put_sst_meta_data(file_id, skip_metadata).await; let mut metrics = MetadataCacheMetrics::default(); assert!( @@ -2461,7 +3057,7 @@ mod tests { ) .unwrap(), ); - cache.put_sst_meta_data(file_id, optional_metadata); + cache.put_sst_meta_data(file_id, optional_metadata).await; let mut metrics = MetadataCacheMetrics::default(); assert!( @@ -2500,12 +3096,19 @@ mod tests { let cached = Arc::new( CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(), ); - let key = SstMetaKey(region_metadata.region_id, FileId::random()); + let key = SstDecodedMetaKey::new( + RegionFileId::new(region_metadata.region_id, FileId::random()), + PageIndexPolicy::Skip, + ); + let value = SstMetaLoadResult::new(cached.clone(), false, SstMetaResolution::default()); assert!(cached.region_metadata_weight > json_len); assert_eq!( - decoded_meta_cache_weight(&key, &cached) as usize, - key.estimated_size() + cached.parquet_metadata_size + cached.region_metadata_weight + decoded_meta_cache_weight(&key, &value) as usize, + size_of_val(&key) + + size_of::() + + cached.parquet_metadata_size + + cached.region_metadata_weight ); assert_eq!( cached.parquet_metadata_size, @@ -2521,9 +3124,13 @@ mod tests { CachedSstMeta::try_new("test.parquet", Arc::unwrap_or_clone(metadata)).unwrap(); cached.region_metadata_weight = u32::MAX as usize + 1; let cached = Arc::new(cached); - let key = SstMetaKey(region_metadata.region_id, FileId::random()); + let key = SstDecodedMetaKey::new( + RegionFileId::new(region_metadata.region_id, FileId::random()), + PageIndexPolicy::Skip, + ); + let value = SstMetaLoadResult::new(cached.clone(), false, SstMetaResolution::default()); - assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &cached)); + assert_eq!(u32::MAX, decoded_meta_cache_weight(&key, &value)); } #[test] diff --git a/src/mito2/src/cache/test_util.rs b/src/mito2/src/cache/test_util.rs index 82cabe24297..13c711b35a6 100644 --- a/src/mito2/src/cache/test_util.rs +++ b/src/mito2/src/cache/test_util.rs @@ -15,11 +15,13 @@ //! Utilities for testing cache. use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use bytes::Bytes; use datatypes::arrow::array::{ArrayRef, Int64Array}; use datatypes::arrow::record_batch::RecordBatch; use object_store::ObjectStore; +use object_store::layers::mock; use object_store::services::Fs; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; @@ -31,6 +33,33 @@ use store_api::metadata::RegionMetadataRef; use crate::sst::parquet::PARQUET_METADATA_KEY; use crate::test_util::sst_util::sst_region_metadata; +/// Holds storage reads until every concurrent caller has been polled. +pub(crate) struct GatedStorageReader { + pub(crate) inner: mock::oio::Reader, + pub(crate) gate: Arc, + pub(crate) reads: Arc, +} + +impl mock::oio::Read for GatedStorageReader { + async fn open( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, Box)> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.gate.acquire().await.unwrap().forget(); + self.inner.open(range).await + } + + async fn read( + &self, + range: mock::BytesRange, + ) -> object_store::Result<(mock::RpRead, object_store::Buffer)> { + self.reads.fetch_add(1, Ordering::Relaxed); + self.gate.acquire().await.unwrap().forget(); + self.inner.read(range).await + } +} + /// Returns a parquet meta data. pub(crate) fn parquet_meta() -> Arc { let file_data = parquet_file_data(); @@ -60,7 +89,9 @@ fn parquet_file_data() -> Vec { parquet_file_data_inner(None) } -fn parquet_file_data_with_region_metadata(region_metadata: &RegionMetadataRef) -> Vec { +pub(crate) fn parquet_file_data_with_region_metadata( + region_metadata: &RegionMetadataRef, +) -> Vec { let json = region_metadata.to_json().unwrap(); let key_value = KeyValue::new(PARQUET_METADATA_KEY.to_string(), json); parquet_file_data_inner(Some(vec![key_value])) diff --git a/src/mito2/src/compaction/reader.rs b/src/mito2/src/compaction/reader.rs index 0968f0e7a60..0a7f30ee34f 100644 --- a/src/mito2/src/compaction/reader.rs +++ b/src/mito2/src/compaction/reader.rs @@ -21,6 +21,7 @@ use common_time::range::TimestampRange; use common_time::timestamp::TimeUnit; use datafusion_common::ScalarValue; use datafusion_expr::Expr; +use futures::FutureExt; use parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use snafu::OptionExt; use store_api::metadata::RegionMetadataRef; @@ -67,8 +68,10 @@ impl CompactionSstReaderBuilder<'_> { let schema = scan_input.mapper.output_schema(); let schema = rewrite_json2_schema(schema.arrow_schema(), &plans); + // Erase the nested scan future before embedding it in the compaction task. let stream = SeqScan::new(scan_input) .build_flat_reader_for_compaction() + .boxed() .await?; Ok(FlatSource::new_stream(schema, stream)) } diff --git a/src/mito2/src/error.rs b/src/mito2/src/error.rs index a336fb89673..91402711136 100644 --- a/src/mito2/src/error.rs +++ b/src/mito2/src/error.rs @@ -198,6 +198,13 @@ pub enum Error { location: Location, }, + #[snafu(display("Failed to load SST metadata"))] + LoadSstMeta { + source: Arc, + #[snafu(implicit)] + location: Location, + }, + #[snafu(display("Failed to write parquet file"))] WriteParquet { #[snafu(source)] @@ -1455,6 +1462,7 @@ impl Error { match self { Error::OpenDal { error, .. } => error.kind() == ErrorKind::NotFound, Error::ManifestDeltaNotFound { .. } => true, + Error::LoadSstMeta { source, .. } => source.is_object_not_found(), _ => false, } } @@ -1577,7 +1585,7 @@ impl ErrorExt for Error { OpenRegion { source, .. } => source.status_code(), WriteParquet { .. } => StatusCode::StorageUnavailable, - WriteGroup { source, .. } => source.status_code(), + WriteGroup { source, .. } | LoadSstMeta { source, .. } => source.status_code(), InvalidBatch { .. } => StatusCode::InvalidArguments, InvalidRecordBatch { .. } => StatusCode::InvalidArguments, ConvertVector { source, .. } => source.status_code(), @@ -1732,6 +1740,7 @@ impl ErrorExt for Error { | External { source, .. } => source.retry_hint(), OpenRegion { source, .. } + | LoadSstMeta { source, .. } | WriteGroup { source, .. } | FlushRegion { source, .. } | BuildIndexAsync { source, .. } @@ -1861,6 +1870,31 @@ mod tests { } } + #[test] + fn test_shared_sst_metadata_preserves_error_semantics() { + for source in [ + OpenDalSnafu {} + .into_error(object_store::Error::new(ErrorKind::NotFound, "missing SST")), + OpenDalSnafu {}.into_error(object_store::Error::new( + ErrorKind::Unexpected, + "read failed", + )), + RejectWriteSnafu { + region_id: RegionId::new(1, 1), + } + .build(), + ] { + let expected_not_found = source.is_object_not_found(); + let expected_status = source.status_code(); + let expected_retry = source.retry_hint(); + let source = Arc::new(source); + let wrapped = LoadSstMetaSnafu.into_error(source); + assert_eq!(expected_not_found, wrapped.is_object_not_found()); + assert_eq!(expected_status, wrapped.status_code()); + assert_eq!(expected_retry, wrapped.retry_hint()); + } + } + #[test] fn test_manifest_update_persistence() { let rejected_kinds = [ diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index f7af9d54ffc..acb6a9da4c4 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -272,8 +272,44 @@ lazy_static! { ).unwrap(); } +pub(crate) fn cache_loading_metrics( + cache_type: &'static str, +) -> std::sync::Arc { + use common_cache::metrics::{CacheMetrics, OperationMetrics}; + + let operation = |stage| OperationMetrics { + started: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "started"]), + succeeded: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "success"]), + failed: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "error"]), + cancelled: CACHE_NATIVE_LOAD_EVENTS.with_label_values(&[cache_type, stage, "cancelled"]), + active: CACHE_NATIVE_LOAD_ACTIVE.with_label_values(&[cache_type, stage]), + elapsed: CACHE_NATIVE_LOAD_ELAPSED.with_label_values(&[cache_type, stage]), + }; + std::sync::Arc::new(CacheMetrics { + requests: operation("request"), + loads: operation("loader"), + }) +} + // Cache metrics. lazy_static! { + /// Native cache request and actual loader lifecycle; request is not a waiter count. + pub static ref CACHE_NATIVE_LOAD_EVENTS: IntCounterVec = register_int_counter_vec!( + "greptime_mito_cache_native_load_events_total", + "Native cache request and executed loader outcomes; requests include hits and shared loads", + &[TYPE_LABEL, STAGE_LABEL, "event"] + ).unwrap(); + pub static ref CACHE_NATIVE_LOAD_ACTIVE: IntGaugeVec = register_int_gauge_vec!( + "greptime_mito_cache_native_load_active", + "Active native cache requests or executed loaders, not waiting requests", + &[TYPE_LABEL, STAGE_LABEL] + ).unwrap(); + pub static ref CACHE_NATIVE_LOAD_ELAPSED: HistogramVec = register_histogram_vec!( + "greptime_mito_cache_native_load_seconds", + "Native cache request or executed loader wall time including suspension and cancellation, not CPU time", + &[TYPE_LABEL, STAGE_LABEL], exponential_buckets(0.0001, 10.0, 6).unwrap() + ).unwrap(); + /// Cache hit counter. pub static ref CACHE_HIT: IntCounterVec = register_int_counter_vec!( "greptime_mito_cache_hit", @@ -289,12 +325,14 @@ lazy_static! { ) .unwrap(); /// Cache size in bytes. - pub static ref CACHE_BYTES: IntGaugeVec = register_int_gauge_vec!( - "greptime_mito_cache_bytes", - "mito cache bytes", - &[TYPE_LABEL] - ) - .unwrap(); + pub static ref CACHE_BYTES: common_cache::metrics::CacheSizeMetrics = { + let metrics = common_cache::metrics::CacheSizeMetrics::new( + prometheus::Opts::new("greptime_mito_cache_bytes", "mito cache bytes; backend cache weights are approximate"), + &[TYPE_LABEL], + ).unwrap(); + prometheus::register(Box::new(metrics.clone())).unwrap(); + metrics + }; /// Download bytes counter in the write cache. pub static ref WRITE_CACHE_DOWNLOAD_BYTES_TOTAL: IntCounter = register_int_counter!( "mito_write_cache_download_bytes_total", diff --git a/src/mito2/src/read/pruner.rs b/src/mito2/src/read/pruner.rs index ef09f5f48ab..47689db161c 100644 --- a/src/mito2/src/read/pruner.rs +++ b/src/mito2/src/read/pruner.rs @@ -810,7 +810,7 @@ mod tests { use super::*; use crate::read::flat_projection::FlatProjectionMapper; - use crate::read::range::RowGroupIndex; + use crate::read::range::{RangeMeta, RowGroupIndex}; use crate::read::scan_region::{PredicateGroup, ScanInput}; use crate::read::scan_util::PartitionMetrics; use crate::sst::file::{FileHandle, FileMeta}; @@ -853,7 +853,13 @@ mod tests { .with_files(files) .with_append_mode(true) .build(); - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx( + input, + ranges, + query_start, + )); let pruner = Arc::new(Pruner::new_with_options( stream_ctx, 1, @@ -912,7 +918,13 @@ mod tests { .with_predicate(predicate) .with_append_mode(true) .build(); - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx( + input, + ranges, + query_start, + )); let pruner = Arc::new(Pruner::new(stream_ctx, 1)); (env, pruner) } diff --git a/src/mito2/src/read/range.rs b/src/mito2/src/read/range.rs index bf47e045cad..5e6da0a5986 100644 --- a/src/mito2/src/read/range.rs +++ b/src/mito2/src/read/range.rs @@ -15,6 +15,7 @@ //! Structs for partition ranges. use common_time::Timestamp; +use futures::future::BoxFuture; use smallvec::{SmallVec, smallvec}; use store_api::region_engine::PartitionRange; use store_api::storage::TimeSeriesDistribution; @@ -105,20 +106,24 @@ impl RangeMeta { } /// Creates a list of ranges from the `input` for unordered scan. - pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> Vec { - let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len()); - Self::push_unordered_mem_ranges(&input.memtables, &mut ranges); - Self::push_unordered_file_ranges( - input.memtables.len(), - &input.files, - &input.cache_strategy, - &mut ranges, - ); + pub(crate) fn unordered_scan_ranges(input: &ScanInput) -> BoxFuture<'_, Vec> { + // Keep Moka's async lookup type out of callers' nested scanner futures. + Box::pin(async move { + let mut ranges = Vec::with_capacity(input.memtables.len() + input.files.len()); + Self::push_unordered_mem_ranges(&input.memtables, &mut ranges); + Self::push_unordered_file_ranges( + input.memtables.len(), + &input.files, + &input.cache_strategy, + &mut ranges, + ) + .await; - #[cfg(feature = "enterprise")] - Self::push_extension_ranges(input, &mut ranges); + #[cfg(feature = "enterprise")] + Self::push_extension_ranges(input, &mut ranges); - ranges + ranges + }) } /// Returns true if the time range of given `meta` overlaps with the time range of this meta. @@ -203,7 +208,7 @@ impl RangeMeta { } } - fn push_unordered_file_ranges( + async fn push_unordered_file_ranges( num_memtables: usize, files: &[FileHandle], cache: &CacheStrategy, @@ -213,7 +218,9 @@ impl RangeMeta { for (i, file) in files.iter().enumerate() { let file_index = num_memtables + i; // Get parquet meta from the cache. - let parquet_meta = cache.get_parquet_meta_data_from_mem_cache(file.file_id()); + let parquet_meta = cache + .get_parquet_meta_data_from_mem_cache(file.file_id()) + .await; if let Some(parquet_meta) = parquet_meta { // Scans each row group. for row_group_index in 0..file.meta_ref().num_row_groups { diff --git a/src/mito2/src/read/scan_region.rs b/src/mito2/src/read/scan_region.rs index 31db8581b27..6dd6fe81c63 100644 --- a/src/mito2/src/read/scan_region.rs +++ b/src/mito2/src/read/scan_region.rs @@ -435,7 +435,10 @@ impl ScanRegion { #[tracing::instrument(skip_all, fields(region_id = %self.region_id()))] pub(crate) async fn unordered_scan(self) -> Result { let input = self.scan_input().await?; - Ok(UnorderedScan::new(input)) + let query_start = input.query_start.unwrap_or_else(Instant::now); + let ranges = RangeMeta::unordered_scan_ranges(&input).await; + let stream_ctx = StreamContext::unordered_scan_ctx(input, ranges, query_start); + Ok(UnorderedScan::new(stream_ctx)) } /// Scans by series. @@ -2041,9 +2044,11 @@ impl StreamContext { } /// Creates a new [StreamContext] for [UnorderedScan]. - pub(crate) fn unordered_scan_ctx(input: ScanInput) -> Self { - let query_start = input.query_start.unwrap_or_else(Instant::now); - let ranges = RangeMeta::unordered_scan_ranges(&input); + pub(crate) fn unordered_scan_ctx( + input: ScanInput, + ranges: Vec, + query_start: Instant, + ) -> Self { READ_SST_COUNT.observe(input.num_files() as f64); Self { input, diff --git a/src/mito2/src/read/unordered_scan.rs b/src/mito2/src/read/unordered_scan.rs index 05387be0ce1..bda554e4c95 100644 --- a/src/mito2/src/read/unordered_scan.rs +++ b/src/mito2/src/read/unordered_scan.rs @@ -35,7 +35,9 @@ use store_api::region_engine::{ use crate::error::{PartitionOutOfRangeSnafu, Result}; use crate::read::pruner::{PartitionPruner, Pruner}; -use crate::read::scan_region::{ScanInput, StreamContext}; +#[cfg(test)] +use crate::read::scan_region::ScanInput; +use crate::read::scan_region::StreamContext; use crate::read::scan_util::{ PartitionMetrics, PartitionMetricsList, prepare_flat_mem_stream, scan_flat_file_ranges, scan_flat_mem_ranges, @@ -59,7 +61,8 @@ pub struct UnorderedScan { impl UnorderedScan { /// Creates a new [UnorderedScan]. - pub(crate) fn new(input: ScanInput) -> Self { + pub(crate) fn new(stream_ctx: StreamContext) -> Self { + let input = &stream_ctx.input; let mut properties = ScannerProperties::default() .with_append_mode(input.append_mode) .with_total_rows(input.total_rows()) @@ -67,7 +70,7 @@ impl UnorderedScan { if let Some(counters) = input.query_stat_counters.clone() { properties.set_query_stat_counters(counters); } - let stream_ctx = Arc::new(StreamContext::unordered_scan_ctx(input)); + let stream_ctx = Arc::new(stream_ctx); properties.partitions = vec![stream_ctx.partition_ranges()]; // Create the shared pruner with number of workers equal to CPU cores. diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index ba4fe615951..30eb6b99cbc 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -1299,7 +1299,7 @@ async fn preload_parquet_meta_cache_for_files( } match metadata { SstMetaPreparation::Prepared(metadata) => { - cache_manager.put_prepared_sst_meta(file_id, metadata, false); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); loaded += 1; } SstMetaPreparation::DecodedOnly { encoding_error, .. } => warn!( @@ -1342,7 +1342,7 @@ async fn preload_parquet_meta_cache_for_files( Ok(SstMetaPreparation::Prepared(metadata)) => { // Preload the compact canonical entry only. Query traffic promotes decoded // entries into the separate acceleration tier according to actual demand. - cache_manager.put_prepared_sst_meta(file_id, metadata, false); + cache_manager.put_compact_sst_meta(file_id, metadata.compact); loaded += 1; } Ok(SstMetaPreparation::DecodedOnly { encoding_error, .. }) => warn!( @@ -1801,6 +1801,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); @@ -1877,6 +1878,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); @@ -1896,6 +1898,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); } @@ -1950,6 +1953,7 @@ mod tests { assert!( cache_manager .get_parquet_meta_data_from_mem_cache(region_file_id) + .await .is_none() ); diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index a627513d310..6422819607b 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -873,7 +873,9 @@ pub async fn delete_files( if let Some(cache) = &cache_manager { for (file_id, _) in file_ids { cache.invalidate_range_results(*file_id); - cache.remove_parquet_meta_data(RegionFileId::new(region_id, *file_id)); + cache + .remove_parquet_meta_data(RegionFileId::new(region_id, *file_id)) + .await; } } let mut attempted_files = Vec::with_capacity(file_ids.len()); diff --git a/src/mito2/src/sst/index/inverted_index/creator.rs b/src/mito2/src/sst/index/inverted_index/creator.rs index 19c55253b58..28242e79105 100644 --- a/src/mito2/src/sst/index/inverted_index/creator.rs +++ b/src/mito2/src/sst/index/inverted_index/creator.rs @@ -636,7 +636,8 @@ mod tests { move |expr| { let _d = &d; let cache = Arc::new(InvertedIndexCache::new(10, 10, 100)); - let puffin_metadata_cache = Arc::new(PuffinMetadataCache::new(10, &CACHE_BYTES)); + let puffin_metadata_cache = + Arc::new(PuffinMetadataCache::new(10, CACHE_BYTES.gauges())); let applier = InvertedIndexApplierBuilder::new( table_dir.clone(), PathType::Bare, diff --git a/src/mito2/src/sst/parquet/reader.rs b/src/mito2/src/sst/parquet/reader.rs index 1c98444946c..2ddb66a53f8 100644 --- a/src/mito2/src/sst/parquet/reader.rs +++ b/src/mito2/src/sst/parquet/reader.rs @@ -53,7 +53,10 @@ use store_api::storage::{ColumnId, FileId}; use table::predicate::Predicate; use crate::cache::index::result_cache::PredicateKey; -use crate::cache::{CacheStrategy, CachedSstMeta, SstMetaPreparation, prepare_sst_meta}; +use crate::cache::{ + CacheStrategy, CachedSstMeta, SstMetaObservation, SstMetaPreparation, SstMetaResolution, + prepare_sst_meta, +}; use crate::error::{ ParquetToArrowSchemaSnafu, ReadDataPartSnafu, Result, SerializePartitionExprSnafu, UnexpectedSnafu, @@ -822,14 +825,66 @@ impl ParquetReaderBuilder { .with_label_values(&["read_parquet_metadata"]) .start_timer(); + let observation = self + .cache_strategy + .observe_sst_meta(self.file_handle.file_id(), page_index_policy, cache_metrics) + .await; + if let SstMetaObservation::Decoded(metadata) = observation { + cache_metrics.metadata_load_cost += start.elapsed(); + return Ok((metadata, false)); + } + + let mut loader_metrics = MetadataCacheMetrics::default(); + // Keep the shared loader future off the stack of nested series-index builders. + // Decoded metadata hits above do not need this allocation. + let (result, resolution) = Box::pin(self.cache_strategy.get_or_load_observed_sst_meta( + self.file_handle.file_id(), + page_index_policy, + async { + let result = self + .read_parquet_metadata_inner( + file_path, + file_size, + &mut loader_metrics, + page_index_policy, + observation, + ) + .await; + ( + result, + SstMetaResolution { + file_cache_hit: loader_metrics.file_cache_hit, + cache_miss: loader_metrics.cache_miss, + }, + ) + }, + )) + .await; + // These are resolution counts, including a known miss before a failed remote read. + // Only the selected loader contributes actual I/O counts and bytes. + cache_metrics.file_cache_hit += resolution.file_cache_hit; + cache_metrics.cache_miss += resolution.cache_miss; + cache_metrics.num_reads += loader_metrics.num_reads; + cache_metrics.bytes_read += loader_metrics.bytes_read; + cache_metrics.metadata_load_cost += start.elapsed(); + result + } + + async fn read_parquet_metadata_inner( + &self, + file_path: &str, + file_size: u64, + cache_metrics: &mut MetadataCacheMetrics, + page_index_policy: PageIndexPolicy, + observation: SstMetaObservation, + ) -> Result<(Arc, bool)> { let file_id = self.file_handle.file_id(); // Tries to get from cache with metrics tracking. if let Some(metadata) = self .cache_strategy - .get_sst_meta_data(file_id, cache_metrics, page_index_policy) + .resolve_sst_meta(observation, file_id, cache_metrics, page_index_policy) .await { - cache_metrics.metadata_load_cost += start.elapsed(); return Ok((metadata, false)); } @@ -852,7 +907,7 @@ impl ParquetReaderBuilder { match metadata { SstMetaPreparation::Prepared(metadata) => { self.cache_strategy - .put_prepared_sst_meta(file_id, metadata, true); + .put_compact_sst_meta(file_id, metadata.compact); } SstMetaPreparation::DecodedOnly { encoding_error, .. } => { warn!( @@ -872,7 +927,6 @@ impl ParquetReaderBuilder { )?) }; - cache_metrics.metadata_load_cost += start.elapsed(); Ok((decoded, true)) } @@ -1635,13 +1689,13 @@ impl ReaderFilterMetrics { pub struct MetadataCacheMetrics { /// Number of memory cache hits for parquet metadata. pub mem_cache_hit: usize, - /// Number of file cache hits for parquet metadata. + /// Successful file-cache resolutions, including results shared with concurrent callers. pub file_cache_hit: usize, /// Number of cache misses for parquet metadata. pub cache_miss: usize, /// Duration to load parquet metadata. pub metadata_load_cost: Duration, - /// Number of read operations performed. + /// Physical reads executed by this caller; shared waiters do not repeat them. pub num_reads: usize, /// Total bytes read from storage. pub bytes_read: u64, @@ -2470,10 +2524,12 @@ mod tests { use datatypes::extension::json::Json2ExtensionType; use datatypes::prelude::ConcreteDataType; use datatypes::schema::ColumnSchema; + use object_store::layers::mock::MockLayerBuilder; use object_store::services::Memory; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::RowSelector; use parquet::file::properties::WriterProperties; + use std::sync::atomic::{AtomicUsize, Ordering}; use store_api::codec::PrimaryKeyEncoding; use store_api::metadata::{ColumnMetadata, RegionMetadata, RegionMetadataBuilder}; use store_api::region_request::PathType; @@ -2482,12 +2538,276 @@ mod tests { use super::*; use crate::cache::CacheManager; + use crate::cache::test_util::{GatedStorageReader, parquet_file_data_with_region_metadata}; use crate::sst::parquet::metadata::MetadataLoader; use crate::sst::parquet::prefilter::{build_reader_filter_plan, execute_prefilter}; use crate::sst::parquet::read_columns::{ParquetReadColumn, ParquetReadColumns}; use crate::sst::parquet::row_group::ParquetFetchMetrics; use crate::test_util::sst_util::{sst_file_handle, sst_region_metadata}; + #[tokio::test] + async fn test_shared_file_metadata_resolution_counts_without_duplicate_io() { + use crate::cache::file_cache::{FileType, IndexKey, IndexValue}; + use crate::cache::write_cache::WriteCache; + use crate::test_util::TestEnv; + use common_base::readable_size::ReadableSize; + + let env = TestEnv::new().await; + let metadata = Arc::new(sst_region_metadata()); + let bytes = parquet_file_data_with_region_metadata(&metadata); + let handle = sst_file_handle(0, 3); + let reads = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new({ + let reads = reads.clone(); + let gate = gate.clone(); + move |_, _, inner| { + Box::new(GatedStorageReader { + inner, + reads: reads.clone(), + gate: gate.clone(), + }) + } + })) + .build() + .unwrap(); + let local = ObjectStore::new(Memory::default()).unwrap().layer(layer); + let write_cache = Arc::new( + WriteCache::new( + local.clone(), + ReadableSize::mb(10), + None, + None, + false, + env.get_puffin_manager(), + env.get_intermediate_manager(), + None, + ) + .await + .unwrap(), + ); + let file_cache = write_cache.file_cache(); + let key = IndexKey::new( + handle.region_id(), + handle.file_id().file_id(), + FileType::Parquet, + ); + local + .write(&file_cache.cache_file_path(key), bytes.clone()) + .await + .unwrap(); + file_cache + .put( + key, + IndexValue { + file_size: bytes.len() as _, + }, + ) + .await; + let cache = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .write_cache(Some(write_cache)) + .build(), + ); + // The remote store is empty: success must come from the local file cache. + let builder = ParquetReaderBuilder::new( + "file_resolution".to_string(), + PathType::Bare, + handle.clone(), + ObjectStore::new(Memory::default()).unwrap(), + ) + .cache(CacheStrategy::EnableAll(cache)); + let path = handle.file_path("file_resolution", PathType::Bare); + let mut requests = (0..8) + .map(|_| { + Box::pin(async { + let mut metrics = MetadataCacheMetrics::default(); + let result = builder + .read_parquet_metadata( + &path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + ) + .await; + (result, metrics) + }) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request).is_pending()); + } + assert_eq!(1, reads.load(Ordering::Relaxed)); + gate.add_permits(32); + let results = + tokio::time::timeout(Duration::from_secs(10), futures::future::join_all(requests)) + .await + .unwrap(); + assert_eq!(1, reads.load(Ordering::Relaxed)); + assert_eq!( + 1, + results + .iter() + .map(|(_, metrics)| metrics.num_reads) + .sum::() + ); + for (result, metrics) in results { + assert!(!result.unwrap().1); + assert_eq!(0, metrics.mem_cache_hit); + assert_eq!(1, metrics.file_cache_hit); + assert_eq!(0, metrics.cache_miss); + } + let mut warm = MetadataCacheMetrics::default(); + builder + .read_parquet_metadata(&path, bytes.len() as u64, &mut warm, PageIndexPolicy::Skip) + .await + .unwrap(); + assert_eq!(1, warm.mem_cache_hit); + assert_eq!( + 0, warm.file_cache_hit, + "hot hits must not replay the original file source" + ); + assert_eq!(0, warm.num_reads); + } + + #[tokio::test] + #[allow( + clippy::print_stdout, + reason = "records the controlled storage-read comparison" + )] + async fn test_metadata_load_coalescing_storage_reads() { + const CONCURRENCY: usize = 8; + let metadata = Arc::new(sst_region_metadata()); + let bytes = parquet_file_data_with_region_metadata(&metadata); + let file_handle = sst_file_handle(0, 3); + let file_path = file_handle.file_path("metadata_coalescing", PathType::Bare); + // Same bytes, file, cache budget and concurrency; only bypass the coordinator + // in the control arm. Repeat with fresh caches to verify cold behavior. + for _ in 0..3 { + for coalesce in [false, true] { + let reads = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let layer = MockLayerBuilder::default() + .reader_factory(Arc::new({ + let reads = reads.clone(); + let gate = gate.clone(); + move |_, _, inner| { + Box::new(GatedStorageReader { + inner, + reads: reads.clone(), + gate: gate.clone(), + }) + } + })) + .build() + .unwrap(); + let store = ObjectStore::new(Memory::default()).unwrap().layer(layer); + store.write(&file_path, bytes.clone()).await.unwrap(); + let cache = Arc::new( + CacheManager::builder() + .sst_meta_cache_size(1024 * 1024) + .build(), + ); + let builder = ParquetReaderBuilder::new( + "metadata_coalescing".to_string(), + PathType::Bare, + file_handle.clone(), + store, + ) + .cache(CacheStrategy::EnableAll(cache.clone())); + let mut requests = (0..CONCURRENCY) + .map(|_| { + Box::pin(async { + let mut metrics = MetadataCacheMetrics::default(); + let result = if coalesce { + builder + .read_parquet_metadata( + &file_path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + ) + .await + } else { + let observation = builder + .cache_strategy + .observe_sst_meta( + file_handle.file_id(), + PageIndexPolicy::Skip, + &mut metrics, + ) + .await; + builder + .read_parquet_metadata_inner( + &file_path, + bytes.len() as u64, + &mut metrics, + PageIndexPolicy::Skip, + observation, + ) + .await + }; + (result.unwrap(), metrics) + }) + }) + .collect::>(); + for request in &mut requests { + assert!(futures::poll!(request).is_pending()); + } + let expected_reads = if coalesce { 1 } else { CONCURRENCY }; + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + gate.add_permits(CONCURRENCY * 4); + let results = tokio::time::timeout( + Duration::from_secs(10), + futures::future::join_all(requests), + ) + .await + .unwrap(); + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + assert_eq!( + expected_reads, + results.iter().map(|(_, m)| m.num_reads).sum::() + ); + for ((value, miss), request_metrics) in &results { + assert_eq!(1, request_metrics.cache_miss); + assert_eq!(0, request_metrics.mem_cache_hit); + assert_eq!(0, request_metrics.file_cache_hit); + assert!( + *miss, + "each request must preserve background-download eligibility" + ); + assert_eq!(3, value.parquet_metadata().file_metadata().num_rows()); + if coalesce { + assert!(Arc::ptr_eq(value, &results[0].0.0)); + } + } + let mut warm_metrics = MetadataCacheMetrics::default(); + let (_, miss) = tokio::time::timeout( + Duration::from_secs(5), + builder.read_parquet_metadata( + &file_path, + bytes.len() as u64, + &mut warm_metrics, + PageIndexPolicy::Skip, + ), + ) + .await + .unwrap() + .unwrap(); + assert!(!miss); + assert_eq!(1, warm_metrics.mem_cache_hit); + assert_eq!(0, warm_metrics.cache_miss); + assert_eq!(0, warm_metrics.num_reads); + assert_eq!(expected_reads, reads.load(Ordering::Relaxed)); + println!( + "coalesce={coalesce} concurrency={CONCURRENCY} cold_storage_reads={expected_reads} warm_storage_reads=0" + ); + } + } + } + async fn prefilter_test_builder( object_store: ObjectStore, predicate: Predicate,