diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index cf0e8de3458..c4b8ffa0d73 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -19,6 +19,7 @@ pub(crate) mod cache_size; pub(crate) mod file_cache; pub(crate) mod index; pub(crate) mod manifest_cache; +pub(crate) mod range_cache; #[cfg(test)] pub(crate) mod test_util; pub(crate) mod write_cache; @@ -145,6 +146,40 @@ impl RangeResultMemoryLimiter { self.semaphore.available_permits() } + /// Never wait while holding a partial load: several partial loads could + /// otherwise exhaust the budget and wait on each other indefinitely. + pub(crate) fn try_acquire(&self, bytes: usize) -> Option { + let permits = bytes.div_ceil(self.permit_bytes).max(1); + if permits > self.total_permits { + return None; + } + self.semaphore + .clone() + .try_acquire_many_owned(permits.try_into().ok()?) + .ok() + } + + /// Transfers a load's reservation to its completed result without a gap in + /// accounting. Compaction may shrink buffers or add metadata; growing the + /// reservation must remain nonblocking just like loading more input. + pub(crate) fn try_resize( + &self, + permit: Option, + bytes: usize, + ) -> Option { + let Some(mut permit) = permit else { + return self.try_acquire(bytes); + }; + let required = bytes.div_ceil(self.permit_bytes).max(1); + let reserved = permit.num_permits(); + if required > reserved { + permit.merge(self.try_acquire((required - reserved).checked_mul(self.permit_bytes)?)?); + } else { + drop(permit.split(reserved - required)); + } + Some(permit) + } + pub(crate) async fn acquire(&self, bytes: usize) -> Result> { let permits = bytes.div_ceil(self.permit_bytes).max(1); if permits > self.total_permits { @@ -898,6 +933,7 @@ impl CacheStrategy { /// Calls [CacheManager::get_range_result()]. /// It returns None if the strategy is [CacheStrategy::Compaction] or [CacheStrategy::Disabled]. + #[cfg(test)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -910,6 +946,7 @@ impl CacheStrategy { /// Calls [CacheManager::put_range_result()]. /// It does nothing if the strategy isn't [CacheStrategy::EnableAll]. + #[cfg(any(test, feature = "test"))] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, @@ -920,6 +957,13 @@ impl CacheStrategy { } } + pub(crate) fn range_cache_worker(&self) -> Option<&Arc> { + match self { + Self::EnableAll(manager) => manager.range_result_cache.as_ref(), + Self::Compaction(_) | Self::Disabled => None, + } + } + /// Returns true if the range result cache is enabled. pub(crate) fn has_range_result_cache(&self) -> bool { match self { @@ -1389,7 +1433,15 @@ impl CacheManager { } } + /// Retires decoded range results when an SST leaves the current version. + pub(crate) fn invalidate_range_results(&self, file_id: FileId) { + if let Some(worker) = &self.range_result_cache { + worker.invalidate_file(file_id); + } + } + /// Gets cached result for range scan. + #[cfg(test)] pub(crate) fn get_range_result( &self, key: &RangeScanCacheKey, @@ -1400,15 +1452,13 @@ impl CacheManager { } /// Puts range scan result into cache. + #[cfg(any(test, feature = "test"))] pub(crate) fn put_range_result( &self, key: RangeScanCacheKey, result: Arc, ) { if let Some(cache) = &self.range_result_cache { - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .add(range_result_cache_weight(&key, &result).into()); cache.insert(key, result); } } @@ -1649,21 +1699,8 @@ impl CacheManagerBuilder { }) .build() }); - let range_result_cache = (self.range_result_cache_size != 0).then(|| { - Cache::builder() - .max_capacity(self.range_result_cache_size) - .weigher(range_result_cache_weight) - .eviction_listener(move |k, v, cause| { - let size = range_result_cache_weight(&k, &v); - CACHE_BYTES - .with_label_values(&[RANGE_RESULT_TYPE]) - .sub(size.into()); - CACHE_EVICTION - .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)]) - .inc(); - }) - .build() - }); + let range_result_cache = (self.range_result_cache_size != 0) + .then(|| range_cache::RangeCacheWorker::new(self.range_result_cache_size)); CacheManager { sst_meta_cache, sst_decoded_meta_cache, @@ -2093,7 +2130,7 @@ type VectorCache = Cache<(ConcreteDataType, Value), VectorRef>; /// Maps (file id, row group id, time series row selector) to [SelectorResultValue]. type SelectorResultCache = Cache>; /// Maps partition-range scan key to cached flat batches. -type RangeResultCache = Cache>; +type RangeResultCache = Arc; #[cfg(test)] mod tests { @@ -2900,6 +2937,38 @@ mod tests { assert_eq!(limiter.available_permits(), 2); } + #[rstest::rstest] + #[case::empty_result(None, 0, Some(1))] + #[case::new_result(None, 1025, Some(2))] + #[case::unchanged(Some(1024), 1024, Some(1))] + #[case::grow(Some(1024), 2048, Some(2))] + #[case::shrink(Some(3072), 1024, Some(1))] + #[case::exhausted(Some(1024), 4096, None)] + #[case::oversized(Some(1024), usize::MAX, None)] + #[test] + fn range_result_memory_reservation_resizes_without_waiting( + #[case] input_size: Option, + #[case] result_size: usize, + #[case] expected_permits: Option, + ) { + let limiter = RangeResultMemoryLimiter::new(4096, 1024); + let blocker = limiter.try_acquire(1024).unwrap(); + let input_permit = input_size.map(|size| limiter.try_acquire(size).unwrap()); + let result_permit = limiter.try_resize(input_permit, result_size); + assert_eq!( + result_permit.as_ref().map(|permit| permit.num_permits()), + expected_permits + ); + assert_eq!( + limiter.available_permits(), + 3 - expected_permits.unwrap_or(0) + ); + drop(result_permit); + assert_eq!(limiter.available_permits(), 3); + drop(blocker); + assert_eq!(limiter.available_permits(), 4); + } + #[tokio::test] async fn test_evict_puffin_cache_clears_all_entries() { use std::collections::{BTreeMap, HashMap}; diff --git a/src/mito2/src/cache/range_cache.rs b/src/mito2/src/cache/range_cache.rs new file mode 100644 index 00000000000..060af974764 --- /dev/null +++ b/src/mito2/src/cache/range_cache.rs @@ -0,0 +1,733 @@ +// Copyright 2023 Greptime Team +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Coordinates range-result cache loads and SST retirement. No I/O or decoding +//! runs under the coordinator lock; loaders use the shared query runtime. + +use std::collections::HashMap; +use std::future::Future; +use std::sync::{Arc, Mutex, Weak, mpsc}; + +use moka::sync::Cache; +use store_api::storage::FileId; +use tokio::sync::{OwnedSemaphorePermit, watch}; + +use crate::cache::{ + RANGE_RESULT_TYPE, RangeResultMemoryLimiter, range_result_cache_weight, removal_cause_str, +}; +use crate::metrics::{CACHE_BYTES, CACHE_EVICTION, RANGE_CACHE_WORKER_EVENTS}; +use crate::read::range_cache::{RangeScanCacheKey, RangeScanCacheValue}; +use crate::sst::file::FileHandle; + +/// Admission, not a queue: excess distinct misses keep using ordinary scans. +const MAX_IN_FLIGHT_RANGES: usize = 32; + +type SharedResult = Option>; + +/// Keeps completed results charged while subscribers or replay streams own them. +/// Cache entries retain only a weak reference: after the last reader leaves, +/// the resident value is covered solely by the cache's capacity accounting. +pub(crate) struct RangeCacheDelivery { + value: Arc, + _permit: OwnedSemaphorePermit, +} + +impl RangeCacheDelivery { + /// Shares a decoded result together with its memory reservation. + pub(crate) fn new(value: Arc, permit: OwnedSemaphorePermit) -> Arc { + Arc::new(Self { + value, + _permit: permit, + }) + } + + /// Borrows the result without detaching it from the delivery reservation. + pub(crate) fn value(&self) -> &RangeScanCacheValue { + &self.value + } +} + +pub(crate) enum RangeCacheLookup { + Hit(Arc), + Loading(watch::Receiver), + Bypass, +} + +struct Entry { + generation: u64, + value: Arc, + delivery: Mutex>, +} + +impl Entry { + fn deliver(&self, memory: &RangeResultMemoryLimiter) -> SharedResult { + // Serialize cache-hit reservations so readers of the same result share + // one charge, including when the last previous reader just left. + let mut shared_delivery = self + .delivery + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(delivery) = shared_delivery.upgrade() { + return Some(delivery); + } + let permit = memory.try_acquire(self.value.estimated_size())?; + let delivery = RangeCacheDelivery::new(self.value.clone(), permit); + *shared_delivery = Arc::downgrade(&delivery); + Some(delivery) + } +} + +struct Flight { + generation: u64, + sender: watch::Sender, + invalidated: bool, +} + +#[derive(Default)] +struct State { + generation: u64, + flights: HashMap, + /// File IDs are globally unique. Index by physical file, not query region: + /// a copied region can read an SST whose original region differs from its own. + files: HashMap, u64>>, +} + +impl State { + fn next_generation(&mut self) -> u64 { + self.generation += 1; + self.generation + } + + fn forget(&mut self, key: &RangeScanCacheKey, generation: u64) { + for (file, _) in &key.row_groups { + if let Some(keys) = self.files.get_mut(file) { + // Eviction callbacks may arrive after the same key is refilled. + if keys.get(key) == Some(&generation) { + keys.remove(key); + } + if keys.is_empty() { + self.files.remove(file); + } + } + } + } +} + +pub(crate) struct RangeCacheWorker { + cache: Cache>, + state: Mutex, + /// Moka can call listeners during a cache operation. Queue only cleanup + /// metadata instead of reentering `state` from its eviction listener. + evictions: Mutex, u64)>>, + /// Covers loads and completed deliveries, separately from cache residency + /// and the temporary concat-memory limiter. Never wait for this budget. + pub(crate) load_memory: RangeResultMemoryLimiter, +} + +impl RangeCacheWorker { + pub(crate) fn new(capacity: u64) -> Arc { + let (tx, rx) = mpsc::channel(); + let cache = Cache::builder() + .max_capacity(capacity) + .weigher(|k: &RangeScanCacheKey, v: &Arc| range_result_cache_weight(k, &v.value)) + .eviction_listener(move |key, entry: Arc, cause| { + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .sub(range_result_cache_weight(&key, &entry.value).into()); + CACHE_EVICTION + .with_label_values(&[RANGE_RESULT_TYPE, removal_cause_str(cause)]) + .inc(); + let _ = tx.send((key, entry.generation)); + }) + .build(); + Arc::new(Self { + cache, + state: Mutex::new(State::default()), + evictions: Mutex::new(rx), + load_memory: RangeResultMemoryLimiter::new(capacity as usize, 1024), + }) + } + + fn drain_evictions(&self, state: &mut State) { + for (key, generation) in self.evictions.lock().unwrap().try_iter() { + state.forget(&key, generation); + } + } + + #[cfg(test)] + pub(crate) fn get(&self, key: &RangeScanCacheKey) -> Option> { + self.cache.get(key).map(|entry| entry.value.clone()) + } + + fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option { + let entry = self.cache.get(key)?; + // A hit without delivery budget must bypass, not start another load. + Some(match entry.deliver(&self.load_memory) { + Some(delivery) => RangeCacheLookup::Hit(delivery), + None => { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_memory"]) + .inc(); + RangeCacheLookup::Bypass + } + }) + } + + /// For test/benchmark producers that do not hold live SST handles. + #[cfg(any(test, feature = "test"))] + pub(crate) fn insert(&self, key: RangeScanCacheKey, value: Arc) { + let Some(permit) = self.load_memory.try_acquire(value.estimated_size()) else { + return; + }; + let delivery = RangeCacheDelivery::new(value, permit); + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + self.publish(&mut state, key, &delivery); + } + + fn publish( + &self, + state: &mut State, + key: RangeScanCacheKey, + delivery: &Arc, + ) { + let generation = state.next_generation(); + let indexed_key = Arc::new(key.clone()); + for (file, _) in &key.row_groups { + state + .files + .entry(*file) + .or_default() + .insert(indexed_key.clone(), generation); + } + CACHE_BYTES + .with_label_values(&[RANGE_RESULT_TYPE]) + .add(range_result_cache_weight(&key, &delivery.value).into()); + self.cache.insert( + key, + Arc::new(Entry { + generation, + value: delivery.value.clone(), + delivery: Mutex::new(Arc::downgrade(delivery)), + }), + ); + self.drain_evictions(state); + } + + /// Called at logical retirement, before the last snapshot releases the SST. + /// Existing readers may finish, but their results must not refill this cache. + pub(crate) fn invalidate_file(&self, file: FileId) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["invalidate"]) + .inc(); + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + if let Some(keys) = state.files.remove(&file) { + for (key, generation) in keys { + self.cache.invalidate(key.as_ref()); + state.forget(&key, generation); + } + } + for (key, flight) in &mut state.flights { + if key.row_groups.iter().any(|(id, _)| *id == file) { + flight.invalidated = true; + } + } + self.drain_evictions(&mut state); + } + + pub(crate) fn get_or_load( + self: &Arc, + key: RangeScanCacheKey, + files: Vec, + load: F, + ) -> RangeCacheLookup + where + F: Future + Send + 'static, + { + if files.iter().any(FileHandle::is_deleted) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_retired"]) + .inc(); + return RangeCacheLookup::Bypass; + } + if let Some(lookup) = + crate::cache::update_hit_miss(self.get_for_delivery(&key), RANGE_RESULT_TYPE) + { + return lookup; + } + let mut state = self.state.lock().unwrap(); + self.drain_evictions(&mut state); + // Close the miss/register race with completion and invalidation. + if files.iter().any(FileHandle::is_deleted) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_retired"]) + .inc(); + return RangeCacheLookup::Bypass; + } + if let Some(lookup) = self.get_for_delivery(&key) { + return lookup; + } + if let Some(flight) = state.flights.get(&key) { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["coalesced"]) + .inc(); + return RangeCacheLookup::Loading(flight.sender.subscribe()); + } + if state.flights.len() >= MAX_IN_FLIGHT_RANGES { + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["bypass_capacity"]) + .inc(); + return RangeCacheLookup::Bypass; + } + let generation = state.next_generation(); + let (sender, receiver) = watch::channel(None); + state.flights.insert( + key.clone(), + Flight { + generation, + sender: sender.clone(), + invalidated: false, + }, + ); + drop(state); + + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["started"]) + .inc(); + let guard = LoadGuard { + worker: self.clone(), + key, + generation, + }; + common_runtime::spawn_query(async move { + // Losing the first requester does not cancel other subscribers. + // Losing every requester drops the loader and unregisters the flight. + let value = tokio::select! { + biased; + _ = sender.closed() => return, + value = load => value, + }; + guard.complete(value, &files); + }); + RangeCacheLookup::Loading(receiver) + } +} + +/// Also unregisters a cancelled or panicking load. Generation matching keeps a +/// stale task from unregistering a later retry of the same key. +struct LoadGuard { + worker: Arc, + key: RangeScanCacheKey, + generation: u64, +} + +impl LoadGuard { + fn complete(&self, value: SharedResult, files: &[FileHandle]) { + let mut state = self.worker.state.lock().unwrap(); + self.worker.drain_evictions(&mut state); + let Some(flight) = state.flights.get(&self.key) else { + return; + }; + if flight.generation != self.generation { + return; + } + let flight = state.flights.remove(&self.key).unwrap(); + let event = if let Some(value) = &value { + if !flight.invalidated && !files.iter().any(FileHandle::is_deleted) { + self.worker.publish(&mut state, self.key.clone(), value); + "published" + } else { + "late_fill_skipped" + } + } else { + "load_skipped" + }; + RANGE_CACHE_WORKER_EVENTS.with_label_values(&[event]).inc(); + // Fan out independently of Moka admission, including to old snapshots. + // Do not hold the coordinator lock while notifying readers. + drop(state); + let _ = flight.sender.send(value); + } +} + +impl Drop for LoadGuard { + fn drop(&mut self) { + let mut state = self.worker.state.lock().unwrap(); + if state + .flights + .get(&self.key) + .is_some_and(|flight| flight.generation == self.generation) + { + state.flights.remove(&self.key); + RANGE_CACHE_WORKER_EVENTS + .with_label_values(&["abandoned"]) + .inc(); + } + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use tokio::sync::oneshot; + + use super::*; + use crate::cache::test_util::range_cache_key; + use crate::sst::file::FileMeta; + use crate::sst::file_purger::NoopFilePurger; + + fn empty_result() -> Arc { + Arc::new(RangeScanCacheValue::new(vec![], 0)) + } + + fn empty_delivery() -> Arc { + let value = empty_result(); + let memory = RangeResultMemoryLimiter::new(1024, 1024); + let permit = memory.try_acquire(value.estimated_size()).unwrap(); + RangeCacheDelivery::new(value, permit) + } + + fn receiver(lookup: RangeCacheLookup) -> watch::Receiver { + match lookup { + RangeCacheLookup::Loading(receiver) => receiver, + _ => panic!("expected a shared load"), + } + } + + async fn loaded(mut receiver: watch::Receiver) -> SharedResult { + tokio::time::timeout(Duration::from_secs(5), async { + receiver + .wait_for(|value| value.is_some()) + .await + .ok() + .and_then(|value| value.clone()) + }) + .await + .expect("load must finish or cancel") + } + + #[tokio::test] + async fn same_key_misses_share_one_load_even_after_first_reader_leaves() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let count = Arc::new(AtomicUsize::new(0)); + let (release, gate) = oneshot::channel(); + let first_count = count.clone(); + let first = receiver(worker.get_or_load(key.clone(), vec![], async move { + first_count.fetch_add(1, Ordering::Relaxed); + gate.await.unwrap(); + Some(empty_delivery()) + })); + let mut followers = Vec::new(); + for _ in 0..16 { + let count = count.clone(); + followers.push(receiver(worker.get_or_load( + key.clone(), + vec![], + async move { + count.fetch_add(1, Ordering::Relaxed); + Some(empty_delivery()) + }, + ))); + } + drop(first); + release.send(()).unwrap(); + let mut results = Vec::new(); + for follower in followers { + results.push(loaded(follower).await.unwrap()); + } + assert_eq!(count.load(Ordering::Relaxed), 1); + assert!(results.iter().all(|value| Arc::ptr_eq(value, &results[0]))); + assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &results[0].value)); + } + + #[test] + fn retirement_removes_all_dependent_keys_but_preserves_other_files() { + let worker = RangeCacheWorker::new(1024 * 1024); + let a = FileId::random(); + let b = FileId::random(); + let combined = range_cache_key(&[a, b, a]); + let mut copied = range_cache_key(&[a]); + copied.region_id = store_api::storage::RegionId::new(2, 1); + let unaffected = range_cache_key(&[b]); + for key in [&combined, &copied, &unaffected] { + worker.insert(key.clone(), empty_result()); + } + worker.invalidate_file(a); + worker.invalidate_file(a); + assert!(worker.get(&combined).is_none()); + assert!(worker.get(&copied).is_none()); + assert!(worker.get(&unaffected).is_some()); + worker.invalidate_file(b); + assert!(worker.get(&unaffected).is_none()); + assert!(worker.state.lock().unwrap().files.is_empty()); + } + + #[test] + fn delayed_replacement_eviction_cannot_hide_a_new_entry_from_invalidation() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + worker.insert(key.clone(), empty_result()); + worker.cache.run_pending_tasks(); + let replacement = empty_result(); + worker.insert(key.clone(), replacement.clone()); + worker.cache.run_pending_tasks(); + assert!(Arc::ptr_eq(&worker.get(&key).unwrap(), &replacement)); + worker.invalidate_file(file); + assert!(worker.get(&key).is_none()); + } + + #[tokio::test] + async fn retired_inflight_result_serves_old_readers_without_refilling() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + let (release, gate) = oneshot::channel(); + let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { + gate.await.unwrap(); + Some(empty_delivery()) + })); + worker.invalidate_file(file); + release.send(()).unwrap(); + assert!(loaded(pending).await.is_some()); + assert!(worker.get(&key).is_none()); + } + + #[test] + fn old_snapshot_cannot_start_a_fill_after_retirement() { + let worker = RangeCacheWorker::new(1024 * 1024); + let file = FileHandle::new(FileMeta::default(), Arc::new(NoopFilePurger)); + file.mark_deleted(); + let key = range_cache_key(&[file.file_id().file_id()]); + assert!(matches!( + worker.get_or_load(key, vec![file], async { Some(empty_delivery()) }), + RangeCacheLookup::Bypass + )); + } + + #[rstest::rstest] + #[tokio::test] + async fn index_replacement_retirement_invalidates_and_blocks_old_reader_fills( + #[values(false, true)] gc_enabled: bool, + ) { + use common_base::readable_size::ReadableSize; + use store_api::region_engine::RegionEngine; + use store_api::region_request::RegionRequest; + use store_api::storage::RegionId; + + use crate::config::MitoConfig; + use crate::gc::GcConfig; + use crate::sst::file::IndexType; + use crate::sst::version::SstVersion; + use crate::test_util::{CreateRequestBuilder, TestEnv}; + + let mut env = TestEnv::new().await; + let engine = env + .create_engine(MitoConfig { + range_result_cache_size: ReadableSize::mb(1), + gc: GcConfig { + enable: gc_enabled, + ..Default::default() + }, + ..Default::default() + }) + .await; + let region_id = RegionId::new(1, 1); + engine + .handle_request( + region_id, + RegionRequest::Create(CreateRequestBuilder::new().build()), + ) + .await + .unwrap(); + let region = engine.get_region(region_id).unwrap(); + let cache = engine.cache_manager(); + let worker = cache.range_result_cache.as_ref().unwrap(); + let meta = FileMeta { + region_id, + file_id: FileId::random(), + ..Default::default() + }; + let mut version = SstVersion::new(crate::test_util::memtable_util::metadata_for_test()); + version.add_files(region.file_purger.clone(), [meta.clone()].into_iter()); + let old_reader = version.levels()[0].files[&meta.file_id].clone(); + let key = range_cache_key(&[meta.file_id]); + worker.insert(key.clone(), empty_result()); + assert!(worker.get(&key).is_some()); + let loading_key = key.clone().for_cross_range_merge(); + let (release, gate) = oneshot::channel(); + let pending = receiver(worker.get_or_load( + loading_key.clone(), + vec![old_reader.clone()], + async move { + gate.await.unwrap(); + Some(empty_delivery()) + }, + )); + let updated = FileMeta { + index_version: 1, + available_indexes: smallvec::smallvec![IndexType::InvertedIndex], + ..meta + }; + version.add_files(region.file_purger.clone(), [updated.clone()].into_iter()); + version.remove_files([updated].into_iter()); + + // R1-C1: neither a pre-existing flight nor a new request through the old + // index handle may refill a range after the current handle is retired. + assert!(old_reader.is_deleted()); + assert!(worker.get(&key).is_none()); + assert!(matches!( + worker.get_or_load(key, vec![old_reader], async { Some(empty_delivery()) }), + RangeCacheLookup::Bypass + )); + release.send(()).unwrap(); + assert!(loaded(pending).await.is_some()); + assert!(worker.get(&loading_key).is_none()); + } + + #[tokio::test] + async fn admission_rejection_does_not_discard_the_shared_result() { + let worker = RangeCacheWorker::new(1); + let key = range_cache_key(&[FileId::random()]); + let value = empty_result(); + let permit = worker + .load_memory + .try_acquire(value.estimated_size()) + .unwrap(); + let delivery = RangeCacheDelivery::new(value, permit); + let pending = + receiver(worker.get_or_load(key.clone(), vec![], async move { Some(delivery) })); + let delivery = loaded(pending).await.unwrap(); + worker.cache.run_pending_tasks(); + assert!(worker.get(&key).is_none()); + { + let mut state = worker.state.lock().unwrap(); + worker.drain_evictions(&mut state); + assert!(state.files.is_empty()); + } + assert_eq!(worker.load_memory.available_permits(), 0); + drop(delivery); + // Publishing wakes readers before the producer drops its watch sender. + // Its LoadGuard also needs the coordinator lock to finish cleaning up. + tokio::time::timeout(Duration::from_secs(5), async { + while worker.load_memory.available_permits() != 1 { + tokio::task::yield_now().await; + } + }) + .await + .expect("a rejected value must release its reservation after all owners leave"); + } + + #[test] + fn cache_hits_share_delivery_budget_until_the_last_reader_leaves() { + let worker = RangeCacheWorker::new(1024); + let file = FileId::random(); + let key = range_cache_key(&[file]); + worker.insert(key.clone(), empty_result()); + assert_eq!(worker.load_memory.available_permits(), 1); + + let blocker = worker.load_memory.try_acquire(1024).unwrap(); + assert!(matches!( + worker.get_or_load(key.clone(), vec![], async { + panic!("must not load on a hit") + }), + RangeCacheLookup::Bypass + )); + assert!(worker.get(&key).is_some()); + drop(blocker); + + let RangeCacheLookup::Hit(first) = + worker.get_or_load(key.clone(), vec![], async { panic!("must hit the cache") }) + else { + panic!("expected an accounted cache hit"); + }; + let RangeCacheLookup::Hit(second) = + worker.get_or_load(key.clone(), vec![], async { panic!("must reuse delivery") }) + else { + panic!("readers of the same value must share its reservation"); + }; + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(worker.load_memory.available_permits(), 0); + + worker.invalidate_file(file); + worker.cache.run_pending_tasks(); + assert!(worker.get(&key).is_none()); + drop(first); + assert_eq!(worker.load_memory.available_permits(), 0); + drop(second); + assert_eq!(worker.load_memory.available_permits(), 1); + } + + #[tokio::test] + async fn failed_load_releases_followers_and_allows_a_retry() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let failed = receiver(worker.get_or_load(key.clone(), vec![], async { None })); + assert!(loaded(failed).await.is_none()); + let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); + assert!(loaded(retry).await.is_some()); + } + + #[tokio::test] + async fn abandoning_all_subscribers_cancels_and_unregisters_the_load() { + let worker = RangeCacheWorker::new(1024 * 1024); + let key = range_cache_key(&[FileId::random()]); + let (started, start_rx) = oneshot::channel(); + let marker = Arc::new(()); + let weak_marker = Arc::downgrade(&marker); + let pending = receiver(worker.get_or_load(key.clone(), vec![], async move { + let _marker = marker; + started.send(()).unwrap(); + std::future::pending().await + })); + start_rx.await.unwrap(); + drop(pending); + tokio::time::timeout(Duration::from_secs(5), async { + while weak_marker.upgrade().is_some() + || !worker.state.lock().unwrap().flights.is_empty() + { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled loaders must release resources and the flight key"); + assert!(worker.get(&key).is_none()); + let retry = receiver(worker.get_or_load(key, vec![], async { Some(empty_delivery()) })); + assert!(loaded(retry).await.is_some()); + } + + #[tokio::test] + async fn distinct_loads_are_bounded_without_queueing_misses() { + let worker = RangeCacheWorker::new(1024 * 1024); + let mut subscribers = Vec::new(); + for _ in 0..MAX_IN_FLIGHT_RANGES { + subscribers.push(receiver(worker.get_or_load( + range_cache_key(&[FileId::random()]), + vec![], + std::future::pending(), + ))); + } + assert!(matches!( + worker.get_or_load(range_cache_key(&[FileId::random()]), vec![], async { + Some(empty_delivery()) + }), + RangeCacheLookup::Bypass + )); + drop(subscribers); + } +} diff --git a/src/mito2/src/cache/test_util.rs b/src/mito2/src/cache/test_util.rs index 58b8272fbc7..82cabe24297 100644 --- a/src/mito2/src/cache/test_util.rs +++ b/src/mito2/src/cache/test_util.rs @@ -145,3 +145,29 @@ pub(crate) fn assert_parquet_metadata_equal(x: Arc, y: Arc crate::read::range_cache::RangeScanCacheKey { + use crate::read::range_cache::{RangeScanCacheKey, ScanRequestFingerprintBuilder}; + use crate::read::read_columns::ReadColumns; + + RangeScanCacheKey { + region_id: store_api::storage::RegionId::new(1, 1), + row_groups: files.iter().map(|file| (*file, 0)).collect(), + scan: ScanRequestFingerprintBuilder { + read_columns: ReadColumns::new(std::iter::empty()), + read_column_types: vec![], + filters: vec![], + time_filters: vec![], + series_row_selector: None, + append_mode: false, + filter_deleted: true, + merge_mode: crate::region::options::MergeMode::LastRow, + sequence_range: None, + partition_expr_version: 0, + } + .build(), + } +} diff --git a/src/mito2/src/metrics.rs b/src/mito2/src/metrics.rs index e8718f0e758..f7af9d54ffc 100644 --- a/src/mito2/src/metrics.rs +++ b/src/mito2/src/metrics.rs @@ -319,6 +319,14 @@ lazy_static! { "mito upload bytes total", ) .unwrap(); + /// Bounded event labels distinguish lookup misses from actual shared loads. + pub static ref RANGE_CACHE_WORKER_EVENTS: IntCounterVec = register_int_counter_vec!( + "greptime_mito_range_cache_worker_events_total", + "Range cache load coordination and SST retirement events", + &["event"] + ) + .unwrap(); + /// Cache eviction counter, labeled with cache type and eviction reason. pub static ref CACHE_EVICTION: IntCounterVec = register_int_counter_vec!( "greptime_mito_cache_eviction", diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index 081efa7bec8..a3b857fc24d 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -14,6 +14,7 @@ //! Utilities for the partition range scan result cache. +use std::future::Future; use std::mem; use std::sync::Arc; @@ -36,6 +37,7 @@ use table::predicate::is_string_timestamp_literal; use tokio::sync::{mpsc, oneshot}; use crate::cache::CacheStrategy; +use crate::cache::range_cache::{RangeCacheDelivery, RangeCacheLookup, RangeCacheWorker}; use crate::error::{ComputeArrowSnafu, Result}; use crate::read::BoxedRecordBatchStream; use crate::read::read_columns::ReadColumns; @@ -608,9 +610,14 @@ fn build_range_cache_key_inner( } /// Returns a stream that replays cached record batches. -pub(crate) fn cached_flat_range_stream(value: Arc) -> BoxedRecordBatchStream { +pub(crate) fn cached_flat_range_stream( + delivery: Arc, +) -> BoxedRecordBatchStream { Box::pin(try_stream! { - for cached_batch in &value.cached_batches { + // Borrow through the delivery guard so it survives every yield. Batches + // handed to downstream operators are then owned by the query, not by + // this completed-result backlog. + for cached_batch in &delivery.value().cached_batches { let mut offset = 0; for &len in &cached_batch.slice_lengths { yield cached_batch.batch.slice(offset, len); @@ -620,14 +627,147 @@ pub(crate) fn cached_flat_range_stream(value: Arc) -> Boxed }) } +/// Builds only one decoded result per full cache key. A load must be driven +/// independently of any caller: merge readers can wait for another stream with +/// the same key before polling their first stream to completion. +pub(crate) struct CachedRangeReader<'a> { + pub(crate) stream_ctx: &'a StreamContext, + pub(crate) part_range: &'a PartitionRange, + pub(crate) key: Option, + pub(crate) metrics: &'a PartitionMetrics, +} + +impl CachedRangeReader<'_> { + pub(crate) async fn build(self, build: F) -> Result<(BoxedRecordBatchStream, usize)> + where + F: FnOnce() -> Fut + Clone + Send + 'static, + Fut: Future> + Send + 'static, + { + let Self { + stream_ctx, + part_range, + key, + metrics, + } = self; + let strategy = &stream_ctx.input.cache_strategy; + let (Some(key), Some(worker)) = (key, strategy.range_cache_worker()) else { + return build().await; + }; + let files = stream_ctx.ranges[part_range.identifier] + .row_group_indices + .iter() + .map(|index| stream_ctx.input.file_from_index(*index).clone()) + .collect(); + let loader = build.clone(); + let load_worker = worker.clone(); + let load_strategy = strategy.clone(); + let load_metrics = metrics.clone(); + let key_size = key.estimated_size(); + match worker.get_or_load(key, files, async move { + let result = async { + let (stream, _) = loader().await?; + collect_range_result(stream, &load_worker, &load_strategy).await + } + .await; + match result { + Ok(Some(value)) => { + load_metrics.inc_range_cache_size(key_size + value.value().estimated_size()); + Some(value) + } + Ok(None) => None, + Err(error) => { + warn!(error; "Failed to load range cache; falling back to an uncached scan"); + None + } + } + }) { + RangeCacheLookup::Hit(value) => { + metrics.inc_range_cache_hit(); + Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)) + } + RangeCacheLookup::Loading(receiver) => { + metrics.inc_range_cache_miss(); + Ok(( + shared_range_stream(receiver, build), + DEFAULT_READ_BATCH_SIZE, + )) + } + RangeCacheLookup::Bypass => { + metrics.inc_range_cache_miss(); + build().await + } + } + } +} + +/// Waiting is lazy, but the producer is independent of every consumer. +fn shared_range_stream( + mut receiver: tokio::sync::watch::Receiver>>, + build: F, +) -> BoxedRecordBatchStream +where + F: FnOnce() -> Fut + Send + 'static, + Fut: Future> + Send + 'static, +{ + Box::pin(try_stream! { + let cached = receiver.wait_for(|cached| cached.is_some()).await + .ok().and_then(|cached| cached.clone()); + drop(receiver); + let mut stream = if let Some(cached) = cached { + cached_flat_range_stream(cached) + } else { + // A failed/oversized/cancelled fill is never a query result. + // Retry once through the ordinary streaming path, not the worker. + build().await?.0 + }; + while let Some(batch) = stream.try_next().await? { + yield batch; + } + }) +} + +/// The load budget covers retained input buffers and completed deliveries. Acquisition +/// must be nonblocking: waiting with a partially accumulated result deadlocks +/// when several loads consume all permits. Oversized ranges retain streaming +/// query semantics by falling back to an uncached scan. +async fn collect_range_result( + mut stream: BoxedRecordBatchStream, + worker: &RangeCacheWorker, + strategy: &CacheStrategy, +) -> Result>> { + let mut permit: Option = None; + let mut buffer = CacheBatchBuffer::new(strategy); + while let Some(batch) = stream.try_next().await? { + let Some(next_permit) = worker + .load_memory + .try_acquire(batch.get_array_memory_size()) + else { + return Ok(None); + }; + if let Some(permit) = &mut permit { + permit.merge(next_permit); + } else { + permit = Some(next_permit); + } + buffer.push(batch)?; + } + let Some(value) = buffer.finish().await.transpose()? else { + return Ok(None); + }; + let Some(permit) = worker + .load_memory + .try_resize(permit, value.estimated_size()) + else { + return Ok(None); + }; + Ok(Some(RangeCacheDelivery::new(value, permit))) +} + enum CacheConcatCommand { Compact(Vec), Finish { pending: Vec, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - result_tx: Option>>>, + result_tx: oneshot::Sender>>, }, } @@ -708,13 +848,7 @@ async fn run_cache_concat_task( return; } } - CacheConcatCommand::Finish { - pending, - key, - cache_strategy, - part_metrics, - result_tx, - } => { + CacheConcatCommand::Finish { pending, result_tx } => { let compact_result = state .compact(pending, &limiter) .await @@ -722,10 +856,6 @@ async fn run_cache_concat_task( let result = match compact_result { Ok(v) => { let value = Arc::new(v); - part_metrics - .inc_range_cache_size(key.estimated_size() + value.estimated_size()); - cache_strategy.put_range_result(key, value.clone()); - Ok(value) } Err(e) => { @@ -735,9 +865,7 @@ async fn run_cache_concat_task( } }; - if let Some(tx) = result_tx { - let _ = tx.send(result); - } + let _ = result_tx.send(result); break; } @@ -750,8 +878,7 @@ struct CacheBatchBuffer { buffered_rows: usize, buffered_size: usize, sender: Option>, - /// Handle of the background concat task. Cleared once the task owns a finish - /// command, so only unfinished scans cancel it. + /// Kept until concat finishes so abandoning the shared load cancels concat too. concat_task: Option>, } @@ -816,39 +943,26 @@ impl CacheBatchBuffer { } } - fn finish( - mut self, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - result_tx: Option>>>, - ) { - let Some(sender) = self.sender.take() else { - return; - }; - - if sender + async fn finish(mut self) -> Option>> { + let sender = self.sender.take()?; + let (result_tx, result_rx) = oneshot::channel(); + sender .send(CacheConcatCommand::Finish { pending: mem::take(&mut self.buffered_batches), - key, - cache_strategy, - part_metrics, result_tx, }) - .is_ok() - { - // The task now owns the finish command, so it may keep populating the - // cache after the scan stream is dropped. - self.concat_task = None; - } + .ok()?; + // Keep ownership until the result arrives so cancellation also aborts + // concat while it waits for its temporary memory permits. + let result = result_rx.await.ok(); + self.concat_task = None; + result } } impl Drop for CacheBatchBuffer { fn drop(&mut self) { - // Still holding the handle means no finish command was sent: the scan was - // cancelled or failed and the queued batches can never reach the cache. - // Aborting releases them even while the task waits for a memory permit. + // Aborting releases queued batches even while concat waits for a permit. if let Some(task) = &self.concat_task { task.abort(); } @@ -856,6 +970,7 @@ impl Drop for CacheBatchBuffer { } /// Wraps a stream to cache its output for future range cache hits. +#[cfg(feature = "test")] pub(crate) fn cache_flat_range_stream( mut stream: BoxedRecordBatchStream, cache_strategy: CacheStrategy, @@ -869,7 +984,10 @@ pub(crate) fn cache_flat_range_stream( yield batch; } - buffer.finish(key, cache_strategy, part_metrics, None); + if let Some(Ok(value)) = buffer.finish().await { + part_metrics.inc_range_cache_size(key.estimated_size() + value.estimated_size()); + cache_strategy.put_range_result(key, value); + } }) } @@ -936,9 +1054,11 @@ mod tests { use datafusion_expr::{Expr, col, lit}; use smallvec::smallvec; use store_api::storage::{FileId, RegionId}; + use tokio::sync::watch; use super::*; use crate::cache::CacheManager; + use crate::cache::test_util::range_cache_key; use crate::read::flat_projection::FlatProjectionMapper; use crate::read::range::{RangeMeta, RowGroupIndex, SourceIndex}; use crate::read::scan_region::{PredicateGroup, ScanInput}; @@ -994,19 +1114,6 @@ mod tests { (key, part_metrics) } - async fn finish_cache_batch_buffer( - buffer: CacheBatchBuffer, - key: RangeScanCacheKey, - cache_strategy: CacheStrategy, - part_metrics: PartitionMetrics, - ) -> Result> { - let (tx, rx) = oneshot::channel(); - common_telemetry::info!("finish start"); - buffer.finish(key, cache_strategy, part_metrics, Some(tx)); - common_telemetry::info!("finish end"); - rx.await.context(crate::error::RecvSnafu)? - } - async fn new_stream_context( filters: Vec, query_time_range: Option, @@ -1050,6 +1157,295 @@ mod tests { (stream_ctx, partition_range) } + #[rstest::rstest] + #[case(1024 * 1024, true)] + #[case(1, false)] + #[tokio::test] + async fn shared_range_read_preserves_batches_or_falls_back_on_budget_exhaustion( + #[case] capacity: u64, + #[case] cached: bool, + ) { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let (mut ctx, range) = + new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await; + ctx.input.cache_strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(capacity) + .build(), + )); + let key = build_range_cache_key(&ctx, &range).unwrap(); + let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); + let calls = Arc::new(AtomicUsize::new(0)); + let expected = vec![make_batch(&[1, 2]), make_batch(&[3, 4, 5])]; + let batches = expected.clone(); + let count = calls.clone(); + let (stream, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key.clone()), + metrics: &metrics, + } + .build(move || async move { + count.fetch_add(1, Ordering::Relaxed); + Ok(( + Box::pin(futures::stream::iter(batches.into_iter().map(Ok))) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }) + .await + .unwrap(); + let actual: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(actual, expected); + assert_eq!(calls.load(Ordering::Relaxed), if cached { 1 } else { 2 }); + assert_eq!( + ctx.input.cache_strategy.get_range_result(&key).is_some(), + cached + ); + assert_eq!( + ctx.input + .cache_strategy + .range_cache_worker() + .unwrap() + .load_memory + .available_permits(), + (capacity as usize).div_ceil(1024) + ); + } + + fn start_shared_load( + strategy: &CacheStrategy, + key: RangeScanCacheKey, + batches: Vec, + ) -> watch::Receiver>> { + let worker = strategy.range_cache_worker().unwrap(); + let loader_worker = worker.clone(); + let loader_strategy = strategy.clone(); + let lookup = worker.get_or_load(key, vec![], async move { + collect_range_result( + Box::pin(futures::stream::iter(batches.into_iter().map(Ok))), + &loader_worker, + &loader_strategy, + ) + .await + .unwrap() + }); + let RangeCacheLookup::Loading(receiver) = lookup else { + panic!("expected a shared load"); + }; + receiver + } + + async fn wait_for_shared_load( + receiver: &mut watch::Receiver>>, + ) -> Option> { + tokio::time::timeout(std::time::Duration::from_secs(5), async { + receiver + .wait_for(|value| value.is_some()) + .await + .ok() + .and_then(|value| value.clone()) + }) + .await + .expect("a shared load must finish or fall back without waiting for consumers") + } + + #[tokio::test] + async fn completed_load_keeps_memory_until_all_subscribers_release() { + let strategy = test_cache_strategy(); + let worker = strategy.range_cache_worker().unwrap(); + let capacity = worker.load_memory.available_permits(); + let (key, _) = test_cache_context(&strategy); + let mut first = start_shared_load(&strategy, key.clone(), vec![make_batch(&[1, 2, 3])]); + let second = first.clone(); + drop(wait_for_shared_load(&mut first).await.unwrap()); + + // A completed flight no longer counts toward the flight limit, but + // unpolled subscribers must still reserve its decoded result memory. + assert!(worker.load_memory.available_permits() < capacity); + drop(first); + assert!(worker.load_memory.available_permits() < capacity); + drop(second); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("the last subscriber must release the delivery reservation"); + assert!( + worker.get(&key).is_some(), + "the cache still owns the result" + ); + } + + #[rstest::rstest] + #[tokio::test] + async fn successive_loads_with_slow_consumers_exhaust_delivery_budget( + #[values(false, true)] invalidate: bool, + ) { + let strategy = CacheStrategy::EnableAll(Arc::new( + CacheManager::builder() + .range_result_cache_size(8192) + .build(), + )); + let worker = strategy.range_cache_worker().unwrap(); + let capacity = worker.load_memory.available_permits(); + let mut held_streams = Vec::new(); + let mut fallbacks = 0; + + // Each flight finishes before the next starts, exceeding the 32-flight + // limit cumulatively. Only the delivery budget can bound this backlog. + for _ in 0..40 { + let file = FileId::random(); + let key = range_cache_key(&[file]); + let expected = vec![make_large_binary_batch(1, 2048)]; + let mut receiver = start_shared_load(&strategy, key, expected.clone()); + let shared = wait_for_shared_load(&mut receiver).await.is_some(); + if invalidate { + worker.invalidate_file(file); + } + let fallback_batches = expected.clone(); + let stream = shared_range_stream(receiver, move || async move { + Ok(( + Box::pin(futures::stream::iter(fallback_batches.into_iter().map(Ok))) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }); + if shared { + held_streams.push(stream); + } else { + fallbacks += 1; + let actual: Vec<_> = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.try_collect()) + .await + .expect("fallback must not wait for an unpolled range") + .unwrap(); + assert_eq!(actual, expected); + } + } + + assert!(!held_streams.is_empty()); + assert!(held_streams.len() < 32); + assert!( + fallbacks > 0, + "completed results must not escape the budget" + ); + assert!(worker.load_memory.available_permits() < capacity); + + // Transfer ownership out of watch receivers into actively polled replay + // streams, then leave them suspended at their first yield. + for stream in &mut held_streams { + assert!(stream.try_next().await.unwrap().is_some()); + } + assert!(worker.load_memory.available_permits() < capacity); + drop(held_streams); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping slow readers must unblock new shared loads"); + + let mut retry = start_shared_load( + &strategy, + range_cache_key(&[FileId::random()]), + vec![make_large_binary_batch(1, 2048)], + ); + assert!(wait_for_shared_load(&mut retry).await.is_some()); + } + + #[tokio::test] + async fn cancelling_a_partial_load_releases_its_reservation() { + let strategy = test_cache_strategy(); + let worker = strategy.range_cache_worker().unwrap().clone(); + let capacity = worker.load_memory.available_permits(); + let (key, _) = test_cache_context(&strategy); + let loader_worker = worker.clone(); + let lookup = worker.get_or_load(key.clone(), vec![], async move { + let stream = Box::pin(try_stream! { + yield make_batch(&[1, 2, 3]); + std::future::pending::<()>().await; + }); + collect_range_result(stream, &loader_worker, &strategy) + .await + .unwrap() + }); + let RangeCacheLookup::Loading(receiver) = lookup else { + panic!("expected a shared load"); + }; + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() == capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("the partial result must acquire its input reservation"); + drop(receiver); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while worker.load_memory.available_permits() != capacity { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancellation must release the partial result reservation"); + assert!(worker.get(&key).is_none()); + } + + #[tokio::test] + async fn follower_can_finish_without_polling_the_first_range_stream() { + use std::sync::atomic::{AtomicUsize, Ordering}; + + let (ctx, range) = + new_stream_context(vec![col("k0").eq(lit("a"))], None, FileTimeRange::default()).await; + let key = build_range_cache_key(&ctx, &range).unwrap(); + let (_, metrics) = test_cache_context(&ctx.input.cache_strategy); + let gate = Arc::new(tokio::sync::Semaphore::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + let count = calls.clone(); + let loader_gate = gate.clone(); + let build = move || async move { + count.fetch_add(1, Ordering::Relaxed); + let _permit = loader_gate.acquire().await.unwrap(); + Ok(( + Box::pin(futures::stream::iter([Ok(make_batch(&[1, 2, 3]))])) + as BoxedRecordBatchStream, + DEFAULT_READ_BATCH_SIZE, + )) + }; + let (first, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key.clone()), + metrics: &metrics, + } + .build(build.clone()) + .await + .unwrap(); + let (second, _) = CachedRangeReader { + stream_ctx: &ctx, + part_range: &range, + key: Some(key), + metrics: &metrics, + } + .build(build) + .await + .unwrap(); + gate.add_permits(1); + let second_batches: Vec<_> = + tokio::time::timeout(std::time::Duration::from_secs(5), second.try_collect()) + .await + .expect("a follower must not depend on polling the first query") + .unwrap(); + let first_batches: Vec<_> = first.try_collect().await.unwrap(); + assert_eq!(first_batches, second_batches); + assert_eq!(second_batches, vec![make_batch(&[1, 2, 3])]); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + /// Helper to create a timestamp millisecond literal. fn ts_lit(val: i64) -> Expr { lit(ScalarValue::TimestampMillisecond(Some(val), None)) @@ -1600,7 +1996,13 @@ mod tests { make_batch(&[1, 2, 3]).get_array_memory_size(), )); - let replayed = cached_flat_range_stream(value) + let worker = RangeCacheWorker::new(1024 * 1024); + let permit = worker + .load_memory + .try_acquire(value.estimated_size()) + .unwrap(); + let delivery = RangeCacheDelivery::new(value, permit); + let replayed = cached_flat_range_stream(delivery) .try_collect::>() .await .unwrap(); @@ -1652,28 +2054,20 @@ mod tests { let strategy = test_cache_strategy(); let batch = make_batch(&[1, 2, 3]); let expected_size = batch.get_array_memory_size(); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(batch).unwrap(); - let value = finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!(value.cached_batches[0].slice_lengths, vec![3]); assert_eq!(value.estimated_batches_size, expected_size); - assert!(Arc::ptr_eq( - &value, - &strategy.get_range_result(&key).unwrap() - )); } #[tokio::test] async fn cache_batch_buffer_compacts_when_rows_exceed_default_batch_size() { let strategy = test_cache_strategy(); let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(batch.clone()).unwrap(); @@ -1682,9 +2076,7 @@ mod tests { assert_eq!(buffer.buffered_rows, 0); assert!(buffer.buffered_batches.is_empty()); - let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!( value.cached_batches[0].slice_lengths, @@ -1703,7 +2095,6 @@ mod tests { .range_result_cache_size((large_batch.get_array_memory_size() * 3) as u64) .build(), )); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); buffer.push(large_batch.clone()).unwrap(); @@ -1716,9 +2107,7 @@ mod tests { assert_eq!(buffer.buffered_rows, 0); assert!(buffer.buffered_batches.is_empty()); - let value = finish_cache_batch_buffer(buffer, key, strategy, part_metrics) - .await - .unwrap(); + let value = buffer.finish().await.unwrap().unwrap(); assert_eq!(value.cached_batches.len(), 1); assert_eq!( value.cached_batches[0].slice_lengths, @@ -1736,17 +2125,11 @@ mod tests { .range_result_cache_size(budget) .build(), )); - let (key, part_metrics) = test_cache_context(&strategy); let mut buffer = CacheBatchBuffer::new(&strategy); for _ in 0..4 { buffer.push(large_batch.clone()).unwrap(); } - assert!( - finish_cache_batch_buffer(buffer, key.clone(), strategy.clone(), part_metrics) - .await - .is_err() - ); - assert!(strategy.get_range_result(&key).is_none()); + assert!(!matches!(buffer.finish().await, Some(Ok(_)))); } } diff --git a/src/mito2/src/read/seq_scan.rs b/src/mito2/src/read/seq_scan.rs index 42119135805..4b09533e29a 100644 --- a/src/mito2/src/read/seq_scan.rs +++ b/src/mito2/src/read/seq_scan.rs @@ -42,9 +42,7 @@ use crate::read::flat_merge::FlatMergeReader; use crate::read::last_row::FlatLastRowReader; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RangeMeta; -use crate::read::range_cache::{ - build_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_range_cache_key}; use crate::read::scan_region::{ScanInput, StreamContext}; use crate::read::scan_util::{ BATCH_SIZE_THRESHOLD, MIN_COMPACTION_SPLIT_ROWS, PartitionMetrics, PartitionMetricsList, @@ -307,15 +305,39 @@ impl SeqScan { key } }); - - if let Some(key) = cache_key.as_ref() { - if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { - part_metrics.inc_range_cache_hit(); - return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); - } - part_metrics.inc_range_cache_miss(); + let ctx = stream_ctx.clone(); + let range = *part_range; + let metrics = part_metrics.clone(); + CachedRangeReader { + stream_ctx, + part_range, + key: cache_key, + metrics: part_metrics, } + .build(move || async move { + Self::build_uncached_partition_range( + &ctx, + &range, + read_mode, + &metrics, + partition_pruner, + file_scan_semaphore, + merge_semaphore, + ) + .await + }) + .await + } + async fn build_uncached_partition_range( + stream_ctx: &Arc, + part_range: &PartitionRange, + read_mode: FlatReadMode, + part_metrics: &PartitionMetrics, + partition_pruner: Arc, + file_scan_semaphore: Option>, + merge_semaphore: Option>, + ) -> Result<(BoxedRecordBatchStream, usize)> { let mut sources = Vec::new(); let split_batch_size = build_flat_sources( stream_ctx, @@ -339,16 +361,6 @@ impl SeqScan { ) .await?; - let stream = match cache_key { - Some(key) => cache_flat_range_stream( - stream, - stream_ctx.input.cache_strategy.clone(), - key, - part_metrics.clone(), - ), - None => stream, - }; - Ok((stream, estimated_rows_per_batch)) } diff --git a/src/mito2/src/read/series_candidate.rs b/src/mito2/src/read/series_candidate.rs index 3ba5512a7bd..185d3ea0dff 100644 --- a/src/mito2/src/read/series_candidate.rs +++ b/src/mito2/src/read/series_candidate.rs @@ -47,9 +47,7 @@ use crate::error::{ use crate::read::BoxedRecordBatchStream; use crate::read::pruner::{PartitionPruner, Pruner}; use crate::read::range::RowGroupIndex; -use crate::read::range_cache::{ - build_candidate_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_candidate_range_cache_key}; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{PartitionMetrics, new_filter_metrics, scan_flat_mem_ranges}; use crate::series_index::{ @@ -371,14 +369,29 @@ impl SeriesCandidateRangeBuilder { } else { build_candidate_range_cache_key(&self.stream_ctx, &part_range) }; - if let Some(key) = cache_key.as_ref() { - if let Some(value) = self.stream_ctx.input.cache_strategy.get_range_result(key) { - self.part_metrics.inc_range_cache_hit(); - return Ok(cached_flat_range_stream(value)); - } - self.part_metrics.inc_range_cache_miss(); + let builder = self.clone(); + let (stream, _) = CachedRangeReader { + stream_ctx: &self.stream_ctx, + part_range: &part_range, + key: cache_key, + metrics: &self.part_metrics, } + .build(move || async move { + let stream = builder + .build_uncached_range_stream(part_range, merge_partition) + .await?; + Ok((stream, DEFAULT_READ_BATCH_SIZE)) + }) + .await?; + Ok(stream) + } + async fn build_uncached_range_stream( + &self, + part_range: PartitionRange, + merge_partition: usize, + ) -> Result { + let range_meta = &self.stream_ctx.ranges[part_range.identifier]; let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); for index in &range_meta.row_group_indices { let source = self.build_source(*index, range_meta.time_range).await?; @@ -400,15 +413,7 @@ impl SeriesCandidateRangeBuilder { "SeriesCandidateScanner::range_merge", )?; - Ok(match cache_key { - Some(key) => cache_flat_range_stream( - stream, - self.stream_ctx.input.cache_strategy.clone(), - key, - self.part_metrics.clone(), - ), - None => stream, - }) + Ok(stream) } async fn build_source( diff --git a/src/mito2/src/read/series_reader.rs b/src/mito2/src/read/series_reader.rs index 5ac69d191a6..c030b5b3e7e 100644 --- a/src/mito2/src/read/series_reader.rs +++ b/src/mito2/src/read/series_reader.rs @@ -31,9 +31,7 @@ use crate::error::InvalidRequestSnafu; use crate::error::{JoinSnafu, Result, UnexpectedSnafu}; use crate::read::BoxedRecordBatchStream; use crate::read::pruner::PartitionPruner; -use crate::read::range_cache::{ - build_series_range_cache_key, cache_flat_range_stream, cached_flat_range_stream, -}; +use crate::read::range_cache::{CachedRangeReader, build_series_range_cache_key}; use crate::read::scan_region::StreamContext; use crate::read::scan_util::{ PartitionMetrics, SplitRecordBatchStream, compute_average_batch_size, @@ -377,14 +375,36 @@ async fn build_series_partition_range( .enable_range_cache .then(|| build_series_range_cache_key(&stream_ctx, &part_range, range)) .flatten(); - if let Some(key) = cache_key.as_ref() { - if let Some(value) = stream_ctx.input.cache_strategy.get_range_result(key) { - part_metrics.inc_range_cache_hit(); - return Ok((cached_flat_range_stream(value), DEFAULT_READ_BATCH_SIZE)); - } - part_metrics.inc_range_cache_miss(); + let ctx = stream_ctx.clone(); + let metrics = part_metrics.clone(); + CachedRangeReader { + stream_ctx: &stream_ctx, + part_range: &part_range, + key: cache_key, + metrics: &part_metrics, } + .build(move || async move { + build_uncached_series_partition_range( + ctx, + part_range, + filter, + codec, + partition_pruner, + metrics, + ) + .await + }) + .await +} +async fn build_uncached_series_partition_range( + stream_ctx: Arc, + part_range: PartitionRange, + filter: MetricSeriesFilter, + codec: SparsePrimaryKeyCodec, + partition_pruner: Arc, + part_metrics: PartitionMetrics, +) -> Result<(BoxedRecordBatchStream, usize)> { let range_meta = &stream_ctx.ranges[part_range.identifier]; let split_batch_size = should_split_flat_batches_for_merge(&stream_ctx, range_meta); let mut sources = Vec::with_capacity(range_meta.row_group_indices.len()); @@ -480,15 +500,6 @@ async fn build_series_partition_range( compute_parallel_channel_size(estimated_batch_size), ) .await?; - let stream = match cache_key { - Some(key) => cache_flat_range_stream( - stream, - stream_ctx.input.cache_strategy.clone(), - key, - part_metrics, - ), - None => stream, - }; Ok((stream, estimated_batch_size)) } diff --git a/src/mito2/src/sst/file.rs b/src/mito2/src/sst/file.rs index 4c6613dee40..a627513d310 100644 --- a/src/mito2/src/sst/file.rs +++ b/src/mito2/src/sst/file.rs @@ -633,7 +633,12 @@ impl FileHandle { /// Mark the file as deleted and will delete it on drop asynchronously pub fn mark_deleted(&self) { - self.inner.lifetime.deleted.store(true, Ordering::Relaxed); + if !self.inner.lifetime.deleted.swap(true, Ordering::Relaxed) { + self.inner + .lifetime + .file_purger + .invalidate_file(&self.inner.meta); + } } pub fn compacting(&self) -> bool { @@ -867,6 +872,7 @@ pub async fn delete_files( // Remove meta of the file from cache. 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)); } } diff --git a/src/mito2/src/sst/file_purger.rs b/src/mito2/src/sst/file_purger.rs index 727c6c1cc9e..ad44c5bcc05 100644 --- a/src/mito2/src/sst/file_purger.rs +++ b/src/mito2/src/sst/file_purger.rs @@ -33,6 +33,9 @@ pub trait FilePurger: Send + Sync + fmt::Debug { /// Each metadata view has already released its reader reference. fn purge_file(&self, file_meta: FileMeta); + /// Retire cached scan results immediately, without waiting for physical deletion. + fn invalidate_file(&self, _: &FileMeta) {} + /// Returns the object-store reference manager used by SST metadata views. /// Local and no-op purgers do not track GC references. fn file_ref_manager(&self) -> Option { @@ -105,6 +108,7 @@ pub fn create_file_purger( && matches!(path_type, PathType::Data | PathType::Bare) { Arc::new(ObjectStoreFilePurger { + cache_manager, file_ref_manager, scheduler, range_index_deleter, @@ -180,6 +184,12 @@ impl LocalFilePurger { } impl FilePurger for LocalFilePurger { + fn invalidate_file(&self, file_meta: &FileMeta) { + if let Some(cache) = &self.cache_manager { + cache.invalidate_range_results(file_meta.file_id); + } + } + fn purge_file(&self, file_meta: FileMeta) { schedule_range_index_deletion( &self.scheduler, @@ -191,6 +201,7 @@ impl FilePurger for LocalFilePurger { } pub struct ObjectStoreFilePurger { + cache_manager: Option, file_ref_manager: FileReferenceManagerRef, scheduler: SchedulerRef, range_index_deleter: Option, @@ -224,6 +235,12 @@ fn schedule_range_index_deletion( } impl FilePurger for ObjectStoreFilePurger { + fn invalidate_file(&self, file_meta: &FileMeta) { + if let Some(cache) = &self.cache_manager { + cache.invalidate_range_results(file_meta.file_id); + } + } + fn purge_file(&self, file_meta: FileMeta) { schedule_range_index_deletion( &self.scheduler, @@ -407,12 +424,25 @@ mod tests { let owner = RegionId::new(9, 1); let index_path = crate::sst::range_index::range_index_path(owner, sst_file_id.file_id()); index_store.write(&index_path, "range index").await.unwrap(); + let cache = Arc::new( + crate::cache::CacheManager::builder() + .range_result_cache_size(1024 * 1024) + .build(), + ); + let cache_key = crate::cache::test_util::range_cache_key(&[sst_file_id.file_id()]); + cache.put_range_result( + cache_key.clone(), + Arc::new(crate::read::range_cache::RangeScanCacheValue::new( + vec![], + 0, + )), + ); let file_purger = create_file_purger( gc_enabled, PathType::Bare, scheduler.clone(), env.layer.clone(), - None, + Some(cache.clone()), Arc::new(crate::sst::file_ref::FileReferenceManager::new(None)), Some(RangeIndexDeleter::new(index_store.clone(), owner)), ); @@ -443,6 +473,8 @@ mod tests { if is_delete { version.mark_all_deleted(); } + // Invalidate before any old index-version reader releases the SST. + assert_eq!(cache.get_range_result(&cache_key).is_none(), is_delete); drop(version); env.wait_for_purges().await; assert!(object_store.exists(&path).await.unwrap());