From 43ce07fa694ca53c87a2f104bfe264e0da2e96e9 Mon Sep 17 00:00:00 2001 From: dennis zhuang Date: Sun, 20 Sep 2026 04:58:47 +0000 Subject: [PATCH] fix(mito2): cancel cache construction for incomplete scans (#9254) * fix(mito2): cancel cache construction for incomplete scans CacheBatchBuffer spawned the background concat task and dropped its join handle. A scan that was cancelled or failed therefore left the task alive: it kept compacting already queued batches, and while waiting for a range result memory permit it held them, even though without a finish command the result can never be put into the cache. Keep the handle and abort it when the buffer is dropped. The handle is cleared once the task owns the finish command, so a completed scan still populates the cache after its stream is dropped. Abort does not preempt a concat that is already running; it takes effect the next time the task is polled. Signed-off-by: Dennis Zhuang * test(mito2): wait for the permit park before cancelling the buffer An empty buffered_batches only proves the batches were enqueued, so the cancellation test could abort a concat task that had never been polled. Count acquisitions that find too few permits, a test-only signal, and drop the buffer once the task has reached that wait. Nothing awaits between the check and the parking, so an observed increment means the caller is about to wait for permits the test holds for the rest of the case. Signed-off-by: Dennis Zhuang --------- Signed-off-by: Dennis Zhuang --- src/mito2/src/cache.rs | 20 ++++++++ src/mito2/src/read/range_cache.rs | 81 ++++++++++++++++++++++++++----- 2 files changed, 89 insertions(+), 12 deletions(-) diff --git a/src/mito2/src/cache.rs b/src/mito2/src/cache.rs index 2e98ef33454..0d43d31f434 100644 --- a/src/mito2/src/cache.rs +++ b/src/mito2/src/cache.rs @@ -99,6 +99,10 @@ pub(crate) struct RangeResultMemoryLimiter { semaphore: Arc, permit_bytes: usize, total_permits: usize, + /// Number of acquisitions that found too few permits and parked. Test-only + /// signal to synchronize with a caller waiting inside [Self::acquire]. + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize, } impl Default for RangeResultMemoryLimiter { @@ -120,6 +124,8 @@ impl RangeResultMemoryLimiter { semaphore: Arc::new(tokio::sync::Semaphore::new(total_permits)), permit_bytes, total_permits, + #[cfg(test)] + waited_acquires: std::sync::atomic::AtomicUsize::new(0), } } @@ -128,6 +134,12 @@ impl RangeResultMemoryLimiter { self.permit_bytes } + #[cfg(test)] + pub(crate) fn waited_acquires(&self) -> usize { + self.waited_acquires + .load(std::sync::atomic::Ordering::Acquire) + } + #[cfg(test)] pub(crate) fn available_permits(&self) -> usize { self.semaphore.available_permits() @@ -144,6 +156,14 @@ impl RangeResultMemoryLimiter { } .fail(); } + // Nothing awaits between this check and the parking below, so an observed + // increment means the caller is about to wait for the missing permits. + #[cfg(test)] + if self.semaphore.available_permits() < permits { + self.waited_acquires + .fetch_add(1, std::sync::atomic::Ordering::Release); + } + self.semaphore .acquire_many(permits as u32) .await diff --git a/src/mito2/src/read/range_cache.rs b/src/mito2/src/read/range_cache.rs index ddf8c3cc356..1108b074697 100644 --- a/src/mito2/src/read/range_cache.rs +++ b/src/mito2/src/read/range_cache.rs @@ -736,26 +736,33 @@ 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. + concat_task: Option>, } impl CacheBatchBuffer { fn new(cache_strategy: &CacheStrategy) -> Self { - let sender = cache_strategy.range_result_memory_limiter().map(|limiter| { - let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); - let (tx, rx) = mpsc::unbounded_channel(); - common_runtime::spawn_query(run_cache_concat_task( - rx, - limiter.clone(), - skip_threshold_bytes, - )); - tx - }); + let (sender, concat_task) = cache_strategy + .range_result_memory_limiter() + .map(|limiter| { + let skip_threshold_bytes = cache_strategy.range_result_cache_size().unwrap_or(0); + let (tx, rx) = mpsc::unbounded_channel(); + let task = common_runtime::spawn_query(run_cache_concat_task( + rx, + limiter.clone(), + skip_threshold_bytes, + )); + (tx, task) + }) + .unzip(); Self { buffered_batches: Vec::new(), buffered_rows: 0, buffered_size: 0, sender, + concat_task, } } @@ -814,9 +821,22 @@ impl CacheBatchBuffer { part_metrics, result_tx, }) - .is_err() + .is_ok() { - self.sender = None; + // The task now owns the finish command, so it may keep populating the + // cache after the scan stream is dropped. + self.concat_task = None; + } + } +} + +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. + if let Some(task) = &self.concat_task { + task.abort(); } } } @@ -1570,6 +1590,43 @@ mod tests { assert_eq!(replayed[1].num_rows(), 1); } + #[tokio::test] + async fn dropped_cache_buffer_releases_batches_while_waiting_for_memory() { + let strategy = test_cache_strategy(); + let limiter = strategy.range_result_memory_limiter().unwrap(); + // Take every permit so the concat task blocks before it can compact. + let _permit = limiter + .acquire(limiter.available_permits() * limiter.permit_bytes()) + .await + .unwrap(); + let batch = make_batch(&vec![1; DEFAULT_READ_BATCH_SIZE / 2 + 1]); + let weak = Arc::downgrade(batch.column(0)); + let mut buffer = CacheBatchBuffer::new(&strategy); + buffer.push(batch.clone()).unwrap(); + buffer.push(batch).unwrap(); + // Both batches were handed to the task instead of staying in the buffer. + assert!(buffer.buffered_batches.is_empty()); + + // Cancel only once the task is parked on the permit, otherwise the abort + // could land on a task that was never polled. + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while limiter.waited_acquires() == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("concat task must wait for a memory permit"); + drop(buffer); + + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while weak.upgrade().is_some() { + tokio::task::yield_now().await; + } + }) + .await + .expect("cancelled cache work must release input without waiting for a permit"); + } + #[tokio::test] async fn cache_batch_buffer_finishes_pending_batches() { let strategy = test_cache_strategy();