diff --git a/src/client/src/database.rs b/src/client/src/database.rs index e97fff6db00..de3c722a86a 100644 --- a/src/client/src/database.rs +++ b/src/client/src/database.rs @@ -16,7 +16,7 @@ use std::collections::HashMap; use std::pin::Pin; use std::str::FromStr; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{Arc, RwLock}; +use std::sync::{Arc, Mutex, RwLock}; use std::task::{Context, Poll}; use std::time::Duration; @@ -55,6 +55,7 @@ use futures::future; use futures_util::{Stream, StreamExt, TryStreamExt}; use prost::Message; use snafu::{IntoError, ResultExt}; +use tokio::sync::Notify; use tonic::metadata::{AsciiMetadataKey, AsciiMetadataValue, MetadataMap, MetadataValue}; use tonic::transport::Channel; @@ -70,20 +71,47 @@ type FlightDataStream = Pin + Send>>; type DoPutResponseStream = Pin>>>; const HINTS_METADATA_KEY: &str = "x-greptime-hints"; +/// Maximum time to wait for the optional trailing metrics message after +/// affected rows have already been delivered. +const FLIGHT_TRAILING_METRICS_TIMEOUT: Duration = Duration::from_secs(5); /// Terminal metrics associated with a query output. /// /// For streaming outputs, metrics are only final after the stream is fully -/// drained and [`Self::is_ready`] returns `true`. +/// drained and [`Self::is_ready`] returns `true`. Affected-row outputs may +/// briefly await a compatibility trailing metrics message. #[derive(Debug, Clone, Default)] pub struct OutputMetrics { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct OutputMetricsInner { metrics: RwLock>, + completion_error: RwLock>, ready: AtomicBool, + ready_notify: Notify, + compatibility_task: Mutex>, +} + +impl Default for OutputMetricsInner { + fn default() -> Self { + Self { + metrics: RwLock::new(None), + completion_error: RwLock::new(None), + ready: AtomicBool::new(false), + ready_notify: Notify::new(), + compatibility_task: Mutex::new(None), + } + } +} + +impl Drop for OutputMetricsInner { + fn drop(&mut self) { + if let Some(handle) = self.compatibility_task.get_mut().unwrap().take() { + handle.abort(); + } + } } impl OutputMetrics { @@ -98,10 +126,45 @@ impl OutputMetrics { /// Marks the terminal metrics as final for this output. pub fn mark_ready(&self) { - let _ = self + if self .inner .ready - .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire); + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_ok() + { + self.inner.ready_notify.notify_waiters(); + } + } + + /// Waits until terminal metrics are final. + pub async fn wait_ready(&self) { + loop { + let notified = self.inner.ready_notify.notified(); + if self.is_ready() { + return; + } + notified.await; + } + } + + /// Returns an error encountered while completing the output, if any. + pub fn completion_error(&self) -> Option { + self.inner.completion_error.read().unwrap().clone() + } + + fn set_completion_error(&self, error: impl Into) { + *self.inner.completion_error.write().unwrap() = Some(error.into()); + } + + fn set_compatibility_task(&self, handle: tokio::task::AbortHandle) { + let mut task = self.inner.compatibility_task.lock().unwrap(); + if !self.is_ready() { + *task = Some(handle); + } + } + + fn take_compatibility_task(&self) -> Option { + self.inner.compatibility_task.lock().unwrap().take() } /// Returns whether terminal metrics are final. @@ -148,7 +211,8 @@ impl OutputMetrics { /// /// The contained [`OutputMetrics`] lets callers read stream terminal metrics /// after consuming `output`. For non-stream outputs, metrics are ready -/// immediately. +/// immediately. Flight affected-row outputs without inline metrics require +/// [`OutputMetrics::wait_ready`] before compatibility trailing metrics can be read. #[derive(Debug)] pub struct OutputWithMetrics { pub output: Output, @@ -194,6 +258,55 @@ fn parse_terminal_metrics(metrics_json: &str) -> Result { }) } +fn spawn_affected_rows_trailing_metrics_task( + terminal_metrics: &OutputMetrics, + mut reader: FlightMessageReader, +) where + S: Stream> + Send + Unpin + 'static, +{ + let metrics_ref = Arc::downgrade(&terminal_metrics.inner); + let remote_addr = reader.remote_addr().to_string(); + let task = common_runtime::spawn_global(async move { + let result = + tokio::time::timeout(FLIGHT_TRAILING_METRICS_TIMEOUT, reader.read_next()).await; + let Some(inner) = metrics_ref.upgrade() else { + return; + }; + let metrics = OutputMetrics { inner }; + match result { + Ok(Ok(Some(FlightMessage::Metrics(s)))) => match parse_terminal_metrics(&s) { + Ok(metrics_json) => metrics.update(Some(metrics_json)), + Err(error) => { + metrics.set_completion_error(error.to_string()); + warn!( + "Failed to decode trailing Flight metrics from {}: {}", + remote_addr, error + ); + } + }, + Ok(Ok(None)) => {} + Ok(Ok(Some(other))) => { + let error = format!("Unexpected trailing Flight message: {other:?}"); + metrics.set_completion_error(error.clone()); + warn!("{} from {}", error, remote_addr); + } + Ok(Err(error)) => { + let error = flight_stream_error(&remote_addr, error); + metrics.set_completion_error(error.to_string()); + warn!("{}", error); + } + Err(_) => { + let error = "Timed out waiting for trailing Flight metrics"; + metrics.set_completion_error(error); + warn!("{} from {}", error, remote_addr); + } + } + metrics.mark_ready(); + metrics.take_compatibility_task(); + }); + terminal_metrics.set_compatibility_task(task.abort_handle()); +} + struct StreamWithMetrics { stream: common_recordbatch::SendableRecordBatchStream, metrics: OutputMetrics, @@ -233,6 +346,9 @@ impl Stream for StreamWithMetrics { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let polled = Pin::new(&mut self.stream).poll_next(cx); + if let Poll::Ready(Some(Err(error))) = &polled { + self.metrics.set_completion_error(error.to_string()); + } if let Poll::Ready(None) = &polled { self.sync_terminal_metrics(); self.metrics.mark_ready(); @@ -280,32 +396,12 @@ where FlightMessage::AffectedRows { rows, metrics } => { let terminal_metrics = OutputMetrics::new(); if let Some(metrics) = metrics { + // Inline metrics are authoritative. They complete the output + // without touching the rest of the Flight stream. terminal_metrics.update(Some(parse_terminal_metrics(&metrics)?)); - } - let next_message = reader - .read_next() - .await - .map_err(|error| flight_stream_error(reader.remote_addr(), error))?; - match next_message { - None => terminal_metrics.mark_ready(), - Some(FlightMessage::Metrics(s)) if terminal_metrics.get().is_none() => { - terminal_metrics.update(Some(parse_terminal_metrics(&s)?)); - terminal_metrics.mark_ready(); - } - Some(FlightMessage::Metrics(_)) => { - return IllegalFlightMessagesSnafu { - reason: "'AffectedRows' Flight metadata already carries Metrics and cannot be followed by another Metrics message", - } - .fail(); - } - Some(other) => { - return IllegalFlightMessagesSnafu { - reason: format!( - "'AffectedRows' Flight message can only be followed by a Metrics message, got {other:?}" - ), - } - .fail(); - } + terminal_metrics.mark_ready(); + } else { + spawn_affected_rows_trailing_metrics_task(&terminal_metrics, reader); } Ok(OutputWithMetrics { output: Output::new_with_affected_rows(rows), @@ -749,7 +845,9 @@ impl Database { /// Executes a SQL query and returns the output with terminal metrics. /// /// For stream outputs, callers must consume the stream before reading final - /// terminal metrics from [`OutputWithMetrics::metrics`]. + /// terminal metrics from [`OutputWithMetrics::metrics`]. For affected-row + /// outputs without inline metrics, call [`OutputMetrics::wait_ready`] when + /// compatibility trailing metrics are required. pub async fn sql_with_terminal_metrics( &self, sql: S, @@ -986,6 +1084,7 @@ mod tests { use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::Int32Vector; use futures_util::StreamExt; + use tokio::sync::oneshot; use tonic::codegen::http::{HeaderMap, HeaderValue}; use tonic::metadata::MetadataMap; use tonic::{Code, Status}; @@ -1296,11 +1395,6 @@ mod tests { ); } - #[test] - fn test_parse_terminal_metrics_rejects_invalid_json() { - assert!(parse_terminal_metrics("{not-json}").is_err()); - } - #[tokio::test] async fn test_affected_rows_inline_metrics_are_parsed() { let output = output_from_flight_message_stream( @@ -1322,27 +1416,191 @@ mod tests { } #[tokio::test] - async fn test_affected_rows_inline_metrics_rejects_trailing_metrics() { + async fn test_affected_rows_inline_metrics_do_not_poll_trailer() { let metrics_json = terminal_metrics_json(); - let err = output_from_flight_message_stream( + let output = output_from_flight_message_stream( "test-peer".to_string(), futures_util::stream::iter(vec![ Ok(FlightMessage::AffectedRows { rows: 3, - metrics: Some(metrics_json.clone()), + metrics: Some(metrics_json), }), - Ok(FlightMessage::Metrics(metrics_json)), + Ok(FlightMessage::Metrics(terminal_metrics_json_with_seq(99))), ] as Vec>), ) .await - .unwrap_err(); + .unwrap(); - assert!( - err.to_string().contains("already carries Metrics"), - "unexpected error: {err:?}" + assert!(output.metrics.is_ready()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) ); } + #[tokio::test] + async fn test_affected_rows_without_inline_metrics_becomes_ready_after_trailer() { + // Hold the trailer until the not-ready state has been observed. + let (trailer_tx, trailer_rx) = oneshot::channel(); + let trailer = futures_util::stream::once(trailer_rx).map(|ready| { + ready.unwrap(); + Ok(FlightMessage::Metrics(terminal_metrics_json())) + }); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })]) + .chain(trailer), + ) + .await + .unwrap(); + + assert!(!output.metrics.is_ready()); + trailer_tx.send(()).unwrap(); + tokio::time::timeout(Duration::from_secs(1), output.metrics.wait_ready()) + .await + .expect("terminal metrics must become ready once the trailer arrives"); + assert!(output.metrics.completion_error().is_none()); + assert_eq!( + output.metrics.region_watermark_map(), + Some(std::collections::HashMap::from([(7, 42)])) + ); + } + + #[tokio::test] + async fn test_affected_rows_malformed_trailer_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Ok(FlightMessage::Metrics("{not-json}".to_string())), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("Invalid terminal metrics message")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_transport_trailer_error_sets_completion_error_and_ready() { + let output = output_from_flight_message_stream( + "test-peer".to_string(), + futures_util::stream::iter(vec![ + Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + }), + Err(Status::unavailable("trailer read failed").into()), + ] as Vec>), + ) + .await + .unwrap(); + + output.metrics.wait_ready().await; + let error = output.metrics.completion_error().unwrap(); + assert!(error.contains("trailer read failed")); + assert!(output.metrics.is_ready()); + } + + #[tokio::test] + async fn test_affected_rows_compatibility_reader_is_cancelled_after_second_poll_begins() { + struct DropProbe { + first: Option>, + polled: Option>, + dropped: Option>, + } + + impl Stream for DropProbe { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if let Some(message) = self.first.take() { + return Poll::Ready(Some(message)); + } + if let Some(polled) = self.polled.take() { + let _ = polled.send(()); + } + Poll::Pending + } + } + + impl Drop for DropProbe { + fn drop(&mut self) { + if let Some(dropped) = self.dropped.take() { + let _ = dropped.send(()); + } + } + } + + let (polled_tx, polled_rx) = oneshot::channel(); + let (dropped_tx, dropped_rx) = oneshot::channel(); + let output = output_from_flight_message_stream( + "test-peer".to_string(), + DropProbe { + first: Some(Ok(FlightMessage::AffectedRows { + rows: 3, + metrics: None, + })), + polled: Some(polled_tx), + dropped: Some(dropped_tx), + }, + ) + .await + .unwrap(); + polled_rx.await.unwrap(); + drop(output); + // Require cancellation before the five-second compatibility timeout. + tokio::time::timeout(Duration::from_secs(1), dropped_rx) + .await + .expect("the compatibility reader must be dropped once the output is dropped") + .unwrap(); + } + + #[tokio::test] + async fn test_schema_record_batch_yields_before_pending_message_stream() { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "v", + ConcreteDataType::int32_datatype(), + false, + )])); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + let messages = futures_util::stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ] as Vec>) + .chain(futures_util::stream::pending()); + let output = output_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + let OutputData::Stream(mut stream) = output.output.data else { + panic!("expected stream output"); + }; + + let batch = tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + assert!(!output.metrics.is_ready()); + } + #[tokio::test] async fn test_invalid_terminal_metrics_after_record_batch_yields_batch_then_error() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( diff --git a/src/client/src/flight.rs b/src/client/src/flight.rs index 1753fdd5246..3a21e6daa4e 100644 --- a/src/client/src/flight.rs +++ b/src/client/src/flight.rs @@ -12,40 +12,18 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::pin::Pin; - use arrow_flight::FlightData; use common_grpc::flight::{FlightDecoder, FlightMessage}; -use futures_util::stream::Peekable; use futures_util::{Stream, StreamExt}; use snafu::{OptionExt, ResultExt}; use crate::Result; use crate::error::{ConvertFlightDataSnafu, Error, IllegalFlightMessagesSnafu}; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) enum FlightMessageKind { - Schema, - RecordBatch, - AffectedRows, - Metrics, -} - -impl From<&FlightMessage> for FlightMessageKind { - fn from(message: &FlightMessage) -> Self { - match message { - FlightMessage::Schema(_) => Self::Schema, - FlightMessage::RecordBatch(_) => Self::RecordBatch, - FlightMessage::AffectedRows { .. } => Self::AffectedRows, - FlightMessage::Metrics(_) => Self::Metrics, - } - } -} - pub(crate) struct FlightMessageReader { /// Remote Flight peer associated with this response stream. remote_addr: String, - messages: Peekable, + messages: S, } impl FlightMessageReader @@ -55,7 +33,7 @@ where pub(crate) fn new(remote_addr: impl Into, messages: S) -> Self { Self { remote_addr: remote_addr.into(), - messages: messages.peekable(), + messages, } } @@ -72,21 +50,6 @@ where pub(crate) async fn read_next(&mut self) -> Result> { self.messages.next().await.transpose() } - - pub(crate) async fn peek_next_message_kind(&mut self) -> Result> { - match Pin::new(&mut self.messages).peek().await { - Some(Ok(message)) => Ok(Some(message.into())), - None => Ok(None), - Some(Err(_)) => match self.read_next().await { - // `peek` only borrows the error; consume it to preserve the source error. - Err(error) => Err(error), - Ok(_) => IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek".to_string(), - } - .fail(), - }, - } - } } pub(crate) fn decode_flight_data( diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f0c683ec155..fb2e4d10ca1 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -48,7 +48,7 @@ use crate::error::{ self, FlightGetSnafu, IllegalDatabaseResponseSnafu, IllegalFlightMessagesSnafu, MissingFieldSnafu, Result, ServerSnafu, }; -use crate::flight::{FlightMessageKind, FlightMessageReader, decode_flight_data}; +use crate::flight::{FlightMessageReader, decode_flight_data}; use crate::{Client, metrics}; const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); @@ -247,9 +247,7 @@ where "poll_flight_data_stream" )); - let mut stream_ended = false; - - while !stream_ended { + loop { let flight_message = match reader.read_next().await { Ok(Some(message)) => message, Ok(None) => break, @@ -262,59 +260,27 @@ where match flight_message { FlightMessage::RecordBatch(record_batch) => { - let result_to_yield = - RecordBatch::from_df_record_batch(schema_cloned.clone(), record_batch); - - // Metrics follow a batch so MergeScan can observe them before yielding it. - match reader.peek_next_message_kind().await { - Ok(Some(FlightMessageKind::Metrics)) => { - let metrics_message = match reader.read_next().await { - Ok(Some(FlightMessage::Metrics(metrics))) => metrics, - Ok(Some(_) | None) => { - yield IllegalFlightMessagesSnafu { - reason: "Flight stream changed after peek", - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Err(error) => { - yield Err(BoxedError::new(flight_stream_error( - &stream_addr, - error, - ))) - .context(ExternalSnafu); - break; - } - }; - let metrics = serde_json::from_str(&metrics_message).ok().map(Arc::new); - metrics_ref.swap(metrics); - } - Ok(Some(FlightMessageKind::RecordBatch)) => {} - Ok(Some(FlightMessageKind::Schema | FlightMessageKind::AffectedRows)) => { - yield IllegalFlightMessagesSnafu { - reason: "A RecordBatch message can only be succeeded by a Metrics message or another RecordBatch message" - } - .fail() - .map_err(BoxedError::new) - .context(ExternalSnafu); - break; - } - Ok(None) => stream_ended = true, - Err(error) => { - yield Err(BoxedError::new(flight_stream_error(&stream_addr, error))) - .context(ExternalSnafu); - break; - } - } - - yield Ok(result_to_yield); + // Deliver each batch immediately. In particular, do not + // wait for a possible following Metrics message; it is + // consumed on the next poll of this stream. + yield Ok(RecordBatch::from_df_record_batch( + schema_cloned.clone(), + record_batch, + )); } FlightMessage::Metrics(s) => { // Metrics may arrive before the next RecordBatch. - let m = serde_json::from_str(&s).ok().map(Arc::new); - metrics_ref.swap(m); + match serde_json::from_str(&s) { + Ok(metrics) => { + metrics_ref.swap(Some(Arc::new(metrics))); + } + Err(error) => { + common_telemetry::warn!( + "Failed to decode region Flight metrics: {}", + error + ); + } + } continue; } _ => { @@ -719,7 +685,34 @@ mod test { } #[tokio::test] - async fn test_record_batch_stream_updates_following_metrics_before_yielding_batch() { + async fn test_record_batch_is_yielded_without_waiting_for_next_message() { + let schema = test_schema(); + let batch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_slice([1])) as VectorRef], + ) + .unwrap(); + + let messages = stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]) + .chain(stream::pending::>()); + let mut recordbatches = + recordbatches_from_flight_message_stream("test-peer".to_string(), messages) + .await + .unwrap(); + + let batch = tokio::time::timeout(Duration::from_secs(1), recordbatches.next()) + .await + .expect("the first batch must not wait for lookahead") + .unwrap() + .unwrap(); + assert_eq!(batch.num_rows(), 1); + } + + #[tokio::test] + async fn test_malformed_region_metrics_are_non_fatal() { let schema = test_schema(); let batch = RecordBatch::new( schema.clone(), @@ -730,23 +723,19 @@ mod test { "test-peer".to_string(), stream::iter(vec![ Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics("{not-json}".to_string())), Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - Ok(FlightMessage::Metrics(test_metrics_json())), ]), ) .await .unwrap(); - let batch = recordbatches.next().await.unwrap().unwrap(); - assert_eq!(batch.num_rows(), 1); - - let metrics = recordbatches.metrics().unwrap(); - assert_eq!(metrics.elapsed_compute, 7); + assert_eq!(recordbatches.next().await.unwrap().unwrap().num_rows(), 1); assert!(recordbatches.next().await.is_none()); } #[tokio::test] - async fn test_record_batch_stream_preserves_peeked_record_batch() { + async fn test_record_batch_stream_preserves_following_record_batch() { let schema = test_schema(); let first_batch = RecordBatch::new( schema.clone(), diff --git a/src/flow/src/batching_mode/task.rs b/src/flow/src/batching_mode/task.rs index 40020f639c1..c2c3082995a 100644 --- a/src/flow/src/batching_mode/task.rs +++ b/src/flow/src/batching_mode/task.rs @@ -1400,6 +1400,12 @@ impl BatchingTask { match res { Ok(res) => { let (affected_rows, _) = res.output.extract_rows_and_cost(); + if matches!(&res.output.data, common_query::OutputData::AffectedRows(_)) { + res.metrics.wait_ready().await; + } + if let Some(error) = res.metrics.completion_error() { + warn!("Flow {flow_id} completed with terminal metrics error: {error}"); + } debug!( "Flow {flow_id} executed, affected_rows: {affected_rows:?}, elapsed: {:?}, watermark: {:?}", elapsed, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index ef3df79f558..e74d4d5bf70 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -990,7 +990,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); - assert!(result.metrics.is_ready()); + result.metrics.wait_ready().await; assert!(result.region_watermark_map().is_none()); let err = client @@ -1016,6 +1016,7 @@ mod test { panic!("expected affected rows output"); }; assert_eq!(affected_rows, 9); + result.metrics.wait_ready().await; assert_eq!( result.region_watermark_map(), Some(std::collections::HashMap::from([previous_watermark]))