fix(cluster): try fix broken pipe or connect reset

Author: zhang2014

scripts/ci/ci-run-stateless-tests-cluster.sh

@@ -15,5 +15,4 @@ cd "$SCRIPT_PATH/../../tests" || exit
 
 echo "Starting databend-test"
 # 13_0004_q4: https://github.com/datafuselabs/databend/issues/8107
-# 13_0005_q5: https://github.com/datafuselabs/databend/issues/7986
-./databend-test --mode 'cluster' --run-dir 0_stateless --skip '13_0004_q4' --skip '13_0005_q5'
+./databend-test --mode 'cluster' --run-dir 0_stateless --skip '13_0004_q4'

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

@@ -264,6 +264,7 @@ impl Processor for ExchangeTransform {
                 DataPacket::ProgressAndPrecommit { .. } => unreachable!(),
                 DataPacket::FetchProgressAndPrecommit => unreachable!(),
                 DataPacket::FragmentData(v) => self.on_recv_data(v),
+                DataPacket::ClosingClient => Ok(()),
             };
         }
 

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

@@ -130,6 +130,7 @@ impl Processor for ExchangeSourceTransform {
                 DataPacket::FragmentData(v) => self.on_recv_data(v),
                 DataPacket::FetchProgressAndPrecommit => unreachable!(),
                 DataPacket::ProgressAndPrecommit { .. } => unreachable!(),
+                DataPacket::ClosingClient => Ok(()),
             };
         }
 

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

@@ -139,6 +139,7 @@ impl StatisticsReceiver {
     fn recv_data(ctx: &Arc<QueryContext>, recv_data: Result<Option<DataPacket>>) -> Result<bool> {
         match recv_data {
             Ok(None) => Ok(true),
+            Ok(Some(DataPacket::ClosingClient)) => Ok(true),
             Err(transport_error) => Err(transport_error),
             Ok(Some(DataPacket::ErrorCode(error))) => Err(error),
             Ok(Some(DataPacket::FragmentData(_))) => unreachable!(),

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

@@ -90,6 +90,11 @@ impl StatisticsSender {
                         notified = right;
                         recv = Box::pin(flight_exchange.recv());
 
+                        if matches!(command, DataPacket::ClosingClient) {
+                            ctx.get_exchange_manager().shutdown_query(&query_id);
+                            return;
+                        }
+
                         if let Err(_cause) = Self::on_command(&ctx, command, &flight_exchange).await
                         {
                             ctx.get_exchange_manager().shutdown_query(&query_id);
@@ -100,6 +105,10 @@ impl StatisticsSender {
             }
 
             if let Ok(Some(command)) = flight_exchange.recv().await {
+                if matches!(command, DataPacket::ClosingClient) {
+                    return;
+                }
+
                 if let Err(error) = Self::on_command(&ctx, command, &flight_exchange).await {
                     tracing::warn!("Statistics send has error, cause: {:?}.", error);
                 }
@@ -140,6 +149,7 @@ impl StatisticsSender {
                     })
                     .await
             }
+            DataPacket::ClosingClient => unreachable!(),
         }
     }
 

src/query/service/src/api/rpc/flight_client.rs

@@ -135,10 +135,18 @@ impl FlightExchange {
     ) -> FlightExchange {
         let mut streaming = streaming.into_inner();
         let (tx, rx) = async_channel::bounded(1);
+
         common_base::base::tokio::spawn(async move {
             while let Some(message) = streaming.next().await {
-                if let Err(_cause) = tx.send(message).await {
-                    break;
+                match message {
+                    Ok(message) if DataPacket::is_closing_client(&message) => {
+                        break;
+                    }
+                    other => {
+                        if let Err(_c) = tx.send(other).await {
+                            break;
+                        }
+                    }
                 }
             }
         });
@@ -159,8 +167,15 @@ impl FlightExchange {
         let (tx, request_rx) = async_channel::bounded(1);
         common_base::base::tokio::spawn(async move {
             while let Some(message) = streaming.next().await {
-                if let Err(_cause) = tx.send(message).await {
-                    break;
+                match message {
+                    Ok(flight_data) if DataPacket::is_closing_client(&flight_data) => {
+                        break;
+                    }
+                    other => {
+                        if let Err(_cause) = tx.send(other).await {
+                            break;
+                        }
+                    }
                 }
             }
         });
@@ -282,7 +297,9 @@ impl ClientFlightExchange {
         if !self.is_closed_response.fetch_or(true, Ordering::SeqCst)
             && self.state.response_count.fetch_sub(1, Ordering::AcqRel) == 1
         {
-            self.response_tx.close();
+            let _ = self
+                .response_tx
+                .send_blocking(FlightData::from(DataPacket::ClosingClient));
         }
     }
 }
@@ -304,17 +321,8 @@ impl Clone for ClientFlightExchange {
 
 impl Drop for ClientFlightExchange {
     fn drop(&mut self) {
-        if !self.is_closed_request.fetch_or(true, Ordering::SeqCst)
-            && self.state.request_count.fetch_sub(1, Ordering::AcqRel) == 1
-        {
-            self.request_rx.close();
-        }
-
-        if !self.is_closed_response.fetch_or(true, Ordering::SeqCst)
-            && self.state.response_count.fetch_sub(1, Ordering::AcqRel) == 1
-        {
-            self.response_tx.close();
-        }
+        self.close_input();
+        self.close_output();
     }
 }
 
@@ -398,7 +406,9 @@ impl ServerFlightExchange {
         if !self.is_closed_response.fetch_or(true, Ordering::SeqCst)
             && self.state.response_count.fetch_sub(1, Ordering::AcqRel) == 1
         {
-            self.response_tx.close();
+            let _ = self
+                .response_tx
+                .send_blocking(Ok(FlightData::from(DataPacket::ClosingClient)));
         }
     }
 }

src/query/service/src/api/rpc/packets/packet_data.rs

@@ -56,6 +56,13 @@ pub enum DataPacket {
         progress: Vec<ProgressInfo>,
         precommit: Vec<PrecommitBlock>,
     },
+    ClosingClient,
+}
+
+impl DataPacket {
+    pub fn is_closing_client(data: &FlightData) -> bool {
+        data.app_metadata.last() == Some(&0x05)
+    }
 }
 
 impl From<DataPacket> for FlightData {
@@ -103,6 +110,12 @@ impl From<DataPacket> for FlightData {
                     app_metadata: vec![0x04],
                 }
             }
+            DataPacket::ClosingClient => FlightData {
+                data_body: vec![],
+                data_header: vec![],
+                flight_descriptor: None,
+                app_metadata: vec![0x05],
+            },
         }
     }
 }
@@ -155,6 +168,7 @@ impl TryFrom<FlightData> for DataPacket {
                     progress: progress_info,
                 })
             }
+            0x05 => Ok(DataPacket::ClosingClient),
             _ => Err(ErrorCode::BadBytes("Unknown flight data packet type.")),
         }
     }