fix(mito2): coordinate range cache invalidation and singleflight reads (#9935)

* fix(mito2): invalidate retired SST range results and coalesce loads

Coordinate decoded range loads by their full scan key, bound shared work, and prevent late fills after SST retirement. Invalidate range results for local purge and object-store GC while preserving old snapshot reads.

Fixes #9934

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(mito): preserve SST lifetime across index replacement

Share retirement and final purge across immutable index metadata views. Keep version-specific GC references until readers release them and match Parquet liveness by file identity.

Fixes #9944

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* refactor(mito): list-based SST purge and single FilePurger::remove_file

The final purge of a retired SST no longer depends on FileMeta: index
files of every version are discovered by listing
({region_dir}/index/{file_id}. prefix), so even when the last live view
is a stale one, the purge still removes all newer index files.

- SstLifetime is slimmed down to { deleted, file_purger }; the purge now
  runs in the last FileHandleInner::drop (Arc::strong_count == 1, which
  is race-free here since only inners of the same file hold the lifetime).
- FilePurger is merged back into a single remove_file(file_meta,
  is_delete) method; is_delete means the file was retired AND this is the
  last live view.
- LocalFilePurger deletes the SST and all its index files via the new
  delete_file_and_all_indexes helper (list before delete, so a listing
  failure leaves the file untouched).
- Remove the dead index_outdated flag and its accessors, which had no
  callers in the entire repo.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(mito): elect final SST view release atomically

Track metadata views with an atomic counter instead of inspecting the lifetime Arc strong count. Concurrent destructors now elect exactly one final releaser while preserving the single remove_file interface.

Add a deterministic concurrent-drop regression covering retired and live files.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* refactor(mito): restore Arc-owned SST cleanup and clarify reference release

Restore the lifetime design from 835227dcbc, using the final Arc drop and latest metadata instead of manual view counting and per-SST index listing.

Rename FilePurger::remove_file to release_file_ref while retaining purge_file for final cleanup. Preserve and adapt the concurrent view-release regression test.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* refactor(mito): track SST references by file lifetime

Register one reference guard per logical SST lifetime and release it on final drop, including normal region close. Remove per-view release_file_ref and index_outdated handling.

Protect all index generations while their SST is referenced. Preserve the existing GC snapshot format by expanding the recorded index generation upper bound for older version-aware workers.

Update coverage for independent references, normal close, cross-region ownership, concurrent final cleanup and both GC listing modes.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* refactor(mito): release SST references directly on lifetime drop

Remove the redundant FileReferenceGuard and let SstLifetime register and unregister its file reference directly. Preserve unconditional reference release on normal close and retirement-only purge.

Expose the reference manager through FilePurger and adapt reference manager tests without changing GC or snapshot behavior.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* test(mito): verify SST cleanup with real storage and isolate GC policy cases

Replace the recording-only concurrent purge test with filesystem assertions using LocalFilePurger and shared fixtures. Keep intermediate-view retention and retired versus live file cleanup coverage without claiming exhaustive thread interleavings.

Separate GC reference-policy cases from reader-release ordering and reuse their setup. Production code remains unchanged.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* style(mito): format SST lifetime changes

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(mito2): use crate-rooted range cache imports

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

* fix(mito2): account for range cache delivery lifetime

Keep completed range results charged while watch subscribers or replay streams retain them. Cache entries keep weak delivery references, allowing cache-only ownership to release the delivery reservation. Cache hits reuse or nonblockingly acquire the reservation, falling back when the budget is exhausted.

Resize input reservations to completed-result estimates and cover successive slow consumers, cache-only handoff, hit sharing, invalidation, admission rejection, and cancellation with regression tests.

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>

---------

Signed-off-by: Lei, HUANG <ratuthomm@gmail.com>
This commit is contained in:
Lei, HUANG
2026-10-27 13:11:44 +00:00
committed by GitHub
parent 984f80e277
commit 44df6e5ded
10 changed files with 1446 additions and 161 deletions
+88 -19
View File
@@ -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<tokio::sync::OwnedSemaphorePermit> {
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<tokio::sync::OwnedSemaphorePermit>,
bytes: usize,
) -> Option<tokio::sync::OwnedSemaphorePermit> {
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<tokio::sync::SemaphorePermit<'_>> {
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<range_cache::RangeCacheWorker>> {
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<RangeScanCacheValue>,
) {
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<SelectorResultKey, Arc<SelectorResultValue>>;
/// Maps partition-range scan key to cached flat batches.
type RangeResultCache = Cache<RangeScanCacheKey, Arc<RangeScanCacheValue>>;
type RangeResultCache = Arc<range_cache::RangeCacheWorker>;
#[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<usize>,
#[case] result_size: usize,
#[case] expected_permits: Option<usize>,
) {
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};
+733
View File
@@ -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<Arc<RangeCacheDelivery>>;
/// 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<RangeScanCacheValue>,
_permit: OwnedSemaphorePermit,
}
impl RangeCacheDelivery {
/// Shares a decoded result together with its memory reservation.
pub(crate) fn new(value: Arc<RangeScanCacheValue>, permit: OwnedSemaphorePermit) -> Arc<Self> {
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<RangeCacheDelivery>),
Loading(watch::Receiver<SharedResult>),
Bypass,
}
struct Entry {
generation: u64,
value: Arc<RangeScanCacheValue>,
delivery: Mutex<Weak<RangeCacheDelivery>>,
}
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<SharedResult>,
invalidated: bool,
}
#[derive(Default)]
struct State {
generation: u64,
flights: HashMap<RangeScanCacheKey, Flight>,
/// 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<FileId, HashMap<Arc<RangeScanCacheKey>, 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<RangeScanCacheKey, Arc<Entry>>,
state: Mutex<State>,
/// Moka can call listeners during a cache operation. Queue only cleanup
/// metadata instead of reentering `state` from its eviction listener.
evictions: Mutex<mpsc::Receiver<(Arc<RangeScanCacheKey>, 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<Self> {
let (tx, rx) = mpsc::channel();
let cache = Cache::builder()
.max_capacity(capacity)
.weigher(|k: &RangeScanCacheKey, v: &Arc<Entry>| range_result_cache_weight(k, &v.value))
.eviction_listener(move |key, entry: Arc<Entry>, 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<Arc<RangeScanCacheValue>> {
self.cache.get(key).map(|entry| entry.value.clone())
}
fn get_for_delivery(&self, key: &RangeScanCacheKey) -> Option<RangeCacheLookup> {
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<RangeScanCacheValue>) {
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<RangeCacheDelivery>,
) {
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<F>(
self: &Arc<Self>,
key: RangeScanCacheKey,
files: Vec<FileHandle>,
load: F,
) -> RangeCacheLookup
where
F: Future<Output = SharedResult> + 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<RangeCacheWorker>,
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<RangeScanCacheValue> {
Arc::new(RangeScanCacheValue::new(vec![], 0))
}
fn empty_delivery() -> Arc<RangeCacheDelivery> {
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<SharedResult> {
match lookup {
RangeCacheLookup::Loading(receiver) => receiver,
_ => panic!("expected a shared load"),
}
}
async fn loaded(mut receiver: watch::Receiver<SharedResult>) -> 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);
}
}
+26
View File
@@ -145,3 +145,29 @@ pub(crate) fn assert_parquet_metadata_equal(x: Arc<ParquetMetaData>, y: Arc<Parq
let y = normalize_statistics(Arc::unwrap_or_clone(y));
assert_eq!(x, y);
}
/// A file-only range key shared by cache and SST-lifecycle regression tests.
pub(crate) fn range_cache_key(
files: &[store_api::storage::FileId],
) -> 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(),
}
}
+8
View File
@@ -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",
+467 -84
View File
@@ -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<RangeScanCacheValue>) -> BoxedRecordBatchStream {
pub(crate) fn cached_flat_range_stream(
delivery: Arc<RangeCacheDelivery>,
) -> 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<RangeScanCacheValue>) -> 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<RangeScanCacheKey>,
pub(crate) metrics: &'a PartitionMetrics,
}
impl CachedRangeReader<'_> {
pub(crate) async fn build<F, Fut>(self, build: F) -> Result<(BoxedRecordBatchStream, usize)>
where
F: FnOnce() -> Fut + Clone + Send + 'static,
Fut: Future<Output = Result<(BoxedRecordBatchStream, usize)>> + 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<F, Fut>(
mut receiver: tokio::sync::watch::Receiver<Option<Arc<RangeCacheDelivery>>>,
build: F,
) -> BoxedRecordBatchStream
where
F: FnOnce() -> Fut + Send + 'static,
Fut: Future<Output = Result<(BoxedRecordBatchStream, usize)>> + 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<Option<Arc<RangeCacheDelivery>>> {
let mut permit: Option<tokio::sync::OwnedSemaphorePermit> = 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<RecordBatch>),
Finish {
pending: Vec<RecordBatch>,
key: RangeScanCacheKey,
cache_strategy: CacheStrategy,
part_metrics: PartitionMetrics,
result_tx: Option<oneshot::Sender<Result<Arc<RangeScanCacheValue>>>>,
result_tx: oneshot::Sender<Result<Arc<RangeScanCacheValue>>>,
},
}
@@ -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<mpsc::UnboundedSender<CacheConcatCommand>>,
/// 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<common_runtime::JoinHandle<()>>,
}
@@ -816,39 +943,26 @@ impl CacheBatchBuffer {
}
}
fn finish(
mut self,
key: RangeScanCacheKey,
cache_strategy: CacheStrategy,
part_metrics: PartitionMetrics,
result_tx: Option<oneshot::Sender<Result<Arc<RangeScanCacheValue>>>>,
) {
let Some(sender) = self.sender.take() else {
return;
};
if sender
async fn finish(mut self) -> Option<Result<Arc<RangeScanCacheValue>>> {
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<Arc<RangeScanCacheValue>> {
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<Expr>,
query_time_range: Option<TimestampRange>,
@@ -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<RecordBatch>,
) -> watch::Receiver<Option<Arc<RangeCacheDelivery>>> {
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<Arc<RangeCacheDelivery>>>,
) -> Option<Arc<RangeCacheDelivery>> {
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::<Vec<_>>()
.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(_))));
}
}
+32 -20
View File
@@ -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<StreamContext>,
part_range: &PartitionRange,
read_mode: FlatReadMode,
part_metrics: &PartitionMetrics,
partition_pruner: Arc<PartitionPruner>,
file_scan_semaphore: Option<Arc<Semaphore>>,
merge_semaphore: Option<Arc<Semaphore>>,
) -> 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))
}
+23 -18
View File
@@ -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<BoxedRecordBatchStream> {
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(
+29 -18
View File
@@ -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<StreamContext>,
part_range: PartitionRange,
filter: MetricSeriesFilter,
codec: SparsePrimaryKeyCodec,
partition_pruner: Arc<PartitionPruner>,
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))
}
+7 -1
View File
@@ -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));
}
}
+33 -1
View File
@@ -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<FileReferenceManagerRef> {
@@ -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<CacheManagerRef>,
file_ref_manager: FileReferenceManagerRef,
scheduler: SchedulerRef,
range_index_deleter: Option<RangeIndexDeleter>,
@@ -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());