From 7472204db545364cf09a526be1df44492e0f179c Mon Sep 17 00:00:00 2001 From: LFC <990479+MichaelScofield@users.noreply.github.com> Date: Tue, 27 Oct 2026 15:32:46 +0800 Subject: [PATCH] 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 e6b60644b3ce10b96d8be3d8eb11c43c8d18d0ed) Signed-off-by: Dennis Zhuang Co-authored-by: dennis zhuang --- src/mito2/src/read/scan_util.rs | 10 ---- src/mito2/src/read/series_scan.rs | 99 +++++++++++++++++-------------- 2 files changed, 56 insertions(+), 53 deletions(-) diff --git a/src/mito2/src/read/scan_util.rs b/src/mito2/src/read/scan_util.rs index c5eed8853ae..9c0fbc017c2 100644 --- a/src/mito2/src/read/scan_util.rs +++ b/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. diff --git a/src/mito2/src/read/series_scan.rs b/src/mito2/src/read/series_scan.rs index 4328a20f59c..d11d06951e4 100644 --- a/src/mito2/src/read/series_scan.rs +++ b/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],