From 6573fc97e8bac5e45484561208bb8e8b3576170c Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 15 Sep 2026 16:18:21 +0800 Subject: [PATCH] fix(servers): bound ready Flight batch admission and cover lifecycle Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- src/servers/src/grpc/flight/stream.rs | 590 +++++++++++++++++++++----- 1 file changed, 485 insertions(+), 105 deletions(-) diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index c6dd5af9360..a38ee81b19f 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -264,10 +264,8 @@ impl FlightRecordBatchStream { let mut metrics = StreamMetrics::new(should_send_partial_metrics); let mut last_metrics_str = None; - let mut sent_first_batch = false; - let recordbatch_schema = recordbatches.schema(); - let schema = recordbatches.schema().arrow_schema().clone(); + let schema = recordbatch_schema.arrow_schema().clone(); let start = Instant::now(); if let Err(e) = tx.send(Ok(FlightMessage::Schema(schema))).await { warn!(e; "stop sending Flight data"); @@ -275,92 +273,48 @@ impl FlightRecordBatchStream { } metrics.send_schema_duration += start.elapsed(); - loop { - let start = Instant::now(); - let batch_or_err = if should_send_partial_metrics && can_send_metrics_before_batch { - match time::timeout( - FLIGHT_METRICS_HEARTBEAT_INTERVAL, - recordbatches.next().in_current_span(), - ) - .await - { - Ok(result) => result, - Err(_) => { - if let Some(metrics_str) = recordbatches - .metrics() - .and_then(|m| serde_json::to_string(&m).ok()) - && !Self::send_metrics_if_changed( - &mut tx, - &mut metrics, - &mut last_metrics_str, - metrics_str, - ) - .await - { - return; - } - metrics.fetch_content_duration += start.elapsed(); - continue; - } - } - } else { - recordbatches.next().in_current_span().await - }; - metrics.fetch_content_duration += start.elapsed(); - let Some(batch_or_err) = batch_or_err else { - break; - }; - match batch_or_err { - Ok(recordbatch) => { - metrics.total_rows += recordbatch.num_rows(); - metrics.record_batch_count += 1; - metrics.total_bytes += recordbatch.df_record_batch().get_array_memory_size(); - - let mut batches = vec![recordbatch]; - let mut eof = false; - let mut stream_error = None; - if !should_send_partial_metrics && sent_first_batch { - let mut rows = batches[0].num_rows(); - let mut bytes = batches[0].df_record_batch().get_array_memory_size(); - while batches.len() < MAX_BATCHES && rows < MAX_ROWS && bytes < MAX_BYTES { - let start = Instant::now(); - let next = - poll_fn(|cx| Poll::Ready(recordbatches.as_mut().poll_next(cx))) - .await; - metrics.fetch_content_duration += start.elapsed(); - match next { - Poll::Ready(Some(Ok(recordbatch))) => { - rows += recordbatch.num_rows(); - bytes += recordbatch.df_record_batch().get_array_memory_size(); - metrics.total_rows += recordbatch.num_rows(); - metrics.record_batch_count += 1; - metrics.total_bytes += - recordbatch.df_record_batch().get_array_memory_size(); - batches.push(recordbatch); - } - Poll::Ready(Some(e)) => { - stream_error = Some(e); - break; - } - Poll::Ready(None) => { - eof = true; - break; - } - Poll::Pending => break, + if should_send_partial_metrics { + loop { + let start = Instant::now(); + let batch_or_err = if can_send_metrics_before_batch { + match time::timeout( + FLIGHT_METRICS_HEARTBEAT_INTERVAL, + recordbatches.next().in_current_span(), + ) + .await + { + Ok(result) => result, + Err(_) => { + if let Some(metrics_str) = recordbatches + .metrics() + .and_then(|m| serde_json::to_string(&m).ok()) + && !Self::send_metrics_if_changed( + &mut tx, + &mut metrics, + &mut last_metrics_str, + metrics_str, + ) + .await + { + return; } + metrics.fetch_content_duration += start.elapsed(); + continue; } } - sent_first_batch = true; - - let batches = if batches.len() >= 2 { - match merge_record_batches(recordbatch_schema.clone(), &batches) { - Ok(merged) => vec![merged], - Err(_) => batches, - } - } else { - batches - }; - for recordbatch in batches { + } else { + recordbatches.next().in_current_span().await + }; + metrics.fetch_content_duration += start.elapsed(); + let Some(batch_or_err) = batch_or_err else { + break; + }; + match batch_or_err { + Ok(recordbatch) => { + metrics.total_rows += recordbatch.num_rows(); + metrics.record_batch_count += 1; + metrics.total_bytes += + recordbatch.df_record_batch().get_array_memory_size(); let start = Instant::now(); if let Err(e) = tx .send(Ok(FlightMessage::RecordBatch( @@ -372,9 +326,18 @@ impl FlightRecordBatchStream { return; } metrics.send_record_batch_duration += start.elapsed(); + if let Some(metrics_str) = recordbatches + .metrics() + .and_then(|m| serde_json::to_string(&m).ok()) + && { + last_metrics_str = Some(metrics_str.clone()); + !Self::send_metrics(&mut tx, &mut metrics, metrics_str).await + } + { + return; + } } - - if let Some(Err(e)) = stream_error { + Err(e) => { if e.status_code().should_log_error() { error!("{e:?}"); } @@ -384,36 +347,141 @@ impl FlightRecordBatchStream { } return; } - if eof { - break; + } + } + } else { + let mut sent_first_batch = false; + let mut pending = None; + loop { + let (batch_or_err, fetched) = match pending.take() { + Some(recordbatch) => (Ok(recordbatch), false), + None => { + let start = Instant::now(); + let batch_or_err = recordbatches.next().in_current_span().await; + metrics.fetch_content_duration += start.elapsed(); + let Some(batch_or_err) = batch_or_err else { + break; + }; + (batch_or_err, true) } - - if should_send_partial_metrics - && let Some(metrics_str) = recordbatches - .metrics() - .and_then(|m| serde_json::to_string(&m).ok()) - && { - last_metrics_str = Some(metrics_str.clone()); - !Self::send_metrics(&mut tx, &mut metrics, metrics_str).await + }; + let recordbatch = match batch_or_err { + Ok(recordbatch) => recordbatch, + Err(e) => { + if e.status_code().should_log_error() { + error!("{e:?}"); + } + let e = Err(e).context(error::CollectRecordbatchSnafu); + if let Err(e) = tx.send(e.map_err(|x| x.into())).await { + warn!(e; "stop sending Flight data"); } - { return; } + }; + let batch_rows = recordbatch.num_rows(); + let batch_bytes = recordbatch.df_record_batch().get_array_memory_size(); + if fetched { + metrics.total_rows += batch_rows; + metrics.record_batch_count += 1; + metrics.total_bytes += batch_bytes; } - Err(e) => { + + if !sent_first_batch || batch_rows >= MAX_ROWS || batch_bytes >= MAX_BYTES { + let start = Instant::now(); + if let Err(e) = tx + .send(Ok(FlightMessage::RecordBatch( + recordbatch.into_df_record_batch(), + ))) + .await + { + warn!(e; "stop sending Flight data"); + return; + } + metrics.send_record_batch_duration += start.elapsed(); + sent_first_batch = true; + continue; + } + + let mut batches = vec![recordbatch]; + let mut rows = batch_rows; + let mut bytes = batch_bytes; + let mut eof = false; + let mut stream_error = None; + loop { + if batches.len() == MAX_BATCHES || rows >= MAX_ROWS || bytes >= MAX_BYTES { + break; + } + let start = Instant::now(); + let next = + poll_fn(|cx| Poll::Ready(recordbatches.as_mut().poll_next(cx))).await; + metrics.fetch_content_duration += start.elapsed(); + match next { + Poll::Ready(Some(Ok(recordbatch))) => { + let batch_rows = recordbatch.num_rows(); + let batch_bytes = recordbatch.df_record_batch().get_array_memory_size(); + if batch_rows > MAX_ROWS - rows || batch_bytes > MAX_BYTES - bytes { + metrics.total_rows += batch_rows; + metrics.record_batch_count += 1; + metrics.total_bytes += batch_bytes; + pending = Some(recordbatch); + break; + } + metrics.total_rows += batch_rows; + metrics.record_batch_count += 1; + metrics.total_bytes += batch_bytes; + rows += batch_rows; + bytes += batch_bytes; + batches.push(recordbatch); + } + Poll::Ready(Some(Err(e))) => { + stream_error = Some(e); + break; + } + Poll::Ready(None) => { + eof = true; + break; + } + Poll::Pending => break, + } + } + + let batches = if batches.len() >= 2 { + match merge_record_batches(recordbatch_schema.clone(), &batches) { + Ok(merged) => vec![merged], + Err(_) => batches, + } + } else { + batches + }; + for recordbatch in batches { + let start = Instant::now(); + if let Err(e) = tx + .send(Ok(FlightMessage::RecordBatch( + recordbatch.into_df_record_batch(), + ))) + .await + { + warn!(e; "stop sending Flight data"); + return; + } + metrics.send_record_batch_duration += start.elapsed(); + } + if let Some(e) = stream_error { if e.status_code().should_log_error() { error!("{e:?}"); } - let e = Err(e).context(error::CollectRecordbatchSnafu); if let Err(e) = tx.send(e.map_err(|x| x.into())).await { warn!(e; "stop sending Flight data"); } return; } + if eof { + break; + } } } - // make last package to pass metrics + // Make the last package pass metrics exactly once at EOF. if let Some(metrics_str) = recordbatches .metrics() .and_then(|m| serde_json::to_string(&m).ok()) @@ -506,6 +574,7 @@ mod test { enum ScriptedItem { Batch(common_recordbatch::error::Result), Pending, + PersistentPending, } struct ScriptedBatchStream { @@ -514,6 +583,11 @@ mod test { poll_count: Arc, } + struct DropFlagStream { + schema: SchemaRef, + dropped: Arc, + } + fn query_context_with_matching_capability() -> Arc { let query_ctx = QueryContext::arc(); let remote_query_id = query_ctx @@ -588,7 +662,7 @@ mod test { } fn metrics(&self) -> Option { - None + Some(RecordBatchMetrics::default()) } } @@ -604,11 +678,44 @@ mod test { cx.waker().wake_by_ref(); Poll::Pending } + Some(ScriptedItem::PersistentPending) => { + self.items.push_front(ScriptedItem::PersistentPending); + Poll::Pending + } None => Poll::Ready(None), } } } + impl RecordBatchStream for DropFlagStream { + fn schema(&self) -> SchemaRef { + self.schema.clone() + } + + fn output_ordering(&self) -> Option<&[OrderOption]> { + None + } + + fn metrics(&self) -> Option { + None + } + } + + impl Stream for DropFlagStream { + type Item = common_recordbatch::error::Result; + + fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Pending + } + } + + impl Drop for DropFlagStream { + fn drop(&mut self) { + self.dropped + .store(true, std::sync::atomic::Ordering::Relaxed); + } + } + fn int_batch(schema: SchemaRef, values: impl IntoIterator) -> RecordBatch { RecordBatch::new( schema, @@ -617,8 +724,9 @@ mod test { .unwrap() } - async fn flight_messages( + async fn flight_messages_with_context( recordbatches: SendableRecordBatchStream, + query_ctx: Arc, ) -> Vec> { let mut stream = FlightRecordBatchStream::new( FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( @@ -626,7 +734,7 @@ mod test { )), TracingContext::default(), FlightCompression::default(), - QueryContext::arc(), + query_ctx, ); let decoder = &mut FlightDecoder::default(); let mut messages = Vec::new(); @@ -643,6 +751,78 @@ mod test { messages } + async fn flight_messages( + recordbatches: SendableRecordBatchStream, + ) -> Vec> { + flight_messages_with_context(recordbatches, QueryContext::arc()).await + } + + #[tokio::test] + async fn test_drop_cancels_and_releases_upstream_stream() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let dropped = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches(Box::pin( + DropFlagStream { + schema, + dropped: dropped.clone(), + }, + ))), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + drop(stream); + tokio::time::timeout(Duration::from_secs(1), async { + while !dropped.load(std::sync::atomic::Ordering::Relaxed) { + tokio::task::yield_now().await; + } + }) + .await + .expect("dropping Flight stream must release upstream"); + } + + #[tokio::test] + async fn test_first_batch_is_sent_before_any_additional_poll() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let recordbatches = Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::PersistentPending, + ]), + poll_count: poll_count.clone(), + }); + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + let decoder = &mut FlightDecoder::default(); + let schema_data = stream.next().await.unwrap().unwrap(); + assert!(matches!( + decoder.try_decode(&schema_data).unwrap().unwrap(), + FlightMessage::Schema(_) + )); + let first = stream.next().await.unwrap().unwrap(); + assert!(matches!( + decoder.try_decode(&first).unwrap().unwrap(), + FlightMessage::RecordBatch(_) + )); + } + #[tokio::test] async fn test_ready_only_coalesces_later_batches_without_polling_after_first() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -674,6 +854,111 @@ mod test { assert_eq!(merged.num_rows(), 2); } + #[tokio::test] + async fn test_ready_only_exact_row_cap_flushes() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let recordbatches = Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..1023))), + ScriptedItem::PersistentPending, + ]), + poll_count: poll_count.clone(), + }); + let (tx, mut rx) = mpsc::channel::>(1); + let handle = tokio::spawn(FlightRecordBatchStream::flight_data_stream( + recordbatches, + tx, + false, + false, + )); + + assert!(matches!( + rx.next().await.unwrap().unwrap(), + FlightMessage::Schema(_) + )); + tokio::time::timeout(Duration::from_secs(1), async { + while poll_count.load(std::sync::atomic::Ordering::Relaxed) != 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("exact-cap group must be formed"); + assert!(matches!( + rx.next().await.unwrap().unwrap(), + FlightMessage::RecordBatch(_) + )); + let FlightMessage::RecordBatch(batch) = rx.next().await.unwrap().unwrap() else { + panic!("expected the exact-cap group"); + }; + assert_eq!(batch.num_rows(), 1024); + handle.abort(); + } + + #[tokio::test] + async fn test_pending_batch_is_sent_before_upstream_is_polled_again() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let recordbatches = Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), 0..1024))), + ScriptedItem::PersistentPending, + ]), + poll_count: poll_count.clone(), + }); + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchStreamInput::ready(FlightRecordBatchSource::RecordBatches( + recordbatches, + )), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + let decoder = &mut FlightDecoder::default(); + let schema_data = stream.next().await.unwrap().unwrap(); + assert!(matches!( + decoder.try_decode(&schema_data).unwrap().unwrap(), + FlightMessage::Schema(_) + )); + let first = stream.next().await.unwrap().unwrap(); + assert!(matches!( + decoder.try_decode(&first).unwrap().unwrap(), + FlightMessage::RecordBatch(_) + )); + tokio::time::timeout(Duration::from_secs(1), async { + while poll_count.load(std::sync::atomic::Ordering::Relaxed) != 3 { + tokio::task::yield_now().await; + } + }) + .await + .expect("second group must fetch the held batch"); + let held = stream.next().await.unwrap().unwrap(); + assert!(matches!( + decoder.try_decode(&held).unwrap().unwrap(), + FlightMessage::RecordBatch(_) + )); + let pending = stream.next().await.unwrap().unwrap(); + let FlightMessage::RecordBatch(batch) = decoder.try_decode(&pending).unwrap().unwrap() + else { + panic!("expected the held oversized batch"); + }; + assert_eq!(batch.num_rows(), 1024); + } + #[tokio::test] async fn test_ready_only_flushes_before_pending_and_before_error() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -705,6 +990,101 @@ mod test { assert_eq!(messages.len(), 5); } + #[tokio::test] + async fn test_verbose_ready_batches_preserve_batch_metrics_order() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let query_ctx = QueryContext::arc(); + query_ctx.set_explain_verbose(true); + let messages = flight_messages_with_context( + Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))), + ]), + poll_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + }), + query_ctx, + ) + .await; + assert!(matches!(messages[1], Ok(FlightMessage::RecordBatch(_)))); + assert!(matches!(messages[2], Ok(FlightMessage::Metrics(_)))); + assert!(matches!(messages[3], Ok(FlightMessage::RecordBatch(_)))); + assert!(matches!(messages[4], Ok(FlightMessage::Metrics(_)))); + } + + #[tokio::test] + async fn test_ready_only_preserves_oversized_and_empty_batch_caps() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::int32_datatype(), + false, + )])); + let oversized = int_batch(schema.clone(), 0..1024); + let mut items = vec![ + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [1]))), + ScriptedItem::Batch(Ok(int_batch(schema.clone(), [2]))), + ScriptedItem::Batch(Ok(oversized)), + ]; + items.extend( + (0..16).map(|_| ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone())))), + ); + let poll_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let messages = flight_messages(Box::pin(ScriptedBatchStream { + schema, + items: items.into(), + poll_count: poll_count.clone(), + })) + .await; + let batches = messages + .iter() + .filter_map(|message| match message.as_ref().unwrap() { + FlightMessage::RecordBatch(batch) => Some(batch.num_rows()), + _ => None, + }) + .collect::>(); + assert_eq!(batches, vec![1, 1, 1024, 0]); + assert_eq!(poll_count.load(std::sync::atomic::Ordering::Relaxed), 20); + } + + #[tokio::test] + async fn test_ready_only_preserves_byte_oversized_batch() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "a", + ConcreteDataType::string_datatype(), + false, + )])); + let oversized = RecordBatch::new( + schema.clone(), + vec![Arc::new(datatypes::vectors::StringVector::from_slice( + &vec!["x".repeat(300); 1023], + )) as VectorRef], + ) + .unwrap(); + let messages = flight_messages(Box::pin(ScriptedBatchStream { + schema: schema.clone(), + items: VecDeque::from([ + ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone()))), + ScriptedItem::Batch(Ok(RecordBatch::new_empty(schema.clone()))), + ScriptedItem::Batch(Ok(oversized)), + ]), + poll_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + })) + .await; + let batches = messages + .iter() + .filter_map(|message| match message.as_ref().unwrap() { + FlightMessage::RecordBatch(batch) => Some(batch.num_rows()), + _ => None, + }) + .collect::>(); + assert_eq!(batches, vec![0, 0, 1023]); + } + #[tokio::test] async fn test_ready_only_coalesces_dictionary_batches_with_different_mappings() { let arrow_schema = Arc::new(datatypes::arrow::datatypes::Schema::new(vec![