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 <killme2008@gmail.com>

* 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 <killme2008@gmail.com>

---------

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>

src/mito2/src/cache.rs

@@ -99,6 +99,10 @@ pub(crate) struct RangeResultMemoryLimiter {
     semaphore: Arc<tokio::sync::Semaphore>,
     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

src/mito2/src/read/range_cache.rs

@@ -736,26 +736,33 @@ 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.
+    concat_task: Option<common_runtime::JoinHandle<()>>,
 }
 
 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();