From 31e5cbb37e1d88b415e4b0b1a574133df8b6a6ff Mon Sep 17 00:00:00 2001 From: WenyXu Date: Tue, 25 Aug 2026 10:31:06 +0000 Subject: [PATCH] fix(flight): bound DoGet response wait (#8943) * fix(flight): defer datanode query initialization Signed-off-by: WenyXu * fix(client): retain Flight stream peer context Signed-off-by: WenyXu * fix(client): improve Flight stream diagnostics Signed-off-by: WenyXu --------- Signed-off-by: WenyXu (cherry picked from commit 28398138ec03d154524a8217672d463f28e7f466) --- src/client/src/region.rs | 116 +++++++++++++++++++++----- src/datanode/src/region_server.rs | 20 +++-- src/servers/src/grpc/flight.rs | 6 +- src/servers/src/grpc/flight/stream.rs | 104 +++++++++++++++++++---- tests-integration/src/grpc/flight.rs | 102 +++++++++++++++++++++- 5 files changed, 299 insertions(+), 49 deletions(-) diff --git a/src/client/src/region.rs b/src/client/src/region.rs index f1730a93c8..2bb5f7dd0a 100644 --- a/src/client/src/region.rs +++ b/src/client/src/region.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::sync::Arc; +use std::time::Duration; use api::region::RegionResponse; use api::v1::ResponseHeader; @@ -24,7 +25,7 @@ use arc_swap::ArcSwapOption; use arrow_flight::Ticket; use async_stream::stream; use async_trait::async_trait; -use common_error::ext::BoxedError; +use common_error::ext::{BoxedError, ErrorExt}; use common_error::status_code::StatusCode; use common_grpc::flight::{FlightDecoder, FlightMessage}; use common_meta::error::{self as meta_error, Result as MetaResult}; @@ -49,6 +50,8 @@ use crate::error::{ use crate::flight::decode_flight_data; use crate::{Client, metrics}; +const FLIGHT_DO_GET_TIMEOUT: Duration = Duration::from_secs(10); + #[derive(Debug)] pub struct RegionRequester { client: Client, @@ -109,20 +112,24 @@ impl RegionRequester { let mut flight_client = self .client .make_flight_client(self.send_compression, self.accept_compression)?; + // Limit Flight DoGet response time without limiting query stream execution. + let addr = flight_client.addr().to_string(); + let mut request = tonic::Request::new(ticket); + request.set_timeout(FLIGHT_DO_GET_TIMEOUT); let response = flight_client .mut_inner() - .do_get(ticket) + .do_get(request) .await .or_else(|e| { let tonic_code = e.code(); let e: error::Error = e.into(); error!( e; "Failed to do Flight get, addr: {}, code: {}", - flight_client.addr(), + addr, tonic_code ); Err(BoxedError::new(e)).with_context(|_| FlightGetSnafu { - addr: flight_client.addr().to_string(), + addr: addr.clone(), tonic_code, }) })?; @@ -133,7 +140,7 @@ impl RegionRequester { let flight_message_stream = flight_data_stream .filter_map(move |flight_data| decode_flight_data(&mut decoder, flight_data)); - recordbatches_from_flight_message_stream(flight_message_stream).await + recordbatches_from_flight_message_stream(addr, flight_message_stream).await } async fn handle_inner(&self, request: RegionRequest) -> Result { @@ -195,6 +202,7 @@ impl RegionRequester { } async fn recordbatches_from_flight_message_stream( + addr: String, mut flight_message_stream: S, ) -> Result where @@ -204,13 +212,17 @@ where return IllegalFlightMessagesSnafu { reason: "Expect the response not to be empty", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; - let FlightMessage::Schema(schema) = first_flight_message? else { + let FlightMessage::Schema(schema) = + first_flight_message.map_err(|e| flight_stream_error(&addr, e))? + else { return IllegalFlightMessagesSnafu { reason: "Expect schema to be the first flight message", } - .fail(); + .fail() + .map_err(|error| flight_stream_error(&addr, error)); }; let metrics = Arc::new(ArcSwapOption::from(None)); @@ -221,6 +233,7 @@ where let schema = Arc::new(datatypes::schema::Schema::try_from(schema).context(error::ConvertSchemaSnafu)?); let schema_cloned = schema.clone(); + let stream_addr = addr.clone(); let stream = Box::pin(stream!({ let _span = tracing_context.attach(common_telemetry::tracing::info_span!( "poll_flight_data_stream" @@ -240,7 +253,8 @@ where let flight_message = match flight_message_item { Some(Ok(message)) => message, Some(Err(e)) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } None => break, @@ -273,7 +287,8 @@ where break; } Err(e) => { - yield Err(BoxedError::new(e)).context(ExternalSnafu); + yield Err(BoxedError::new(flight_stream_error(&stream_addr, e))) + .context(ExternalSnafu); break; } } @@ -312,6 +327,23 @@ where Ok(Box::pin(record_batch_stream)) } +fn flight_stream_error(addr: &str, error: error::Error) -> error::Error { + let tonic_code = error.tonic_code().unwrap_or(tonic::Code::Unknown); + if error.status_code().should_log_error() { + error!( + error; "Failed to receive Flight data, addr: {}, code: {}", + addr, + tonic_code + ); + } + + error::Error::FlightGet { + addr: addr.to_string(), + tonic_code, + source: BoxedError::new(error), + } +} + pub fn build_remote_dyn_filter_update_request( query_id: impl Into, update: RemoteDynFilterUpdate, @@ -389,6 +421,44 @@ mod test { use super::*; use crate::Error::{self, IllegalDatabaseResponse, Server}; + #[test] + fn test_flight_stream_error_preserves_peer_address() { + let error = flight_stream_error( + "127.0.0.1:4001", + tonic::Status::unavailable("datanode unavailable").into(), + ); + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unavailable, + .. + } if addr == "127.0.0.1:4001" + )); + } + + #[tokio::test] + async fn test_empty_flight_stream_preserves_peer_address() { + let Err(error) = recordbatches_from_flight_message_stream( + "127.0.0.1:4001".to_string(), + stream::empty::>(), + ) + .await + else { + panic!("expected empty Flight stream to fail"); + }; + + assert!(matches!( + error, + error::Error::FlightGet { + addr, + tonic_code: tonic::Code::Unknown, + .. + } if addr == "127.0.0.1:4001" + )); + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![ColumnSchema::new( "v", @@ -509,11 +579,14 @@ mod test { ) .unwrap(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Ok(FlightMessage::RecordBatch(batch.into_df_record_batch())), + ]), + ) .await .unwrap(); @@ -528,11 +601,14 @@ mod test { #[tokio::test] async fn test_record_batch_stream_exposes_error_after_pre_batch_metrics() { let schema = test_schema(); - let mut recordbatches = recordbatches_from_flight_message_stream(stream::iter(vec![ - Ok(FlightMessage::Schema(schema.arrow_schema().clone())), - Ok(FlightMessage::Metrics(test_metrics_json())), - Err(Error::from(Status::internal("boom after metrics"))), - ])) + let mut recordbatches = recordbatches_from_flight_message_stream( + "test-peer".to_string(), + stream::iter(vec![ + Ok(FlightMessage::Schema(schema.arrow_schema().clone())), + Ok(FlightMessage::Metrics(test_metrics_json())), + Err(Error::from(Status::internal("boom after metrics"))), + ]), + ) .await .unwrap(); diff --git a/src/datanode/src/region_server.rs b/src/datanode/src/region_server.rs index 57facfd53b..0d9398b088 100644 --- a/src/datanode/src/region_server.rs +++ b/src/datanode/src/region_server.rs @@ -66,7 +66,9 @@ use servers::error::{ self as servers_error, ExecuteGrpcRequestSnafu, Result as ServerResult, SuspendedSnafu, }; use servers::grpc::FlightCompression; -use servers::grpc::flight::{FlightCraft, FlightRecordBatchStream, TonicStream}; +use servers::grpc::flight::{ + FlightCraft, FlightRecordBatchSource, FlightRecordBatchStream, TonicStream, +}; use servers::grpc::region_server::RegionServerHandler; use session::context::{ FLIGHT_METRICS_HEARTBEAT_INTERVAL, QueryContext, QueryContextBuilder, QueryContextRef, @@ -976,13 +978,19 @@ impl FlightCraft for RegionServer { .map(|h| Arc::new(QueryContext::from(h))) .unwrap_or(QueryContext::arc()); - let result = self - .handle_remote_read(request, query_ctx.clone()) - .trace(tracing_context.attach(info_span!("RegionServer::handle_read"))) - .await?; + let region_server = self.clone(); + let initializer_query_ctx = query_ctx.clone(); + let initializer_tracing_context = tracing_context.clone(); + let initializer = async move { + region_server + .handle_remote_read(request, initializer_query_ctx) + .trace(initializer_tracing_context.attach(info_span!("RegionServer::handle_read"))) + .await + .map_err(Into::into) + }; let stream = Box::pin(FlightRecordBatchStream::new( - result, + FlightRecordBatchSource::initializer(initializer), tracing_context, self.flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight.rs b/src/servers/src/grpc/flight.rs index 97a579d952..e6f35e17b2 100644 --- a/src/servers/src/grpc/flight.rs +++ b/src/servers/src/grpc/flight.rs @@ -53,7 +53,7 @@ use tokio_stream::wrappers::ReceiverStream; use tonic::{Request, Response, Status, Streaming}; use crate::error::{InvalidParameterSnafu, Result, ToJsonSnafu}; -pub use crate::grpc::flight::stream::FlightRecordBatchStream; +pub use crate::grpc::flight::stream::{FlightRecordBatchSource, FlightRecordBatchStream}; use crate::grpc::greptime_handler::{ GreptimeRequestHandler, create_query_context, get_request_type, }; @@ -583,7 +583,7 @@ fn to_flight_data_stream( match output.data { OutputData::Stream(stream) => { let stream = FlightRecordBatchStream::new( - stream, + FlightRecordBatchSource::RecordBatches(stream), tracing_context, flight_compression, query_ctx, @@ -592,7 +592,7 @@ fn to_flight_data_stream( } OutputData::RecordBatches(x) => { let stream = FlightRecordBatchStream::new( - x.as_stream(), + FlightRecordBatchSource::RecordBatches(x.as_stream()), tracing_context, flight_compression, query_ctx, diff --git a/src/servers/src/grpc/flight/stream.rs b/src/servers/src/grpc/flight/stream.rs index c222c72e1c..319a797d9c 100644 --- a/src/servers/src/grpc/flight/stream.rs +++ b/src/servers/src/grpc/flight/stream.rs @@ -13,6 +13,7 @@ // limitations under the License. use std::collections::VecDeque; +use std::future::Future; use std::pin::Pin; use std::task::{Context, Poll}; use std::time::{Duration, Instant}; @@ -40,6 +41,30 @@ use crate::error; use crate::grpc::FlightCompression; use crate::grpc::flight::TonicResult; +pub type FlightRecordBatchStreamInitializer = Pin< + Box< + dyn Future> + + Send + + 'static, + >, +>; + +pub enum FlightRecordBatchSource { + RecordBatches(SendableRecordBatchStream), + Initializer(FlightRecordBatchStreamInitializer), +} + +impl FlightRecordBatchSource { + pub fn initializer(initializer: F) -> Self + where + F: Future> + + Send + + 'static, + { + Self::Initializer(Box::pin(initializer)) + } +} + /// Metrics collector for Flight stream with RAII logging pattern struct StreamMetrics { send_schema_duration: Duration, @@ -136,7 +161,7 @@ impl FlightRecordBatchStream { } pub fn new( - recordbatches: SendableRecordBatchStream, + source: FlightRecordBatchSource, tracing_context: TracingContext, compression: FlightCompression, query_ctx: QueryContextRef, @@ -148,17 +173,43 @@ impl FlightRecordBatchStream { .remote_query_id() .zip(query_ctx.extension(SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY)) .is_some_and(|(remote_query_id, capability)| capability == remote_query_id); - let (tx, rx) = mpsc::channel::>(1); - let join_handle = common_runtime::spawn_global(async move { - Self::flight_data_stream( - recordbatches, - tx, - should_send_partial_metrics, - can_send_metrics_before_batch, - ) - .trace(tracing_context.attach(info_span!("flight_data_stream"))) - .await - }); + let (mut tx, rx) = mpsc::channel::>(1); + let source_type = match &source { + FlightRecordBatchSource::RecordBatches(_) => "record_batches", + FlightRecordBatchSource::Initializer(_) => "initializer", + }; + let initializer_tracing_context = tracing_context.clone(); + let join_handle = common_runtime::spawn_global( + async move { + let recordbatches = async move { + match source { + FlightRecordBatchSource::RecordBatches(recordbatches) => Ok(recordbatches), + FlightRecordBatchSource::Initializer(initializer) => initializer.await, + } + } + .trace( + initializer_tracing_context + .attach(info_span!("flight_data_stream_init", source_type)), + ) + .await; + + match recordbatches { + Ok(recordbatches) => { + Self::flight_data_stream( + recordbatches, + tx, + should_send_partial_metrics, + can_send_metrics_before_batch, + ) + .await; + } + Err(status) => { + let _ = tx.send(Err(status)).await; + } + } + } + .trace(tracing_context.attach(info_span!("flight_data_stream"))), + ); let encoder = if compression.arrow_compression() { FlightEncoder::default() } else { @@ -434,7 +485,7 @@ mod test { .unwrap() .as_stream(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), QueryContext::arc(), @@ -468,6 +519,23 @@ mod test { } } + #[tokio::test] + async fn test_flight_record_batch_stream_forwards_initializer_error() { + let mut stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + Err(tonic::Status::unavailable( + "remote read initialization failed", + )) + }), + TracingContext::default(), + FlightCompression::default(), + QueryContext::arc(), + ); + + let error = stream.next().await.unwrap().unwrap_err(); + assert_eq!(tonic::Code::Unavailable, error.code()); + assert!(stream.next().await.is_none()); + } #[tokio::test] async fn test_flight_record_batch_stream_emits_metrics_while_pending() { let schema = Arc::new(Schema::new(vec![ColumnSchema::new( @@ -486,7 +554,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -541,7 +609,7 @@ mod test { let query_ctx = query_context_with_live_metrics_and_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -595,7 +663,7 @@ mod test { let query_ctx = query_context_with_matching_capability(); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -638,7 +706,7 @@ mod test { let query_ctx = Arc::new(query_ctx); query_ctx.set_explain_verbose(true); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, @@ -677,7 +745,7 @@ mod test { }); let query_ctx = query_context_with_matching_capability(); let mut stream = FlightRecordBatchStream::new( - recordbatches, + FlightRecordBatchSource::RecordBatches(recordbatches), TracingContext::default(), FlightCompression::default(), query_ctx, diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 068cb98142..3dee75e066 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/tests-integration/src/grpc/flight.rs @@ -17,34 +17,132 @@ mod test { use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; + use std::time::{Duration, Instant}; use api::v1::auth_header::AuthScheme; use api::v1::query_request::Query; use api::v1::{Basic, ColumnDataType, ColumnDef, CreateTableExpr, QueryRequest, SemanticType}; - use arrow_flight::FlightDescriptor; + use arrow_flight::flight_service_server::FlightServiceServer; + use arrow_flight::{FlightData, FlightDescriptor, Ticket}; use auth::user_provider_from_option; use client::{Client, Database}; use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME}; + use common_grpc::channel_manager::{ChannelConfig, ChannelManager}; use common_grpc::flight::do_put::DoPutMetadata; use common_grpc::flight::{FlightEncoder, FlightMessage}; use common_query::OutputData; - use common_recordbatch::RecordBatch; use common_recordbatch::adapter::RegionWatermarkEntry; + use common_recordbatch::{RecordBatch, RecordBatches, SendableRecordBatchStream}; + use common_telemetry::tracing_context::TracingContext; use datatypes::prelude::{ConcreteDataType, ScalarVector, VectorRef}; use datatypes::schema::{ColumnSchema, Schema}; use datatypes::vectors::{Int32Vector, StringVector, TimestampMillisecondVector}; use futures_util::StreamExt; + use hyper_util::rt::TokioIo; use itertools::Itertools; use servers::grpc::builder::GrpcServerBuilder; + use servers::grpc::flight::{ + FlightCraft, FlightCraftWrapper, FlightRecordBatchSource, FlightRecordBatchStream, + TonicStream, + }; use servers::grpc::greptime_handler::GreptimeRequestHandler; use servers::grpc::{FlightCompression, GrpcServerConfig}; use servers::server::Server; + use tonic::Response; + use tonic::transport::Server as TonicServer; + use tower::service_fn; use crate::cluster::GreptimeDbClusterBuilder; use crate::grpc::query_and_expect; use crate::test_util::{StorageType, setup_grpc_server}; use crate::tests::test_util::MockInstance; + struct SlowFlightCraft; + + fn slow_recordbatch_stream() -> SendableRecordBatchStream { + let schema = Arc::new(Schema::new(vec![ColumnSchema::new( + "value", + ConcreteDataType::int32_datatype(), + false, + )])); + let recordbatch = RecordBatch::new( + schema.clone(), + vec![Arc::new(Int32Vector::from_vec(vec![1])) as VectorRef], + ) + .unwrap(); + + RecordBatches::try_new(schema, vec![recordbatch]) + .unwrap() + .as_stream() + } + + #[async_trait::async_trait] + impl FlightCraft for SlowFlightCraft { + async fn do_get( + &self, + _: tonic::Request, + ) -> std::result::Result>, tonic::Status> { + let stream = FlightRecordBatchStream::new( + FlightRecordBatchSource::initializer(async { + tokio::time::sleep(Duration::from_secs(2)).await; + Ok(slow_recordbatch_stream()) + }), + TracingContext::default(), + FlightCompression::default(), + session::context::QueryContext::arc(), + ); + + Ok(Response::new(Box::pin(stream))) + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_do_get_timeout_does_not_cancel_slow_flight_stream() { + let (client_io, server_io) = tokio::io::duplex(1024); + tokio::spawn(async move { + TonicServer::builder() + .add_service(FlightServiceServer::new(FlightCraftWrapper( + SlowFlightCraft, + ))) + .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>( + server_io, + )])) + .await + .unwrap(); + }); + + let channel_manager = ChannelManager::with_config(ChannelConfig::new().timeout(None), None); + let mut client_io = Some(client_io); + channel_manager + .reset_with_connector( + "slow-flight", + service_fn(move |_| { + let client_io = client_io.take(); + + async move { + client_io + .map(TokioIo::new) + .ok_or_else(|| std::io::Error::other("Client already taken")) + } + }), + ) + .unwrap(); + let client = Client::with_manager_and_urls(channel_manager, ["slow-flight"]); + let mut flight_client = client.make_flight_client(false, false).unwrap(); + + let start = Instant::now(); + let mut request = tonic::Request::new(Ticket::default()); + request.set_timeout(Duration::from_secs(1)); + let response = flight_client.mut_inner().do_get(request).await.unwrap(); + assert!(start.elapsed() < Duration::from_secs(1)); + let mut stream = response.into_inner(); + + let start = Instant::now(); + assert!(stream.message().await.unwrap().is_some()); + assert!(start.elapsed() >= Duration::from_secs(1)); + + assert!(stream.message().await.unwrap().is_some()); + } #[tokio::test(flavor = "multi_thread")] async fn test_standalone_flight_do_put() { common_telemetry::init_default_ut_logging();