refactor(query): support do_exchange reconnection

src/query/service/src/servers/flight/flight_client.rs

@@ -44,10 +44,11 @@ use crate::servers::flight::request_builder::RequestBuilder;
 use crate::servers::flight::v1::packets::DataPacket;
 
 /// Parameters for a do_exchange RPC call, serialized as JSON in metadata.
-#[derive(Serialize, Deserialize)]
+#[derive(Clone, Serialize, Deserialize)]
 pub struct DoExchangeParams {
     pub query_id: String,
     pub exchange_id: String,
+    pub source_id: String,
     pub num_threads: usize,
 }
 

src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs

@@ -81,8 +81,12 @@ use crate::servers::flight::v1::actions::init_query_fragments;
 use crate::servers::flight::v1::exchange::DataExchange;
 use crate::servers::flight::v1::exchange::DefaultExchangeInjector;
 use crate::servers::flight::v1::exchange::ExchangeInjector;
+use crate::servers::flight::v1::network::DoExchangeConnection;
+use crate::servers::flight::v1::network::DoExchangeConnector;
+use crate::servers::flight::v1::network::DoExchangeRetryPolicy;
 use crate::servers::flight::v1::network::NetworkInboundChannelSet;
 use crate::servers::flight::v1::network::NetworkInboundSender;
+use crate::servers::flight::v1::network::NetworkInboundSession;
 use crate::servers::flight::v1::network::PingPongExchange;
 use crate::servers::flight::v1::packets::Edge;
 use crate::servers::flight::v1::packets::QueryEnv;
@@ -299,6 +303,8 @@ impl DataExchangeManager {
             None => env.settings.clone(),
         };
         let keep_alive = settings.get_flight_keep_alive_params()?;
+        let retry_times = settings.get_flight_max_retry_times()?;
+        let retry_interval = settings.get_flight_retry_interval()?;
 
         let mut request_exchanges = HashMap::new();
         let mut targets_exchanges = HashMap::<String, Vec<FlightExchange>>::new();
@@ -385,43 +391,70 @@ impl DataExchangeManager {
                         let query_id = env.query_id.clone();
                         let address = target.flight_address.clone();
                         let keep_alive_params = keep_alive;
+                        let retry_policy = DoExchangeRetryPolicy {
+                            max_retries: retry_times,
+                            retry_interval: Duration::from_secs(retry_interval),
+                        };
                         let num_threads = channels.len();
+                        let source_id = config.query.node_id.clone();
                         warn!(
                             "do_exchange: node={} -> target={}, exchange_id={}, num_threads={}",
                             config.query.node_id, target_id, exchange_id, num_threads
                         );
 
                         flight_exchanges.push(Box::pin(async move {
-                            let (send_tx, response_stream) = {
-                                let mut flight_client = create_flight_client(
-                                    target_id.clone(),
-                                    address,
-                                    with_cur_rt,
-                                    keep_alive_params,
-                                )
-                                .await?;
-
-                                let (send_tx, send_rx) = async_channel::bounded(1);
-                                let response_stream = flight_client
-                                    .do_exchange(send_rx, DoExchangeParams {
-                                        query_id,
-                                        num_threads,
-                                        exchange_id: exchange_id.clone(),
+                            let params = DoExchangeParams {
+                                query_id,
+                                num_threads,
+                                source_id,
+                                exchange_id: exchange_id.clone(),
+                            };
+                            let connector_target_id = target_id.clone();
+                            let connector: DoExchangeConnector = Arc::new(move || {
+                                let address = address.clone();
+                                let params = params.clone();
+                                let target_id = connector_target_id.clone();
+                                Box::pin(async move {
+                                    let mut flight_client = create_flight_client(
+                                        target_id,
+                                        address,
+                                        with_cur_rt,
+                                        keep_alive_params,
+                                    )
+                                    .await
+                                    .map_err(|error| Status::unavailable(error.to_string()))?;
+
+                                    let (request_tx, request_rx) = async_channel::bounded(1);
+                                    let response_stream = flight_client
+                                        .do_exchange(request_rx, params)
+                                        .await
+                                        .map_err(|error| Status::unavailable(error.to_string()))?;
+                                    Ok(DoExchangeConnection {
+                                        request_tx,
+                                        response_stream: Box::pin(response_stream),
                                     })
-                                    .await?;
-                                Ok::<_, ErrorCode>((send_tx, response_stream))
-                            }?;
+                                })
+                            });
+
+                            let exchange = PingPongExchange::connect(
+                                num_threads,
+                                connector,
+                                retry_policy,
+                                local_node_id,
+                                target_id.clone(),
+                            )
+                            .await
+                            .map_err(|error| {
+                                ErrorCode::Internal(format!(
+                                    "PingPong connect failed after retries: {}",
+                                    error
+                                ))
+                            })?;
 
                             Ok::<QueryExchange, ErrorCode>(QueryExchange::PingPong {
                                 target_id: target_id.clone(),
                                 exchange_id,
-                                exchange: PingPongExchange::from_parts(
-                                    num_threads,
-                                    send_tx,
-                                    response_stream,
-                                    local_node_id,
-                                    target_id.clone(),
-                                ),
+                                exchange,
                             })
                         }));
                     }
@@ -703,24 +736,28 @@ impl DataExchangeManager {
     /// `NetworkInboundChannelSet` for the given channel_id. The caller (flight_service)
     /// uses the sender to push incoming FlightData into per-tid queues.
     #[fastrace::trace]
-    pub fn handle_do_exchange(
+    pub(crate) fn handle_do_exchange(
         &self,
         query_id: &str,
         channel_id: &str,
+        source_id: &str,
         num_threads: usize,
-    ) -> Result<NetworkInboundSender> {
+    ) -> Result<Arc<NetworkInboundSession>> {
         warn!(
-            "handle_do_exchange: query_id={}, channel_id={}, num_threads={}",
-            query_id, channel_id, num_threads
+            "handle_do_exchange: query_id={}, channel_id={}, source_id={}, num_threads={}",
+            query_id, channel_id, source_id, num_threads
         );
         let queries_coordinator_guard = self.queries_coordinator.lock();
         let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
 
         match queries_coordinator.entry(query_id.to_string()) {
-            Entry::Occupied(mut v) => v.get_mut().create_inbound_sender(channel_id, num_threads),
+            Entry::Occupied(mut v) => {
+                v.get_mut()
+                    .get_or_create_inbound_session(channel_id, source_id, num_threads)
+            }
             Entry::Vacant(v) => v
                 .insert(QueryCoordinator::create())
-                .create_inbound_sender(channel_id, num_threads),
+                .get_or_create_inbound_session(channel_id, source_id, num_threads),
         }
     }
 
@@ -1035,6 +1072,7 @@ pub(crate) struct QueryCoordinator {
     flight_data_senders: HashMap<String, Vec<FlightSender>>,
     flight_data_receivers: HashMap<String, Vec<FlightReceiver>>,
     inbound_channel_sets: HashMap<String, Arc<NetworkInboundChannelSet>>,
+    inbound_sessions: HashMap<(String, String), Arc<NetworkInboundSession>>,
     ping_pong_exchanges: HashMap<String, HashMap<String, PingPongExchange>>,
 }
 
@@ -1048,6 +1086,7 @@ impl QueryCoordinator {
             statistics_exchanges: HashMap::new(),
             fragments_coordinator: HashMap::new(),
             inbound_channel_sets: HashMap::new(),
+            inbound_sessions: HashMap::new(),
             ping_pong_exchanges: HashMap::new(),
         }
     }
@@ -1144,23 +1183,39 @@ impl QueryCoordinator {
         }
     }
 
-    /// Create a NetworkInboundSender for a new do_exchange connection.
-    ///
-    /// The `num_threads` value is provided by the coordinator via DoExchangeParams.
-    fn create_inbound_sender(
+    fn get_or_create_inbound_session(
         &mut self,
         channel_id: &str,
+        source_id: &str,
         num_threads: usize,
-    ) -> Result<NetworkInboundSender> {
-        let channel_set = self.get_or_create_inbound_channel_set(channel_id, num_threads)?;
+    ) -> Result<Arc<NetworkInboundSession>> {
+        let session_key = (channel_id.to_string(), source_id.to_string());
+        if let Some(session) = self.inbound_sessions.get(&session_key) {
+            return Ok(session.clone());
+        }
+
+        let channel_set = self
+            .inbound_channel_sets
+            .entry(channel_id.to_string())
+            .or_insert_with(|| Arc::new(NetworkInboundChannelSet::new(num_threads)))
+            .clone();
 
-        // TODO: get max_bytes_per_connection from query settings
-        let max_bytes_per_connection = 20 * 1024 * 1024; // 20MB
+        if channel_set.channels.len() != num_threads {
+            return Err(ErrorCode::Internal(format!(
+                "do_exchange num_threads mismatch for channel {}: expected {}, received {}",
+                channel_id,
+                channel_set.channels.len(),
+                num_threads
+            )));
+        }
 
-        Ok(NetworkInboundSender::new(
+        let max_bytes_per_connection = 20 * 1024 * 1024;
+        let session = Arc::new(NetworkInboundSession::new(NetworkInboundSender::new(
             &channel_set,
             max_bytes_per_connection,
-        ))
+        )));
+        self.inbound_sessions.insert(session_key, session.clone());
+        Ok(session)
     }
 
     fn get_or_create_inbound_channel_set(

src/query/service/src/servers/flight/v1/flight_service.rs

@@ -154,9 +154,10 @@ impl FlightService for DatabendQueryFlightService {
             Status::invalid_argument(format!("Failed to parse DoExchangeParams: {}", e))
         })?;
 
-        let sender = DataExchangeManager::instance().handle_do_exchange(
+        let session = DataExchangeManager::instance().handle_do_exchange(
             &params.query_id,
             &params.exchange_id,
+            &params.source_id,
             params.num_threads,
         )?;
 
@@ -165,20 +166,23 @@ impl FlightService for DatabendQueryFlightService {
 
         GlobalIORuntime::instance().spawn(async move {
             while let Some(result) = stream.next().await {
-                let Ok(flight_data) = result else {
-                    break;
+                let flight_data = match result {
+                    Ok(flight_data) => flight_data,
+                    Err(_) => break,
                 };
 
-                if sender.add_data(flight_data).await.is_err() {
-                    break; // Receiver closed
-                }
-
-                // Send pong (empty response signals readiness for next ping)
-                if let Err(_cause) = tx.try_send(Ok(FlightData::default())) {
-                    break;
+                match session.handle_frame(flight_data).await {
+                    Ok(response) => {
+                        if tx.send(Ok(response.data)).await.is_err() || response.terminal {
+                            break;
+                        }
+                    }
+                    Err(status) => {
+                        let _ = tx.send(Err(status)).await;
+                        break;
+                    }
                 }
             }
-            // sender is dropped here → closes sub-queues, notifies processors
         });
 
         Ok(RawResponse::new(Box::pin(rx)))