Merge pull request #10532 from zhang2014/fix/cluster_hang

fix(cluster): using global io spawn flight client

Author: zhang2014

src/query/service/src/api/rpc/exchange/exchange_sink.rs

@@ -79,6 +79,8 @@ impl ExchangeSink {
                     Ok(ProcessorPtr::create(ExchangeWriterSink::create(
                         input,
                         flight_exchange.clone(),
+                        params.query_id.to_string(),
+                        params.fragment_id,
                     )))
                 })
             }
@@ -87,7 +89,11 @@ impl ExchangeSink {
 
                 // exchange writer sink
                 let len = pipeline.output_len();
-                let items = create_writer_items(flight_exchanges);
+                let items = create_writer_items(
+                    flight_exchanges,
+                    params.query_id.clone(),
+                    params.fragment_id,
+                );
                 pipeline.add_pipe(Pipe::create(len, 0, items));
                 Ok(())
             }

src/query/service/src/api/rpc/exchange/exchange_sink_writer.rs

@@ -24,17 +24,27 @@ use common_pipeline_core::processors::processor::ProcessorPtr;
 use common_pipeline_core::processors::Processor;
 use common_pipeline_sinks::AsyncSink;
 use common_pipeline_sinks::AsyncSinker;
+use tracing::info;
 
 use crate::api::rpc::exchange::serde::exchange_serializer::ExchangeSerializeMeta;
 use crate::api::rpc::flight_client::FlightExchangeRef;
 
 pub struct ExchangeWriterSink {
+    query_id: String,
+    fragment: usize,
     exchange: FlightExchangeRef,
 }
 
 impl ExchangeWriterSink {
-    pub fn create(input: Arc<InputPort>, flight_exchange: FlightExchangeRef) -> Box<dyn Processor> {
+    pub fn create(
+        input: Arc<InputPort>,
+        flight_exchange: FlightExchangeRef,
+        query_id: String,
+        fragment: usize,
+    ) -> Box<dyn Processor> {
         AsyncSinker::create(input, ExchangeWriterSink {
+            query_id,
+            fragment,
             exchange: flight_exchange,
         })
     }
@@ -45,12 +55,20 @@ impl AsyncSink for ExchangeWriterSink {
     const NAME: &'static str = "ExchangeWriterSink";
 
     async fn on_start(&mut self) -> Result<()> {
-        self.exchange.close_input().await;
+        let res = self.exchange.close_input().await;
+        info!(
+            "Started query:{:?}, fragment:{:?} exchange write. {}",
+            self.query_id, self.fragment, res
+        );
         Ok(())
     }
 
     async fn on_finish(&mut self) -> Result<()> {
-        self.exchange.close_output().await;
+        let res = self.exchange.close_output().await;
+        info!(
+            "Finished query:{:?}, fragment:{:?} exchange write. {}",
+            self.query_id, self.fragment, res
+        );
         Ok(())
     }
 
@@ -79,15 +97,34 @@ impl AsyncSink for ExchangeWriterSink {
     }
 }
 
-pub fn create_writer_item(exchange: FlightExchangeRef) -> PipeItem {
+pub fn create_writer_item(
+    exchange: FlightExchangeRef,
+    query_id: String,
+    fragment: usize,
+) -> PipeItem {
     let input = InputPort::create();
     PipeItem::create(
-        ProcessorPtr::create(ExchangeWriterSink::create(input.clone(), exchange)),
+        ProcessorPtr::create(ExchangeWriterSink::create(
+            input.clone(),
+            exchange,
+            query_id,
+            fragment,
+        )),
         vec![input],
         vec![],
     )
 }
 
-pub fn create_writer_items(exchanges: Vec<FlightExchangeRef>) -> Vec<PipeItem> {
-    exchanges.into_iter().map(create_writer_item).collect()
+pub fn create_writer_items(
+    exchanges: Vec<FlightExchangeRef>,
+    query_id: String,
+    fragment: usize,
+) -> Vec<PipeItem> {
+    let mut items = Vec::with_capacity(exchanges.len());
+
+    for exchange in exchanges {
+        items.push(create_writer_item(exchange, query_id.clone(), fragment));
+    }
+
+    items
 }

src/query/service/src/api/rpc/exchange/exchange_source.rs

@@ -54,7 +54,13 @@ pub fn via_exchange_source(
 
     let last_output_len = pipeline.output_len();
     let flight_exchanges_len = flight_exchanges.len();
-    exchange_source_reader::via_reader(last_output_len, flight_exchanges, pipeline);
+    exchange_source_reader::via_reader(
+        last_output_len,
+        flight_exchanges,
+        pipeline,
+        params.query_id.clone(),
+        params.fragment_id,
+    );
 
     injector.apply_merge_deserializer(flight_exchanges_len, params, pipeline)
 }

src/query/service/src/api/rpc/exchange/exchange_source_reader.rs

@@ -25,6 +25,7 @@ use common_pipeline_core::processors::processor::ProcessorPtr;
 use common_pipeline_core::processors::Processor;
 use common_pipeline_core::Pipeline;
 use common_pipeline_transforms::processors::transforms::TransformDummy;
+use tracing::info;
 
 use crate::api::rpc::exchange::serde::exchange_deserializer::ExchangeDeserializeMeta;
 use crate::api::rpc::flight_client::FlightExchangeRef;
@@ -33,19 +34,28 @@ use crate::api::DataPacket;
 pub struct ExchangeSourceReader {
     finished: bool,
     initialized: bool,
+    query_id: String,
+    fragment: usize,
     output: Arc<OutputPort>,
     output_data: Option<DataPacket>,
     flight_exchange: FlightExchangeRef,
 }
 
 impl ExchangeSourceReader {
-    pub fn create(output: Arc<OutputPort>, flight_exchange: FlightExchangeRef) -> ProcessorPtr {
+    pub fn create(
+        output: Arc<OutputPort>,
+        flight_exchange: FlightExchangeRef,
+        query_id: String,
+        fragment: usize,
+    ) -> ProcessorPtr {
         ProcessorPtr::create(Box::new(ExchangeSourceReader {
             output,
             flight_exchange,
             finished: false,
             output_data: None,
             initialized: false,
+            fragment,
+            query_id,
         }))
     }
 }
@@ -90,7 +100,16 @@ impl Processor for ExchangeSourceReader {
     async fn async_process(&mut self) -> common_exception::Result<()> {
         if !self.initialized {
             self.initialized = true;
-            self.flight_exchange.close_output().await;
+            info!(
+                "Start query:{:?}, fragment:{:?} exchange read.",
+                self.query_id, self.fragment
+            );
+            let res = self.flight_exchange.close_output().await;
+
+            info!(
+                "Started query:{:?}, fragment:{:?} exchange read. {}",
+                self.query_id, self.fragment, res
+            );
         }
 
         if self.output_data.is_none() {
@@ -102,14 +121,29 @@ impl Processor for ExchangeSourceReader {
 
         if !self.finished {
             self.finished = true;
-            self.flight_exchange.close_input().await;
+            info!(
+                "Finish query:{:?}, fragment:{:?} exchange read.",
+                self.query_id, self.fragment
+            );
+            let res = self.flight_exchange.close_input().await;
+
+            info!(
+                "Finished query:{:?}, fragment:{:?} exchange read. {}",
+                self.query_id, self.fragment, res
+            );
         }
 
         Ok(())
     }
 }
 
-pub fn via_reader(prefix_size: usize, exchanges: Vec<FlightExchangeRef>, pipeline: &mut Pipeline) {
+pub fn via_reader(
+    prefix_size: usize,
+    exchanges: Vec<FlightExchangeRef>,
+    pipeline: &mut Pipeline,
+    query_id: String,
+    fragment: usize,
+) {
     let mut items = Vec::with_capacity(prefix_size + exchanges.len());
 
     for _index in 0..prefix_size {
@@ -126,7 +160,12 @@ pub fn via_reader(prefix_size: usize, exchanges: Vec<FlightExchangeRef>, pipelin
     for flight_exchange in exchanges {
         let output = OutputPort::create();
         items.push(PipeItem::create(
-            ExchangeSourceReader::create(output.clone(), flight_exchange),
+            ExchangeSourceReader::create(
+                output.clone(),
+                flight_exchange,
+                query_id.clone(),
+                fragment,
+            ),
             vec![],
             vec![output],
         ));
@@ -135,10 +174,14 @@ pub fn via_reader(prefix_size: usize, exchanges: Vec<FlightExchangeRef>, pipelin
     pipeline.add_pipe(Pipe::create(prefix_size, items.len(), items));
 }
 
-pub fn create_reader_item(exchange: FlightExchangeRef) -> PipeItem {
+pub fn create_reader_item(
+    exchange: FlightExchangeRef,
+    query_id: String,
+    fragment: usize,
+) -> PipeItem {
     let output = OutputPort::create();
     PipeItem::create(
-        ExchangeSourceReader::create(output.clone(), exchange),
+        ExchangeSourceReader::create(output.clone(), exchange, query_id, fragment),
         vec![],
         vec![output],
     )

src/query/service/src/api/rpc/exchange/exchange_transform.rs

@@ -59,7 +59,11 @@ impl ExchangeTransform {
                     items.push(match destination_id == &params.executor_id {
                         true if max_threads == 1 => create_dummy_item(),
                         true => create_resize_item(1, max_threads),
-                        false => create_writer_item(exchange),
+                        false => create_writer_item(
+                            exchange,
+                            params.query_id.clone(),
+                            params.fragment_id,
+                        ),
                     });
                 }
 
@@ -68,7 +72,11 @@ impl ExchangeTransform {
                 for (destination_id, exchange) in params.destination_ids.iter().zip(exchanges) {
                     if destination_id != &params.executor_id {
                         nodes_source += 1;
-                        items.push(create_reader_item(exchange));
+                        items.push(create_reader_item(
+                            exchange,
+                            params.query_id.clone(),
+                            params.fragment_id,
+                        ));
                     }
                 }
 

src/query/service/src/api/rpc/exchange/statistics_sender.rs

@@ -18,7 +18,7 @@ use std::sync::Arc;
 
 use async_channel::Receiver;
 use async_channel::Sender;
-use common_base::runtime::GlobalQueryRuntime;
+use common_base::base::tokio::task::JoinHandle;
 use common_base::runtime::TrySpawn;
 use common_catalog::table_context::TableContext;
 use common_exception::ErrorCode;
@@ -38,6 +38,7 @@ pub struct StatisticsSender {
     shutdown_flag: Arc<AtomicBool>,
     shutdown_flag_sender: Sender<Option<ErrorCode>>,
     shutdown_flag_receiver: Receiver<Option<ErrorCode>>,
+    join_handle: Option<JoinHandle<()>>,
 }
 
 impl StatisticsSender {
@@ -54,6 +55,7 @@ impl StatisticsSender {
             exchange: Some(exchange),
             query_id: query_id.to_string(),
             shutdown_flag: Arc::new(AtomicBool::new(false)),
+            join_handle: None,
         }
     }
 
@@ -64,8 +66,8 @@ impl StatisticsSender {
         let shutdown_flag = self.shutdown_flag.clone();
         let shutdown_flag_receiver = self.shutdown_flag_receiver.clone();
 
-        let spawner = GlobalQueryRuntime::instance();
-        spawner.runtime().spawn(async move {
+        let spawner = ctx.clone();
+        self.join_handle = Some(spawner.spawn(async move {
             let mut recv = Box::pin(flight_exchange.recv());
             let mut notified = Box::pin(shutdown_flag_receiver.recv());
 
@@ -113,13 +115,14 @@ impl StatisticsSender {
 
             flight_exchange.close_input().await;
             flight_exchange.close_output().await;
-        });
+        }));
     }
 
     pub fn shutdown(&mut self, error: Option<ErrorCode>) {
         self.shutdown_flag.store(true, Ordering::Release);
         let shutdown_flag_sender = self.shutdown_flag_sender.clone();
 
+        let join_handle = self.join_handle.take();
         futures::executor::block_on(async move {
             if let Err(error_code) = shutdown_flag_sender.send(error).await {
                 tracing::warn!(
@@ -129,6 +132,10 @@ impl StatisticsSender {
             }
 
             shutdown_flag_sender.close();
+
+            if let Some(join_handle) = join_handle {
+                let _ = join_handle.await;
+            }
         });
     }