diff --git a/src/query/service/src/servers/flight/flight_client.rs b/src/query/service/src/servers/flight/flight_client.rs index 540a86a204336..e80ea13411682 100644 --- a/src/query/service/src/servers/flight/flight_client.rs +++ b/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, } diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs index a0e80fc40f15b..c77606e00358d 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_manager.rs +++ b/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::>::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::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 { + ) -> Result> { 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>, flight_data_receivers: HashMap>, inbound_channel_sets: HashMap>, + inbound_sessions: HashMap<(String, String), Arc>, ping_pong_exchanges: HashMap>, } @@ -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 { - let channel_set = self.get_or_create_inbound_channel_set(channel_id, num_threads)?; + ) -> Result> { + 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( diff --git a/src/query/service/src/servers/flight/v1/flight_service.rs b/src/query/service/src/servers/flight/v1/flight_service.rs index 73bd91f77d01f..073fb6ba34f5b 100644 --- a/src/query/service/src/servers/flight/v1/flight_service.rs +++ b/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( ¶ms.query_id, ¶ms.exchange_id, + ¶ms.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))) diff --git a/src/query/service/src/servers/flight/v1/network/do_exchange_protocol.rs b/src/query/service/src/servers/flight/v1/network/do_exchange_protocol.rs new file mode 100644 index 0000000000000..c5ec56f59999e --- /dev/null +++ b/src/query/service/src/servers/flight/v1/network/do_exchange_protocol.rs @@ -0,0 +1,192 @@ +// Copyright 2021 Datafuse Labs +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use arrow_flight::FlightData; +use bytes::BufMut; +use bytes::Bytes; +use bytes::BytesMut; +use tonic::Status; + +const MAGIC: &[u8; 4] = b"DX01"; +const HEADER_LEN: usize = MAGIC.len() + 1 + std::mem::size_of::(); + +#[derive(Clone, Copy, Debug)] +#[repr(u8)] +enum FrameKind { + Data = 1, + Finish = 2, + Ack = 3, + ReceiverFinished = 4, + FinishAck = 5, +} + +impl TryFrom for FrameKind { + type Error = Status; + + fn try_from(value: u8) -> Result { + match value { + value if value == FrameKind::Data as u8 => Ok(FrameKind::Data), + value if value == FrameKind::Finish as u8 => Ok(FrameKind::Finish), + value if value == FrameKind::Ack as u8 => Ok(FrameKind::Ack), + value if value == FrameKind::ReceiverFinished as u8 => Ok(FrameKind::ReceiverFinished), + value if value == FrameKind::FinishAck as u8 => Ok(FrameKind::FinishAck), + value => Err(Status::invalid_argument(format!( + "unknown do_exchange frame kind: {}", + value + ))), + } + } +} + +#[derive(Debug)] +pub enum DoExchangeFrame { + Data { sequence: u64, data: FlightData }, + Ack { sequence: u64 }, + Finish { sequence: u64 }, + FinishAck { sequence: u64 }, + ReceiverFinished, +} + +impl From for FlightData { + fn from(frame: DoExchangeFrame) -> Self { + let (kind, sequence, mut data) = match frame { + DoExchangeFrame::Data { sequence, data } => (FrameKind::Data, sequence, data), + DoExchangeFrame::Ack { sequence } => (FrameKind::Ack, sequence, FlightData::default()), + DoExchangeFrame::Finish { sequence } => { + (FrameKind::Finish, sequence, FlightData::default()) + } + DoExchangeFrame::FinishAck { sequence } => { + (FrameKind::FinishAck, sequence, FlightData::default()) + } + DoExchangeFrame::ReceiverFinished => { + (FrameKind::ReceiverFinished, 0, FlightData::default()) + } + }; + + data.app_metadata = encode_header(kind, sequence, &data.app_metadata); + data + } +} + +impl TryFrom for DoExchangeFrame { + type Error = Status; + + fn try_from(mut data: FlightData) -> Result { + let (kind, sequence, payload) = decode_header(&data)?; + match kind { + FrameKind::Data => { + data.app_metadata = payload; + Ok(DoExchangeFrame::Data { sequence, data }) + } + FrameKind::Ack => { + ensure_empty_payload(kind, &payload)?; + Ok(DoExchangeFrame::Ack { sequence }) + } + FrameKind::Finish => { + ensure_empty_payload(kind, &payload)?; + Ok(DoExchangeFrame::Finish { sequence }) + } + FrameKind::FinishAck => { + ensure_empty_payload(kind, &payload)?; + Ok(DoExchangeFrame::FinishAck { sequence }) + } + FrameKind::ReceiverFinished => { + ensure_empty_payload(kind, &payload)?; + Ok(DoExchangeFrame::ReceiverFinished) + } + } + } +} + +fn encode_header(kind: FrameKind, sequence: u64, payload: &[u8]) -> Bytes { + let mut metadata = BytesMut::with_capacity(HEADER_LEN + payload.len()); + metadata.put_slice(MAGIC); + metadata.put_u8(kind as u8); + metadata.put_u64_le(sequence); + metadata.put_slice(payload); + metadata.freeze() +} + +fn decode_header(data: &FlightData) -> Result<(FrameKind, u64, Bytes), Status> { + if data.app_metadata.len() < HEADER_LEN || &data.app_metadata[..MAGIC.len()] != MAGIC { + return Err(Status::invalid_argument( + "invalid do_exchange protocol header", + )); + } + + let kind = FrameKind::try_from(data.app_metadata[MAGIC.len()])?; + let sequence_offset = MAGIC.len() + 1; + let sequence = u64::from_le_bytes( + data.app_metadata[sequence_offset..HEADER_LEN] + .try_into() + .expect("do_exchange sequence has a fixed size"), + ); + + Ok((kind, sequence, data.app_metadata.slice(HEADER_LEN..))) +} + +fn ensure_empty_payload(kind: FrameKind, payload: &Bytes) -> Result<(), Status> { + if payload.is_empty() { + return Ok(()); + } + + Err(Status::invalid_argument(format!( + "do_exchange {:?} frame must not contain a payload", + kind + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_data_round_trip_preserves_metadata() { + let data = FlightData { + app_metadata: Bytes::from_static(b"payload"), + data_body: Bytes::from_static(b"body"), + ..Default::default() + }; + + let encoded: FlightData = DoExchangeFrame::Data { sequence: 42, data }.into(); + let DoExchangeFrame::Data { sequence, data } = DoExchangeFrame::try_from(encoded).unwrap() + else { + panic!("expected DATA frame"); + }; + + assert_eq!(sequence, 42); + assert_eq!(data.app_metadata, Bytes::from_static(b"payload")); + assert_eq!(data.data_body, Bytes::from_static(b"body")); + } + + #[test] + fn test_control_frame_round_trip() { + let cases = [ + DoExchangeFrame::Ack { sequence: 7 }, + DoExchangeFrame::ReceiverFinished, + DoExchangeFrame::FinishAck { sequence: 9 }, + ]; + + for expected in cases { + let encoded: FlightData = expected.into(); + let decoded = DoExchangeFrame::try_from(encoded).unwrap(); + match decoded { + DoExchangeFrame::Ack { sequence } => assert_eq!(sequence, 7), + DoExchangeFrame::ReceiverFinished => {} + DoExchangeFrame::FinishAck { sequence } => assert_eq!(sequence, 9), + frame => panic!("unexpected control frame: {:?}", frame), + } + } + } +} diff --git a/src/query/service/src/servers/flight/v1/network/inbound_channel.rs b/src/query/service/src/servers/flight/v1/network/inbound_channel.rs index 2a3ef86a1dedd..ccf5d5fd9f6fb 100644 --- a/src/query/service/src/servers/flight/v1/network/inbound_channel.rs +++ b/src/query/service/src/servers/flight/v1/network/inbound_channel.rs @@ -14,6 +14,7 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -30,8 +31,11 @@ use databend_common_expression::DataBlock; use databend_common_expression::DataSchemaRef; use databend_common_io::prelude::BinaryRead; use databend_common_io::prelude::bincode_deserialize_from_stream; +use tokio::sync::Mutex; use tokio::sync::Semaphore; +use tonic::Status; +use super::do_exchange_protocol::DoExchangeFrame; use super::inbound_quota::QueueItem; use super::inbound_quota::SubQueue; @@ -81,17 +85,14 @@ impl NetworkInboundChannelSet { } } -/// Network-side handle. Each do_exchange connection gets one. -/// -/// When dropped, closes this connection's sub-queues and notifies processors. +/// Network-side handle for one logical sender across all of its reconnects. pub struct NetworkInboundSender { - /// This connection's sub-queue in each tid's NetworkInboundChannel. sub_queues: Vec>, + finished: AtomicBool, } impl NetworkInboundSender { - /// Create a new sender for a connection. - /// Adds a sub-queue to each NetworkInboundChannel for this connection. + /// Adds one sub-queue to every NetworkInboundChannel for this logical sender. pub fn new(channel_set: &NetworkInboundChannelSet, max_bytes_per_connection: usize) -> Self { let semaphore = Arc::new(Semaphore::new(max_bytes_per_connection)); let mut sub_queues = Vec::with_capacity(channel_set.channels.len()); @@ -110,7 +111,10 @@ impl NetworkInboundSender { sub_queues.push(sub_queue); } - Self { sub_queues } + Self { + sub_queues, + finished: AtomicBool::new(false), + } } /// Add data to the inbound channel. @@ -156,10 +160,12 @@ impl NetworkInboundSender { pub fn all_receivers_closed(&self) -> bool { self.sub_queues.iter().all(|q| q.sender.is_closed()) } -} -impl Drop for NetworkInboundSender { - fn drop(&mut self) { + fn finish(&self) { + if self.finished.swap(true, Ordering::AcqRel) { + return; + } + for sub_queue in &self.sub_queues { if sub_queue.sender_count.fetch_sub(1, Ordering::AcqRel) == 1 { sub_queue.sender.close(); @@ -168,6 +174,141 @@ impl Drop for NetworkInboundSender { } } +impl Drop for NetworkInboundSender { + fn drop(&mut self) { + self.finish(); + } +} + +#[derive(Clone, Copy)] +enum InboundSessionTerminal { + SenderFinished, + ReceiverFinished, +} + +struct InboundSessionState { + next_sequence: u64, + terminal: Option, +} + +pub(crate) struct InboundExchangeResponse { + pub data: FlightData, + pub terminal: bool, +} + +/// Receiver-side state for one logical do_exchange sender. +/// +/// The state outlives individual gRPC streams. A retried sequence is acknowledged +/// without being delivered twice, while FINISH is the only sender-side event that +/// completes the logical inbound channel. +pub(crate) struct NetworkInboundSession { + sender: NetworkInboundSender, + state: Mutex, +} + +impl NetworkInboundSession { + pub(crate) fn new(sender: NetworkInboundSender) -> Self { + Self { + sender, + state: Mutex::new(InboundSessionState { + next_sequence: 0, + terminal: None, + }), + } + } + + pub(crate) async fn handle_frame( + &self, + data: FlightData, + ) -> Result { + match DoExchangeFrame::try_from(data)? { + DoExchangeFrame::Data { sequence, data } => { + let mut state = self.state.lock().await; + + if state.terminal.is_some() { + return Ok(InboundExchangeResponse { + data: DoExchangeFrame::ReceiverFinished.into(), + terminal: true, + }); + } + + if sequence < state.next_sequence { + // ACK may be lost, so acknowledge replays without enqueueing again. + return Ok(InboundExchangeResponse { + data: DoExchangeFrame::Ack { sequence }.into(), + terminal: false, + }); + } + + if sequence > state.next_sequence { + // Only one frame is in flight, so a forward gap is invalid. + return Err(Status::failed_precondition(format!( + "do_exchange sequence gap: expected {}, received {}", + state.next_sequence, sequence + ))); + } + + if self.sender.add_data(data).await.is_err() { + // LIMIT can close all receivers; report this as normal completion. + self.sender.finish(); + state.terminal = Some(InboundSessionTerminal::ReceiverFinished); + return Ok(InboundExchangeResponse { + data: DoExchangeFrame::ReceiverFinished.into(), + terminal: true, + }); + } + + state.next_sequence += 1; + Ok(InboundExchangeResponse { + data: DoExchangeFrame::Ack { sequence }.into(), + terminal: false, + }) + } + DoExchangeFrame::Finish { sequence } => { + let mut state = self.state.lock().await; + + match state.terminal { + Some(InboundSessionTerminal::ReceiverFinished) => { + return Ok(InboundExchangeResponse { + data: DoExchangeFrame::ReceiverFinished.into(), + terminal: true, + }); + } + Some(InboundSessionTerminal::SenderFinished) => { + // FINISH_ACK may be lost, so acknowledge a replayed FINISH again. + return Ok(InboundExchangeResponse { + data: DoExchangeFrame::FinishAck { + sequence: state.next_sequence, + } + .into(), + terminal: true, + }); + } + None => {} + } + + if sequence != state.next_sequence { + return Err(Status::failed_precondition(format!( + "do_exchange FINISH sequence mismatch: expected {}, received {}", + state.next_sequence, sequence + ))); + } + + self.sender.finish(); + state.terminal = Some(InboundSessionTerminal::SenderFinished); + Ok(InboundExchangeResponse { + data: DoExchangeFrame::FinishAck { sequence }.into(), + terminal: true, + }) + } + frame => Err(Status::invalid_argument(format!( + "unexpected inbound do_exchange frame: {:?}", + frame + ))), + } + } +} + /// Trait for receiving data blocks from the network. #[async_trait::async_trait] pub trait InboundChannel: Send + Sync { @@ -371,9 +512,102 @@ mod tests { use bytes::BytesMut; use super::*; + use crate::servers::flight::v1::network::do_exchange_protocol::DoExchangeFrame; const BATCH_MARKER: u8 = 0x02; + fn session_data(tid: u16) -> FlightData { + let mut app_metadata = BytesMut::new(); + app_metadata.put_u16_le(tid); + app_metadata.put_u8(0x01); + FlightData { + app_metadata: app_metadata.freeze(), + data_body: Bytes::from_static(b"data"), + ..Default::default() + } + } + + #[tokio::test] + async fn test_inbound_session_deduplicates_retried_sequence_and_finishes_explicitly() { + let channel_set = NetworkInboundChannelSet::new(1); + let session = NetworkInboundSession::new(NetworkInboundSender::new(&channel_set, 1024)); + + let response = session + .handle_frame( + DoExchangeFrame::Data { + sequence: 0, + data: session_data(0), + } + .into(), + ) + .await + .unwrap(); + assert!(matches!( + DoExchangeFrame::try_from(response.data).unwrap(), + DoExchangeFrame::Ack { sequence: 0 } + )); + + let duplicate = session + .handle_frame( + DoExchangeFrame::Data { + sequence: 0, + data: session_data(0), + } + .into(), + ) + .await + .unwrap(); + assert!(matches!( + DoExchangeFrame::try_from(duplicate.data).unwrap(), + DoExchangeFrame::Ack { sequence: 0 } + )); + assert_eq!( + channel_set.channels[0].receiver.len(), + 1, + "a retransmitted sequence must not be delivered twice" + ); + + let finish = session + .handle_frame(DoExchangeFrame::Finish { sequence: 1 }.into()) + .await + .unwrap(); + assert!(finish.terminal); + assert!(matches!( + DoExchangeFrame::try_from(finish.data).unwrap(), + DoExchangeFrame::FinishAck { sequence: 1 } + )); + + let _ = channel_set.channels[0].receiver.recv().await.unwrap(); + assert!( + channel_set.channels[0].receiver.recv().await.is_err(), + "FINISH must close the logical inbound sender" + ); + } + + #[tokio::test] + async fn test_inbound_session_reports_receiver_finished_without_network_error() { + let channel_set = NetworkInboundChannelSet::new(1); + let session = NetworkInboundSession::new(NetworkInboundSender::new(&channel_set, 1024)); + channel_set.channels[0].receiver.close(); + + let response = session + .handle_frame( + DoExchangeFrame::Data { + sequence: 0, + data: session_data(0), + } + .into(), + ) + .await + .unwrap(); + + assert!(response.terminal); + assert!(matches!( + DoExchangeFrame::try_from(response.data).unwrap(), + DoExchangeFrame::ReceiverFinished + )); + } + /// Build a FlightData with tid prefix in app_metadata. fn make_item(tid: u16, inner_meta: &[u8], header: &[u8], body: &[u8]) -> FlightData { let mut app_metadata = BytesMut::with_capacity(2 + inner_meta.len()); diff --git a/src/query/service/src/servers/flight/v1/network/mod.rs b/src/query/service/src/servers/flight/v1/network/mod.rs index ca4c6b2e5c7e6..7a073acf79883 100644 --- a/src/query/service/src/servers/flight/v1/network/mod.rs +++ b/src/query/service/src/servers/flight/v1/network/mod.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +pub(crate) mod do_exchange_protocol; pub mod inbound_channel; pub mod inbound_quota; pub mod local_channel; @@ -25,6 +26,7 @@ pub use inbound_channel::InboundChannel; pub use inbound_channel::NetworkInboundChannelSet; pub use inbound_channel::NetworkInboundReceiver; pub use inbound_channel::NetworkInboundSender; +pub(crate) use inbound_channel::NetworkInboundSession; pub use local_channel::LocalOutboundChannel; pub use local_channel::create_local_channels; pub use outbound_buffer::ExchangeBufferConfig; @@ -33,6 +35,9 @@ pub use outbound_channel::DummyOutboundChannel; pub use outbound_channel::OutboundChannel; pub use outbound_channel::RemoteChannel; pub use outbound_channel::RoundRobinChannel; +pub(crate) use outbound_transport::DoExchangeConnection; +pub(crate) use outbound_transport::DoExchangeConnector; +pub(crate) use outbound_transport::DoExchangeRetryPolicy; pub use outbound_transport::PingPongCallback; pub use outbound_transport::PingPongExchange; pub use outbound_transport::PingPongExchangeInner; diff --git a/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs b/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs index 974cd1fd675f3..c8a470e230eef 100644 --- a/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs +++ b/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs @@ -30,7 +30,6 @@ use tonic::Status; use super::outbound_transport::PingPongCallback; use super::outbound_transport::PingPongExchange; use super::outbound_transport::PingPongResponse; -use super::outbound_transport::REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE; use crate::servers::flight::FlightOperation; use crate::servers::flight::add_flight_error_context; use crate::servers::flight::v1::network::inbound_quota::RemoteQueueItem; @@ -131,14 +130,11 @@ struct RemoteInstanceState { channels: Vec, /// Last error from the exchange, returned on next poll_send call last_error: Option, + receiver_finished: bool, } impl RemoteInstanceState { - fn close(&mut self, error: ErrorCode) { - if self.last_error.is_none() { - self.last_error = Some(error); - } - + fn close_queues(&mut self) { for channel in &self.channels { channel.pending_queue.close(); } @@ -147,6 +143,19 @@ impl RemoteInstanceState { while channel.pending_queue.pop().is_ok() {} } } + + fn close(&mut self, error: ErrorCode) { + if self.last_error.is_none() { + self.last_error = Some(error); + } + + self.close_queues(); + } + + fn receiver_finished(&mut self) { + self.receiver_finished = true; + self.close_queues(); + } } /// Per-destination remote instance containing multiple sink channels. @@ -164,6 +173,7 @@ impl RemoteInstance { state: Mutex::new(RemoteInstanceState { channels, last_error: None, + receiver_finished: false, }), } } @@ -190,7 +200,7 @@ struct ExchangeSinkBufferInner { impl Drop for ExchangeSinkBufferInner { fn drop(&mut self) { for remote in &self.state.remotes { - remote.exchange.shutdown.notify_waiters(); + remote.exchange.finish(); } } } @@ -262,15 +272,10 @@ impl PingPongCallback for SinkBufferCallback { .try_flush_remote(self.dest_idx, response.data.err()); } - fn on_closed(&self) { + fn on_receiver_finished(&self) { let remote = &self.buffer.remotes[self.dest_idx]; let mut state = remote.state.lock(); - state.close(add_flight_error_context( - ErrorCode::AbortedQuery(REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE), - FlightOperation::DoExchange, - remote.exchange.local_node_id(), - remote.exchange.remote_node_id(), - )); + state.receiver_finished(); } } @@ -330,6 +335,9 @@ impl ExchangeSinkBuffer { if let Some(status) = state.last_error.clone() { return Err(status); } + if state.receiver_finished { + return Ok(()); + } } // Try to send directly first @@ -357,6 +365,9 @@ impl ExchangeSinkBuffer { if let Some(status) = state.last_error.clone() { return Err(status); } + if state.receiver_finished { + return Ok(()); + } // Try to send again match remote.exchange.try_send(data) { @@ -394,7 +405,7 @@ impl ExchangeSinkBuffer { pub fn is_closed(&self, dest_idx: usize) -> bool { let remote = &self.inner.state.remotes[dest_idx]; let state = remote.state.lock(); - state.last_error.is_some() + state.last_error.is_some() || state.receiver_finished } } @@ -409,7 +420,9 @@ mod tests { use tonic::Status; use super::*; + use crate::servers::flight::v1::network::do_exchange_protocol::DoExchangeFrame; use crate::servers::flight::v1::network::outbound_transport::PingPongExchange; + use crate::servers::flight::v1::network::outbound_transport::REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE; fn test_runtime() -> Arc { Arc::new(Runtime::with_worker_threads(2, None).unwrap()) @@ -700,6 +713,36 @@ mod tests { ); } + #[tokio::test(flavor = "multi_thread")] + async fn test_receiver_finished_closes_buffer_without_error() { + let rt = test_runtime(); + let (exchange, send_rx, pong_tx) = create_mock_exchange(1); + let buffer = + ExchangeSinkBuffer::create(vec![exchange], ExchangeBufferConfig::default(), &rt) + .unwrap(); + + buffer.add_data(0, 0, make_flight_data(1)).await.unwrap(); + let _ = send_rx.recv().await.unwrap(); + pong_tx + .send(Ok(DoExchangeFrame::ReceiverFinished.into())) + .await + .unwrap(); + + tokio::time::timeout(Duration::from_secs(2), async { + while !buffer.is_closed(0) { + tokio::task::yield_now().await; + } + }) + .await + .expect("receiver-finished response should close the outbound channel"); + + buffer + .add_data(0, 0, make_flight_data(2)) + .await + .expect("receiver-initiated close is not a network error"); + assert!(send_rx.try_recv().is_err()); + } + #[tokio::test(flavor = "multi_thread")] async fn test_force_send_error_drains_remaining_pending_queue() { let rt = test_runtime(); @@ -781,7 +824,7 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn test_direct_send_closed_returns_aborted_query() { + async fn test_request_stream_close_is_reported_as_aborted_query() { let rt = test_runtime(); let (exchange, send_rx, _pong_tx) = create_mock_exchange(1); let buffer = @@ -790,8 +833,17 @@ mod tests { drop(send_rx); + buffer.add_data(0, 0, make_flight_data(1)).await.unwrap(); + tokio::time::timeout(Duration::from_secs(2), async { + while !buffer.is_closed(0) { + tokio::task::yield_now().await; + } + }) + .await + .expect("transport worker should report the closed request stream"); + let error = buffer - .add_data(0, 0, make_flight_data(1)) + .add_data(0, 0, make_flight_data(2)) .await .expect_err("closed request channel should abort the exchange"); assert_eq!(error.code(), ErrorCode::ABORTED_QUERY); @@ -809,7 +861,7 @@ mod tests { } #[tokio::test(flavor = "multi_thread")] - async fn test_force_send_error_unblocks_waiting_add_data() { + async fn test_request_stream_error_unblocks_waiting_add_data() { let rt = test_runtime(); let (exchange, send_rx, pong_tx) = create_mock_exchange(1); let buffer = Arc::new( @@ -846,7 +898,20 @@ mod tests { .await .expect("blocked add_data should be released after remote close") .unwrap(); - let error = result.expect_err("released add_data should observe exchange abort"); + let _ = result; + + tokio::time::timeout(Duration::from_secs(2), async { + while !buffer.is_closed(0) { + tokio::task::yield_now().await; + } + }) + .await + .expect("request stream failure should close the remote"); + + let error = buffer + .add_data(0, 0, make_flight_data(4)) + .await + .expect_err("later data should observe the exchange abort"); assert_eq!(error.code(), ErrorCode::ABORTED_QUERY); assert_eq!( buffer.semaphore.available_permits(), diff --git a/src/query/service/src/servers/flight/v1/network/outbound_channel.rs b/src/query/service/src/servers/flight/v1/network/outbound_channel.rs index b721b6f86161b..1d7646e0a385c 100644 --- a/src/query/service/src/servers/flight/v1/network/outbound_channel.rs +++ b/src/query/service/src/servers/flight/v1/network/outbound_channel.rs @@ -290,6 +290,7 @@ mod tests { use tonic::Status; use super::*; + use crate::servers::flight::v1::network::do_exchange_protocol::DoExchangeFrame; use crate::servers::flight::v1::network::inbound_channel::deserialize_flight_data; use crate::servers::flight::v1::network::inbound_channel::strip_tid; use crate::servers::flight::v1::network::outbound_buffer::ExchangeBufferConfig; @@ -336,7 +337,11 @@ mod tests { channel.add_block(make_block(10)).await.unwrap(); - let received = send_rx.recv().await.unwrap(); + let DoExchangeFrame::Data { data: received, .. } = + DoExchangeFrame::try_from(send_rx.recv().await.unwrap()).unwrap() + else { + panic!("expected DATA frame"); + }; let meta = received.app_metadata.to_vec(); // First 2 bytes: tid=0 as u16 LE assert_eq!(&meta[..2], &[0, 0]); @@ -376,7 +381,11 @@ mod tests { channel.add_block(make_block(1)).await.unwrap(); - let received = send_rx.recv().await.unwrap(); + let DoExchangeFrame::Data { data: received, .. } = + DoExchangeFrame::try_from(send_rx.recv().await.unwrap()).unwrap() + else { + panic!("expected DATA frame"); + }; let meta = received.app_metadata.to_vec(); // First 2 bytes: tid=5 as u16 LE assert_eq!(&meta[..2], &[5, 0]); @@ -450,7 +459,12 @@ mod tests { channel.add_block(original.clone()).await.unwrap(); // Receive, strip tid, deserialize - let flight_data = send_rx.recv().await.unwrap(); + let DoExchangeFrame::Data { + data: flight_data, .. + } = DoExchangeFrame::try_from(send_rx.recv().await.unwrap()).unwrap() + else { + panic!("expected DATA frame"); + }; let stripped = strip_tid(flight_data); let schema = Arc::new(original.infer_schema()); let arrow_schema = Arc::new(ArrowSchema::from(schema.as_ref())); diff --git a/src/query/service/src/servers/flight/v1/network/outbound_transport.rs b/src/query/service/src/servers/flight/v1/network/outbound_transport.rs index 538f77ff60bc9..ac14680a0f9f0 100644 --- a/src/query/service/src/servers/flight/v1/network/outbound_transport.rs +++ b/src/query/service/src/servers/flight/v1/network/outbound_transport.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; @@ -19,6 +21,7 @@ use std::time::Duration; use std::time::Instant; use arrow_flight::FlightData; +use async_channel::Receiver; use async_channel::Sender; use async_channel::TrySendError; use databend_common_base::base::WatchNotify; @@ -27,24 +30,50 @@ use futures::StreamExt; use futures::future::Either; use futures::future::select; use futures::stream::BoxStream; +use log::warn; use parking_lot::Mutex; use tokio::task::JoinHandle; +use tonic::Code; use tonic::Status; +use super::do_exchange_protocol::DoExchangeFrame; + pub const REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE: &str = "Aborted query, because the remote flight channel is closed."; -/// Response from a ping-pong exchange. pub struct PingPongResponse { pub data: Result, pub rtt: Duration, } -/// Callback trait for handling ping-pong responses. pub trait PingPongCallback: Send + Sync + 'static { fn has_pending(&self) -> bool; fn on_response(&self, response: PingPongResponse); - fn on_closed(&self); + fn on_receiver_finished(&self); +} + +pub struct DoExchangeConnection { + pub request_tx: Sender, + pub response_stream: BoxStream<'static, Result>, +} + +pub type DoExchangeConnectFuture = + Pin> + Send + 'static>>; +pub type DoExchangeConnector = Arc DoExchangeConnectFuture + Send + Sync + 'static>; + +#[derive(Clone, Copy)] +pub struct DoExchangeRetryPolicy { + pub max_retries: u64, + pub retry_interval: Duration, +} + +impl Default for DoExchangeRetryPolicy { + fn default() -> Self { + Self { + max_retries: 0, + retry_interval: Duration::ZERO, + } + } } pub struct PingPongExchangeInner { @@ -53,23 +82,22 @@ pub struct PingPongExchangeInner { send_tx: Sender, local_node_id: String, remote_node_id: String, + finishing: AtomicBool, pub shutdown: WatchNotify, } -/// A non-blocking ping-pong style flight exchange. -/// -/// This exchange guarantees that at most one request is in-flight at any time. -/// When a request is sent, subsequent sends will return the data back to the caller -/// until a response is received. pub struct PingPongExchange { pub num_threads: usize, inner: Arc, - response_stream: Mutex>>>, + send_rx: Mutex>>, + initial_connection: Mutex>, + connector: DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, } impl Drop for PingPongExchange { fn drop(&mut self) { - self.inner.shutdown.notify_waiters(); + self.finish(); } } @@ -89,12 +117,7 @@ impl PingPongExchangeInner { .map(|t| t.elapsed()) .unwrap_or_default() } - /// Try to send data through the exchange. - /// - /// Returns: - /// - `Ok(None)`: Data was sent successfully - /// - `Ok(Some(data))`: A request is already in-flight, data is returned to caller - /// - `Err(status)`: The exchange is closed + pub fn try_send(&self, data: FlightData) -> Result, Status> { if self.in_flight.fetch_or(true, Ordering::SeqCst) { return Ok(Some(data)); @@ -126,9 +149,56 @@ impl PingPongExchangeInner { } impl PingPongExchange { + pub async fn connect( + num_threads: usize, + connector: DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, + local_node_id: impl Into, + remote_node_id: impl Into, + ) -> Result { + let initial_connection = connect_with_retry(&connector, retry_policy).await?; + Ok(Self::from_connection( + num_threads, + initial_connection, + connector, + retry_policy, + local_node_id.into(), + remote_node_id.into(), + )) + } + + fn from_connection( + num_threads: usize, + initial_connection: DoExchangeConnection, + connector: DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, + local_node_id: String, + remote_node_id: String, + ) -> Self { + let (send_tx, send_rx) = async_channel::bounded(1); + let inner = Arc::new(PingPongExchangeInner { + in_flight: AtomicBool::new(false), + send_time: Mutex::new(None), + send_tx, + local_node_id, + remote_node_id, + finishing: AtomicBool::new(false), + shutdown: WatchNotify::new(), + }); + + Self { + inner, + num_threads, + send_rx: Mutex::new(Some(send_rx)), + initial_connection: Mutex::new(Some(initial_connection)), + connector, + retry_policy, + } + } + pub fn from_parts( num_threads: usize, - send_tx: async_channel::Sender, + send_tx: Sender, response_stream: tonic::Streaming, local_node_id: String, remote_node_id: String, @@ -144,25 +214,43 @@ impl PingPongExchange { pub fn from_stream( num_threads: usize, - send_tx: async_channel::Sender, + request_tx: Sender, stream: impl futures::Stream> + Send + 'static, local_node_id: impl Into, remote_node_id: impl Into, ) -> Self { - let inner = Arc::new(PingPongExchangeInner { - in_flight: AtomicBool::new(false), - send_time: Mutex::new(None), - send_tx, - local_node_id: local_node_id.into(), - remote_node_id: remote_node_id.into(), - shutdown: WatchNotify::new(), + let mut sequence = 0; + let stream = stream.map(move |result| match result { + Ok(response) if DoExchangeFrame::try_from(response.clone()).is_ok() => Ok(response), + Ok(_) => { + let response = DoExchangeFrame::Ack { sequence }.into(); + sequence += 1; + Ok(response) + } + Err(status) => Err(status), }); - - Self { - inner, + let connector: DoExchangeConnector = Arc::new(|| { + Box::pin(async { Err(Status::aborted(REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE)) }) + }); + Self::from_connection( num_threads, - response_stream: Mutex::new(Some(Box::pin(stream))), - } + DoExchangeConnection { + request_tx, + response_stream: Box::pin(stream), + }, + connector, + DoExchangeRetryPolicy { + max_retries: 1, + retry_interval: Duration::ZERO, + }, + local_node_id.into(), + remote_node_id.into(), + ) + } + + pub fn finish(&self) { + self.inner.finishing.store(true, Ordering::Release); + self.inner.shutdown.notify_waiters(); } pub fn local_node_id(&self) -> &str { @@ -184,57 +272,249 @@ impl PingPongExchange { callback: Arc, runtime: &Runtime, ) -> Result, Status> { - let Some(mut stream) = self.response_stream.lock().take() else { - return Err(Status::already_exists("Receiver already started")); + let Some(send_rx) = self.send_rx.lock().take() else { + return Err(Status::already_exists("do_exchange sender already started")); + }; + let Some(initial_connection) = self.initial_connection.lock().take() else { + return Err(Status::already_exists( + "do_exchange connection already started", + )); }; let inner = self.inner.clone(); + let connector = self.connector.clone(); + let retry_policy = self.retry_policy; Ok(runtime.spawn(async move { - let mut finished = false; - let mut shutdown_fut = Box::pin(inner.shutdown.notified()); + run_exchange( + inner, + send_rx, + initial_connection, + connector, + retry_policy, + callback, + ) + .await; + })) + } +} - loop { - if finished && !callback.has_pending() { - inner.send_tx.close(); +fn is_retriable(status: &Status) -> bool { + matches!( + status.code(), + Code::Unavailable | Code::Cancelled | Code::DeadlineExceeded | Code::Unknown + ) +} + +async fn connect_with_retry( + connector: &DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, +) -> Result { + let mut retries = 0; + loop { + match connector().await { + Ok(connection) => return Ok(connection), + Err(status) if is_retriable(&status) && retries < retry_policy.max_retries => { + retries += 1; + warn!( + "do_exchange connect failed, retry {}/{}: {}", + retries, retry_policy.max_retries, status + ); + tokio::time::sleep(retry_policy.retry_interval).await; + } + Err(status) => return Err(status), + } + } +} + +async fn round_trip( + connection: &mut Option, + connector: &DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, + frame: FlightData, +) -> Result { + let mut retries = 0; + + loop { + if connection.is_none() { + match connector().await { + Ok(new_connection) => { + *connection = Some(new_connection); + } + Err(status) if is_retriable(&status) && retries < retry_policy.max_retries => { + retries += 1; + warn!( + "do_exchange reconnect failed, retry {}/{}: {}", + retries, retry_policy.max_retries, status + ); + tokio::time::sleep(retry_policy.retry_interval).await; + continue; + } + Err(status) => return Err(status), + } + } + + let result = { + let connection = connection + .as_mut() + .expect("do_exchange connection must be initialized"); + + match connection.request_tx.send(frame.clone()).await { + Ok(()) => match connection.response_stream.next().await { + Some(Ok(response)) => DoExchangeFrame::try_from(response), + Some(Err(status)) => Err(status), + None => Err(Status::unavailable( + "do_exchange response stream closed without a terminal frame", + )), + }, + Err(_) => Err(Status::unavailable( + "do_exchange request stream closed while sending", + )), + } + }; + + match result { + Ok(response) => return Ok(response), + Err(status) if is_retriable(&status) && retries < retry_policy.max_retries => { + retries += 1; + warn!( + "do_exchange communication failed, reconnect {}/{}: {}", + retries, retry_policy.max_retries, status + ); + connection.take(); + tokio::time::sleep(retry_policy.retry_interval).await; + } + Err(status) => return Err(status), + } + } +} + +async fn run_exchange( + inner: Arc, + send_rx: Receiver, + initial_connection: DoExchangeConnection, + connector: DoExchangeConnector, + retry_policy: DoExchangeRetryPolicy, + callback: Arc, +) { + let mut connection = Some(initial_connection); + let mut next_sequence = 0; + let mut finishing = false; + let mut shutdown_fut = Box::pin(inner.shutdown.notified()); + + loop { + finishing |= inner.finishing.load(Ordering::Acquire); + if finishing && !inner.in_flight.load(Ordering::Acquire) && !callback.has_pending() { + match round_trip( + &mut connection, + &connector, + retry_policy, + DoExchangeFrame::Finish { + sequence: next_sequence, } + .into(), + ) + .await + { + Ok(DoExchangeFrame::FinishAck { sequence }) if sequence == next_sequence => {} + Ok(DoExchangeFrame::ReceiverFinished) => callback.on_receiver_finished(), + Ok(response) => callback.on_response(PingPongResponse { + data: Err(Status::failed_precondition(format!( + "unexpected do_exchange FINISH response: {:?}", + response + ))), + rtt: Duration::ZERO, + }), + Err(status) => callback.on_response(PingPongResponse { + data: Err(status), + rtt: Duration::ZERO, + }), + } + inner.send_tx.close(); + break; + } - match select(shutdown_fut, stream.next()).await { - Either::Left(_) => { - finished = true; - shutdown_fut = Box::pin(inner.shutdown.notified()); - } - Either::Right((None, _)) => { - callback.on_closed(); - break; - } - Either::Right((Some(Ok(data)), next_shutdown)) => { - shutdown_fut = next_shutdown; - let rtt = inner.get_rtt(); - callback.on_response(PingPongResponse { - data: Ok(data), - rtt, - }); - } - Either::Right((Some(Err(status)), _)) => { - let rtt = inner.get_rtt(); - callback.on_response(PingPongResponse { - data: Err(status), - rtt, - }); - break; - } + let recv_fut = Box::pin(send_rx.recv()); + let data = match select(shutdown_fut, recv_fut).await { + Either::Left((_, _)) => { + finishing = true; + shutdown_fut = Box::pin(inner.shutdown.notified()); + continue; + } + Either::Right((Ok(data), next_shutdown)) => { + shutdown_fut = next_shutdown; + data + } + Either::Right((Err(_), _)) => break, + }; + + let exchange_fut = Box::pin(round_trip( + &mut connection, + &connector, + retry_policy, + DoExchangeFrame::Data { + sequence: next_sequence, + data, + } + .into(), + )); + let response = if finishing { + exchange_fut.await + } else { + match select(shutdown_fut, exchange_fut).await { + Either::Left((_, pending_exchange)) => { + finishing = true; + shutdown_fut = Box::pin(inner.shutdown.notified()); + pending_exchange.await + } + Either::Right((response, next_shutdown)) => { + shutdown_fut = next_shutdown; + response } } - })) + }; + + let rtt = inner.get_rtt(); + match response { + Ok(DoExchangeFrame::Ack { sequence }) if sequence == next_sequence => { + next_sequence += 1; + callback.on_response(PingPongResponse { + data: Ok(FlightData::default()), + rtt, + }); + } + Ok(DoExchangeFrame::ReceiverFinished) => { + callback.on_receiver_finished(); + inner.ready_send(); + inner.send_tx.close(); + break; + } + Ok(response) => { + callback.on_response(PingPongResponse { + data: Err(Status::failed_precondition(format!( + "unexpected do_exchange DATA response: {:?}", + response + ))), + rtt, + }); + inner.send_tx.close(); + break; + } + Err(status) => { + callback.on_response(PingPongResponse { + data: Err(status), + rtt, + }); + inner.send_tx.close(); + break; + } + } } } #[cfg(test)] mod tests { - use arrow_flight::FlightData; - use tonic::Status; - use super::*; + use crate::servers::flight::v1::network::do_exchange_protocol::DoExchangeFrame; fn create_mock_exchange( num_threads: usize, @@ -262,56 +542,103 @@ mod tests { } } - #[tokio::test] - async fn test_ping_pong_basic_send_recv() { - let (exchange, send_rx, _pong_tx) = create_mock_exchange(2); - - // First send should succeed - assert!(exchange.try_send(make_flight_data(10)).unwrap().is_none()); - - // Data should arrive on send_rx - let received = send_rx.recv().await.unwrap(); - assert_eq!(received.data_body.len(), 10); - - // Simulate pong by clearing in_flight - exchange.ready_send(); + #[test] + fn test_ping_pong_in_flight_returns_data() { + let (exchange, _send_rx, _pong_tx) = create_mock_exchange(2); - // Second send should succeed - assert!(exchange.try_send(make_flight_data(20)).unwrap().is_none()); - let received = send_rx.recv().await.unwrap(); - assert_eq!(received.data_body.len(), 20); + assert!(exchange.try_send(make_flight_data(1)).unwrap().is_none()); + let returned = exchange.try_send(make_flight_data(2)).unwrap(); + assert_eq!(returned.unwrap().data_body.len(), 2); } #[tokio::test] - async fn test_ping_pong_in_flight_returns_data() { - let (exchange, send_rx, _pong_tx) = create_mock_exchange(2); + async fn test_reconnect_resends_unacknowledged_sequence() { + let attempts = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let received_sequences = Arc::new(Mutex::new(Vec::new())); + + let connector: DoExchangeConnector = { + let attempts = attempts.clone(); + let received_sequences = received_sequences.clone(); + Arc::new(move || { + let attempt = attempts.fetch_add(1, Ordering::SeqCst); + let received_sequences = received_sequences.clone(); + Box::pin(async move { + let (request_tx, request_rx) = async_channel::bounded(1); + let (response_tx, response_rx) = async_channel::bounded(1); + databend_common_base::runtime::spawn(async move { + let request = request_rx.recv().await.unwrap(); + let DoExchangeFrame::Data { sequence, .. } = + DoExchangeFrame::try_from(request).unwrap() + else { + panic!("expected DATA"); + }; + received_sequences.lock().push(sequence); + + if attempt == 0 { + response_tx + .send(Err(Status::unavailable("connection lost"))) + .await + .unwrap(); + } else { + response_tx + .send(Ok(DoExchangeFrame::Ack { sequence }.into())) + .await + .unwrap(); + } + }); + + Ok(DoExchangeConnection { + request_tx, + response_stream: Box::pin(response_rx), + }) + }) + }) + }; - // First send succeeds - assert!(exchange.try_send(make_flight_data(1)).unwrap().is_none()); + let exchange = PingPongExchange::connect( + 1, + connector, + DoExchangeRetryPolicy { + max_retries: 1, + retry_interval: Duration::ZERO, + }, + "query-node-0", + "query-node-1", + ) + .await + .unwrap(); - // Second send returns data back (in-flight) - let returned = exchange.try_send(make_flight_data(2)).unwrap(); - assert!(returned.is_some()); - assert_eq!(returned.unwrap().data_body.len(), 2); + struct Callback(Arc); + impl PingPongCallback for Callback { + fn has_pending(&self) -> bool { + false + } - // Drain the channel and simulate pong - let _ = send_rx.recv().await.unwrap(); - exchange.ready_send(); + fn on_response(&self, response: PingPongResponse) { + assert!(response.data.is_ok()); + self.0.store(true, Ordering::SeqCst); + } - // Now send should succeed again - assert!(exchange.try_send(make_flight_data(3)).unwrap().is_none()); - } + fn on_receiver_finished(&self) { + panic!("receiver must remain active"); + } + } - #[tokio::test] - async fn test_ping_pong_closed_returns_error() { - let (exchange, send_rx, _pong_tx) = create_mock_exchange(1); + let completed = Arc::new(AtomicBool::new(false)); + let runtime = Runtime::with_worker_threads(2, None).unwrap(); + exchange + .start(Arc::new(Callback(completed.clone())), &runtime) + .unwrap(); + exchange.try_send(make_flight_data(1)).unwrap(); - // Drop the receiver to close the channel - drop(send_rx); + tokio::time::timeout(Duration::from_secs(2), async { + while !completed.load(Ordering::SeqCst) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); - let result = exchange.try_send(make_flight_data(1)); - let status = result.expect_err("closed request channel should abort the exchange"); - assert_eq!(status.code(), tonic::Code::Aborted); - assert_eq!(status.message(), REMOTE_FLIGHT_CHANNEL_CLOSED_MESSAGE); + assert_eq!(*received_sequences.lock(), vec![0, 0]); } }