fix(mito2): stop the series distributor from stalling on unread partitions (#10003)

I hereby agree to the terms of the [GreptimeDB
CLA](https://github.com/GreptimeTeam/.github/blob/main/CLA.md).

## Refer to a related PR or issue link (optional)

Cherry-pick of #9991 (merged into `main`) onto `narthex`.

## What's changed and what's your intention?

`FusedSeriesExec` with two inputs reads every scan partition one after
another. On `narthex` this hits the series distributor's round-robin
`send_timeout(100us)`. Tokio timers tick at 1 ms, so after a merge (one
batch per series) each series batch costs about 1 ms per unpolled
partition. With 16 partitions, two-selector native queries over regions
with several SSTs spent ~470 ms to 12.7 s in `SeriesScanExec`, while
legacy took ~10 ms.

#9991 waits on all senders at once and sends to the first free slot. The
cherry-pick applies cleanly and the diff is identical to #9991.

## PR Checklist
Please convert it to a draft if some of the following conditions are not
met.

- [x] I have written the necessary rustdoc comments.
- [x] I have added the necessary unit tests and integration tests.
- [ ] This PR requires documentation updates.
- [x] API changes are backward compatible.
- [x] Schema or data changes are backward compatible.
- [ ] This PR needs to be backported to release branches, and I have
added the `backport-<target>` labels (e.g. `backport-v1.3` targets
`release/v1.3`).

src/mito2/src/read/scan_util.rs

@@ -202,8 +202,6 @@ pub(crate) struct ScanMetricsSet {
     /// Elapsed time before the first poll operation.
     first_poll: Duration,
 
-    /// Number of send timeout in SeriesScan.
-    num_series_send_timeout: usize,
     /// Number of send full in SeriesScan.
     num_series_send_full: usize,
     /// Number of rows the series distributor scanned.
@@ -325,7 +323,6 @@ impl fmt::Debug for ScanMetricsSet {
             num_sst_batches,
             num_sst_rows,
             first_poll,
-            num_series_send_timeout,
             num_series_send_full,
             num_distributor_rows,
             num_distributor_batches,
@@ -459,9 +456,6 @@ impl fmt::Debug for ScanMetricsSet {
         }
 
         // Write non-zero distributor metrics
-        if *num_series_send_timeout > 0 {
-            write!(f, ", \"num_series_send_timeout\":{num_series_send_timeout}")?;
-        }
         if *num_series_send_full > 0 {
             write!(f, ", \"num_series_send_full\":{num_series_send_full}")?;
         }
@@ -775,7 +769,6 @@ impl ScanMetricsSet {
     /// Sets distributor metrics.
     fn set_distributor_metrics(&mut self, distributor_metrics: &SeriesDistributorMetrics) {
         let SeriesDistributorMetrics {
-            num_series_send_timeout,
             num_series_send_full,
             num_rows,
             num_batches,
@@ -784,7 +777,6 @@ impl ScanMetricsSet {
             divider_cost,
         } = distributor_metrics;
 
-        self.num_series_send_timeout += *num_series_send_timeout;
         self.num_series_send_full += *num_series_send_full;
         self.num_distributor_rows += *num_rows;
         self.num_distributor_batches += *num_batches;
@@ -1144,8 +1136,6 @@ impl fmt::Debug for PartitionMetrics {
 /// Metrics for the series distributor.
 #[derive(Default)]
 pub(crate) struct SeriesDistributorMetrics {
-    /// Number of send timeout in SeriesScan.
-    pub(crate) num_series_send_timeout: usize,
     /// Number of send full in SeriesScan.
     pub(crate) num_series_send_full: usize,
     /// Number of rows the series distributor scanned.

src/mito2/src/read/series_scan.rs

@@ -21,7 +21,7 @@ use std::collections::HashSet;
 use std::fmt;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
-use std::time::{Duration, Instant};
+use std::time::Instant;
 
 use api::v1::SemanticType;
 use async_stream::try_stream;
@@ -55,7 +55,7 @@ use store_api::region_engine::{
 use store_api::storage::consts::{ReservedColumnId, is_internal_column};
 use table::predicate::Predicate;
 use tokio::sync::Semaphore;
-use tokio::sync::mpsc::error::{SendTimeoutError, TrySendError};
+use tokio::sync::mpsc::error::TrySendError;
 use tokio::sync::mpsc::{self, Receiver, Sender};
 
 use crate::error::{
@@ -76,9 +76,6 @@ use crate::read::stream::{ConvertBatchStream, ScanBatch, ScanBatchStream};
 use crate::sst::parquet::flat_format::primary_key_column_index;
 use crate::sst::parquet::format::PrimaryKeyArray;
 
-/// Timeout to send a batch to a sender.
-const SEND_TIMEOUT: Duration = Duration::from_micros(100);
-
 /// Maximum number of candidate series retained before distributing assignments.
 const CANDIDATE_SERIES_ASSIGNMENT_THRESHOLD: usize = 1_000_000;
 
@@ -1268,7 +1265,6 @@ impl SeriesDistributor {
         }
 
         metrics.scan_cost += fetch_start.elapsed();
-        metrics.num_series_send_timeout = self.senders.num_timeout;
         metrics.num_series_send_full = self.senders.num_full;
         part_metrics.set_distributor_metrics(&metrics);
 
@@ -1423,8 +1419,6 @@ struct SenderList {
     num_nones: usize,
     /// Index of the current partition to send.
     sender_idx: usize,
-    /// Number of timeout.
-    num_timeout: usize,
     /// Number of full senders.
     num_full: usize,
 }
@@ -1436,7 +1430,6 @@ impl SenderList {
             senders,
             num_nones,
             sender_idx: 0,
-            num_timeout: 0,
             num_full: 0,
         }
     }
@@ -1472,46 +1465,35 @@ impl SenderList {
     }
 
     /// Finds a partition and sends the batch to the partition.
-    async fn send_batch(&mut self, mut batch: SeriesBatch) -> Result<()> {
-        // Sends the batch without blocking first.
-        match self.try_send_batch(batch)? {
-            Some(b) => {
-                // Unable to send batch to partition.
-                batch = b;
-            }
-            None => {
-                return Ok(());
-            }
-        }
+    async fn send_batch(&mut self, batch: SeriesBatch) -> Result<()> {
+        let Some(batch) = self.try_send_batch(batch)? else {
+            return Ok(());
+        };
 
         loop {
             ensure!(self.num_nones < self.senders.len(), InvalidSenderSnafu);
 
-            let sender_idx = self.fetch_add_sender_idx();
-            let Some(sender) = &self.senders[sender_idx] else {
-                continue;
+            // Waits on every partition at once and sends to whichever frees a slot first. A
+            // consumer may leave partitions unpolled for a long time, e.g. one that drains them one
+            // after another, and waiting on such a partition would stall the scan.
+            let closed = {
+                let reservations = self.senders.iter().enumerate().filter_map(|(idx, sender)| {
+                    sender
+                        .as_ref()
+                        .map(|sender| Box::pin(async move { (idx, sender.reserve().await) }))
+                });
+                let ((idx, reserved), _, _) = futures::future::select_all(reservations).await;
+                match reserved {
+                    Ok(permit) => {
+                        permit.send(Ok(batch));
+                        return Ok(());
+                    }
+                    Err(_) => idx,
+                }
             };
-            // Adds a timeout to avoid blocking indefinitely and sending
-            // the batch in a round-robin fashion when some partitions
-            // don't poll their inputs. This may happen if we have a
-            // node like sort merging. But it is rare when we are using SeriesScan.
-            match sender.send_timeout(Ok(batch), SEND_TIMEOUT).await {
-                Ok(()) => break,
-                Err(SendTimeoutError::Timeout(res)) => {
-                    self.num_timeout += 1;
-                    // Safety: we send Ok.
-                    batch = res.unwrap();
-                }
-                Err(SendTimeoutError::Closed(res)) => {
-                    self.senders[sender_idx] = None;
-                    self.num_nones += 1;
-                    // Safety: we send Ok.
-                    batch = res.unwrap();
-                }
-            }
+            self.senders[closed] = None;
+            self.num_nones += 1;
         }
-
-        Ok(())
     }
 
     async fn send_error(&self, error: Error) {
@@ -1721,6 +1703,37 @@ mod tests {
         assert_eq!(0, active_receivers.load(Ordering::Relaxed));
     }
 
+    #[tokio::test]
+    async fn distributor_does_not_stall_on_partitions_read_one_after_another() {
+        const NUM_BATCHES: usize = 2000;
+        let (mut senders, receivers) = new_legacy_channel_list(16);
+        let producer = tokio::spawn(async move {
+            for _ in 0..NUM_BATCHES {
+                senders
+                    .send_batch(SeriesBatch::Flat(FlatSeriesBatch::default()))
+                    .await
+                    .unwrap();
+            }
+        });
+        // Reads partition 0 to its end before touching the next one, so all other partitions stay
+        // full while the producer still has batches to send.
+        let consume = async {
+            let mut received = 0;
+            for mut receiver in receivers.into_iter().flatten() {
+                while let Some(batch) = receiver.recv().await {
+                    batch.unwrap();
+                    received += 1;
+                }
+            }
+            received
+        };
+        let received = tokio::time::timeout(std::time::Duration::from_secs(1), consume)
+            .await
+            .expect("the distributor stalled on partitions the consumer was not reading");
+        producer.await.unwrap();
+        assert_eq!(NUM_BATCHES, received);
+    }
+
     fn new_test_record_batch(
         primary_keys: &[&[u8]],
         timestamps: &[i64],