fix(mito2): stop the series distributor from stalling on unread partitions [Backport release/v1.3] (#9994)

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

When every partition channel is full, the distributor walked the senders
round-robin with a 100us send timeout each. Tokio timers fire on a 1ms
tick, so a consumer that leaves some partitions unpolled costs about one
millisecond per unpolled partition for every series batch.
SeriesScan::build_stream and scan_all_partitions drain partitions one
after another, and SortPreservingMergeExec above a multi-partition
SeriesScan only pulls the partition holding the current minimum key.
With 16 partitions and one batch per series after a merge, a 40-series
scan read partition by partition took ~480ms (measured on the native
PromQL two-input route).

Wait on all senders at once and send to the first one that frees a slot.
The send-timeout metric goes away with the timeout.


(cherry picked from commit e6b60644b3)

Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
Co-authored-by: dennis zhuang <killme2008@gmail.com>

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

@@ -17,7 +17,7 @@
 use std::fmt;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::{Arc, Mutex};
-use std::time::{Duration, Instant};
+use std::time::Instant;
 
 use async_stream::try_stream;
 use common_error::ext::BoxedError;
@@ -38,7 +38,7 @@ use store_api::region_engine::{
     PartitionRange, PrepareRequest, QueryScanContext, RegionScanner, ScannerProperties,
 };
 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::{
@@ -59,9 +59,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;
 
@@ -830,7 +827,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);
 
@@ -982,8 +978,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,
 }
@@ -995,7 +989,6 @@ impl SenderList {
             senders,
             num_nones,
             sender_idx: 0,
-            num_timeout: 0,
             num_full: 0,
         }
     }
@@ -1031,46 +1024,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) {
@@ -1280,6 +1262,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],