fix(flight): bound DoGet response wait (#8943)

* fix(flight): defer datanode query initialization

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(client): retain Flight stream peer context

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix(client): improve Flight stream diagnostics

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
(cherry picked from commit 28398138ec)

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<RegionResponse> {
@@ -195,6 +202,7 @@ impl RegionRequester {
 }
 
 async fn recordbatches_from_flight_message_stream<S>(
+    addr: String,
     mut flight_message_stream: S,
 ) -> Result<SendableRecordBatchStream>
 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<String>,
     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::<Result<FlightMessage>>(),
+        )
+        .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<Schema> {
         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();
 

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,

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,

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<Output = std::result::Result<SendableRecordBatchStream, tonic::Status>>
+            + Send
+            + 'static,
+    >,
+>;
+
+pub enum FlightRecordBatchSource {
+    RecordBatches(SendableRecordBatchStream),
+    Initializer(FlightRecordBatchStreamInitializer),
+}
+
+impl FlightRecordBatchSource {
+    pub fn initializer<F>(initializer: F) -> Self
+    where
+        F: Future<Output = std::result::Result<SendableRecordBatchStream, tonic::Status>>
+            + 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::<TonicResult<FlightMessage>>(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::<TonicResult<FlightMessage>>(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,

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<Ticket>,
+        ) -> std::result::Result<Response<TonicStream<FlightData>>, 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();