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();