diff --git a/.github/actions/test_query_flight_reconnect_tpch/action.yml b/.github/actions/test_query_flight_reconnect_tpch/action.yml new file mode 100644 index 00000000000..22219e3b7f0 --- /dev/null +++ b/.github/actions/test_query_flight_reconnect_tpch/action.yml @@ -0,0 +1,14 @@ +name: "Test Query Flight Reconnect with TPC-H" +description: "Verify New Flight reconnects after a random temporary network partition" +runs: + using: "composite" + steps: + - uses: ./.github/actions/setup_test + with: + artifacts: meta,query,sqllogictests + + - uses: ./.github/actions/setup_minio + + - name: Run TPC-H Flight reconnect test + shell: bash + run: ./scripts/ci/ci-run-query-flight-reconnect-tpch.sh diff --git a/.github/workflows/reuse.linux.yml b/.github/workflows/reuse.linux.yml index 5197fc7488e..d2cca5d375a 100644 --- a/.github/workflows/reuse.linux.yml +++ b/.github/workflows/reuse.linux.yml @@ -423,6 +423,24 @@ jobs: with: name: test-stateful-cluster-linux + test_query_flight_reconnect_tpch: + needs: [build, check] + runs-on: + - self-hosted + - "${{ inputs.runner_arch }}" + - Linux + - 2c + - "${{ inputs.runner_provider }}" + steps: + - uses: actions/checkout@v6 + - uses: ./.github/actions/test_query_flight_reconnect_tpch + timeout-minutes: 20 + - name: Upload failure + if: failure() || cancelled() + uses: ./.github/actions/artifact_failure + with: + name: test-query-flight-reconnect-tpch + test_stateful_large_data: if: contains(github.event.pull_request.labels.*.name, 'ci-largedata') needs: [build, check] diff --git a/scripts/ci/ci-run-query-flight-reconnect-tpch.sh b/scripts/ci/ci-run-query-flight-reconnect-tpch.sh new file mode 100755 index 00000000000..823f466a9ce --- /dev/null +++ b/scripts/ci/ci-run-query-flight-reconnect-tpch.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# Copyright 2020-2026 The Databend Authors. +# SPDX-License-Identifier: Apache-2.0. + +set -euo pipefail + +export STORAGE_TYPE=s3 +export STORAGE_S3_BUCKET=testbucket +export STORAGE_S3_ROOT=admin +export STORAGE_S3_ENDPOINT_URL=http://127.0.0.1:9900 +export STORAGE_S3_ACCESS_KEY_ID=minioadmin +export STORAGE_S3_SECRET_ACCESS_KEY=minioadmin +export STORAGE_ALLOW_INSECURE=true + +readonly BUILD_PROFILE="${BUILD_PROFILE:-debug}" +readonly SCRIPT_PATH="$(cd "$(dirname "$0")" >/dev/null 2>&1 && pwd)" +readonly REPO_PATH="$(cd "$SCRIPT_PATH/../.." >/dev/null 2>&1 && pwd)" +readonly TPCH_DATA_PATH=/tmp/tpch_1 + +python3 -m pip install --quiet mysql-connector-python requests +sudo apt-get update -yq +sudo apt-get install -yq iproute2 iptables lsof + +cd "$REPO_PATH" +./scripts/ci/deploy/databend-query-cluster-3-nodes.sh + +rm -rf -- "$TPCH_DATA_PATH" +bash tests/sqllogictests/scripts/prepare_tpch_data.sh tpch_test 1 + +python3 tests/query-flight-reconnect/test_tpch_reconnect.py \ + --sqllogictests "target/${BUILD_PROFILE}/databend-sqllogictests" \ + --tpch-suite tests/sqllogictests/suites/tpch/queries.test \ + --operation-log .databend/tpch-flight-reconnect/operations.log \ + --repo-dir "$REPO_PATH" diff --git a/src/query/service/src/schedulers/fragments/query_fragment_actions.rs b/src/query/service/src/schedulers/fragments/query_fragment_actions.rs index 79564e26c56..94c5c520bfc 100644 --- a/src/query/service/src/schedulers/fragments/query_fragment_actions.rs +++ b/src/query/service/src/schedulers/fragments/query_fragment_actions.rs @@ -230,6 +230,8 @@ impl QueryFragmentsActions { /// unique map(target, map(source, vec(fragment_id))) fn fragments_connections(&self, builder: &mut DataflowDiagramBuilder) -> Result<()> { + let new_flight = self.ctx.get_settings().get_enable_experiment_new_flight()?; + for fragment_actions in &self.fragments_actions { if let Some(exchange) = &fragment_actions.data_exchange { let destinations = exchange.get_destinations(); @@ -249,7 +251,11 @@ impl QueryFragmentsActions { )?; } else { for channel in exchange.get_channels(destination) { - builder.add_data_edge(&source, destination, &channel)?; + if new_flight && matches!(exchange, DataExchange::Merge(_)) { + builder.add_merge_edge(&source, destination, &channel)?; + } else { + builder.add_data_edge(&source, destination, &channel)?; + } } } } diff --git a/src/query/service/src/servers/flight/flight_client.rs b/src/query/service/src/servers/flight/flight_client.rs index 540a86a2043..80dc19dadf4 100644 --- a/src/query/service/src/servers/flight/flight_client.rs +++ b/src/query/service/src/servers/flight/flight_client.rs @@ -20,8 +20,6 @@ use arrow_flight::FlightData; use arrow_flight::Ticket; use arrow_flight::flight_service_client::FlightServiceClient; use async_channel::Receiver; -use async_channel::Sender; -use databend_common_base::runtime::drop_guard; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use fastrace::Span; @@ -33,7 +31,6 @@ use serde::Deserialize; use serde::Serialize; use tokio::time::Duration; use tonic::Request; -use tonic::Status; use tonic::Streaming; use tonic::metadata::AsciiMetadataKey; use tonic::metadata::AsciiMetadataValue; @@ -41,14 +38,121 @@ use tonic::transport::channel::Channel; use crate::pipelines::executor::WatchNotify; use crate::servers::flight::request_builder::RequestBuilder; -use crate::servers::flight::v1::packets::DataPacket; +use crate::servers::flight::v1::transport::legacy::LegacyInbound; /// 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 num_threads: usize, + /// Present only for New Flight streams. Absent keeps the existing Flight wire shape, so a + /// node running either transport can parse the other's requests. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub new_flight: Option, +} + +/// Identifies a New Flight logical stream and how long its receiver waits for a replacement +/// connection. Keeping these together makes a partially specified attachment unrepresentable. +#[derive(Clone, Serialize, Deserialize)] +pub struct NewFlightAttachment { + pub source_id: String, + pub receiver_lease_secs: u64, + pub stream: NewFlightStream, +} + +#[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum NewFlightStream { + Fragment, + Exchange, + Merge, + Statistics, +} + +impl DoExchangeParams { + pub fn create(query_id: String, exchange_id: String, num_threads: usize) -> Self { + Self { + query_id, + exchange_id, + num_threads, + new_flight: None, + } + } + + pub fn new_flight_fragment( + query_id: String, + exchange_id: String, + source_id: String, + num_threads: usize, + receiver_lease_secs: u64, + ) -> Self { + Self { + query_id, + exchange_id, + num_threads, + new_flight: Some(NewFlightAttachment { + source_id, + receiver_lease_secs, + stream: NewFlightStream::Fragment, + }), + } + } + + pub fn new_flight_merge( + query_id: String, + exchange_id: String, + source_id: String, + receiver_lease_secs: u64, + ) -> Self { + Self { + query_id, + exchange_id, + num_threads: 1, + new_flight: Some(NewFlightAttachment { + source_id, + receiver_lease_secs, + stream: NewFlightStream::Merge, + }), + } + } + + pub fn new_flight_exchange( + query_id: String, + exchange_id: String, + source_id: String, + num_threads: usize, + receiver_lease_secs: u64, + ) -> Self { + Self { + query_id, + exchange_id, + num_threads, + new_flight: Some(NewFlightAttachment { + source_id, + receiver_lease_secs, + stream: NewFlightStream::Exchange, + }), + } + } + + pub fn new_flight_statistics( + query_id: String, + source_id: String, + receiver_lease_secs: u64, + ) -> Self { + Self { + // Statistics is a single stream per source, so it needs no exchange id or thread fan-out. + query_id, + exchange_id: String::new(), + num_threads: 1, + new_flight: Some(NewFlightAttachment { + source_id, + receiver_lease_secs, + stream: NewFlightStream::Statistics, + }), + } + } } pub struct FlightClient { @@ -196,7 +300,7 @@ impl FlightClient { &mut self, query_id: &str, target: &str, - ) -> Result { + ) -> Result { let streaming = self .get_streaming( RequestBuilder::create(Ticket::default()) @@ -212,12 +316,12 @@ impl FlightClient { self.local_node_id.clone(), self.remote_node_id.clone(), ); - Ok(FlightExchange::create_receiver(notify, rx)) + Ok(LegacyInbound::create(notify, rx)) } #[async_backtrace::framed] #[fastrace::trace] - pub async fn do_get(&mut self, query_id: &str, channel_id: &str) -> Result { + pub async fn do_get(&mut self, query_id: &str, channel_id: &str) -> Result { let request = RequestBuilder::create(Ticket::default()) .with_metadata("x-type", "exchange_fragment")? .with_metadata("x-query-id", query_id)? @@ -232,7 +336,7 @@ impl FlightClient { self.local_node_id.clone(), self.remote_node_id.clone(), ); - Ok(FlightExchange::create_receiver(notify, rx)) + Ok(LegacyInbound::create(notify, rx)) } fn streaming_receiver( @@ -331,109 +435,3 @@ impl FlightClient { }) } } - -pub struct FlightReceiver { - notify: Arc, - rx: Receiver>, -} - -impl Drop for FlightReceiver { - fn drop(&mut self) { - drop_guard(move || { - self.close(); - }) - } -} - -impl FlightReceiver { - pub fn create(rx: Receiver>) -> FlightReceiver { - FlightReceiver { - rx, - notify: Arc::new(WatchNotify::new()), - } - } - - #[async_backtrace::framed] - pub async fn recv(&self) -> Result> { - match self.rx.recv().await { - Err(_) => Ok(None), - Ok(Err(error)) => Err(error), - Ok(Ok(message)) => Ok(Some(DataPacket::try_from(message)?)), - } - } - - pub fn close(&self) { - self.rx.close(); - self.notify.notify_waiters(); - } -} - -pub struct FlightSender { - tx: Sender>, -} - -impl FlightSender { - pub fn create(tx: Sender>) -> FlightSender { - FlightSender { tx } - } - - pub fn is_closed(&self) -> bool { - self.tx.is_closed() - } - - #[async_backtrace::framed] - pub async fn send(&self, data: DataPacket) -> Result<()> { - if let Err(_cause) = self.tx.send(Ok(FlightData::try_from(data)?)).await { - return Err(ErrorCode::AbortedQuery( - "Aborted query, because the remote flight channel is closed.", - )); - } - - Ok(()) - } - - pub fn close(&self) { - self.tx.close(); - } -} - -pub enum FlightExchange { - Dummy, - Receiver { - notify: Arc, - receiver: Receiver>, - }, - Sender(Sender>), -} - -impl FlightExchange { - pub fn create_sender( - sender: Sender>, - ) -> FlightExchange { - FlightExchange::Sender(sender) - } - - pub fn create_receiver( - notify: Arc, - receiver: Receiver>, - ) -> FlightExchange { - FlightExchange::Receiver { notify, receiver } - } - - pub fn convert_to_sender(self) -> FlightSender { - match self { - FlightExchange::Sender(tx) => FlightSender { tx }, - _ => unreachable!(), - } - } - - pub fn convert_to_receiver(self) -> FlightReceiver { - match self { - FlightExchange::Receiver { notify, receiver } => FlightReceiver { - notify, - rx: receiver, - }, - _ => unreachable!(), - } - } -} diff --git a/src/query/service/src/servers/flight/mod.rs b/src/query/service/src/servers/flight/mod.rs index f6ce7e521c2..e1631e4e285 100644 --- a/src/query/service/src/servers/flight/mod.rs +++ b/src/query/service/src/servers/flight/mod.rs @@ -20,9 +20,8 @@ pub mod v1; pub use flight_client::DoExchangeParams; pub use flight_client::FlightClient; -pub use flight_client::FlightExchange; pub(crate) use flight_client::FlightOperation; -pub use flight_client::FlightReceiver; -pub use flight_client::FlightSender; +pub use flight_client::NewFlightAttachment; +pub use flight_client::NewFlightStream; pub(crate) use flight_client::add_flight_error_context; pub use flight_service::FlightService; diff --git a/src/query/service/src/servers/flight/v1/exchange/broadcast_recv_transform.rs b/src/query/service/src/servers/flight/v1/exchange/broadcast_recv_transform.rs index 143f606ea3c..41665f1992c 100644 --- a/src/query/service/src/servers/flight/v1/exchange/broadcast_recv_transform.rs +++ b/src/query/service/src/servers/flight/v1/exchange/broadcast_recv_transform.rs @@ -29,11 +29,11 @@ use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::core::SyncTaskHandle; +use databend_common_pipeline::core::SyncTaskSet; use petgraph::graph::NodeIndex; -use crate::servers::flight::v1::network::InboundChannel; -use crate::servers::flight::v1::network::SyncTaskHandle; -use crate::servers::flight::v1::network::SyncTaskSet; +use crate::servers::flight::v1::exchange::exchange_packet_receiver::InboundChannel; pub struct ExchangeRecvTransform { input: Arc, diff --git a/src/query/service/src/servers/flight/v1/exchange/broadcast_send_transform.rs b/src/query/service/src/servers/flight/v1/exchange/broadcast_send_transform.rs index 96f7a98fe6f..833cb5c0e92 100644 --- a/src/query/service/src/servers/flight/v1/exchange/broadcast_send_transform.rs +++ b/src/query/service/src/servers/flight/v1/exchange/broadcast_send_transform.rs @@ -25,12 +25,12 @@ use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::core::SyncTaskSet; use petgraph::prelude::NodeIndex; use super::outbound_send_channels::OutboundSendChannels; use super::outbound_send_channels::OutboundSendHandle; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::SyncTaskSet; +use super::outbound_send_channels::SharedOutboundChannels; pub struct BroadcastSendTransform { id: NodeIndex, @@ -47,7 +47,7 @@ impl BroadcastSendTransform { pub fn create_item( worker_id: usize, local_pos: usize, - channels: Vec>, + channels: SharedOutboundChannels, waker: Arc, ) -> PipeItem { let input = InputPort::create(); @@ -70,6 +70,12 @@ impl BroadcastSendTransform { fn no_active_downstream(&self) -> bool { self.output.is_finished() && self.channels.all_closed_except(self.local_pos) } + + fn finish_processor(&mut self) -> Result { + self.input.finish(); + self.output.finish(); + self.channels.poll_complete_event(&self.tasks, self.id) + } } impl Processor for BroadcastSendTransform { @@ -88,8 +94,7 @@ impl Processor for BroadcastSendTransform { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -100,8 +105,7 @@ impl Processor for BroadcastSendTransform { } if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } if self.input.has_data() { @@ -140,8 +144,7 @@ impl Processor for BroadcastSendTransform { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -156,12 +159,8 @@ impl Processor for BroadcastSendTransform { } } - // Input finished → close channels if self.input.is_finished() { - self.output.finish(); - - self.channels.close_all(); - return Ok(Event::Finished); + return self.finish_processor(); } self.input.set_need_data(); @@ -170,12 +169,10 @@ impl Processor for BroadcastSendTransform { fn details_status(&self) -> Option { Some(format!( - "handle_pending={}, local_pos={}, closed_channels={}/{}, closed={:?}", + "handle_pending={}, local_pos={}, closed_channels={}", self.handle.is_some(), self.local_pos, - self.channels.closed_count(), - self.channels.len(), - self.channels.closed_status(), + self.channels.closed_summary(), )) } 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 a0e80fc40f1..3b1d373ce38 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 @@ -40,11 +40,13 @@ use databend_common_pipeline::core::always_callback; use databend_common_pipeline::core::basic_callback; use databend_common_settings::FlightKeepAliveParams; use fastrace::prelude::*; +use futures::StreamExt; use log::warn; use parking_lot::Mutex; use parking_lot::ReentrantMutex; use petgraph::Direction; use petgraph::prelude::EdgeRef; +use tokio::sync::Semaphore; use tokio::sync::oneshot; use tonic::Status; @@ -55,6 +57,8 @@ use super::exchange_params::MergeExchangeParams; use super::exchange_params::ShuffleExchangeParams; use super::exchange_sink::ExchangeSink; use super::exchange_transform::ExchangeTransform; +use super::packet_receiver::PacketReceiver; +use super::reliable_delivery::StatisticsDelivery; use super::statistics_receiver::StatisticsReceiver; use super::statistics_sender::StatisticsSender; use crate::clusters::ClusterHelper; @@ -69,10 +73,8 @@ use crate::pipelines::executor::PlanNodeMemoryUsage; use crate::schedulers::QueryFragmentsActions; use crate::servers::flight::DoExchangeParams; use crate::servers::flight::FlightClient; -use crate::servers::flight::FlightExchange; use crate::servers::flight::FlightOperation; -use crate::servers::flight::FlightReceiver; -use crate::servers::flight::FlightSender; +use crate::servers::flight::NewFlightStream; use crate::servers::flight::add_flight_error_context; use crate::servers::flight::keep_alive::build_keep_alive_config; use crate::servers::flight::v1::actions::INIT_QUERY_FRAGMENTS; @@ -81,27 +83,53 @@ 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::NetworkInboundChannelSet; -use crate::servers::flight::v1::network::NetworkInboundSender; -use crate::servers::flight::v1::network::PingPongExchange; +use crate::servers::flight::v1::exchange::exchange_packet_receiver::ExchangePacketReceiverSet; +use crate::servers::flight::v1::exchange::exchange_packet_receiver::NetworkInboundSender; use crate::servers::flight::v1::packets::Edge; use crate::servers::flight::v1::packets::QueryEnv; use crate::servers::flight::v1::packets::QueryFragment; use crate::servers::flight::v1::packets::QueryFragments; +use crate::servers::flight::v1::transport::InboundDelivery; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::legacy::LegacyInbound; +use crate::servers::flight::v1::transport::legacy::LegacyOutbound; +use crate::servers::flight::v1::transport::legacy::PingPongExchange; +use crate::servers::flight::v1::transport::reliable::DoExchangeConnector; +use crate::servers::flight::v1::transport::reliable::DoExchangeTransport; +use crate::servers::flight::v1::transport::reliable::FlightReconnectPolicy; +use crate::servers::flight::v1::transport::reliable::PendingReliableOutbound; +use crate::servers::flight::v1::transport::reliable::ReliableInboundConnection; +use crate::servers::flight::v1::transport::reliable::ReliableInboundSource; use crate::sessions::QueryContext; use crate::sessions::TableContextCluster; use crate::sessions::TableContextPerf; use crate::sessions::TableContextQueryIdentity; use crate::sessions::TableContextSettings; +/// Inbound queue quota per do_exchange connection. +// TODO: get max_bytes_per_connection from query settings +const MAX_INBOUND_BYTES_PER_CONNECTION: usize = 20 * 1024 * 1024; + +/// Queued statistics packets allowed before the source blocks. +const STATISTICS_QUEUE_CAPACITY: usize = 8; + enum QueryExchange { Fragment { channel: String, - exchange: FlightExchange, + exchange: LegacyInbound, }, Statistics { source: String, - exchange: FlightExchange, + exchange: LegacyInbound, + }, + StatisticsSender { + target: String, + sender: OutboundStreamRef, + }, + NewFlightFragmentOutbound { + exchange_id: String, + target_id: String, + outbound: PendingReliableOutbound, }, PingPong { exchange_id: String, @@ -110,6 +138,30 @@ enum QueryExchange { }, } +fn create_do_exchange_connector( + target_id: String, + address: String, + use_current_rt: bool, + keep_alive: FlightKeepAliveParams, + params: DoExchangeParams, +) -> DoExchangeConnector { + Arc::new(move || { + let target_id = target_id.clone(); + let address = address.clone(); + let params = params.clone(); + Box::pin(async move { + let mut client = + create_flight_client(target_id, address, use_current_rt, keep_alive).await?; + let (send_tx, send_rx) = async_channel::bounded(1); + let response_stream = client.do_exchange(send_rx, params).await?; + Ok(DoExchangeTransport { + send_tx, + response_stream: response_stream.boxed(), + }) + }) + }) +} + async fn create_flight_client( remote_node_id: String, address: String, @@ -299,9 +351,19 @@ impl DataExchangeManager { None => env.settings.clone(), }; let keep_alive = settings.get_flight_keep_alive_params()?; + let new_flight = FlightReconnectPolicy::from_settings(&settings)?; + log::info!( + "Flight transport selected: query_id={}, mode={}", + env.query_id, + if new_flight.is_some() { + "new flight" + } else { + "legacy flight" + } + ); let mut request_exchanges = HashMap::new(); - let mut targets_exchanges = HashMap::>::new(); + let mut targets_exchanges = HashMap::>::new(); for index in env.dataflow_diagram.node_indices() { if env.dataflow_diagram[index].id == config.query.node_id { @@ -311,7 +373,6 @@ impl DataExchangeManager { >, > = vec![]; - // Process incoming edges: do_get for Fragment, skip ExchangeFragment let incoming_edges = env .dataflow_diagram .edges_directed(index, Direction::Incoming); @@ -320,14 +381,13 @@ impl DataExchangeManager { let source = env.dataflow_diagram[edge.source()].clone(); let target = env.dataflow_diagram[edge.target()].clone(); let edge = edge.weight().clone(); - let query_id = env.query_id.clone(); let address = source.flight_address.clone(); let source_id = source.id.clone(); - let keep_alive_params = keep_alive; - match edge { - Edge::Fragment(channel) => { + + match (new_flight, edge) { + (None, Edge::Fragment(channel)) | (None, Edge::Merge(channel)) => { flight_exchanges.push(Box::pin(async move { let mut flight_client = Self::create_client( &source_id, @@ -342,7 +402,7 @@ impl DataExchangeManager { }) })); } - Edge::Statistics => { + (None, Edge::Statistics) => { flight_exchanges.push(Box::pin(async move { let mut flight_client = Self::create_client( &source_id, @@ -359,14 +419,10 @@ impl DataExchangeManager { }) })); } - Edge::ExchangeFragment { .. } => { - // Skip: remote sender will call do_exchange on us, - // handled by handle_do_exchange → NetworkInboundSender - } + _ => {} } } - // Process outgoing edges: do_exchange for ExchangeFragment let outgoing_edges = env .dataflow_diagram .edges_directed(index, Direction::Outgoing); @@ -375,24 +431,71 @@ impl DataExchangeManager { let target = env.dataflow_diagram[edge.target()].clone(); let edge = edge.weight().clone(); - if let Edge::ExchangeFragment { - exchange_id, - channels, - } = edge - { + if let (Some(reconnect), Edge::Statistics) = (new_flight, &edge) { let target_id = target.id.clone(); - let local_node_id = config.query.node_id.clone(); let query_id = env.query_id.clone(); + let source_id = config.query.node_id.clone(); let address = target.flight_address.clone(); let keep_alive_params = keep_alive; - let num_threads = channels.len(); - 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 params = DoExchangeParams::new_flight_statistics( + query_id, + source_id.clone(), + reconnect.receiver_lease_secs(), + ); + let connector = create_do_exchange_connector( + target_id.clone(), + address, + with_cur_rt, + keep_alive_params, + params, + ); + let outbound = PendingReliableOutbound::connect( + 1, + connector, + reconnect, + source_id, + target_id.clone(), + ) + .await?; + let slots = Arc::new(Semaphore::new(STATISTICS_QUEUE_CAPACITY)); + let sender = + Arc::new(outbound.start(slots, None, &GlobalIORuntime::instance())) + as OutboundStreamRef; + Ok::(QueryExchange::StatisticsSender { + target: target_id, + sender, + }) + })); + continue; + } + + let (exchange_id, channels, stream) = match (new_flight, edge) { + (Some(_), Edge::Merge(channel)) => { + (channel.clone(), vec![channel], NewFlightStream::Merge) + } + (Some(_), Edge::Fragment(channel)) => { + (channel.clone(), vec![channel], NewFlightStream::Fragment) + } + ( + _, + Edge::ExchangeFragment { + exchange_id, + channels, + }, + ) => (exchange_id, channels, NewFlightStream::Exchange), + _ => continue, + }; + + let target_id = target.id.clone(); + let local_node_id = config.query.node_id.clone(); + let query_id = env.query_id.clone(); + let address = target.flight_address.clone(); + let keep_alive_params = keep_alive; + let num_threads = channels.len(); + flight_exchanges.push(Box::pin(async move { + match new_flight { + None => { let mut flight_client = create_flight_client( target_id.clone(), address, @@ -400,102 +503,158 @@ impl DataExchangeManager { 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(), - }) + .do_exchange( + send_rx, + DoExchangeParams::create( + query_id, + exchange_id.clone(), + num_threads, + ), + ) .await?; - Ok::<_, ErrorCode>((send_tx, response_stream)) - }?; - - Ok::(QueryExchange::PingPong { - target_id: target_id.clone(), - exchange_id, - exchange: PingPongExchange::from_parts( + 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, + ), + }) + } + Some(reconnect) => { + let params = match stream { + NewFlightStream::Merge => DoExchangeParams::new_flight_merge( + query_id, + exchange_id.clone(), + local_node_id.clone(), + reconnect.receiver_lease_secs(), + ), + NewFlightStream::Fragment => { + DoExchangeParams::new_flight_fragment( + query_id, + exchange_id.clone(), + local_node_id.clone(), + num_threads, + reconnect.receiver_lease_secs(), + ) + } + NewFlightStream::Exchange => { + DoExchangeParams::new_flight_exchange( + query_id, + exchange_id.clone(), + local_node_id.clone(), + num_threads, + reconnect.receiver_lease_secs(), + ) + } + NewFlightStream::Statistics => unreachable!( + "statistics streams are installed before fragment streams" + ), + }; + let connector = create_do_exchange_connector( + target_id.clone(), + address, + with_cur_rt, + keep_alive_params, + params, + ); + let outbound = PendingReliableOutbound::connect( num_threads, - send_tx, - response_stream, + connector, + reconnect, local_node_id, target_id.clone(), - ), - }) - })); - } + ) + .await?; + Ok::( + QueryExchange::NewFlightFragmentOutbound { + target_id, + exchange_id, + outbound, + }, + ) + } + } + })); } let flight_exchanges = futures::future::try_join_all(flight_exchanges).await?; - + let mut new_flight_fragment_outbounds = + HashMap::>::new(); let mut ping_pong_exchanges = HashMap::>::new(); + let mut statistics_senders = HashMap::::new(); for flight_exchange in flight_exchanges { match flight_exchange { QueryExchange::Fragment { channel, exchange } => { - match targets_exchanges.entry(channel) { - Entry::Occupied(mut v) => v.get_mut().push(exchange), - Entry::Vacant(v) => { - v.insert(vec![exchange]); - } - } + targets_exchanges.entry(channel).or_default().push(exchange); } QueryExchange::Statistics { source, exchange } => { request_exchanges.insert(source, exchange); } + QueryExchange::StatisticsSender { target, sender } => { + statistics_senders.insert(target, sender); + } + QueryExchange::NewFlightFragmentOutbound { + exchange_id, + target_id, + outbound, + } => { + new_flight_fragment_outbounds + .entry(exchange_id) + .or_default() + .insert(target_id, outbound); + } QueryExchange::PingPong { exchange_id, exchange, target_id, } => { - match ping_pong_exchanges.entry(exchange_id) { - Entry::Occupied(mut v) => { - v.get_mut().insert(target_id, exchange); - } - Entry::Vacant(v) => { - v.insert(HashMap::from([(target_id, exchange)])); - } - }; + ping_pong_exchanges + .entry(exchange_id) + .or_default() + .insert(target_id, exchange); } - }; + } } let mut query_info = Self::create_info(ctx)?; - if let Some(query_info) = query_info.as_mut() { let query_id = env.query_id.clone(); query_info.remove_leak_query_worker = Some(GlobalIORuntime::instance().spawn(async move { - let _ = tokio::time::sleep(Duration::from_secs(180)).await; + tokio::time::sleep(Duration::from_secs(180)).await; DataExchangeManager::instance().remove_if_leak_query(query_id); })); } let queries_coordinator_guard = self.queries_coordinator.lock(); let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; - - match queries_coordinator.entry(env.query_id.clone()) { - Entry::Occupied(mut v) => { - let query_coordinator = v.get_mut(); - query_coordinator.info = query_info; - query_coordinator.is_request_server = - GlobalConfig::instance().query.node_id == env.request_server_id; - query_coordinator.register_flight_channel_receiver(targets_exchanges)?; - query_coordinator.register_ping_pong_exchanges(ping_pong_exchanges); - query_coordinator.add_statistics_exchanges(request_exchanges)?; - } - Entry::Vacant(v) => { - let query_coordinator = v.insert(QueryCoordinator::create()); - query_coordinator.info = query_info; - query_coordinator.is_request_server = - GlobalConfig::instance().query.node_id == env.request_server_id; - query_coordinator.register_flight_channel_receiver(targets_exchanges)?; - query_coordinator.register_ping_pong_exchanges(ping_pong_exchanges); - query_coordinator.add_statistics_exchanges(request_exchanges)?; + let query_coordinator = queries_coordinator + .entry(env.query_id.clone()) + .or_insert_with(QueryCoordinator::create); + query_coordinator.info = query_info; + query_coordinator.is_request_server = + GlobalConfig::instance().query.node_id == env.request_server_id; + query_coordinator.register_fragment_receivers(targets_exchanges)?; + for (exchange_id, outbounds) in new_flight_fragment_outbounds { + for (target_id, outbound) in outbounds { + query_coordinator.register_new_flight_fragment_outbound( + exchange_id.clone(), + target_id, + outbound, + ); } - }; + } + query_coordinator.register_ping_pong_exchanges(ping_pong_exchanges); + query_coordinator.add_statistics_exchanges(request_exchanges)?; + query_coordinator.add_statistics_senders(statistics_senders)?; return Ok(()); } @@ -700,7 +859,7 @@ impl DataExchangeManager { /// Handle a do_exchange request from a remote node. /// /// Creates a `NetworkInboundSender` for this connection, bound to the - /// `NetworkInboundChannelSet` for the given channel_id. The caller (flight_service) + /// `ExchangePacketReceiverSet` 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( @@ -724,6 +883,33 @@ impl DataExchangeManager { } } + /// Admits one New Flight `do_exchange` connection, creating the query coordinator if this is + /// the first stream to arrive for the query. + #[fastrace::trace] + pub fn handle_new_flight_do_exchange( + &self, + query_id: &str, + channel_id: &str, + source_id: &str, + num_threads: usize, + stream: NewFlightStream, + receiver_lease: Duration, + ) -> Result { + let queries_coordinator_guard = self.queries_coordinator.lock(); + let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; + + queries_coordinator + .entry(query_id.to_string()) + .or_insert_with(QueryCoordinator::create) + .open_new_flight_inbound_connection( + channel_id, + source_id, + num_threads, + stream, + receiver_lease, + ) + } + /// Get the NetworkInboundReceivers for a given query and channel. /// /// Returns one `Arc` per tid, for building @@ -732,7 +918,7 @@ impl DataExchangeManager { &self, query_id: &str, channel_id: &str, - ) -> Result> { + ) -> Result> { let queries_coordinator_guard = self.queries_coordinator.lock(); let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; @@ -743,7 +929,7 @@ impl DataExchangeManager { ))), Some(coordinator) => match coordinator.inbound_channel_sets.get(channel_id) { None => Err(ErrorCode::Internal(format!( - "NetworkInboundChannelSet not found for channel {}", + "ExchangePacketReceiverSet not found for channel {}", channel_id ))), Some(channel_set) => Ok(channel_set.clone()), @@ -758,7 +944,7 @@ impl DataExchangeManager { query_id: &str, channel_id: &str, num_threads: usize, - ) -> Result> { + ) -> Result> { let queries_coordinator_guard = self.queries_coordinator.lock(); let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; @@ -773,10 +959,26 @@ impl DataExchangeManager { } } - /// Take the PingPongExchanges for a given query and channel. - /// - /// Returns the exchanges that were created during init_query_env. - /// The exchanges are removed from the coordinator (taken, not borrowed). + pub fn take_new_flight_fragment_outbounds( + &self, + query_id: &str, + exchange_id: &str, + ) -> Result> { + let queries_coordinator_guard = self.queries_coordinator.lock(); + let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; + + match queries_coordinator.get_mut(query_id) { + None => Err(ErrorCode::Internal(format!( + "Query {} not found in cluster.", + query_id + ))), + Some(coordinator) => Ok(coordinator + .new_flight_fragment_outbounds + .remove(exchange_id) + .unwrap_or_default()), + } + } + pub fn take_ping_pong_exchanges( &self, query_id: &str, @@ -891,21 +1093,21 @@ impl DataExchangeManager { match queries_coordinator.get_mut(&query_id) { None => Err(ErrorCode::Internal("Query not exists.")), Some(query_coordinator) => { - if !query_coordinator.flight_data_senders.is_empty() { + if !query_coordinator.fragment_outbounds.is_empty() { unreachable!( "query_coordinator.fragment_senders is not empty: {:?}", query_coordinator - .flight_data_senders + .fragment_outbounds .keys() .collect::>() ); } - if !query_coordinator.flight_data_receivers.is_empty() { + if !query_coordinator.fragment_receivers.is_empty() { unreachable!( "query_coordinator.fragment_receivers is not empty: {:?}", query_coordinator - .flight_data_receivers + .fragment_receivers .keys() .collect::>() ); @@ -932,22 +1134,31 @@ impl DataExchangeManager { .extend(sub_build_res.sources_pipelines); } - let exchanges = std::mem::take(&mut query_coordinator.statistics_exchanges); - let statistics_receiver = StatisticsReceiver::spawn_receiver(&ctx, exchanges)?; + let receivers = std::mem::take(&mut query_coordinator.statistics_receivers); + let statistics_receiver = StatisticsReceiver::spawn_receiver(&ctx, receivers)?; let statistics_receiver: Mutex = Mutex::new(statistics_receiver); - // Interrupting the execution of finished callback if network error + // Keep the query coordinator, including reliable inbound sources, alive until all + // statistics streams have reached a terminal state. A statistics sender may still + // be reconnecting after its pipeline finishes; removing the coordinator first + // would admit the replay as a new logical source with sequence zero. build_res.main_pipeline.set_on_finished(basic_callback( move |info: &ExecutionInfo| { let query_id = ctx.get_id(); let mut statistics_receiver = statistics_receiver.lock(); statistics_receiver.shutdown(info.res.is_err()); + let result = statistics_receiver.wait_shutdown(); + let cause = info + .res + .clone() + .err() + .or_else(|| result.as_ref().err().cloned()); ctx.get_exchange_manager() - .on_finished_query(&query_id, info.res.clone().err()); - statistics_receiver.wait_shutdown() + .on_finished_query(&query_id, cause); + result }, )); @@ -964,28 +1175,29 @@ impl DataExchangeManager { } } - pub fn get_flight_sender( + pub fn take_fragment_outbound_streams( &self, params: &ExchangeParams, - ) -> Result> { + ) -> Result)>> { let queries_coordinator_guard = self.queries_coordinator.lock(); let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; match queries_coordinator.get_mut(¶ms.get_query_id()) { None => Err(ErrorCode::Internal("Query not exists.")), - Some(coordinator) => params.take_flight_sender(&mut coordinator.flight_data_senders), + Some(coordinator) => params.take_outbound_streams(&mut coordinator.fragment_outbounds), } } - pub fn get_flight_receiver(&self, params: &ExchangeParams) -> Result> { + pub(super) fn take_packet_receivers( + &self, + params: &ExchangeParams, + ) -> Result> { let queries_coordinator_guard = self.queries_coordinator.lock(); let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() }; match queries_coordinator.get_mut(¶ms.get_query_id()) { None => Err(ErrorCode::Internal("Query not exists.")), - Some(coordinator) => { - params.take_flight_receiver(&mut coordinator.flight_data_receivers) - } + Some(coordinator) => params.take_packet_receivers(&mut coordinator.fragment_receivers), } } @@ -1031,10 +1243,15 @@ pub(crate) struct QueryCoordinator { /// so execute_pipeline() must not start a second one. is_request_server: bool, - statistics_exchanges: HashMap, - flight_data_senders: HashMap>, - flight_data_receivers: HashMap>, - inbound_channel_sets: HashMap>, + statistics_senders: HashMap, + statistics_receivers: HashMap, + fragment_outbounds: HashMap>, + fragment_receivers: HashMap>, + inbound_channel_sets: HashMap>, + /// Logical inbound streams, keyed by the channel id and source node that opened them. + /// Statistics streams use an empty channel id, since there is one per source node. + new_flight_inbound_sources: HashMap<(String, String), Arc>, + new_flight_fragment_outbounds: HashMap>, ping_pong_exchanges: HashMap>, } @@ -1043,11 +1260,14 @@ impl QueryCoordinator { QueryCoordinator { info: None, is_request_server: false, - flight_data_senders: HashMap::new(), - flight_data_receivers: HashMap::new(), - statistics_exchanges: HashMap::new(), + fragment_outbounds: HashMap::new(), + fragment_receivers: HashMap::new(), + statistics_senders: HashMap::new(), + statistics_receivers: HashMap::new(), fragments_coordinator: HashMap::new(), inbound_channel_sets: HashMap::new(), + new_flight_inbound_sources: HashMap::new(), + new_flight_fragment_outbounds: HashMap::new(), ping_pong_exchanges: HashMap::new(), } } @@ -1058,8 +1278,8 @@ impl QueryCoordinator { ) -> Result>> { let (tx, rx) = async_channel::bounded(8); match self - .statistics_exchanges - .insert(target, FlightExchange::create_sender(tx)) + .statistics_senders + .insert(target, LegacyOutbound::create(tx)) { None => Ok(rx), Some(_) => Err(ErrorCode::Internal( @@ -1070,10 +1290,14 @@ impl QueryCoordinator { pub fn add_statistics_exchanges( &mut self, - exchanges: HashMap, + exchanges: HashMap, ) -> Result<()> { - for (source, exchange) in exchanges.into_iter() { - if self.statistics_exchanges.insert(source, exchange).is_some() { + for (source, exchange) in exchanges { + if self + .statistics_receivers + .insert(source, PacketReceiver::from_legacy(exchange)) + .is_some() + { return Err(ErrorCode::Internal( "Internal error, statistics exchange can only have one.", )); @@ -1083,42 +1307,52 @@ impl QueryCoordinator { Ok(()) } + pub fn add_statistics_senders( + &mut self, + senders: HashMap, + ) -> Result<()> { + for (target, sender) in senders { + if self.statistics_senders.insert(target, sender).is_some() { + return Err(ErrorCode::Internal( + "Internal error, statistics exchange can only have one.", + )); + } + } + Ok(()) + } + pub fn register_flight_channel_sender( &mut self, channel_id: String, ) -> Result>> { let (tx, rx) = async_channel::bounded(8); - match self.flight_data_senders.entry(channel_id) { + match self.fragment_outbounds.entry(channel_id) { Entry::Occupied(mut v) => { - v.get_mut() - .push(FlightExchange::create_sender(tx).convert_to_sender()); + v.get_mut().push(LegacyOutbound::create(tx)); } Entry::Vacant(v) => { - v.insert(vec![FlightExchange::create_sender(tx).convert_to_sender()]); + v.insert(vec![LegacyOutbound::create(tx)]); } } Ok(rx) } - pub fn register_flight_channel_receiver( + pub fn register_fragment_receivers( &mut self, - channels: HashMap>, + channels: HashMap>, ) -> Result<()> { for (id, exchanges) in channels.into_iter() { - match self.flight_data_receivers.entry(id) { + match self.fragment_receivers.entry(id) { Entry::Occupied(mut v) => { - v.get_mut().extend( - exchanges - .into_iter() - .map(FlightExchange::convert_to_receiver), - ); + v.get_mut() + .extend(exchanges.into_iter().map(PacketReceiver::from_legacy)); } Entry::Vacant(v) => { v.insert( exchanges .into_iter() - .map(FlightExchange::convert_to_receiver) + .map(PacketReceiver::from_legacy) .collect(), ); } @@ -1128,6 +1362,18 @@ impl QueryCoordinator { Ok(()) } + fn register_new_flight_fragment_outbound( + &mut self, + exchange_id: String, + target_id: String, + outbound: PendingReliableOutbound, + ) { + self.new_flight_fragment_outbounds + .entry(exchange_id) + .or_default() + .insert(target_id, outbound); + } + pub fn register_ping_pong_exchanges( &mut self, exchanges: HashMap>, @@ -1154,12 +1400,111 @@ impl QueryCoordinator { ) -> Result { let channel_set = self.get_or_create_inbound_channel_set(channel_id, num_threads)?; - // TODO: get max_bytes_per_connection from query settings - let max_bytes_per_connection = 20 * 1024 * 1024; // 20MB - Ok(NetworkInboundSender::new( &channel_set, - max_bytes_per_connection, + MAX_INBOUND_BYTES_PER_CONNECTION, + )) + } + + /// Opens (or re-attaches to) the logical inbound stream for one source node. + /// + /// A reconnecting source finds the existing stream and resumes it, so the delivery target is + /// created only on first attach. + fn open_new_flight_inbound_connection( + &mut self, + channel_id: &str, + source_id: &str, + num_threads: usize, + stream: NewFlightStream, + receiver_lease: Duration, + ) -> Result { + let key = (channel_id.to_string(), source_id.to_string()); + let source = match self.new_flight_inbound_sources.entry(key) { + Entry::Occupied(entry) => entry.get().clone(), + Entry::Vacant(entry) => { + let delivery = match stream { + NewFlightStream::Fragment => { + if num_threads != 1 { + return Err(ErrorCode::Internal(format!( + "New Flight fragment stream {} has {} lanes, expected 1", + channel_id, num_threads + ))); + } + let channel_set = Arc::new(ExchangePacketReceiverSet::new(1)); + let receiver = channel_set.receivers[0].clone(); + self.fragment_receivers + .entry(channel_id.to_string()) + .or_default() + .push(PacketReceiver::from_inbound_queue(receiver)); + Arc::new(NetworkInboundSender::new( + &channel_set, + MAX_INBOUND_BYTES_PER_CONNECTION, + )) as Arc + } + NewFlightStream::Exchange => { + let channel_set = Self::inbound_channel_set( + &mut self.inbound_channel_sets, + channel_id, + num_threads, + )?; + Arc::new(NetworkInboundSender::new( + &channel_set, + MAX_INBOUND_BYTES_PER_CONNECTION, + )) as Arc + } + NewFlightStream::Merge => { + if num_threads != 1 { + return Err(ErrorCode::Internal(format!( + "New Flight merge stream {} has {} lanes, expected 1", + channel_id, num_threads + ))); + } + let channel_set = Arc::new(ExchangePacketReceiverSet::new(1)); + let receiver = channel_set.receivers[0].clone(); + self.fragment_receivers + .entry(channel_id.to_string()) + .or_default() + .push(PacketReceiver::from_inbound_queue(receiver)); + Arc::new(NetworkInboundSender::new( + &channel_set, + MAX_INBOUND_BYTES_PER_CONNECTION, + )) as Arc + } + NewFlightStream::Statistics => { + let (delivery, receiver) = + StatisticsDelivery::create(STATISTICS_QUEUE_CAPACITY); + if self + .statistics_receivers + .insert( + source_id.to_string(), + PacketReceiver::from_result_queue(receiver), + ) + .is_some() + { + return Err(ErrorCode::Internal( + "statistics exchange source was admitted twice", + )); + } + delivery + } + }; + + entry + .insert(Arc::new(ReliableInboundSource::new( + delivery, + receiver_lease, + format!("channel_id={}, source_id={}", channel_id, source_id), + ))) + .clone() + } + }; + + Ok(source.connect( + GlobalIORuntime::instance(), + ErrorCode::CannotConnectNode(format!( + "New Flight source {} for channel {} did not reconnect before its lease expired", + source_id, channel_id + )), )) } @@ -1167,17 +1512,25 @@ impl QueryCoordinator { &mut self, channel_id: &str, num_threads: usize, - ) -> Result> { - let channel_set = self - .inbound_channel_sets + ) -> Result> { + Self::inbound_channel_set(&mut self.inbound_channel_sets, channel_id, num_threads) + } + + /// Takes the map rather than `&mut self` so callers can hold a borrow on another field. + fn inbound_channel_set( + channel_sets: &mut HashMap>, + channel_id: &str, + num_threads: usize, + ) -> Result> { + let channel_set = channel_sets .entry(channel_id.to_string()) - .or_insert_with(|| Arc::new(NetworkInboundChannelSet::new(num_threads))) + .or_insert_with(|| Arc::new(ExchangePacketReceiverSet::new(num_threads))) .clone(); - if channel_set.channels.len() != num_threads { + if channel_set.receivers.len() != num_threads { return Err(ErrorCode::Internal(format!( - "NetworkInboundChannelSet {} has {} channels, expected {}", + "ExchangePacketReceiverSet {} has {} channels, expected {}", channel_id, - channel_set.channels.len(), + channel_set.receivers.len(), num_threads ))); } @@ -1257,6 +1610,12 @@ impl QueryCoordinator { } pub fn shutdown_query(&mut self, cause: Option) { + if let Some(cause) = &cause { + for source in self.new_flight_inbound_sources.values() { + source.fail(cause.clone()); + } + } + self.new_flight_fragment_outbounds.clear(); if let Some(query_info) = &mut self.info { if let Some(query_executor) = &query_info.query_executor { query_executor.finish(cause); @@ -1353,27 +1712,27 @@ impl QueryCoordinator { let settings = ExecutorSettings::try_create(info.query_ctx.clone())?; let executor = PipelineCompleteExecutor::from_pipelines(pipelines, settings)?; - assert!(self.flight_data_senders.is_empty() && self.flight_data_receivers.is_empty()); + assert!(self.fragment_outbounds.is_empty() && self.fragment_receivers.is_empty()); let info_mut = self.info.as_mut().expect("Query info is None"); info_mut.query_executor = Some(executor.clone()); let query_id = info_mut.query_id.clone(); let query_ctx = info_mut.query_ctx.clone(); query_ctx.set_executor(executor.get_inner())?; - let request_server_exchanges = std::mem::take(&mut self.statistics_exchanges); + let request_server_senders = std::mem::take(&mut self.statistics_senders); - if request_server_exchanges.len() != 1 { + if request_server_senders.len() != 1 { return Err(ErrorCode::Internal( "Request server must less than 1 if is not request server.", )); } let ctx = query_ctx.clone(); - let (_, request_server_exchange) = request_server_exchanges.into_iter().next().unwrap(); + let (_, request_server_sender) = request_server_senders.into_iter().next().unwrap(); let mut statistics_sender = StatisticsSender::spawn( &query_id, ctx, - request_server_exchange, + request_server_sender, executor.get_inner(), perf_guard, finished_profiling_rx, diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_packet_receiver.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_packet_receiver.rs new file mode 100644 index 00000000000..06e2b8853f9 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_packet_receiver.rs @@ -0,0 +1,413 @@ +// 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 std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +use arrow_flight::FlightData; +use arrow_flight::utils::flight_data_to_arrow_batch; +use arrow_schema::Schema as ArrowSchema; +use async_channel::Receiver; +use async_channel::Sender; +use databend_common_base::base::WatchNotify; +use databend_common_exception::ErrorCode; +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 futures::FutureExt; +use futures::future::BoxFuture; +use parking_lot::Mutex; +use tokio::sync::Semaphore; + +use super::inbound_quota::QueueItem; +use super::inbound_quota::SubQueue; +use crate::servers::flight::v1::transport::DeliveryOutcome; +use crate::servers::flight::v1::transport::InboundDelivery; +use crate::servers::flight::v1::transport::take_lane; + +pub struct ExchangePacketReceiver { + pub sender: Sender, + pub receiver: Receiver, + + pub sender_count: Arc, + pub(crate) closed_notified: Arc, + /// The first terminal error from any logical source feeding this channel. + /// After queued data drains, a closed channel returns this cause, or clean EOF when it is `None`. + pub close_cause: Arc>>, +} + +impl ExchangePacketReceiver { + pub fn create() -> Self { + let (tx, rx) = async_channel::unbounded(); + Self { + sender: tx, + receiver: rx, + sender_count: Arc::new(AtomicUsize::new(0)), + closed_notified: Arc::new(WatchNotify::new()), + close_cause: Arc::new(Mutex::new(None)), + } + } + + pub fn close(&self) { + self.receiver.close(); + while self.receiver.try_recv().is_ok() {} + self.closed_notified.notify_waiters(); + } + + pub async fn recv_raw(&self) -> Result, ErrorCode> { + if let Ok(item) = self.receiver.try_recv() { + return Ok(Some(item)); + } + + match self.receiver.recv().await { + Ok(item) => Ok(Some(item)), + Err(_) => match self.close_cause.lock().clone() { + Some(cause) => Err(cause), + None => Ok(None), + }, + } + } +} + +/// The receivers for one exchange channel id. +pub struct ExchangePacketReceiverSet { + pub receivers: Arc>>, +} + +impl ExchangePacketReceiverSet { + pub fn new(num_threads: usize) -> Self { + let receivers = (0..num_threads) + .map(|_| Arc::new(ExchangePacketReceiver::create())) + .collect(); + Self { + receivers: Arc::new(receivers), + } + } + + pub fn create_receiver(&self, t_idx: usize, schema: &DataSchemaRef) -> Arc { + NetworkInboundReceiver::create(schema, self.receivers[t_idx].clone()) + } +} + +/// Network-side handle. Each do_exchange connection gets one. +/// +/// Routes incoming payloads into the per-tid sub-queues of one channel and releases them once the +/// connection ends. Both transports share this: the existing Flight path drops the handle when its +/// stream ends, while New Flight drives it through [`InboundDelivery`] so a terminal error +/// can be propagated to waiting processors. +pub struct NetworkInboundSender { + destinations: Vec, + /// Guards `release` so a `terminate` followed by `Drop` decrements `sender_count` only once. + released: AtomicBool, +} + +struct InboundDestination { + /// This connection's sub-queue in one tid's ExchangePacketReceiver. + queue: Arc, + close_cause: Arc>>, + closed_notified: Arc, +} + +/// Whether a payload reached a sub-queue, or every consumer on the channel is gone. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum InboundOutcome { + Accepted, + AllReceiversClosed, +} + +impl NetworkInboundSender { + /// Create a new sender for a connection. + /// Adds a sub-queue to each ExchangePacketReceiver for this connection. + pub fn new(channel_set: &ExchangePacketReceiverSet, max_bytes_per_connection: usize) -> Self { + let semaphore = Arc::new(Semaphore::new(max_bytes_per_connection)); + let destinations = channel_set + .receivers + .iter() + .map(|channel| { + channel.sender_count.fetch_add(1, Ordering::AcqRel); + InboundDestination { + queue: Arc::new(SubQueue { + max_bytes_per_connection, + sender: channel.sender.clone(), + semaphore: semaphore.clone(), + sender_count: channel.sender_count.clone(), + }), + close_cause: channel.close_cause.clone(), + closed_notified: channel.closed_notified.clone(), + } + }) + .collect(); + + Self { + destinations, + released: AtomicBool::new(false), + } + } + + /// Add data to the inbound channel. + /// + /// Extracts tid from the FlightData, pushes to the appropriate sub-queue, + /// and waits for backpressure to clear. + /// + /// Returns `Err(())` only when ALL receivers are closed (network should disconnect). + /// If only the target tid's receiver is closed, discards the data and returns `Ok(())`. + pub async fn add_data(&self, data: FlightData) -> Result<(), ()> { + // The existing Flight path has no way to report a protocol error back to the peer, so a + // malformed tid closes the connection just like an exhausted set of receivers. + let (lane, data) = take_lane(data).map_err(|_| ())?; + match self.deliver_one(lane, data).await { + Ok(InboundOutcome::Accepted) => Ok(()), + Ok(InboundOutcome::AllReceiversClosed) | Err(_) => Err(()), + } + } + + async fn deliver_one( + &self, + lane: usize, + data: FlightData, + ) -> Result { + let Some(destination) = self.destinations.get(lane) else { + return Err(ErrorCode::BadBytes(format!( + "do_exchange thread id {} is out of range for {} channels", + lane, + self.destinations.len() + ))); + }; + + match destination.queue.add_data(data).await { + Ok(()) => Ok(InboundOutcome::Accepted), + Err(()) if self.all_receivers_closed() => Ok(InboundOutcome::AllReceiversClosed), + // Only this tid's consumer is gone, so drop the payload and keep the connection. + Err(()) => Ok(InboundOutcome::Accepted), + } + } + + /// Check if all channels are closed by receivers. + pub fn all_receivers_closed(&self) -> bool { + self.destinations + .iter() + .all(|destination| destination.queue.sender.is_closed()) + } + + /// Releases every sub-queue once. `cause` fails the downstream receivers instead of letting + /// them observe a clean end of stream. + fn release(&self, cause: Option) { + if self.released.swap(true, Ordering::AcqRel) { + return; + } + + for destination in &self.destinations { + if let Some(cause) = &cause { + let mut close_cause = destination.close_cause.lock(); + if close_cause.is_none() { + *close_cause = Some(cause.clone()); + } + drop(close_cause); + destination.queue.sender.close(); + destination.queue.semaphore.close(); + } + + if destination + .queue + .sender_count + .fetch_sub(1, Ordering::AcqRel) + == 1 + { + destination.queue.sender.close(); + } + } + } +} + +#[async_trait::async_trait] +impl InboundDelivery for NetworkInboundSender { + async fn deliver( + &self, + lane: usize, + data: FlightData, + ) -> std::result::Result { + Ok(match self.deliver_one(lane, data).await? { + InboundOutcome::Accepted => DeliveryOutcome::Accepted, + InboundOutcome::AllReceiversClosed => DeliveryOutcome::ConsumerClosed, + }) + } + + fn is_closed(&self) -> bool { + self.all_receivers_closed() + } + + fn consumer_closed(&self) -> Option> { + let notifications = self + .destinations + .iter() + .map(|destination| destination.closed_notified.clone()) + .collect::>(); + + Some( + async move { + futures::future::join_all( + notifications + .into_iter() + .map(|notification| async move { notification.notified().await }), + ) + .await; + } + .boxed(), + ) + } + + fn terminate(&self, cause: Option) { + self.release(cause); + } +} + +impl Drop for NetworkInboundSender { + fn drop(&mut self) { + self.release(None); + } +} + +/// Trait for receiving data blocks from the network. +#[async_trait::async_trait] +pub trait InboundChannel: Send + Sync { + fn close(&self); + + fn is_closed(&self) -> bool; + + async fn recv(&self) -> Result, ErrorCode>; +} + +pub struct NetworkInboundReceiver { + channel: Arc, + schema: DataSchemaRef, + arrow_schema: Arc, +} + +impl NetworkInboundReceiver { + pub fn create( + schema: &DataSchemaRef, + channel: Arc, + ) -> Arc { + Arc::new(Self { + channel, + arrow_schema: Arc::new(ArrowSchema::from(schema.as_ref())), + schema: schema.clone(), + }) + } +} + +#[async_trait::async_trait] +impl InboundChannel for NetworkInboundReceiver { + fn close(&self) { + self.channel.close(); + } + + fn is_closed(&self) -> bool { + self.channel.receiver.is_empty() && self.channel.receiver.is_closed() + } + + async fn recv(&self) -> Result, ErrorCode> { + match self.channel.recv_raw().await? { + None => Ok(None), + Some(QueueItem::LocalData(v)) => Ok(Some(v.into_data())), + Some(QueueItem::RemoteData(r)) => Ok(Some(deserialize_flight_data( + r.into_data(), + &self.schema, + &self.arrow_schema, + )?)), + } + } +} + +/// Compute the byte size of a FlightData for quota accounting. +pub fn flight_data_size(data: &FlightData) -> usize { + data.data_body.len() +} + +/// Deserialize a transport-neutral FlightData back into a DataBlock. +/// +/// Format of `app_metadata`: +/// - Fragment (last byte 0x01): `[row_count: u32][block_meta: bincode][0x01]` +/// - Dictionary (last byte 0x05): dictionary IPC data (currently unsupported) +pub(crate) fn deserialize_flight_data( + flight_data: FlightData, + schema: &DataSchemaRef, + arrow_schema: &Arc, +) -> Result { + let meta_bytes = &flight_data.app_metadata; + if meta_bytes.is_empty() { + return Err(ErrorCode::BadBytes("empty app_metadata in FlightData")); + } + + let marker = meta_bytes[meta_bytes.len() - 1]; + if marker == 0x05 { + return Err(ErrorCode::Unimplemented( + "dictionary FlightData not yet supported in broadcast exchange", + )); + } + + if marker != 0x01 { + return Err(ErrorCode::BadBytes(format!( + "unknown FlightData marker: 0x{:02x}", + marker + ))); + } + + // Parse metadata (excluding the trailing 0x01 marker) + let meta = &meta_bytes[..meta_bytes.len() - 1]; + const ROW_HEADER_SIZE: usize = std::mem::size_of::(); + + let mut cursor = &meta[..ROW_HEADER_SIZE]; + let row_count: u32 = cursor + .read_scalar() + .map_err(|e| ErrorCode::BadBytes(format!("failed to read row_count: {}", e)))?; + + let mut remaining = &meta[ROW_HEADER_SIZE..]; + let block_meta: Option = + bincode_deserialize_from_stream(&mut remaining) + .map_err(|e| ErrorCode::BadBytes(format!("failed to deserialize block_meta: {}", e)))?; + + if row_count == 0 { + return Ok(DataBlock::new_with_meta(vec![], 0, block_meta)); + } + + let mut schema = schema.clone(); + let mut arrow_schema = arrow_schema.clone(); + + if let Some(meta) = &block_meta { + if let Some(dynamic_schema) = meta.override_block_schema() { + arrow_schema = Arc::new(ArrowSchema::from(dynamic_schema.as_ref())); + schema = dynamic_schema; + } + } + + let batch = flight_data_to_arrow_batch(&flight_data, arrow_schema, &HashMap::new()) + .map_err(|e| ErrorCode::BadBytes(format!("failed to decode arrow batch: {}", e)))?; + + let block = DataBlock::from_record_batch(&schema, &batch)?; + + if block.num_columns() == 0 { + return Ok(DataBlock::new_with_meta( + vec![], + row_count as usize, + block_meta, + )); + } + + block.add_meta(block_meta) +} diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_packet_sink.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_packet_sink.rs new file mode 100644 index 00000000000..fb0e00c6314 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_packet_sink.rs @@ -0,0 +1,94 @@ +// 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 std::sync::Arc; + +use arrow_flight::FlightData; +use databend_common_base::runtime::profile::Profile; +use databend_common_base::runtime::profile::ProfileStatisticsName; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_expression::BlockMetaInfoDowncast; +use databend_common_expression::DataBlock; +use databend_common_pipeline::core::InputPort; +use databend_common_pipeline::core::PipeItem; +use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::sinks::AsyncSink; +use databend_common_pipeline::sinks::AsyncSinker; + +use super::serde::ExchangeSerializeMeta; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::StreamSendOutcome; + +pub struct ExchangePacketSink { + stream: OutboundStreamRef, + ignore_exchange: bool, +} + +impl ExchangePacketSink { + fn create( + input: Arc, + stream: OutboundStreamRef, + ignore_exchange: bool, + ) -> ProcessorPtr { + ProcessorPtr::create(AsyncSinker::create(input, Self { + stream, + ignore_exchange, + })) + } +} + +#[async_trait::async_trait] +impl AsyncSink for ExchangePacketSink { + const NAME: &'static str = "ExchangePacketSink"; + + async fn on_finish(&mut self) -> Result<()> { + self.stream.finish().await + } + + async fn consume(&mut self, mut data_block: DataBlock) -> Result { + if self.ignore_exchange { + return Ok(false); + } + + let serialize_meta = data_block + .take_meta() + .and_then(ExchangeSerializeMeta::downcast_from) + .ok_or_else(|| { + ErrorCode::Internal("ExchangePacketSink only accepts ExchangeSerializeMeta") + })?; + + let mut bytes = 0; + for packet in serialize_meta.packet { + bytes += packet.bytes_size(); + let flight_data = FlightData::try_from(packet)?; + + if self.stream.send(0, flight_data).await? == StreamSendOutcome::ConsumerClosed { + return Ok(true); + } + } + + Profile::record_usize_profile(ProfileStatisticsName::ExchangeBytes, bytes); + Ok(false) + } +} + +pub fn create_packet_writer_item(stream: OutboundStreamRef, ignore_exchange: bool) -> PipeItem { + let input = InputPort::create(); + PipeItem::create( + ExchangePacketSink::create(input.clone(), stream, ignore_exchange), + vec![input], + vec![], + ) +} diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs index f9bb7e15ac7..6ab3702ffc9 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_params.rs @@ -20,10 +20,10 @@ use databend_common_exception::Result; use databend_common_expression::DataSchemaRef; use databend_common_expression::RemoteExpr; -use crate::servers::flight::FlightReceiver; -use crate::servers::flight::FlightSender; use crate::servers::flight::v1::exchange::ExchangeInjector; +use crate::servers::flight::v1::exchange::packet_receiver::PacketReceiver; use crate::servers::flight::v1::scatter::FlightScatter; +use crate::servers::flight::v1::transport::OutboundStreamRef; #[derive(Clone)] pub struct ShuffleExchangeParams { @@ -96,36 +96,36 @@ impl ExchangeParams { } } - pub fn take_flight_sender( + pub fn take_outbound_streams( &self, - senders: &mut HashMap>, - ) -> Result> { + senders: &mut HashMap>, + ) -> Result)>> { match self { - ExchangeParams::MergeExchange(params) => params.take_flight_sender(senders), - ExchangeParams::BroadcastExchange(params) => params.take_flight_sender(senders), - ExchangeParams::NodeShuffleExchange(params) => params.take_flight_sender(senders), + ExchangeParams::MergeExchange(params) => params.take_outbound_streams(senders), + ExchangeParams::BroadcastExchange(params) => params.take_outbound_streams(senders), + ExchangeParams::NodeShuffleExchange(params) => params.take_outbound_streams(senders), ExchangeParams::GlobalShuffleExchange(_params) => Ok(vec![]), } } - pub fn take_flight_receiver( + pub(super) fn take_packet_receivers( &self, - receivers: &mut HashMap>, - ) -> Result> { + receivers: &mut HashMap>, + ) -> Result> { match self { - ExchangeParams::MergeExchange(params) => params.take_flight_receiver(receivers), - ExchangeParams::BroadcastExchange(params) => params.take_flight_receiver(receivers), - ExchangeParams::NodeShuffleExchange(params) => params.take_flight_receiver(receivers), + ExchangeParams::MergeExchange(params) => params.take_packet_receivers(receivers), + ExchangeParams::BroadcastExchange(params) => params.take_packet_receivers(receivers), + ExchangeParams::NodeShuffleExchange(params) => params.take_packet_receivers(receivers), ExchangeParams::GlobalShuffleExchange(_params) => Ok(vec![]), } } } impl MergeExchangeParams { - fn take_flight_sender( + fn take_outbound_streams( &self, - senders: &mut HashMap>, - ) -> Result> { + senders: &mut HashMap>, + ) -> Result)>> { let Some(sender) = senders.remove(&self.channel_id) else { return Err(ErrorCode::UnknownFragmentExchange(format!( "Unknown fragment exchange channel, {}, {}", @@ -135,14 +135,14 @@ impl MergeExchangeParams { Ok(sender .into_iter() - .map(|x| (self.destination_id.clone(), x)) + .map(|x| (self.destination_id.clone(), Some(x))) .collect()) } - fn take_flight_receiver( + fn take_packet_receivers( &self, - receivers: &mut HashMap>, - ) -> Result> { + receivers: &mut HashMap>, + ) -> Result> { let Some(receivers) = receivers.remove(&self.channel_id) else { return Err(ErrorCode::UnknownFragmentExchange(format!( "Unknown fragment flight receiver, {}, {}", @@ -155,19 +155,16 @@ impl MergeExchangeParams { } impl BroadcastExchangeParams { - fn take_flight_sender( + fn take_outbound_streams( &self, - senders: &mut HashMap>, - ) -> Result> { + senders: &mut HashMap>, + ) -> Result)>> { let mut exchanges = Vec::with_capacity(self.destination_channels.len()); for (destination, channels) in &self.destination_channels { for channel in channels { if destination == &self.executor_id { - exchanges.push(( - destination.clone(), - FlightSender::create(async_channel::bounded(1).0), - )); + exchanges.push((destination.clone(), None)); continue; } @@ -179,17 +176,17 @@ impl BroadcastExchangeParams { ))); }; - exchanges.extend(senders.into_iter().map(|x| (destination.clone(), x))); + exchanges.extend(senders.into_iter().map(|x| (destination.clone(), Some(x)))); } } Ok(exchanges) } - fn take_flight_receiver( + fn take_packet_receivers( &self, - receivers: &mut HashMap>, - ) -> Result> { + receivers: &mut HashMap>, + ) -> Result> { let mut exchanges = Vec::with_capacity(self.destination_channels.len()); for (destination, channels) in &self.destination_channels { @@ -211,19 +208,16 @@ impl BroadcastExchangeParams { } impl ShuffleExchangeParams { - fn take_flight_sender( + fn take_outbound_streams( &self, - senders: &mut HashMap>, - ) -> Result> { + senders: &mut HashMap>, + ) -> Result)>> { let mut exchanges = Vec::with_capacity(self.destination_ids.len()); for (destination, channels) in &self.destination_channels { for channel in channels { if destination == &self.executor_id { - exchanges.push(( - destination.clone(), - FlightSender::create(async_channel::bounded(1).0), - )); + exchanges.push((destination.clone(), None)); continue; } @@ -235,17 +229,17 @@ impl ShuffleExchangeParams { ))); }; - exchanges.extend(senders.into_iter().map(|x| (destination.clone(), x))); + exchanges.extend(senders.into_iter().map(|x| (destination.clone(), Some(x)))); } } Ok(exchanges) } - fn take_flight_receiver( + fn take_packet_receivers( &self, - receivers: &mut HashMap>, - ) -> Result> { + receivers: &mut HashMap>, + ) -> Result> { let mut exchanges = Vec::with_capacity(self.destination_channels.len()); for (destination, channels) in &self.destination_channels { diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_sink.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_sink.rs index c7d44a906f5..91e8921639c 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_sink.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_sink.rs @@ -23,29 +23,46 @@ use databend_common_pipeline::core::Pipe; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Pipeline; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline_transforms::processors::create_dummy_item; +use tokio::sync::Semaphore; +use super::exchange_packet_sink::create_packet_writer_item; use super::exchange_params::BroadcastExchangeParams; use super::exchange_params::ExchangeParams; use super::exchange_params::GlobalExchangeParams; -use super::exchange_sink_writer::create_writer_item; use super::exchange_sorting::ExchangeSorting; use super::exchange_sorting::TransformExchangeSorting; use super::exchange_transform_shuffle::exchange_shuffle; use super::hash_send_sink::HashSendSink; +use super::outbound_send_channels::SharedOutboundChannels; +use super::outbound_send_channels::fail_streams_on_pipeline_error; use super::serde::ExchangeSerializeMeta; use crate::clusters::ClusterHelper; use crate::servers::flight::v1::exchange::DataExchangeManager; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::RemoteChannel; -use crate::servers::flight::v1::network::RoundRobinChannel; -use crate::servers::flight::v1::network::create_local_channels; -use crate::servers::flight::v1::network::outbound_buffer::ExchangeBufferConfig; -use crate::servers::flight::v1::network::outbound_buffer::ExchangeSinkBuffer; +use crate::servers::flight::v1::exchange::local_channel::create_local_channels; +use crate::servers::flight::v1::exchange::outbound_channel::OutboundChannel; +use crate::servers::flight::v1::exchange::outbound_channel::RemoteOutboundChannel; +use crate::servers::flight::v1::exchange::outbound_channel::RoundRobinChannel; use crate::servers::flight::v1::scatter::HashFlightScatter; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::legacy::ExchangeBufferConfig; +use crate::servers::flight::v1::transport::legacy::ExchangeSinkBuffer; +use crate::servers::flight::v1::transport::reliable::PendingReliableOutbound; use crate::sessions::QueryContext; use crate::sessions::TableContextCluster; use crate::sessions::TableContextSettings; +const QUEUE_CAPACITY: usize = 64; +const MAX_BATCH_BYTES: usize = 256 * 1024; + +fn start_reliable_outbound(pending: PendingReliableOutbound) -> OutboundStreamRef { + Arc::new(pending.start( + Arc::new(Semaphore::new(QUEUE_CAPACITY)), + Some(MAX_BATCH_BYTES), + &GlobalIORuntime::instance(), + )) +} + pub struct ExchangeSink; impl ExchangeSink { @@ -88,17 +105,40 @@ impl ExchangeSink { } let exchange_manager = ctx.get_exchange_manager(); - let senders = exchange_manager - .get_flight_sender(&ExchangeParams::MergeExchange(params.clone()))?; - - let output = senders.len(); + let items = if ctx.get_settings().get_enable_experiment_new_flight()? { + let mut pending = exchange_manager + .take_new_flight_fragment_outbounds(¶ms.query_id, ¶ms.channel_id)?; + let outbound = pending.remove(¶ms.destination_id).ok_or_else(|| { + ErrorCode::Internal(format!( + "New Flight outbound not found for target {}", + params.destination_id + )) + })?; + let stream = start_reliable_outbound(outbound); + fail_streams_on_pipeline_error(std::slice::from_ref(&stream), pipeline); + vec![create_packet_writer_item(stream, params.ignore_exchange)] + } else { + let streams = exchange_manager + .take_fragment_outbound_streams(&ExchangeParams::MergeExchange( + params.clone(), + ))? + .into_iter() + .map(|(_, stream)| { + stream.ok_or_else(|| { + ErrorCode::Internal("Merge exchange cannot target the local node") + }) + }) + .collect::>>()?; + build_legacy_packet_sinks( + streams.into_iter().map(Some).collect(), + params.ignore_exchange, + || unreachable!("merge exchange cannot target the local node"), + ) + }; + + let output = items.len(); pipeline.try_resize(output)?; - let items = senders - .into_iter() - .map(|(_, sender)| create_writer_item(sender, params.ignore_exchange)) - .collect::>(); - pipeline.add_pipe(Pipe::create(output, 0, items)); Ok(()) } @@ -109,16 +149,21 @@ impl ExchangeSink { exchange_shuffle(ctx, params, pipeline)?; let exchange_manager = ctx.get_exchange_manager(); - let senders = exchange_manager - .get_flight_sender(&ExchangeParams::NodeShuffleExchange(params.clone()))?; // exchange writer sink let len = pipeline.output_len(); - - let items = senders - .into_iter() - .map(|(_, sender)| create_writer_item(sender, false)) - .collect::>(); + let items = if ctx.get_settings().get_enable_experiment_new_flight()? { + build_node_shuffle_packet_sinks(ctx, params, pipeline, 1)? + } else { + let streams = exchange_manager.take_fragment_outbound_streams( + &ExchangeParams::NodeShuffleExchange(params.clone()), + )?; + build_legacy_packet_sinks( + streams.into_iter().map(|(_, stream)| stream).collect(), + false, + create_dummy_item, + ) + }; pipeline.add_pipe(Pipe::create(len, 0, items)); Ok(()) @@ -147,6 +192,7 @@ impl ExchangeSink { } let compression = ctx.get_settings().get_query_flight_compression()?; + let new_flight = ctx.get_settings().get_enable_experiment_new_flight()?; let rows_threshold = ctx.get_settings().get_hash_shuffle_rows_threshold()?; let bytes_threshold = ctx.get_settings().get_hash_shuffle_bytes_threshold()?; let waker = pipeline.get_waker(); @@ -158,10 +204,12 @@ impl ExchangeSink { let exchange_manager = DataExchangeManager::instance(); let channel_set = exchange_manager.get_exchange_channel_set(query_id, exchange_id)?; - assert_eq!(channel_set.channels.len(), local_threads); + assert_eq!(channel_set.receivers.len(), local_threads); let local_outbound = create_local_channels(&channel_set); - let remote_outbound = build_hash_outbound_channels(params, local_outbound, compression)?; + let remote_outbound = + build_hash_outbound_channels(params, local_outbound, compression, new_flight)?; + remote_outbound.install_failure_handler(pipeline); let scatter = Arc::new(HashFlightScatter::try_create( ctx.get_function_context()?, @@ -187,6 +235,74 @@ impl ExchangeSink { } } +/// Builds legacy packet sinks without sharing completion across independent do_get streams. +pub(super) fn build_legacy_packet_sinks( + streams: Vec>, + ignore_exchange: bool, + mut local_item: impl FnMut() -> PipeItem, +) -> Vec { + streams + .into_iter() + .map(|stream| match stream { + None => local_item(), + Some(stream) => create_packet_writer_item(stream, ignore_exchange), + }) + .collect() +} + +pub(super) fn build_node_shuffle_packet_sinks( + ctx: &Arc, + params: &crate::servers::flight::v1::exchange::ShuffleExchangeParams, + pipeline: &mut Pipeline, + local_output_parallelism: usize, +) -> Result> { + let exchange_manager = ctx.get_exchange_manager(); + let mut pending_outbounds = Vec::new(); + + for (destination, channels) in ¶ms.destination_channels { + if destination == ¶ms.executor_id { + continue; + } + let [channel] = channels.as_slice() else { + return Err(ErrorCode::Internal(format!( + "node shuffle target {} has {} channels, expected one", + destination, + channels.len() + ))); + }; + let mut pending = exchange_manager + .take_new_flight_fragment_outbounds(¶ms.query_id, channel.as_str())?; + pending_outbounds.push(pending.remove(destination).ok_or_else(|| { + ErrorCode::Internal(format!( + "New Flight outbound not found for target {}", + destination + )) + })?); + } + + let streams = pending_outbounds + .into_iter() + .map(start_reliable_outbound) + .collect::>(); + fail_streams_on_pipeline_error(&streams, pipeline); + let mut streams = streams.into_iter(); + + let mut items = Vec::with_capacity(params.destination_channels.len()); + for (destination, _) in ¶ms.destination_channels { + if destination == ¶ms.executor_id { + items.push(if local_output_parallelism == 1 { + databend_common_pipeline_transforms::processors::create_dummy_item() + } else { + databend_common_pipeline::basic::create_resize_item(1, local_output_parallelism) + }); + } else { + items.push(create_packet_writer_item(streams.next().unwrap(), false)); + } + } + debug_assert!(streams.next().is_none()); + Ok(items) +} + struct SinkExchangeSorting; impl SinkExchangeSorting { @@ -216,11 +332,59 @@ pub(super) fn build_broadcast_outbound_channels( params: &BroadcastExchangeParams, local_outbound_channels: Vec>, compression: Option, -) -> Result>> { + new_flight: bool, +) -> Result { let query_id = ¶ms.query_id; let exchange_id = ¶ms.exchange_id; let exchange_manager = DataExchangeManager::instance(); + if new_flight { + let mut pending = + exchange_manager.take_new_flight_fragment_outbounds(query_id, exchange_id)?; + let mut remote_outbounds = Vec::with_capacity(pending.len()); + for (target_id, _) in ¶ms.destination_channels { + if target_id != ¶ms.executor_id { + remote_outbounds.push(pending.remove(target_id).ok_or_else(|| { + ErrorCode::Internal(format!( + "New Flight outbound not found for target {}", + target_id + )) + })?); + } + } + + let streams = remote_outbounds + .into_iter() + .map(start_reliable_outbound) + .collect::>(); + let num_producers = local_outbound_channels.len(); + let local_channel = RoundRobinChannel::create(local_outbound_channels); + let mut remote_idx = 0; + let mut channels = Vec::with_capacity(params.destination_channels.len()); + for (target_id, threads) in ¶ms.destination_channels { + if target_id == ¶ms.executor_id { + channels.push(local_channel.clone()); + continue; + } + + let mut remote_channels = Vec::with_capacity(threads.len()); + for thread_idx in 0..threads.len() { + remote_channels.push(RemoteOutboundChannel::create( + thread_idx, + streams[remote_idx].clone(), + compression, + )?); + } + channels.push(RoundRobinChannel::create(remote_channels)); + remote_idx += 1; + } + return Ok(SharedOutboundChannels::reliable( + channels, + streams, + num_producers, + )); + } + let mut exchanges = exchange_manager.take_ping_pong_exchanges(query_id, exchange_id)?; let mut exchanges_seq = Vec::with_capacity(exchanges.len()); @@ -257,10 +421,9 @@ pub(super) fn build_broadcast_outbound_channels( let mut remote_channels = Vec::with_capacity(threads.len()); for thread_idx in 0..threads.len() { - remote_channels.push(RemoteChannel::create( - remote_idx, + remote_channels.push(RemoteOutboundChannel::create( thread_idx, - shared_buffer.clone(), + shared_buffer.destination(remote_idx), compression, )?); } @@ -269,7 +432,7 @@ pub(super) fn build_broadcast_outbound_channels( remote_idx += 1; } - Ok(channels) + Ok(SharedOutboundChannels::immediate(channels)) } /// Build per-thread OutboundChannels for hash exchange. @@ -277,11 +440,55 @@ pub(super) fn build_hash_outbound_channels( params: &GlobalExchangeParams, mut local_outbound_channels: Vec>, compression: Option, -) -> Result>> { + new_flight: bool, +) -> Result { let num_threads = local_outbound_channels.len(); let query_id = ¶ms.query_id; let exchange_id = ¶ms.exchange_id; let exchange_manager = DataExchangeManager::instance(); + + if new_flight { + let mut pending = + exchange_manager.take_new_flight_fragment_outbounds(query_id, exchange_id)?; + let mut remote_outbounds = Vec::with_capacity(pending.len()); + for (target_id, _) in ¶ms.destination_channels { + if target_id != ¶ms.executor_id { + remote_outbounds.push(pending.remove(target_id).ok_or_else(|| { + ErrorCode::Internal(format!( + "New Flight outbound not found for target {}", + target_id + )) + })?); + } + } + + let streams = remote_outbounds + .into_iter() + .map(start_reliable_outbound) + .collect::>(); + let mut remote_idx = 0; + let mut channels = Vec::with_capacity(params.destination_channels.len() * num_threads); + for (target_id, threads) in ¶ms.destination_channels { + if target_id == ¶ms.executor_id { + channels.extend(std::mem::take(&mut local_outbound_channels)); + continue; + } + for thread_idx in 0..threads.len() { + channels.push(RemoteOutboundChannel::create( + thread_idx, + streams[remote_idx].clone(), + compression, + )?); + } + remote_idx += 1; + } + return Ok(SharedOutboundChannels::reliable( + channels, + streams, + num_threads, + )); + } + let mut exchanges = exchange_manager.take_ping_pong_exchanges(query_id, exchange_id)?; let mut exchanges_seq = Vec::with_capacity(exchanges.len()); @@ -316,10 +523,9 @@ pub(super) fn build_hash_outbound_channels( } for t_idx in 0..threads.len() { - channels.push(RemoteChannel::create( - remote_idx, + channels.push(RemoteOutboundChannel::create( t_idx, - shared_buffer.clone(), + shared_buffer.destination(remote_idx), compression, )?); } @@ -327,5 +533,5 @@ pub(super) fn build_hash_outbound_channels( remote_idx += 1; } - Ok(channels) + Ok(SharedOutboundChannels::immediate(channels)) } diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_sink_writer.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_sink_writer.rs deleted file mode 100644 index 72d484e12c7..00000000000 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_sink_writer.rs +++ /dev/null @@ -1,119 +0,0 @@ -// 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 std::sync::Arc; - -use databend_common_base::runtime::profile::Profile; -use databend_common_base::runtime::profile::ProfileStatisticsName; -use databend_common_exception::ErrorCode; -use databend_common_exception::Result; -use databend_common_expression::BlockMetaInfoDowncast; -use databend_common_expression::DataBlock; -use databend_common_pipeline::core::InputPort; -use databend_common_pipeline::core::PipeItem; -use databend_common_pipeline::core::Processor; -use databend_common_pipeline::core::ProcessorPtr; -use databend_common_pipeline::sinks::AsyncSink; -use databend_common_pipeline::sinks::AsyncSinker; -use databend_common_pipeline::sinks::Sink; -use databend_common_pipeline::sinks::Sinker; - -use crate::servers::flight::FlightSender; -use crate::servers::flight::v1::exchange::serde::ExchangeSerializeMeta; - -pub struct ExchangeWriterSink { - flight_sender: FlightSender, -} - -impl ExchangeWriterSink { - pub fn create(input: Arc, flight_sender: FlightSender) -> Box { - AsyncSinker::create(input, ExchangeWriterSink { flight_sender }) - } -} - -#[async_trait::async_trait] -impl AsyncSink for ExchangeWriterSink { - const NAME: &'static str = "ExchangeWriterSink"; - - #[async_backtrace::framed] - async fn on_finish(&mut self) -> Result<()> { - self.flight_sender.close(); - Ok(()) - } - - #[async_backtrace::framed] - async fn consume(&mut self, mut data_block: DataBlock) -> Result { - let serialize_meta = match data_block.take_meta() { - None => Err(ErrorCode::Internal( - "ExchangeWriterSink only recv ExchangeSerializeMeta, but got none.", - )), - Some(block_meta) => ExchangeSerializeMeta::downcast_from(block_meta).ok_or_else(|| { - ErrorCode::Internal("ExchangeWriterSink only recv ExchangeSerializeMeta") - }), - }?; - - let mut bytes = 0; - for packet in serialize_meta.packet { - bytes += packet.bytes_size(); - if let Err(error) = self.flight_sender.send(packet).await { - if error.code() == ErrorCode::ABORTED_QUERY { - return Ok(true); - } - - return Err(error); - } - } - - { - Profile::record_usize_profile(ProfileStatisticsName::ExchangeBytes, bytes); - } - - Ok(false) - } -} - -pub struct IgnoreExchangeSink { - flight_sender: FlightSender, -} - -impl IgnoreExchangeSink { - pub fn create(input: Arc, flight_sender: FlightSender) -> Box { - Sinker::create(input, IgnoreExchangeSink { flight_sender }) - } -} - -impl Sink for IgnoreExchangeSink { - const NAME: &'static str = "ExchangeWriterSink"; - - fn on_finish(&mut self) -> Result<()> { - self.flight_sender.close(); - Ok(()) - } - - fn consume(&mut self, _: DataBlock) -> Result<()> { - Ok(()) - } -} - -pub fn create_writer_item(exchange: FlightSender, ignore: bool) -> PipeItem { - let input = InputPort::create(); - PipeItem::create( - match ignore { - true => ProcessorPtr::create(IgnoreExchangeSink::create(input.clone(), exchange)), - false => ProcessorPtr::create(ExchangeWriterSink::create(input.clone(), exchange)), - }, - vec![input], - vec![], - ) -} diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_source.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_source.rs index f10ab34d9bf..b3f3ae9e580 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_source.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_source.rs @@ -57,12 +57,23 @@ pub fn via_exchange_source( ))); } - let exchange_params = ExchangeParams::MergeExchange(params.clone()); let exchange_manager = ctx.get_exchange_manager(); - let flight_receivers = exchange_manager.get_flight_receiver(&exchange_params)?; + let exchange_params = ExchangeParams::MergeExchange(params.clone()); + let remote_items = exchange_manager + .take_packet_receivers(&exchange_params)? + .into_iter() + .map(|flight_exchange| { + let output = OutputPort::create(); + PipeItem::create( + ExchangeSourceReader::create(output.clone(), flight_exchange), + vec![], + vec![output], + ) + }) + .collect::>(); let last_output_len = pipeline.output_len(); - let mut items = Vec::with_capacity(last_output_len + flight_receivers.len()); + let mut items = Vec::with_capacity(last_output_len + remote_items.len()); for _index in 0..last_output_len { let input = InputPort::create(); @@ -75,14 +86,7 @@ pub fn via_exchange_source( )); } - for flight_exchange in flight_receivers { - let output = OutputPort::create(); - items.push(PipeItem::create( - ExchangeSourceReader::create(output.clone(), flight_exchange), - vec![], - vec![output], - )); - } + items.extend(remote_items); pipeline.add_pipe(Pipe::create(last_output_len, items.len(), items)); @@ -110,7 +114,7 @@ pub fn via_hash_exchange_source( let waker = pipeline.get_waker(); let last_output_len = pipeline.output_len(); - let num_receivers = channel_set.channels.len(); + let num_receivers = channel_set.receivers.len(); let mut items = Vec::with_capacity(last_output_len + num_receivers); for _index in 0..last_output_len { diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_source_reader.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_source_reader.rs index 67e0fa9a57b..9d9c42616f9 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_source_reader.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_source_reader.rs @@ -27,7 +27,7 @@ use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; use log::info; -use crate::servers::flight::FlightReceiver; +use super::packet_receiver::PacketReceiver; use crate::servers::flight::v1::exchange::serde::ExchangeDeserializeMeta; use crate::servers::flight::v1::packets::DataPacket; @@ -35,14 +35,14 @@ pub struct ExchangeSourceReader { finished: AtomicBool, output: Arc, output_data: Vec, - flight_receiver: FlightReceiver, + receiver: PacketReceiver, } impl ExchangeSourceReader { - pub fn create(output: Arc, flight_receiver: FlightReceiver) -> ProcessorPtr { + pub fn create(output: Arc, receiver: PacketReceiver) -> ProcessorPtr { ProcessorPtr::create(Box::new(ExchangeSourceReader { output, - flight_receiver, + receiver, finished: AtomicBool::new(false), output_data: vec![], })) @@ -67,7 +67,7 @@ impl Processor for ExchangeSourceReader { if self.output.is_finished() { if !self.finished.swap(true, Ordering::SeqCst) { - self.flight_receiver.close(); + self.receiver.close(); } return Ok(Event::Finished); @@ -92,7 +92,7 @@ impl Processor for ExchangeSourceReader { if self.output.is_finished() { info!("un_reacted output finished, id {}", id); if !self.finished.swap(true, Ordering::SeqCst) { - self.flight_receiver.close(); + self.receiver.close(); } } } @@ -104,7 +104,7 @@ impl Processor for ExchangeSourceReader { async fn async_process(&mut self) -> Result<()> { if self.output_data.is_empty() { let mut dictionaries = Vec::new(); - while let Some(output_data) = self.flight_receiver.recv().await? { + while let Some(output_data) = self.receiver.recv().await? { if !matches!(&output_data, DataPacket::Dictionary(_)) { dictionaries.push(output_data); self.output_data = dictionaries; @@ -118,17 +118,17 @@ impl Processor for ExchangeSourceReader { } if !self.finished.swap(true, Ordering::SeqCst) { - self.flight_receiver.close(); + self.receiver.close(); } Ok(()) } } -pub fn create_reader_item(flight_receiver: FlightReceiver) -> PipeItem { +pub fn create_reader_item(receiver: PacketReceiver) -> PipeItem { let output = OutputPort::create(); PipeItem::create( - ExchangeSourceReader::create(output.clone(), flight_receiver), + ExchangeSourceReader::create(output.clone(), receiver), vec![], vec![output], ) diff --git a/src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs b/src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs index 171f12f7f8c..029859dc4a1 100644 --- a/src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs +++ b/src/query/service/src/servers/flight/v1/exchange/exchange_transform.rs @@ -25,7 +25,6 @@ use super::broadcast_send_transform::BroadcastSendTransform; use super::exchange_params::BroadcastExchangeParams; use super::exchange_params::ExchangeParams; use super::exchange_params::GlobalExchangeParams; -use super::exchange_sink_writer::create_writer_item; use super::exchange_source::via_exchange_source; use super::exchange_source_reader::create_reader_item; use super::exchange_transform_shuffle::exchange_shuffle; @@ -36,7 +35,9 @@ use crate::servers::flight::v1::exchange::ExchangeInjector; use crate::servers::flight::v1::exchange::ShuffleExchangeParams; use crate::servers::flight::v1::exchange::exchange_sink::build_broadcast_outbound_channels; use crate::servers::flight::v1::exchange::exchange_sink::build_hash_outbound_channels; -use crate::servers::flight::v1::network::create_local_channels; +use crate::servers::flight::v1::exchange::exchange_sink::build_legacy_packet_sinks; +use crate::servers::flight::v1::exchange::exchange_sink::build_node_shuffle_packet_sinks; +use crate::servers::flight::v1::exchange::local_channel::create_local_channels; use crate::servers::flight::v1::scatter::HashFlightScatter; use crate::sessions::QueryContext; use crate::sessions::TableContextSettings; @@ -85,29 +86,39 @@ impl ExchangeTransform { }; let mut items = Vec::with_capacity(len); - let exchange_params = ExchangeParams::NodeShuffleExchange(params.clone()); let exchange_manager = ctx.get_exchange_manager(); - let flight_senders = exchange_manager.get_flight_sender(&exchange_params)?; - - for (destination_id, sender) in flight_senders { - items.push(match destination_id == params.executor_id { - true => { - if local_pipe == 1 { - create_dummy_item() - } else { - create_resize_item(1, local_pipe) - } + let new_flight = ctx.get_settings().get_enable_experiment_new_flight()?; + + let nodes_source = if new_flight { + items.extend(build_node_shuffle_packet_sinks( + ctx, params, pipeline, local_pipe, + )?); + + let exchange_params = ExchangeParams::NodeShuffleExchange(params.clone()); + let receivers = exchange_manager.take_packet_receivers(&exchange_params)?; + let nodes_source = receivers.len(); + items.extend(receivers.into_iter().map(create_reader_item)); + nodes_source + } else { + let exchange_params = ExchangeParams::NodeShuffleExchange(params.clone()); + let streams = exchange_manager + .take_fragment_outbound_streams(&exchange_params)? + .into_iter() + .map(|(_, stream)| stream) + .collect(); + items.extend(build_legacy_packet_sinks(streams, false, || { + if local_pipe == 1 { + create_dummy_item() + } else { + create_resize_item(1, local_pipe) } - false => create_writer_item(sender, false), - }); - } + })); - let mut nodes_source = 0; - let receivers = exchange_manager.get_flight_receiver(&exchange_params)?; - for receiver in receivers { - nodes_source += 1; - items.push(create_reader_item(receiver)); - } + let receivers = exchange_manager.take_packet_receivers(&exchange_params)?; + let nodes_source = receivers.len(); + items.extend(receivers.into_iter().map(create_reader_item)); + nodes_source + }; let new_outputs = local_pipe + nodes_source; pipeline.add_pipe(Pipe::create(len, new_outputs, items)); @@ -136,6 +147,7 @@ impl ExchangeTransform { } let compression = ctx.get_settings().get_query_flight_compression()?; + let new_flight = ctx.get_settings().get_enable_experiment_new_flight()?; let waker = pipeline.get_waker(); pipeline.resize(local_threads, false)?; @@ -150,10 +162,12 @@ impl ExchangeTransform { local_threads, )?; - assert_eq!(channel_set.channels.len(), local_threads); + assert_eq!(channel_set.receivers.len(), local_threads); let local_outbound = create_local_channels(&channel_set); - let channels = build_broadcast_outbound_channels(params, local_outbound, compression)?; + let channels = + build_broadcast_outbound_channels(params, local_outbound, compression, new_flight)?; + channels.install_failure_handler(pipeline); let mut items = Vec::with_capacity(local_threads); @@ -169,7 +183,7 @@ impl ExchangeTransform { pipeline.add_pipe(Pipe::create(local_threads, local_threads, items)); let mut items = Vec::with_capacity(local_threads); - for idx in 0..channel_set.channels.len() { + for idx in 0..channel_set.receivers.len() { items.push(BroadcastRecvTransform::create_item( idx, channel_set.create_receiver(idx, ¶ms.schema), @@ -200,6 +214,7 @@ impl ExchangeTransform { let waker = pipeline.get_waker(); let compression = ctx.get_settings().get_query_flight_compression()?; + let new_flight = ctx.get_settings().get_enable_experiment_new_flight()?; let rows_threshold = ctx.get_settings().get_hash_shuffle_rows_threshold()?; let bytes_threshold = ctx.get_settings().get_hash_shuffle_bytes_threshold()?; @@ -214,10 +229,12 @@ impl ExchangeTransform { exchange_id, local_threads, )?; - assert_eq!(channel_set.channels.len(), local_threads); + assert_eq!(channel_set.receivers.len(), local_threads); let local_outbound = create_local_channels(&channel_set); - let remote_outbound = build_hash_outbound_channels(params, local_outbound, compression)?; + let remote_outbound = + build_hash_outbound_channels(params, local_outbound, compression, new_flight)?; + remote_outbound.install_failure_handler(pipeline); let scatter = Arc::new(HashFlightScatter::try_create( ctx.get_function_context()?, @@ -242,7 +259,7 @@ impl ExchangeTransform { pipeline.add_pipe(Pipe::create(local_threads, local_threads, items)); let mut items = Vec::with_capacity(local_threads); - for idx in 0..channel_set.channels.len() { + for idx in 0..channel_set.receivers.len() { items.push(ExchangeRecvTransform::create_item( idx, channel_set.create_receiver(idx, ¶ms.schema), diff --git a/src/query/service/src/servers/flight/v1/exchange/hash_send_sink.rs b/src/query/service/src/servers/flight/v1/exchange/hash_send_sink.rs index 88bed40de08..f847bbdce19 100644 --- a/src/query/service/src/servers/flight/v1/exchange/hash_send_sink.rs +++ b/src/query/service/src/servers/flight/v1/exchange/hash_send_sink.rs @@ -25,12 +25,12 @@ use databend_common_pipeline::core::InputPort; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::core::SyncTaskSet; use petgraph::graph::NodeIndex; use super::outbound_send_channels::OutboundSendChannels; use super::outbound_send_channels::OutboundSendHandle; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::SyncTaskSet; +use super::outbound_send_channels::SharedOutboundChannels; use crate::servers::flight::v1::scatter::FlightScatter; pub struct HashSendSink { @@ -47,7 +47,7 @@ impl HashSendSink { pub fn create_item( worker_id: usize, scatter: Arc>, - channels: Vec>, + channels: SharedOutboundChannels, waker: Arc, rows_threshold: usize, bytes_threshold: usize, @@ -71,6 +71,11 @@ impl HashSendSink { PipeItem::create(processor, vec![input], vec![]) } + + fn finish_processor(&mut self) -> Result { + self.input.finish(); + self.channels.poll_complete_event(&self.tasks, self.id) + } } impl Processor for HashSendSink { @@ -89,8 +94,7 @@ impl Processor for HashSendSink { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.channels.all_closed() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -130,8 +134,7 @@ impl Processor for HashSendSink { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.channels.all_closed() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -164,8 +167,7 @@ impl Processor for HashSendSink { } if futures.is_empty() { - self.channels.close_all(); - return Ok(Event::Finished); + return self.finish_processor(); } let joined = Box::pin(futures::future::join_all(futures)); @@ -179,8 +181,7 @@ impl Processor for HashSendSink { } } - self.channels.close_all(); - return Ok(Event::Finished); + return self.finish_processor(); } self.input.set_need_data(); @@ -189,11 +190,9 @@ impl Processor for HashSendSink { fn details_status(&self) -> Option { Some(format!( - "handle_pending={}, closed_channels={}/{}, closed={:?}, buffered_partitions={:?}", + "handle_pending={}, closed_channels={}, buffered_partitions={:?}", self.handle.is_some(), - self.channels.closed_count(), - self.channels.len(), - self.channels.closed_status(), + self.channels.closed_summary(), self.partition_stream.partition_ids(), )) } diff --git a/src/query/service/src/servers/flight/v1/exchange/hash_send_source.rs b/src/query/service/src/servers/flight/v1/exchange/hash_send_source.rs index 8968b6710ba..aea39a7caa4 100644 --- a/src/query/service/src/servers/flight/v1/exchange/hash_send_source.rs +++ b/src/query/service/src/servers/flight/v1/exchange/hash_send_source.rs @@ -25,11 +25,11 @@ use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::core::SyncTaskHandle; +use databend_common_pipeline::core::SyncTaskSet; use petgraph::graph::NodeIndex; -use crate::servers::flight::v1::network::InboundChannel; -use crate::servers::flight::v1::network::SyncTaskHandle; -use crate::servers::flight::v1::network::SyncTaskSet; +use crate::servers::flight::v1::exchange::exchange_packet_receiver::InboundChannel; pub struct HashSendSource { id: NodeIndex, diff --git a/src/query/service/src/servers/flight/v1/exchange/hash_send_transform.rs b/src/query/service/src/servers/flight/v1/exchange/hash_send_transform.rs index 283601d6159..2e75c62e22e 100644 --- a/src/query/service/src/servers/flight/v1/exchange/hash_send_transform.rs +++ b/src/query/service/src/servers/flight/v1/exchange/hash_send_transform.rs @@ -26,12 +26,12 @@ use databend_common_pipeline::core::OutputPort; use databend_common_pipeline::core::PipeItem; use databend_common_pipeline::core::Processor; use databend_common_pipeline::core::ProcessorPtr; +use databend_common_pipeline::core::SyncTaskSet; use petgraph::graph::NodeIndex; use super::outbound_send_channels::OutboundSendChannels; use super::outbound_send_channels::OutboundSendHandle; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::SyncTaskSet; +use super::outbound_send_channels::SharedOutboundChannels; use crate::servers::flight::v1::scatter::FlightScatter; pub struct HashSendTransform { @@ -51,7 +51,7 @@ impl HashSendTransform { worker_id: usize, local_pos: usize, scatter: Arc>, - channels: Vec>, + channels: SharedOutboundChannels, waker: Arc, rows_threshold: usize, bytes_threshold: usize, @@ -82,6 +82,12 @@ impl HashSendTransform { fn no_active_downstream(&self) -> bool { self.output.is_finished() && self.channels.all_closed_except(self.local_pos) } + + fn finish_processor(&mut self) -> Result { + self.input.finish(); + self.output.finish(); + self.channels.poll_complete_event(&self.tasks, self.id) + } } impl Processor for HashSendTransform { @@ -100,8 +106,7 @@ impl Processor for HashSendTransform { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -112,8 +117,7 @@ impl Processor for HashSendTransform { } if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } if self.input.has_data() { @@ -160,8 +164,7 @@ impl Processor for HashSendTransform { Poll::Ready(results) => { self.channels.handle_send_results(results)?; if self.no_active_downstream() { - self.input.finish(); - return Ok(Event::Finished); + return self.finish_processor(); } } Poll::Pending => { @@ -178,8 +181,6 @@ impl Processor for HashSendTransform { } if self.input.is_finished() { - self.output.finish(); - let mut futures = Vec::new(); for partition_id in 0..self.channels.len() { @@ -200,8 +201,7 @@ impl Processor for HashSendTransform { } if futures.is_empty() { - self.channels.close_all(); - return Ok(Event::Finished); + return self.finish_processor(); } let joined = Box::pin(futures::future::join_all(futures)); @@ -215,8 +215,7 @@ impl Processor for HashSendTransform { } } - self.channels.close_all(); - return Ok(Event::Finished); + return self.finish_processor(); } self.input.set_need_data(); @@ -225,12 +224,10 @@ impl Processor for HashSendTransform { fn details_status(&self) -> Option { Some(format!( - "handle_pending={}, local_pos={}, closed_channels={}/{}, closed={:?}, buffered_partitions={:?}", + "handle_pending={}, local_pos={}, closed_channels={}, buffered_partitions={:?}", self.handle.is_some(), self.local_pos, - self.channels.closed_count(), - self.channels.len(), - self.channels.closed_status(), + self.channels.closed_summary(), self.partition_stream.partition_ids(), )) } diff --git a/src/query/service/src/servers/flight/v1/network/inbound_quota.rs b/src/query/service/src/servers/flight/v1/exchange/inbound_quota.rs similarity index 92% rename from src/query/service/src/servers/flight/v1/network/inbound_quota.rs rename to src/query/service/src/servers/flight/v1/exchange/inbound_quota.rs index a0043330fba..a01707420bf 100644 --- a/src/query/service/src/servers/flight/v1/network/inbound_quota.rs +++ b/src/query/service/src/servers/flight/v1/exchange/inbound_quota.rs @@ -16,13 +16,12 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use arrow_flight::FlightData; -use async_channel::Receiver; use async_channel::Sender; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; -use super::inbound_channel::flight_data_size; -use crate::servers::flight::v1::network::local_channel::LocalQueueItem; +use super::exchange_packet_receiver::flight_data_size; +use crate::servers::flight::v1::exchange::local_channel::LocalQueueItem; /// Item stored in a per-connection sub-queue. pub struct RemoteQueueItem { @@ -56,7 +55,7 @@ impl RemoteQueueItem { } } -/// Per-connection sub-queue within a NetworkInboundChannel. +/// Per-connection sub-queue within a ExchangePacketReceiver. pub struct SubQueue { /// The owning connection's quota (for priority comparison and release). pub semaphore: Arc, @@ -65,8 +64,6 @@ pub struct SubQueue { pub sender: Sender, - pub receiver: Receiver, - pub sender_count: Arc, } diff --git a/src/query/service/src/servers/flight/v1/network/local_channel.rs b/src/query/service/src/servers/flight/v1/exchange/local_channel.rs similarity index 67% rename from src/query/service/src/servers/flight/v1/network/local_channel.rs rename to src/query/service/src/servers/flight/v1/exchange/local_channel.rs index 53d370793ad..3084054e8e6 100644 --- a/src/query/service/src/servers/flight/v1/network/local_channel.rs +++ b/src/query/service/src/servers/flight/v1/exchange/local_channel.rs @@ -23,9 +23,10 @@ use databend_common_expression::DataBlock; use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; -use crate::servers::flight::v1::network::NetworkInboundChannelSet; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::inbound_quota::QueueItem; +use crate::servers::flight::v1::exchange::exchange_packet_receiver::ExchangePacketReceiverSet; +use crate::servers::flight::v1::exchange::inbound_quota::QueueItem; +use crate::servers::flight::v1::exchange::outbound_channel::OutboundChannel; +use crate::servers::flight::v1::transport::StreamSendOutcome; pub struct LocalQueueItem { data: DataBlock, @@ -61,38 +62,30 @@ impl OutboundChannel for LocalOutboundChannel { self.sender.is_closed() } - async fn add_block(&self, block: DataBlock) -> Result<()> { + async fn add_block(&self, block: DataBlock) -> Result { let size = block.memory_size(); let size = std::cmp::min(size, self.max_bytes_local_channel); let semaphore = self.semaphore.clone(); - if let Ok(x) = semaphore.try_acquire_many_owned(size as u32) { - let item = LocalQueueItem::create(block, x); - - if let Err(cause) = self.sender.try_send(item) { - if cause.is_full() { - unreachable!("Logical error, local channel quota queue is full"); - } + let permit = match semaphore.try_acquire_many_owned(size as u32) { + Ok(permit) => permit, + Err(_) => { + let semaphore = self.semaphore.clone(); + let Ok(permit) = semaphore.acquire_many_owned(size as u32).await else { + return Err(ErrorCode::Internal( + "Logical error, inbound quota semaphore is closed.", + )); + }; + permit } - - return Ok(()); - } - - let semaphore = self.semaphore.clone(); - let Ok(x) = semaphore.acquire_many_owned(size as u32).await else { - return Err(ErrorCode::Internal( - "Logical error, inbound quota semaphore is closed.", - )); }; - let item = LocalQueueItem::create(block, x); - if let Err(cause) = self.sender.try_send(item) { - if cause.is_full() { - unreachable!("Logical error, local channel quota queue is full"); - } + let item = LocalQueueItem::create(block, permit); + match self.sender.try_send(item) { + Ok(()) => Ok(StreamSendOutcome::Accepted), + Err(cause) if cause.is_closed() => Ok(StreamSendOutcome::ConsumerClosed), + Err(_) => unreachable!("Logical error, local channel quota queue is full"), } - - Ok(()) } } @@ -108,12 +101,12 @@ impl Drop for LocalOutboundChannel { pub const LOCAL_CHANNEL_MAX_BYTES: usize = 20 * 1024 * 1024; pub fn create_local_channels( - channel_set: &NetworkInboundChannelSet, + channel_set: &ExchangePacketReceiverSet, ) -> Vec> { let semaphore = Arc::new(Semaphore::new(LOCAL_CHANNEL_MAX_BYTES)); - let mut outbound = Vec::>::with_capacity(channel_set.channels.len()); - for channel in channel_set.channels.iter() { + let mut outbound = Vec::>::with_capacity(channel_set.receivers.len()); + for channel in channel_set.receivers.iter() { channel.sender_count.fetch_add(1, Ordering::AcqRel); outbound.push(Arc::new(LocalOutboundChannel { diff --git a/src/query/service/src/servers/flight/v1/exchange/mod.rs b/src/query/service/src/servers/flight/v1/exchange/mod.rs index e79e78ffda2..8b4e6aa873b 100644 --- a/src/query/service/src/servers/flight/v1/exchange/mod.rs +++ b/src/query/service/src/servers/flight/v1/exchange/mod.rs @@ -17,9 +17,10 @@ mod broadcast_send_transform; mod data_exchange; mod exchange_injector; mod exchange_manager; +mod exchange_packet_receiver; +mod exchange_packet_sink; mod exchange_params; mod exchange_sink; -mod exchange_sink_writer; mod exchange_sorting; mod exchange_source; mod exchange_source_reader; @@ -29,7 +30,12 @@ mod exchange_transform_shuffle; mod hash_send_sink; mod hash_send_source; mod hash_send_transform; +mod inbound_quota; +mod local_channel; +mod outbound_channel; mod outbound_send_channels; +mod packet_receiver; +mod reliable_delivery; mod statistics_receiver; mod statistics_sender; diff --git a/src/query/service/src/servers/flight/v1/network/outbound_channel.rs b/src/query/service/src/servers/flight/v1/exchange/outbound_channel.rs similarity index 84% rename from src/query/service/src/servers/flight/v1/network/outbound_channel.rs rename to src/query/service/src/servers/flight/v1/exchange/outbound_channel.rs index b721b6f8616..b4332c2d982 100644 --- a/src/query/service/src/servers/flight/v1/network/outbound_channel.rs +++ b/src/query/service/src/servers/flight/v1/exchange/outbound_channel.rs @@ -29,13 +29,15 @@ use arrow_schema::Schema as ArrowSchema; use bytes::Bytes; use databend_common_base::runtime::profile::Profile; use databend_common_base::runtime::profile::ProfileStatisticsName; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; use databend_common_expression::DataBlock; use databend_common_io::prelude::BinaryWrite; use databend_common_io::prelude::bincode_serialize_into_buf; use databend_common_settings::FlightCompression; -use super::outbound_buffer::ExchangeSinkBuffer; +use crate::servers::flight::v1::transport::OutboundStream; +use crate::servers::flight::v1::transport::StreamSendOutcome; /// Outbound channel trait for sending data blocks. /// Supports both local (zero-copy) and remote (serialized) channels. @@ -45,7 +47,7 @@ pub trait OutboundChannel: Send + Sync { fn is_closed(&self) -> bool; - async fn add_block(&self, block: DataBlock) -> Result<()>; + async fn add_block(&self, block: DataBlock) -> Result; } // --------------------------------------------------------------------------- @@ -150,69 +152,62 @@ pub fn serialize_block( } // --------------------------------------------------------------------------- -// RemoteChannel — sends via ExchangeSinkBuffer + PingPongExchange +// RemoteOutboundChannel // --------------------------------------------------------------------------- -/// Remote exchange channel that serializes DataBlock to FlightData -/// and sends through ExchangeSinkBuffer. -pub struct RemoteChannel { - dest_idx: usize, +/// Remote exchange channel that serializes `DataBlock` into opaque transport payloads. +pub struct RemoteOutboundChannel { channel_id: usize, - buffer: Arc, + stream: Arc, ipc_options: IpcWriteOptions, } -impl RemoteChannel { +impl RemoteOutboundChannel { pub fn create( - dest_idx: usize, channel_id: usize, - buffer: Arc, + stream: Arc, compression: Option, ) -> Result> { Ok(Arc::new(Self { - dest_idx, channel_id, - buffer, + stream, ipc_options: make_ipc_options(compression)?, })) } } #[async_trait::async_trait] -impl OutboundChannel for RemoteChannel { +impl OutboundChannel for RemoteOutboundChannel { fn close(&self) {} fn is_closed(&self) -> bool { - self.buffer.is_closed(self.dest_idx) + self.stream.is_closed() } - async fn add_block(&self, block: DataBlock) -> Result<()> { + async fn add_block(&self, block: DataBlock) -> Result { Profile::record_usize_profile(ProfileStatisticsName::ExchangeRows, block.num_rows()); Profile::record_usize_profile(ProfileStatisticsName::ExchangeBytes, block.memory_size()); let flight_data_list = serialize_block(block, &self.ipc_options, None)?; - let tid_prefix = (self.channel_id as u16).to_le_bytes(); - for flight_data in flight_data_list { - let mut metadata = tid_prefix.to_vec(); - metadata.extend_from_slice(&flight_data.app_metadata); - let flight_data = FlightData { - app_metadata: metadata.into(), - ..flight_data + let outcome = match self.stream.send(self.channel_id, flight_data).await { + Err(cause) if cause.code() == ErrorCode::ABORTED_QUERY => { + StreamSendOutcome::ConsumerClosed + } + result => result?, }; - - self.buffer - .add_data(self.channel_id, self.dest_idx, flight_data) - .await?; + if outcome == StreamSendOutcome::ConsumerClosed { + return Ok(outcome); + } } - Ok(()) + Ok(StreamSendOutcome::Accepted) } } // --------------------------------------------------------------------------- -// RoundRobinChannel — round-robin across multiple RemoteChannels for one node +// RoundRobinChannel — round-robin across multiple OutboundChannels for one node // --------------------------------------------------------------------------- /// Wraps multiple OutboundChannels (one per thread on a remote node) @@ -243,9 +238,9 @@ impl OutboundChannel for RoundRobinChannel { self.channels.iter().all(|ch| ch.is_closed()) } - async fn add_block(&self, block: DataBlock) -> Result<()> { + async fn add_block(&self, block: DataBlock) -> Result { if self.channels.is_empty() { - return Ok(()); + return Ok(StreamSendOutcome::Accepted); } let idx = self.next_idx.fetch_add(1, Ordering::Relaxed) % self.channels.len(); @@ -269,8 +264,8 @@ impl OutboundChannel for DummyOutboundChannel { true } - async fn add_block(&self, _block: DataBlock) -> Result<()> { - Ok(()) + async fn add_block(&self, _block: DataBlock) -> Result { + Ok(StreamSendOutcome::ConsumerClosed) } } @@ -290,11 +285,11 @@ mod tests { use tonic::Status; use super::*; - 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; - use crate::servers::flight::v1::network::outbound_buffer::ExchangeSinkBuffer; - use crate::servers::flight::v1::network::outbound_transport::PingPongExchange; + use crate::servers::flight::v1::exchange::exchange_packet_receiver::deserialize_flight_data; + use crate::servers::flight::v1::transport::legacy::ExchangeBufferConfig; + use crate::servers::flight::v1::transport::legacy::ExchangeSinkBuffer; + use crate::servers::flight::v1::transport::legacy::PingPongExchange; + use crate::servers::flight::v1::transport::take_lane; fn test_runtime() -> Arc { Arc::new(Runtime::with_worker_threads(2, None).unwrap()) @@ -332,7 +327,7 @@ mod tests { ExchangeSinkBuffer::create(vec![exchange], ExchangeBufferConfig::default(), &rt) .unwrap(), ); - let channel = RemoteChannel::create(0, 0, buffer, None).unwrap(); + let channel = RemoteOutboundChannel::create(0, buffer.destination(0), None).unwrap(); channel.add_block(make_block(10)).await.unwrap(); @@ -356,7 +351,7 @@ mod tests { ExchangeSinkBuffer::create(vec![exchange], ExchangeBufferConfig::default(), &rt) .unwrap(), ); - let channel = RemoteChannel::create(0, 0, buffer, None).unwrap(); + let channel = RemoteOutboundChannel::create(0, buffer.destination(0), None).unwrap(); // Empty block with no meta should produce no flight data channel.add_block(DataBlock::empty()).await.unwrap(); @@ -372,7 +367,7 @@ mod tests { .unwrap(), ); // tid=5 - let channel = RemoteChannel::create(0, 5, buffer, None).unwrap(); + let channel = RemoteOutboundChannel::create(5, buffer.destination(0), None).unwrap(); channel.add_block(make_block(1)).await.unwrap(); @@ -410,9 +405,9 @@ mod tests { fn is_closed(&self) -> bool { false } - async fn add_block(&self, _block: DataBlock) -> Result<()> { + async fn add_block(&self, _block: DataBlock) -> Result { self.count.fetch_add(1, Ordering::SeqCst); - Ok(()) + Ok(StreamSendOutcome::Accepted) } } @@ -442,16 +437,17 @@ mod tests { ExchangeSinkBuffer::create(vec![exchange], ExchangeBufferConfig::default(), &rt) .unwrap(), ); - let channel = RemoteChannel::create(0, 3, buffer, None).unwrap(); + let channel = RemoteOutboundChannel::create(3, buffer.destination(0), None).unwrap(); // Send a block with known data let col = Int32Type::from_data(vec![10i32, 20, 30, 40, 50]); let original = DataBlock::new_from_columns(vec![col]); channel.add_block(original.clone()).await.unwrap(); - // Receive, strip tid, deserialize + // Receive, remove transport framing, deserialize let flight_data = send_rx.recv().await.unwrap(); - let stripped = strip_tid(flight_data); + let (lane, stripped) = take_lane(flight_data).unwrap(); + assert_eq!(lane, 3); let schema = Arc::new(original.infer_schema()); let arrow_schema = Arc::new(ArrowSchema::from(schema.as_ref())); let decoded = deserialize_flight_data(stripped, &schema, &arrow_schema).unwrap(); diff --git a/src/query/service/src/servers/flight/v1/exchange/outbound_send_channels.rs b/src/query/service/src/servers/flight/v1/exchange/outbound_send_channels.rs index fba419cee7b..426009d7ea4 100644 --- a/src/query/service/src/servers/flight/v1/exchange/outbound_send_channels.rs +++ b/src/query/service/src/servers/flight/v1/exchange/outbound_send_channels.rs @@ -13,25 +13,135 @@ // limitations under the License. use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; -use databend_common_exception::ErrorCode; +use databend_common_base::runtime::GlobalIORuntime; use databend_common_exception::Result; - -use crate::servers::flight::v1::network::DummyOutboundChannel; -use crate::servers::flight::v1::network::OutboundChannel; -use crate::servers::flight::v1::network::SyncTaskHandle; - -pub(super) type OutboundSendResult = (usize, Result<()>); +use databend_common_pipeline::core::Event; +use databend_common_pipeline::core::ExecutionInfo; +use databend_common_pipeline::core::Pipeline; +use databend_common_pipeline::core::SyncTaskHandle; +use databend_common_pipeline::core::SyncTaskSet; +use databend_common_pipeline::core::basic_callback; +use futures::future::BoxFuture; +use petgraph::prelude::NodeIndex; + +use crate::servers::flight::v1::exchange::outbound_channel::DummyOutboundChannel; +use crate::servers::flight::v1::exchange::outbound_channel::OutboundChannel; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::StreamSendOutcome; + +pub(super) type OutboundSendResult = (usize, Result); pub(super) type OutboundSendResults = Vec; pub(super) type OutboundSendHandle = SyncTaskHandle<'static, OutboundSendResults>; +type OutboundFinishHandle = SyncTaskHandle<'static, Result<()>>; + +pub(super) fn fail_streams_on_pipeline_error( + streams: &[OutboundStreamRef], + pipeline: &mut Pipeline, +) { + let streams = streams.to_vec(); + pipeline.lift_on_finished(basic_callback(move |info: &ExecutionInfo| { + if let Err(cause) = &info.res { + let cause = cause.clone(); + let streams = streams.clone(); + GlobalIORuntime::instance().spawn(async move { + futures::future::join_all(streams.iter().map(|stream| stream.fail(cause.clone()))) + .await; + }); + } + Ok(()) + })); +} + +struct ReliableCompletion { + streams: Vec, + remaining_producers: AtomicUsize, +} + +impl ReliableCompletion { + fn finish(self: &Arc) -> BoxFuture<'static, Result<()>> { + let completion = self.clone(); + Box::pin(async move { + if completion + .remaining_producers + .fetch_sub(1, Ordering::AcqRel) + != 1 + { + return Ok(()); + } + for result in + futures::future::join_all(completion.streams.iter().map(|stream| stream.finish())) + .await + { + result?; + } + Ok(()) + }) + } + + fn install_failure_handler(&self, pipeline: &mut Pipeline) { + fail_streams_on_pipeline_error(&self.streams, pipeline); + } +} + +#[derive(Clone)] +pub struct SharedOutboundChannels { + channels: Vec>, + completion: Option>, +} + +impl SharedOutboundChannels { + pub fn reliable( + channels: Vec>, + streams: Vec, + num_producers: usize, + ) -> Self { + Self { + channels, + completion: Some(Arc::new(ReliableCompletion { + streams, + remaining_producers: AtomicUsize::new(num_producers), + })), + } + } + + /// Channels that need no completion handshake, such as a purely local exchange. + pub fn immediate(channels: Vec>) -> Self { + Self { + channels, + completion: None, + } + } + + pub(super) fn len(&self) -> usize { + self.channels.len() + } + + pub fn install_failure_handler(&self, pipeline: &mut Pipeline) { + if let Some(completion) = &self.completion { + completion.install_failure_handler(pipeline); + } + } +} pub(super) struct OutboundSendChannels { channels: Vec>, + completion: Option>, + finished: bool, + finish_handle: Option, } impl OutboundSendChannels { - pub(super) fn create(channels: Vec>) -> Self { - Self { channels } + pub(super) fn create(channels: SharedOutboundChannels) -> Self { + Self { + channels: channels.channels, + completion: channels.completion, + finished: false, + finish_handle: None, + } } pub(super) fn len(&self) -> usize { @@ -61,12 +171,20 @@ impl OutboundSendChannels { .all(|(idx, ch)| idx == except_idx || ch.is_closed()) } - pub(super) fn closed_status(&self) -> Vec { - self.channels.iter().map(|ch| ch.is_closed()).collect() - } - - pub(super) fn closed_count(&self) -> usize { - self.channels.iter().filter(|ch| ch.is_closed()).count() + /// Renders closed-channel counts for `details_status`, e.g. `2/4, closed=[true, false, ...]`. + pub(super) fn closed_summary(&self) -> String { + let closed = self + .channels + .iter() + .map(|ch| ch.is_closed()) + .collect::>(); + + format!( + "{}/{}, closed={:?}", + closed.iter().filter(|closed| **closed).count(), + closed.len(), + closed + ) } pub(super) fn close(&mut self, idx: usize) { @@ -77,19 +195,49 @@ impl OutboundSendChannels { } } - pub(super) fn close_all(&mut self) { - for idx in 0..self.channels.len() { - self.close(idx); + pub(super) fn poll_complete_event( + &mut self, + tasks: &SyncTaskSet, + id: NodeIndex, + ) -> Result { + if self.finished { + return Ok(Event::Finished); + } + + let mut handle = match self.finish_handle.take() { + Some(handle) => handle, + None => { + for idx in 0..self.channels.len() { + self.close(idx); + } + + // Closing the channels is the whole completion for local and legacy groups. + let Some(completion) = &self.completion else { + self.finished = true; + return Ok(Event::Finished); + }; + + tasks.spawn(id, completion.finish()) + } + }; + + match handle.poll(false) { + Poll::Ready(result) => { + self.finished = true; + result.map(|_| Event::Finished) + } + Poll::Pending => { + self.finish_handle = Some(handle); + Ok(Event::NeedConsume) + } } } pub(super) fn handle_send_results(&mut self, results: OutboundSendResults) -> Result<()> { for (idx, result) in results { match result { - Ok(()) => {} - Err(cause) if cause.code() == ErrorCode::ABORTED_QUERY => { - self.close(idx); - } + Ok(StreamSendOutcome::Accepted) => {} + Ok(StreamSendOutcome::ConsumerClosed) => self.close(idx), Err(cause) => return Err(cause), } } diff --git a/src/query/service/src/servers/flight/v1/exchange/packet_receiver.rs b/src/query/service/src/servers/flight/v1/exchange/packet_receiver.rs new file mode 100644 index 00000000000..79bf8a1949f --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/packet_receiver.rs @@ -0,0 +1,87 @@ +// 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 std::sync::Arc; + +use arrow_flight::FlightData; +use async_channel::Receiver; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; + +use super::super::packets::DataPacket; +use super::super::transport::legacy::LegacyInbound; +use super::exchange_packet_receiver::ExchangePacketReceiver; +use super::inbound_quota::QueueItem; + +enum PacketInput { + Legacy(LegacyInbound), + ResultQueue(Receiver>), + InboundQueue(Arc), +} + +pub(super) struct PacketReceiver { + input: PacketInput, +} + +impl PacketReceiver { + pub(super) fn from_legacy(input: LegacyInbound) -> Self { + Self { + input: PacketInput::Legacy(input), + } + } + + pub(super) fn from_result_queue( + receiver: Receiver>, + ) -> Self { + Self { + input: PacketInput::ResultQueue(receiver), + } + } + + pub(super) fn from_inbound_queue(input: Arc) -> Self { + Self { + input: PacketInput::InboundQueue(input), + } + } + + pub(super) async fn recv(&self) -> Result> { + let data = match &self.input { + PacketInput::Legacy(input) => input.recv().await?, + PacketInput::ResultQueue(receiver) => match receiver.recv().await { + Err(_) => None, + Ok(result) => Some(result?), + }, + PacketInput::InboundQueue(receiver) => match receiver.recv_raw().await? { + None => None, + Some(QueueItem::RemoteData(item)) => Some(item.into_data()), + Some(QueueItem::LocalData(_)) => { + return Err(ErrorCode::Internal( + "PacketReceiver received a local block on a network receiver", + )); + } + }, + }; + data.map(DataPacket::try_from).transpose() + } + + pub(super) fn close(&self) { + match &self.input { + PacketInput::Legacy(input) => input.close(), + PacketInput::ResultQueue(receiver) => { + receiver.close(); + } + PacketInput::InboundQueue(receiver) => receiver.close(), + } + } +} diff --git a/src/query/service/src/servers/flight/v1/exchange/reliable_delivery.rs b/src/query/service/src/servers/flight/v1/exchange/reliable_delivery.rs new file mode 100644 index 00000000000..e0bdb372184 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/exchange/reliable_delivery.rs @@ -0,0 +1,68 @@ +// 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 std::sync::Arc; + +use arrow_flight::FlightData; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use futures::future::BoxFuture; + +use crate::servers::flight::v1::transport::DeliveryOutcome; +use crate::servers::flight::v1::transport::InboundDelivery; + +/// Statistics destination for a reliable logical stream. +/// +/// Fragment traffic routes by thread id, so it uses `NetworkInboundSender`. Statistics is a single +/// ordered stream consumed by `StatisticsReceiver`, so it only needs one queue. +pub struct StatisticsDelivery { + sender: async_channel::Sender>, +} + +impl StatisticsDelivery { + pub fn create( + queue_capacity: usize, + ) -> ( + Arc, + async_channel::Receiver>, + ) { + let (sender, receiver) = async_channel::bounded(queue_capacity); + (Arc::new(Self { sender }), receiver) + } +} + +#[async_trait::async_trait] +impl InboundDelivery for StatisticsDelivery { + async fn deliver(&self, _lane: usize, data: FlightData) -> Result { + match self.sender.send(Ok(data)).await { + Ok(()) => Ok(DeliveryOutcome::Accepted), + Err(_) => Ok(DeliveryOutcome::ConsumerClosed), + } + } + + fn is_closed(&self) -> bool { + self.sender.is_closed() + } + + fn consumer_closed(&self) -> Option> { + None + } + + fn terminate(&self, cause: Option) { + if let Some(cause) = cause { + let _ = self.sender.force_send(Err(cause)); + } + self.sender.close(); + } +} diff --git a/src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs b/src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs index f2503f40f6f..49142097b26 100644 --- a/src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs +++ b/src/query/service/src/servers/flight/v1/exchange/statistics_receiver.rs @@ -18,13 +18,14 @@ use std::sync::atomic::Ordering; use databend_common_base::JoinHandle; use databend_common_base::runtime::Runtime; +use databend_common_exception::ErrorCode; use databend_common_exception::Result; use futures_util::future::Either; use futures_util::future::select; use tokio::sync::broadcast::Sender; use tokio::sync::broadcast::channel; -use crate::servers::flight::FlightExchange; +use super::packet_receiver::PacketReceiver; use crate::servers::flight::v1::packets::DataPacket; use crate::servers::flight::v1::packets::ProgressInfo; use crate::sessions::MemoryUpdater; @@ -32,6 +33,7 @@ use crate::sessions::QueryContext; use crate::sessions::TableContext; use crate::sessions::TableContextPartitionStats; use crate::sessions::TableContextPerf; +use crate::sessions::TableContextQueryIdentity; use crate::sessions::TableContextQueryProfile; use crate::sessions::TableContextTelemetry; @@ -44,14 +46,13 @@ pub struct StatisticsReceiver { impl StatisticsReceiver { pub fn spawn_receiver( ctx: &Arc, - statistics_exchanges: HashMap, + statistics_receivers: HashMap, ) -> Result { let (shutdown_tx, _shutdown_rx) = channel(2); - let mut exchange_handler = Vec::with_capacity(statistics_exchanges.len()); + let mut exchange_handler = Vec::with_capacity(statistics_receivers.len()); let runtime = Runtime::with_worker_threads(2, Some(String::from("StatisticsReceiver")))?; - for (source_target, exchange) in statistics_exchanges.into_iter() { - let rx = exchange.convert_to_receiver(); + for (source_target, rx) in statistics_receivers { exchange_handler.push(runtime.spawn({ let ctx = ctx.clone(); let shutdown_rx = shutdown_tx.subscribe(); @@ -79,8 +80,7 @@ impl StatisticsReceiver { return Ok(()); } Err(cause) => { - ctx.get_current_session().force_kill_query(cause.clone()); - return Err(cause); + return Err(Self::fail_query(&ctx, cause)); } _ => loop { match StatisticsReceiver::recv_data( @@ -93,9 +93,7 @@ impl StatisticsReceiver { return Ok(()); } Err(cause) => { - ctx.get_current_session() - .force_kill_query(cause.clone()); - return Err(cause); + return Err(Self::fail_query(&ctx, cause)); } _ => {} } @@ -117,8 +115,7 @@ impl StatisticsReceiver { recv = Box::pin(rx.recv()); } Err(cause) => { - ctx.get_current_session().force_kill_query(cause.clone()); - return Err(cause); + return Err(Self::fail_query(&ctx, cause)); } }; } @@ -135,6 +132,16 @@ impl StatisticsReceiver { }) } + /// Tears down the query after a statistics stream fails, returning `cause` for the caller to + /// propagate. Shutting down the exchange first releases peers still waiting on this node. + fn fail_query(ctx: &Arc, cause: ErrorCode) -> ErrorCode { + let query_id = ctx.get_id(); + ctx.get_exchange_manager() + .shutdown_query(&query_id, Some(cause.clone())); + ctx.get_current_session().force_kill_query(cause.clone()); + cause + } + fn recv_data( ctx: &Arc, source_target: &str, diff --git a/src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs b/src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs index 5f9b3f92160..052d125aec8 100644 --- a/src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs +++ b/src/query/service/src/servers/flight/v1/exchange/statistics_sender.rs @@ -16,6 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; +use arrow_flight::FlightData; use async_channel::Sender; use databend_common_base::JoinHandle; use databend_common_base::runtime::MemStat; @@ -32,10 +33,10 @@ use tokio::sync::oneshot; use tokio::time::sleep; use crate::pipelines::executor::PipelineExecutor; -use crate::servers::flight::FlightExchange; -use crate::servers::flight::FlightSender; use crate::servers::flight::v1::packets::DataPacket; use crate::servers::flight::v1::packets::ProgressInfo; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::StreamSendOutcome; use crate::sessions::QueryContext; use crate::sessions::TableContext; use crate::sessions::TableContextPartitionStats; @@ -53,13 +54,12 @@ impl StatisticsSender { pub fn spawn( query_id: &str, ctx: Arc, - exchange: FlightExchange, + tx: OutboundStreamRef, executor: Arc, perf_guard: Option, profile_rx: oneshot::Receiver>, ) -> Self { let spawner = ctx.clone(); - let tx = exchange.convert_to_sender(); let (shutdown_flag_sender, shutdown_flag_receiver) = async_channel::bounded(1); let handle = spawner @@ -86,14 +86,7 @@ impl StatisticsSender { warn!("IoStats send has error, cause: {:?}.", error); } - let data = DataPacket::ErrorCode(error_code); - if let Err(error_code) = tx.send(data).await { - warn!( - "Cannot send data via flight exchange, cause: {:?}", - error_code - ); - } - + tx.fail(error_code).await; return; } Either::Left((_, right)) => { @@ -103,7 +96,8 @@ impl StatisticsSender { if let Err(cause) = Self::send_progress(&ctx, &mem_stat, &tx).await { ctx.get_exchange_manager() - .shutdown_query(&query_id, Some(cause)); + .shutdown_query(&query_id, Some(cause.clone())); + tx.fail(cause).await; return; } @@ -111,46 +105,42 @@ impl StatisticsSender { if cnt % 5 == 0 { // send profiles per 500 millis - if let Err(error) = + if let Err(cause) = Self::send_profile(&executor, &tx, false).await { - warn!("Profiles send has error, cause: {:?}.", error); + ctx.get_exchange_manager() + .shutdown_query(&query_id, Some(cause.clone())); + tx.fail(cause).await; + return; } } } } } - if let Err(error) = Self::send_final_profile(profile_rx, &tx).await { - warn!("Final profiles send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_copy_status(&ctx, &tx).await { - warn!("CopyStatus send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_mutation_status(&ctx, &tx).await { - warn!("MutationStatus send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_progress(&ctx, &mem_stat, &tx).await { - warn!("Statistics send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_perf(&perf_guard, &tx).await { - warn!("Perf send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_perf_counters(&ctx, &executor, &tx).await { - warn!("PerfCounters send has error, cause: {:?}.", error); + let final_result = async { + Self::send_final_profile(profile_rx, &tx).await?; + Self::send_copy_status(&ctx, &tx).await?; + Self::send_mutation_status(&ctx, &tx).await?; + Self::send_progress(&ctx, &mem_stat, &tx).await?; + Self::send_perf(&perf_guard, &tx).await?; + Self::send_perf_counters(&ctx, &executor, &tx).await?; + Self::send_part_statistics(&ctx, &tx).await?; + Self::send_io_stats(&tx).await } + .await; - if let Err(error) = Self::send_part_statistics(&ctx, &tx).await { - warn!("PartStatistics send has error, cause: {:?}.", error); - } - - if let Err(error) = Self::send_io_stats(&tx).await { - warn!("IoStats send has error, cause: {:?}.", error); + match final_result { + Ok(()) => { + if let Err(error) = tx.finish().await { + warn!("Statistics sender finish has error, cause: {:?}.", error); + } + } + Err(cause) => { + ctx.get_exchange_manager() + .shutdown_query(&query_id, Some(cause.clone())); + tx.fail(cause).await; + } } } })) @@ -185,7 +175,7 @@ impl StatisticsSender { async fn send_progress( ctx: &Arc, mem_stat: &Option>, - tx: &FlightSender, + tx: &OutboundStreamRef, ) -> Result<()> { let mut progress = Self::fetch_progress(ctx); @@ -199,15 +189,18 @@ impl StatisticsSender { } let data_packet = DataPacket::SerializeProgress(progress); - tx.send(data_packet).await + Self::send_packet(tx, data_packet).await } #[async_backtrace::framed] - async fn send_copy_status(ctx: &Arc, flight_sender: &FlightSender) -> Result<()> { + async fn send_copy_status( + ctx: &Arc, + flight_sender: &OutboundStreamRef, + ) -> Result<()> { let copy_status = ctx.copy_state().copy_status(); if !copy_status.files.is_empty() { let data_packet = DataPacket::CopyStatus(copy_status.as_ref().to_owned()); - flight_sender.send(data_packet).await?; + Self::send_packet(flight_sender, data_packet).await?; } Ok(()) } @@ -215,7 +208,7 @@ impl StatisticsSender { #[async_backtrace::framed] async fn send_mutation_status( ctx: &Arc, - flight_sender: &FlightSender, + flight_sender: &OutboundStreamRef, ) -> Result<()> { let mutation_status = { let binding = ctx.mutation_state().mutation_status(); @@ -227,33 +220,33 @@ impl StatisticsSender { } }; let data_packet = DataPacket::MutationStatus(mutation_status); - flight_sender.send(data_packet).await?; + Self::send_packet(flight_sender, data_packet).await?; Ok(()) } #[async_backtrace::framed] async fn send_profile( executor: &PipelineExecutor, - tx: &FlightSender, + tx: &OutboundStreamRef, collect_metrics: bool, ) -> Result<()> { let plans_profile = executor.fetch_profiling(collect_metrics); if !plans_profile.is_empty() { let data_packet = DataPacket::QueryProfiles(plans_profile); - tx.send(data_packet).await?; + Self::send_packet(tx, data_packet).await?; } Ok(()) } #[async_backtrace::framed] - async fn send_part_statistics(ctx: &Arc, tx: &FlightSender) -> Result<()> { + async fn send_part_statistics(ctx: &Arc, tx: &OutboundStreamRef) -> Result<()> { let part_stats = ctx.get_pruned_partitions_stats(); if !part_stats.is_empty() { let data_packet = DataPacket::PartStatistics(part_stats); - tx.send(data_packet).await?; + Self::send_packet(tx, data_packet).await?; } Ok(()) @@ -262,7 +255,7 @@ impl StatisticsSender { #[async_backtrace::framed] async fn send_final_profile( mut rx: oneshot::Receiver>, - tx: &FlightSender, + tx: &OutboundStreamRef, ) -> Result<()> { // The plans_profile comes from the executor's on_finish callback. // We use try_recv() instead of blocking recv() because the execution order @@ -270,7 +263,7 @@ impl StatisticsSender { if let Ok(plans_profile) = rx.try_recv() { if !plans_profile.is_empty() { let data_packet = DataPacket::QueryProfiles(plans_profile); - tx.send(data_packet).await?; + Self::send_packet(tx, data_packet).await?; } } @@ -280,21 +273,21 @@ impl StatisticsSender { #[async_backtrace::framed] async fn send_scan_cache_metrics( ctx: &Arc, - flight_sender: &FlightSender, + flight_sender: &OutboundStreamRef, ) -> Result<()> { let data_cache_metrics = ctx.get_data_cache_metrics(); let data_packet = DataPacket::DataCacheMetrics(data_cache_metrics.as_values()); - flight_sender.send(data_packet).await + Self::send_packet(flight_sender, data_packet).await } async fn send_perf( perf_guard: &Option, - flight_sender: &FlightSender, + flight_sender: &OutboundStreamRef, ) -> Result<()> { if let Some((_flag_guard, profiler_guard)) = perf_guard { let dumped = QueryPerf::dump(profiler_guard)?; let data_packet = DataPacket::QueryPerf(dumped); - flight_sender.send(data_packet).await?; + Self::send_packet(flight_sender, data_packet).await?; } Ok(()) } @@ -302,25 +295,35 @@ impl StatisticsSender { async fn send_perf_counters( ctx: &Arc, executor: &Arc, - flight_sender: &FlightSender, + flight_sender: &OutboundStreamRef, ) -> Result<()> { if ctx.get_perf_config().has_hw_counters() { let counters = executor.fetch_perf_counters(); if !counters.counters.is_empty() { let data_packet = DataPacket::QueryPerfCounters(counters); - flight_sender.send(data_packet).await?; + Self::send_packet(flight_sender, data_packet).await?; } } Ok(()) } - async fn send_io_stats(flight_sender: &FlightSender) -> Result<()> { + async fn send_io_stats(flight_sender: &OutboundStreamRef) -> Result<()> { let Some(stats) = ThreadTracker::io_stats() else { return Ok(()); }; let data_packet = DataPacket::IoStats(stats.snapshot()); - flight_sender.send(data_packet).await + Self::send_packet(flight_sender, data_packet).await + } + + async fn send_packet(tx: &OutboundStreamRef, packet: DataPacket) -> Result<()> { + let data = FlightData::try_from(packet)?; + if tx.send(0, data).await? == StreamSendOutcome::ConsumerClosed { + return Err(ErrorCode::AbortedQuery( + "Aborted query, because the remote statistics stream is closed.", + )); + } + Ok(()) } fn fetch_progress(ctx: &Arc) -> Vec { 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 73bd91f77d0..dbfd80de7de 100644 --- a/src/query/service/src/servers/flight/v1/flight_service.rs +++ b/src/query/service/src/servers/flight/v1/flight_service.rs @@ -46,6 +46,7 @@ use crate::servers::flight::request_builder::RequestGetter; use crate::servers::flight::v1::actions::FlightActions; use crate::servers::flight::v1::actions::flight_actions; use crate::servers::flight::v1::exchange::DataExchangeManager; +use crate::servers::flight::v1::transport::batch; pub type FlightStream = Pin> + Send + Sync + 'static>>; @@ -154,32 +155,52 @@ impl FlightService for DatabendQueryFlightService { Status::invalid_argument(format!("Failed to parse DoExchangeParams: {}", e)) })?; - let sender = DataExchangeManager::instance().handle_do_exchange( - ¶ms.query_id, - ¶ms.exchange_id, - params.num_threads, - )?; - - let mut stream = req.into_inner(); + let stream = req.into_inner(); let (tx, rx) = async_channel::bounded(1); - GlobalIORuntime::instance().spawn(async move { - while let Some(result) = stream.next().await { - let Ok(flight_data) = result else { - 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 params.new_flight { + Some(attachment) => { + let connection = DataExchangeManager::instance().handle_new_flight_do_exchange( + ¶ms.query_id, + ¶ms.exchange_id, + &attachment.source_id, + params.num_threads, + attachment.stream, + std::time::Duration::from_secs(attachment.receiver_lease_secs), + )?; + GlobalIORuntime::instance().spawn(async move { + connection.serve(stream, tx).await; + }); } - // sender is dropped here → closes sub-queues, notifies processors - }); + None => { + let sender = DataExchangeManager::instance().handle_do_exchange( + ¶ms.query_id, + ¶ms.exchange_id, + params.num_threads, + )?; + GlobalIORuntime::instance().spawn(async move { + let mut stream = stream; + while let Some(result) = stream.next().await { + let Ok(flight_data) = result else { + break; + }; + let payloads = if batch::is_batch(&flight_data) { + batch::split(flight_data) + } else { + vec![flight_data] + }; + for payload in payloads { + if sender.add_data(payload).await.is_err() { + return; + } + } + if tx.try_send(Ok(FlightData::default())).is_err() { + break; + } + } + }); + } + } Ok(RawResponse::new(Box::pin(rx))) } diff --git a/src/query/service/src/servers/flight/v1/mod.rs b/src/query/service/src/servers/flight/v1/mod.rs index 8e2c6d2b378..024c46f3fc6 100644 --- a/src/query/service/src/servers/flight/v1/mod.rs +++ b/src/query/service/src/servers/flight/v1/mod.rs @@ -14,9 +14,9 @@ pub mod actions; pub mod exchange; -pub mod network; pub mod packets; pub mod scatter; +pub mod transport; mod flight_service; 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 deleted file mode 100644 index 2a3ef86a1de..00000000000 --- a/src/query/service/src/servers/flight/v1/network/inbound_channel.rs +++ /dev/null @@ -1,481 +0,0 @@ -// 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 std::collections::HashMap; -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; - -use arrow_flight::FlightData; -use arrow_flight::utils::flight_data_to_arrow_batch; -use arrow_schema::Schema as ArrowSchema; -use async_channel::Receiver; -use async_channel::Sender; -use bytes::Buf; -use bytes::BufMut; -use bytes::BytesMut; -use databend_common_exception::ErrorCode; -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::Semaphore; - -use super::inbound_quota::QueueItem; -use super::inbound_quota::SubQueue; - -pub struct NetworkInboundChannel { - pub sender: Sender, - pub receiver: Receiver, - - pub sender_count: Arc, -} - -impl NetworkInboundChannel { - pub fn create() -> Self { - let (tx, rx) = async_channel::unbounded(); - Self { - sender: tx, - receiver: rx, - sender_count: Arc::new(AtomicUsize::new(0)), - } - } - - pub async fn recv_raw(&self) -> Option { - if let Ok(item) = self.receiver.try_recv() { - return Some(item); - } - - self.receiver.recv().await.ok() - } -} - -/// The set of NetworkInboundChannels for one channel_id. -pub struct NetworkInboundChannelSet { - pub channels: Arc>>, -} - -impl NetworkInboundChannelSet { - pub fn new(num_threads: usize) -> Self { - let channels = (0..num_threads) - .map(|_| Arc::new(NetworkInboundChannel::create())) - .collect(); - Self { - channels: Arc::new(channels), - } - } - - pub fn create_receiver(&self, t_idx: usize, schema: &DataSchemaRef) -> Arc { - NetworkInboundReceiver::create(schema, self.channels[t_idx].clone()) - } -} - -/// Network-side handle. Each do_exchange connection gets one. -/// -/// When dropped, closes this connection's sub-queues and notifies processors. -pub struct NetworkInboundSender { - /// This connection's sub-queue in each tid's NetworkInboundChannel. - sub_queues: Vec>, -} - -impl NetworkInboundSender { - /// Create a new sender for a connection. - /// Adds a sub-queue to each NetworkInboundChannel for this connection. - 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()); - - for channel in channel_set.channels.iter() { - channel.sender_count.fetch_add(1, Ordering::AcqRel); - - let sub_queue = Arc::new(SubQueue { - max_bytes_per_connection, - sender: channel.sender.clone(), - receiver: channel.receiver.clone(), - semaphore: semaphore.clone(), - sender_count: channel.sender_count.clone(), - }); - - sub_queues.push(sub_queue); - } - - Self { sub_queues } - } - - /// Add data to the inbound channel. - /// - /// Extracts tid from the FlightData, pushes to the appropriate sub-queue, - /// and waits for backpressure to clear. - /// - /// Returns `Err(())` only when ALL receivers are closed (network should disconnect). - /// If only the target tid's receiver is closed, discards the data and returns `Ok(())`. - pub async fn add_data(&self, data: FlightData) -> Result<(), ()> { - if is_batch(&data) { - return self.add_batch_data(data).await; - } - - let tid = extract_tid(&data); - - match self.sub_queues[tid].add_data(data).await { - Ok(()) => Ok(()), - Err(()) => match self.all_receivers_closed() { - true => Err(()), - false => Ok(()), - }, - } - } - - async fn add_batch_data(&self, data: FlightData) -> Result<(), ()> { - let items = split_batch_flight_data(data); - for item in items { - let tid = extract_tid(&item); - match self.sub_queues[tid].add_data(item).await { - Ok(()) => {} - Err(()) => { - if self.all_receivers_closed() { - return Err(()); - } - } - } - } - Ok(()) - } - - /// Check if all channels are closed by receivers. - 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) { - for sub_queue in &self.sub_queues { - if sub_queue.sender_count.fetch_sub(1, Ordering::AcqRel) == 1 { - sub_queue.sender.close(); - } - } - } -} - -/// Trait for receiving data blocks from the network. -#[async_trait::async_trait] -pub trait InboundChannel: Send + Sync { - fn close(&self); - - fn is_closed(&self) -> bool; - - async fn recv(&self) -> Result, ErrorCode>; -} - -pub struct NetworkInboundReceiver { - channel: Arc, - schema: DataSchemaRef, - arrow_schema: Arc, -} - -impl NetworkInboundReceiver { - pub fn create( - schema: &DataSchemaRef, - channel: Arc, - ) -> Arc { - Arc::new(Self { - channel, - arrow_schema: Arc::new(ArrowSchema::from(schema.as_ref())), - schema: schema.clone(), - }) - } -} - -#[async_trait::async_trait] -impl InboundChannel for NetworkInboundReceiver { - fn close(&self) { - self.channel.receiver.close(); - - while self.channel.receiver.try_recv().is_ok() {} - } - - fn is_closed(&self) -> bool { - self.channel.receiver.is_empty() && self.channel.receiver.is_closed() - } - - async fn recv(&self) -> Result, ErrorCode> { - match self.channel.recv_raw().await { - None => Ok(None), - Some(QueueItem::LocalData(v)) => Ok(Some(v.into_data())), - Some(QueueItem::RemoteData(r)) => { - let flight_data = strip_tid(r.into_data()); - Ok(Some(deserialize_flight_data( - flight_data, - &self.schema, - &self.arrow_schema, - )?)) - } - } - } -} - -/// Compute the byte size of a FlightData for quota accounting. -pub fn flight_data_size(data: &FlightData) -> usize { - data.data_body.len() -} - -/// Extract tid from FlightData app_metadata (first 2 bytes, little-endian u16). -pub fn extract_tid(data: &FlightData) -> usize { - if data.app_metadata.len() >= 2 { - u16::from_le_bytes([data.app_metadata[0], data.app_metadata[1]]) as usize - } else { - 0 - } -} - -/// Strip the tid prefix (first 2 bytes) from FlightData app_metadata. -/// Returns the FlightData in its original format (without tid encoding). -pub fn strip_tid(mut data: FlightData) -> FlightData { - if data.app_metadata.len() >= 2 { - data.app_metadata = data.app_metadata.slice(2..); - } - data -} - -/// Detect a batch FlightData by checking for the BATCH_MARKER (0x02) as the last byte. -fn is_batch(data: &FlightData) -> bool { - const BATCH_MARKER: u8 = 0x02; - data.app_metadata.len() >= 5 && data.app_metadata[data.app_metadata.len() - 1] == BATCH_MARKER -} - -/// Split a batch FlightData back into individual FlightData items. -/// Uses zero-copy `Bytes::split_to()` for the large data_header and data_body fields. -fn split_batch_flight_data(data: FlightData) -> Vec { - let meta = &data.app_metadata; - let tid_bytes: [u8; 2] = [meta[0], meta[1]]; - let num_items = u16::from_le_bytes([meta[2], meta[3]]) as usize; - - let mut buf = data.data_body; // Bytes implements Buf - let mut items = Vec::with_capacity(num_items); - - for _ in 0..num_items { - let meta_len = buf.get_u32_le() as usize; - let inner_meta = buf.split_to(meta_len); - - let header_len = buf.get_u32_le() as usize; - let data_header = buf.split_to(header_len); - - let body_len = buf.get_u32_le() as usize; - let data_body = buf.split_to(body_len); - - // Only app_metadata needs a small copy to prepend tid - let mut app_metadata = BytesMut::with_capacity(2 + meta_len); - app_metadata.put_slice(&tid_bytes); - app_metadata.extend_from_slice(&inner_meta); - - items.push(FlightData { - flight_descriptor: None, - app_metadata: app_metadata.freeze(), - data_header, - data_body, - }); - } - - items -} - -/// Deserialize a FlightData (after tid stripping) back into a DataBlock. -/// -/// Format of `app_metadata`: -/// - Fragment (last byte 0x01): `[row_count: u32][block_meta: bincode][0x01]` -/// - Dictionary (last byte 0x05): dictionary IPC data (currently unsupported) -pub(crate) fn deserialize_flight_data( - flight_data: FlightData, - schema: &DataSchemaRef, - arrow_schema: &Arc, -) -> Result { - let meta_bytes = &flight_data.app_metadata; - if meta_bytes.is_empty() { - return Err(ErrorCode::BadBytes("empty app_metadata in FlightData")); - } - - let marker = meta_bytes[meta_bytes.len() - 1]; - if marker == 0x05 { - return Err(ErrorCode::Unimplemented( - "dictionary FlightData not yet supported in broadcast exchange", - )); - } - - if marker != 0x01 { - return Err(ErrorCode::BadBytes(format!( - "unknown FlightData marker: 0x{:02x}", - marker - ))); - } - - // Parse metadata (excluding the trailing 0x01 marker) - let meta = &meta_bytes[..meta_bytes.len() - 1]; - const ROW_HEADER_SIZE: usize = std::mem::size_of::(); - - let mut cursor = &meta[..ROW_HEADER_SIZE]; - let row_count: u32 = cursor - .read_scalar() - .map_err(|e| ErrorCode::BadBytes(format!("failed to read row_count: {}", e)))?; - - let mut remaining = &meta[ROW_HEADER_SIZE..]; - let block_meta: Option = - bincode_deserialize_from_stream(&mut remaining) - .map_err(|e| ErrorCode::BadBytes(format!("failed to deserialize block_meta: {}", e)))?; - - if row_count == 0 { - return Ok(DataBlock::new_with_meta(vec![], 0, block_meta)); - } - - let mut schema = schema.clone(); - let mut arrow_schema = arrow_schema.clone(); - - if let Some(meta) = &block_meta { - if let Some(dynamic_schema) = meta.override_block_schema() { - arrow_schema = Arc::new(ArrowSchema::from(dynamic_schema.as_ref())); - schema = dynamic_schema; - } - } - - let batch = flight_data_to_arrow_batch(&flight_data, arrow_schema, &HashMap::new()) - .map_err(|e| ErrorCode::BadBytes(format!("failed to decode arrow batch: {}", e)))?; - - let block = DataBlock::from_record_batch(&schema, &batch)?; - - if block.num_columns() == 0 { - return Ok(DataBlock::new_with_meta( - vec![], - row_count as usize, - block_meta, - )); - } - - block.add_meta(block_meta) -} - -#[cfg(test)] -mod tests { - use arrow_flight::FlightData; - use bytes::BufMut; - use bytes::Bytes; - use bytes::BytesMut; - - use super::*; - - const BATCH_MARKER: u8 = 0x02; - - /// 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()); - app_metadata.put_u16_le(tid); - app_metadata.put_slice(inner_meta); - FlightData { - flight_descriptor: None, - app_metadata: app_metadata.freeze(), - data_header: Bytes::copy_from_slice(header), - data_body: Bytes::copy_from_slice(body), - } - } - - /// Build a batch FlightData by hand (mirrors merge_flight_data_batch logic). - fn build_batch(tid: u16, items: &[FlightData]) -> FlightData { - let mut app_metadata = BytesMut::with_capacity(5); - app_metadata.put_u16_le(tid); - app_metadata.put_u16_le(items.len() as u16); - app_metadata.put_u8(BATCH_MARKER); - - let mut body = BytesMut::new(); - for item in items { - let inner_meta = &item.app_metadata[2..]; - body.put_u32_le(inner_meta.len() as u32); - body.put_slice(inner_meta); - body.put_u32_le(item.data_header.len() as u32); - body.put_slice(&item.data_header); - body.put_u32_le(item.data_body.len() as u32); - body.put_slice(&item.data_body); - } - - FlightData { - flight_descriptor: None, - app_metadata: app_metadata.freeze(), - data_header: Bytes::new(), - data_body: body.freeze(), - } - } - - #[test] - fn test_is_batch_detection() { - // A proper batch: 5 bytes with BATCH_MARKER at end - let batch = build_batch(0, &[make_item(0, &[0x01], &[], &[1, 2, 3])]); - assert!(is_batch(&batch)); - - // A normal single item (3 bytes, marker 0x01) - let single = make_item(0, &[0x01], &[], &[1, 2, 3]); - assert!(!is_batch(&single)); - - // Too short to be a batch - let short = FlightData { - app_metadata: Bytes::from_static(&[0x00, 0x00, 0x02]), - ..Default::default() - }; - assert!(!is_batch(&short)); - } - - #[test] - fn test_split_batch_roundtrip() { - let items = vec![ - make_item(7, &[0xAA, 0x01], &[10, 20], &[1, 2, 3, 4, 5]), - make_item(7, &[0xBB, 0x01], &[], &[6, 7, 8]), - make_item(7, &[0xCC, 0x01], &[30], &[]), - ]; - - let batch = build_batch(7, &items); - assert!(is_batch(&batch)); - - let split = split_batch_flight_data(batch); - assert_eq!(split.len(), 3); - - for (original, restored) in items.iter().zip(split.iter()) { - assert_eq!(restored.app_metadata, original.app_metadata); - assert_eq!(restored.data_header, original.data_header); - assert_eq!(restored.data_body, original.data_body); - } - } - - #[test] - fn test_split_batch_single_item() { - let items = vec![make_item(0, &[0x01], &[1, 2], &[3, 4, 5])]; - let batch = build_batch(0, &items); - let split = split_batch_flight_data(batch); - assert_eq!(split.len(), 1); - assert_eq!(split[0].app_metadata, items[0].app_metadata); - assert_eq!(split[0].data_header, items[0].data_header); - assert_eq!(split[0].data_body, items[0].data_body); - } - - #[test] - fn test_split_preserves_tid() { - let tid: u16 = 42; - let items = vec![ - make_item(tid, &[0x01], &[], &[1]), - make_item(tid, &[0x02], &[], &[2]), - ]; - let batch = build_batch(tid, &items); - let split = split_batch_flight_data(batch); - - for item in &split { - let restored_tid = u16::from_le_bytes([item.app_metadata[0], item.app_metadata[1]]); - assert_eq!(restored_tid, tid); - } - } -} diff --git a/src/query/service/src/servers/flight/v1/network/mod.rs b/src/query/service/src/servers/flight/v1/network/mod.rs deleted file mode 100644 index ca4c6b2e5c7..00000000000 --- a/src/query/service/src/servers/flight/v1/network/mod.rs +++ /dev/null @@ -1,39 +0,0 @@ -// 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. - -pub mod inbound_channel; -pub mod inbound_quota; -pub mod local_channel; -pub mod outbound_buffer; -pub mod outbound_channel; -pub mod outbound_transport; - -pub use databend_common_pipeline::core::SyncTaskHandle; -pub use databend_common_pipeline::core::SyncTaskSet; -pub use inbound_channel::InboundChannel; -pub use inbound_channel::NetworkInboundChannelSet; -pub use inbound_channel::NetworkInboundReceiver; -pub use inbound_channel::NetworkInboundSender; -pub use local_channel::LocalOutboundChannel; -pub use local_channel::create_local_channels; -pub use outbound_buffer::ExchangeBufferConfig; -pub use outbound_buffer::ExchangeSinkBuffer; -pub use outbound_channel::DummyOutboundChannel; -pub use outbound_channel::OutboundChannel; -pub use outbound_channel::RemoteChannel; -pub use outbound_channel::RoundRobinChannel; -pub use outbound_transport::PingPongCallback; -pub use outbound_transport::PingPongExchange; -pub use outbound_transport::PingPongExchangeInner; -pub use outbound_transport::PingPongResponse; diff --git a/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs b/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs index 61f02acc37a..fe764469e85 100644 --- a/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs +++ b/src/query/service/src/servers/flight/v1/packets/packet_publisher.rs @@ -46,8 +46,10 @@ use crate::sessions::TableContextQueryIdentity; #[derive(Debug, Clone, Serialize, Deserialize)] pub enum Edge { Statistics, - /// do_get based channel (unidirectional: receiver pulls from sender) + /// A single-lane fragment channel. Legacy Flight receives it with do_get. Fragment(String), + /// Merge channel whose receiver must retain one ordered input per source node. + Merge(String), /// do_exchange based channel (bidirectional: sender pushes via ping-pong) /// One edge per node pair, identified by exchange_id, carrying all channel_ids. ExchangeFragment { @@ -101,6 +103,10 @@ impl DataflowDiagramBuilder { self.add_edge_inner(source, destination, Edge::Fragment(channel.to_string())) } + pub fn add_merge_edge(&mut self, source: &str, destination: &str, channel: &str) -> Result<()> { + self.add_edge_inner(source, destination, Edge::Merge(channel.to_string())) + } + pub fn add_exchange_edge( &mut self, source: &str, diff --git a/src/query/service/src/servers/flight/v1/transport/batch.rs b/src/query/service/src/servers/flight/v1/transport/batch.rs new file mode 100644 index 00000000000..5ac3103cedb --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/batch.rs @@ -0,0 +1,123 @@ +// 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. + +//! Wire format for merging several same-thread payloads into one `FlightData`. +//! +//! Both transports batch on the sending side and split on the receiving side, so the codec lives +//! here rather than in either transport. A batch carries its thread id once in `app_metadata` and +//! concatenates the original items into `data_body`: +//! +//! ```text +//! app_metadata: [tid: u16 le][item_count: u16 le][BATCH_MARKER] +//! data_body: ([meta_len: u32 le][meta][header_len: u32 le][header][body_len: u32 le][body])* +//! ``` +//! +//! Each item's `app_metadata` is stored without its 2-byte tid prefix, which `split` restores. + +use arrow_flight::FlightData; +use bytes::Buf; +use bytes::BufMut; +use bytes::Bytes; +use bytes::BytesMut; + +/// Trailing `app_metadata` byte marking a merged batch. Distinct from the fragment (0x01) and +/// dictionary (0x05) markers a single payload carries. +pub const BATCH_MARKER: u8 = 0x02; + +const TID_LEN: usize = 2; +const BATCH_HEADER_LEN: usize = 5; +/// Per-item overhead in `data_body`: three u32 length prefixes. +const ITEM_LENGTH_PREFIXES: usize = 12; + +/// Detects a merged batch by its trailing marker. +pub fn is_batch(data: &FlightData) -> bool { + data.app_metadata.len() >= BATCH_HEADER_LEN + && data.app_metadata[data.app_metadata.len() - 1] == BATCH_MARKER +} + +/// Merges same-thread items into one batch payload. +/// +/// The thread id is taken from the first item; callers must only merge items sharing a tid. +/// Panics if `items` is empty or its first item has no tid prefix. +pub fn merge(items: Vec) -> FlightData { + let mut app_metadata = BytesMut::with_capacity(BATCH_HEADER_LEN); + app_metadata.put_slice(&items[0].app_metadata[..TID_LEN]); + app_metadata.put_u16_le(items.len() as u16); + app_metadata.put_u8(BATCH_MARKER); + + let estimated = items + .iter() + .map(|item| { + ITEM_LENGTH_PREFIXES + + (item.app_metadata.len() - TID_LEN) + + item.data_header.len() + + item.data_body.len() + }) + .sum(); + + let mut body = BytesMut::with_capacity(estimated); + for item in items { + let metadata = &item.app_metadata[TID_LEN..]; + body.put_u32_le(metadata.len() as u32); + body.put_slice(metadata); + body.put_u32_le(item.data_header.len() as u32); + body.put_slice(&item.data_header); + body.put_u32_le(item.data_body.len() as u32); + body.put_slice(&item.data_body); + } + + FlightData { + flight_descriptor: None, + app_metadata: app_metadata.freeze(), + data_header: Bytes::new(), + data_body: body.freeze(), + } +} + +/// Splits a batch back into its individual items, restoring each tid prefix. +/// +/// `data_header` and `data_body` are sliced without copying; only the small `app_metadata` is +/// rebuilt. Callers must check [`is_batch`] first. +pub fn split(data: FlightData) -> Vec { + let meta = &data.app_metadata; + let tid_bytes: [u8; TID_LEN] = [meta[0], meta[1]]; + let num_items = u16::from_le_bytes([meta[2], meta[3]]) as usize; + + let mut buf = data.data_body; + let mut items = Vec::with_capacity(num_items); + + for _ in 0..num_items { + let meta_len = buf.get_u32_le() as usize; + let inner_meta = buf.split_to(meta_len); + + let header_len = buf.get_u32_le() as usize; + let data_header = buf.split_to(header_len); + + let body_len = buf.get_u32_le() as usize; + let data_body = buf.split_to(body_len); + + let mut app_metadata = BytesMut::with_capacity(TID_LEN + meta_len); + app_metadata.put_slice(&tid_bytes); + app_metadata.extend_from_slice(&inner_meta); + + items.push(FlightData { + flight_descriptor: None, + app_metadata: app_metadata.freeze(), + data_header, + data_body, + }); + } + + items +} diff --git a/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs b/src/query/service/src/servers/flight/v1/transport/legacy/buffer.rs similarity index 89% rename from src/query/service/src/servers/flight/v1/network/outbound_buffer.rs rename to src/query/service/src/servers/flight/v1/transport/legacy/buffer.rs index 974cd1fd675..f8d23a4b9b3 100644 --- a/src/query/service/src/servers/flight/v1/network/outbound_buffer.rs +++ b/src/query/service/src/servers/flight/v1/transport/legacy/buffer.rs @@ -15,25 +15,44 @@ use std::sync::Arc; use arrow_flight::FlightData; -use bytes::BufMut; -use bytes::Bytes; -use bytes::BytesMut; use concurrent_queue::ConcurrentQueue; use databend_common_base::runtime::Runtime; use databend_common_exception::ErrorCode; use databend_common_exception::Result; use parking_lot::Mutex; +use tokio::sync::OwnedSemaphorePermit; use tokio::sync::Semaphore; use tonic::Code; 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 super::ping_pong::PingPongCallback; +use super::ping_pong::PingPongExchange; +use super::ping_pong::PingPongResponse; +use super::ping_pong::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; +use crate::servers::flight::v1::transport::OutboundStream; +use crate::servers::flight::v1::transport::StreamSendOutcome; +use crate::servers::flight::v1::transport::batch; +use crate::servers::flight::v1::transport::frame_lane; + +struct PendingFlightData { + data: FlightData, + _permit: OwnedSemaphorePermit, +} + +impl PendingFlightData { + fn new(data: FlightData, permit: OwnedSemaphorePermit) -> Self { + Self { + data, + _permit: permit, + } + } + + fn into_data(self) -> FlightData { + self.data + } +} /// Configuration for ExchangeSinkBuffer. #[derive(Clone)] @@ -55,7 +74,7 @@ impl Default for ExchangeBufferConfig { /// Per-sink channel containing its own pending queue. struct Channel { - pending_queue: ConcurrentQueue, + pending_queue: ConcurrentQueue, } impl Channel { @@ -85,46 +104,11 @@ impl Channel { match items.len() { 0 => None, 1 => Some(items.into_iter().next().unwrap()), - _ => { - let tid_bytes: [u8; 2] = [items[0].app_metadata[0], items[0].app_metadata[1]]; - Some(merge_flight_data_batch(tid_bytes, items)) - } + _ => Some(batch::merge(items)), } } } -const BATCH_MARKER: u8 = 0x02; - -fn merge_flight_data_batch(tid_bytes: [u8; 2], items: Vec) -> FlightData { - let mut app_metadata = BytesMut::with_capacity(5); - app_metadata.put_slice(&tid_bytes); - app_metadata.put_u16_le(items.len() as u16); - app_metadata.put_u8(BATCH_MARKER); - - let estimated: usize = items - .iter() - .map(|i| 12 + (i.app_metadata.len() - 2) + i.data_header.len() + i.data_body.len()) - .sum(); - - let mut body = BytesMut::with_capacity(estimated); - for item in items { - let inner_meta = &item.app_metadata[2..]; // strip tid - body.put_u32_le(inner_meta.len() as u32); - body.put_slice(inner_meta); - body.put_u32_le(item.data_header.len() as u32); - body.put_slice(&item.data_header); - body.put_u32_le(item.data_body.len() as u32); - body.put_slice(&item.data_body); - } - - FlightData { - flight_descriptor: None, - app_metadata: app_metadata.freeze(), - data_header: Bytes::new(), - data_body: body.freeze(), - } -} - /// Mutable state within RemoteInstance, protected by its own lock. struct RemoteInstanceState { /// Pre-allocated channels, indexed by channel_id @@ -279,6 +263,12 @@ pub struct ExchangeSinkBuffer { inner: Arc, } +struct LegacyPingPongOutbound { + semaphore: Arc, + inner: Arc, + destination: usize, +} + impl ExchangeSinkBuffer { /// Create a new ExchangeSinkBuffer. /// @@ -322,8 +312,20 @@ impl ExchangeSinkBuffer { }) } - pub async fn add_data(&self, tid: usize, dest_idx: usize, data: FlightData) -> Result<()> { - let remote = &self.inner.state.remotes[dest_idx]; + pub fn destination(&self, destination: usize) -> Arc { + Arc::new(LegacyPingPongOutbound { + semaphore: self.semaphore.clone(), + inner: self.inner.clone(), + destination, + }) + } +} + +#[async_trait::async_trait] +impl OutboundStream for LegacyPingPongOutbound { + async fn send(&self, lane: usize, data: FlightData) -> Result { + let remote = &self.inner.state.remotes[self.destination]; + let data = frame_lane(lane, data)?; { let state = remote.state.lock(); @@ -334,7 +336,7 @@ impl ExchangeSinkBuffer { // Try to send directly first let data = match remote.exchange.try_send(data) { - Ok(None) => return Ok(()), + Ok(None) => return Ok(StreamSendOutcome::Accepted), Ok(Some(data)) => data, Err(status) => { let error = ExchangeSinkBufferSharedState::status_to_error( @@ -342,7 +344,7 @@ impl ExchangeSinkBuffer { remote.exchange.local_node_id(), remote.exchange.remote_node_id(), ); - return Err(self.close_remote(dest_idx, error)); + return Err(self.close(error)); } }; @@ -362,8 +364,8 @@ impl ExchangeSinkBuffer { match remote.exchange.try_send(data) { Ok(None) => {} Ok(Some(data)) => { - let item = RemoteQueueItem::new(data, owned_semaphore_permit); - let _ = state.channels[tid].pending_queue.push(item); + let item = PendingFlightData::new(data, owned_semaphore_permit); + let _ = state.channels[lane].pending_queue.push(item); } Err(status) => { let error = ExchangeSinkBufferSharedState::status_to_error( @@ -377,11 +379,34 @@ impl ExchangeSinkBuffer { } } + Ok(StreamSendOutcome::Accepted) + } + + async fn finish(&self) -> Result<()> { + // Existing ping-pong streams complete when their shared buffer is dropped. Closing an + // individual destination here would truncate sibling block producers. Ok(()) } - fn close_remote(&self, dest_idx: usize, error: ErrorCode) -> ErrorCode { - let remote = &self.inner.state.remotes[dest_idx]; + async fn fail(&self, cause: ErrorCode) { + self.close(cause); + } + + fn abort(&self) { + self.close(ErrorCode::AbortedQuery( + "legacy ping-pong stream was cancelled", + )); + } + + fn is_closed(&self) -> bool { + let remote = &self.inner.state.remotes[self.destination]; + remote.state.lock().last_error.is_some() + } +} + +impl LegacyPingPongOutbound { + fn close(&self, error: ErrorCode) -> ErrorCode { + let remote = &self.inner.state.remotes[self.destination]; let mut state = remote.state.lock(); if state.last_error.is_none() { @@ -390,11 +415,19 @@ impl ExchangeSinkBuffer { state.last_error.clone().unwrap_or(error) } +} + +#[cfg(test)] +impl ExchangeSinkBuffer { + async fn add_data(&self, lane: usize, destination: usize, data: FlightData) -> Result<()> { + self.destination(destination) + .send(lane, data) + .await + .map(|_| ()) + } - 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() + fn is_closed(&self, destination: usize) -> bool { + self.destination(destination).is_closed() } } @@ -404,12 +437,16 @@ mod tests { use std::time::Duration; use arrow_flight::FlightData; + use bytes::BufMut; + use bytes::Bytes; + use bytes::BytesMut; use databend_common_base::runtime::Runtime; use databend_common_base::runtime::spawn; use tonic::Status; + use super::PingPongExchange; use super::*; - use crate::servers::flight::v1::network::outbound_transport::PingPongExchange; + use crate::servers::flight::v1::transport::batch::BATCH_MARKER; fn test_runtime() -> Arc { Arc::new(Runtime::with_worker_threads(2, None).unwrap()) @@ -461,7 +498,7 @@ mod tests { let permit = permit.try_acquire_many_owned(1).unwrap(); let _ = channel .pending_queue - .push(RemoteQueueItem::new(data, permit)); + .push(PendingFlightData::new(data, permit)); let result = channel.pop_front(256 * 1024).unwrap(); // Single item: no batch wrapping, original data returned as-is @@ -478,7 +515,7 @@ mod tests { for i in 0..3 { let data = make_flight_data_with_tid(5, &[i, 0x01], 50); let p = permit.clone().try_acquire_many_owned(1).unwrap(); - let _ = channel.pending_queue.push(RemoteQueueItem::new(data, p)); + let _ = channel.pending_queue.push(PendingFlightData::new(data, p)); } let result = channel.pop_front(256 * 1024).unwrap(); @@ -506,7 +543,7 @@ mod tests { for _ in 0..10 { let data = make_flight_data_with_tid(0, &[0x01], 100); let p = permit.clone().try_acquire_many_owned(1).unwrap(); - let _ = channel.pending_queue.push(RemoteQueueItem::new(data, p)); + let _ = channel.pending_queue.push(PendingFlightData::new(data, p)); } // Budget of 250 bytes: should pop 3 items (100+100+100 >= 250) diff --git a/src/query/service/src/servers/flight/v1/transport/legacy/inbound.rs b/src/query/service/src/servers/flight/v1/transport/legacy/inbound.rs new file mode 100644 index 00000000000..d94fe215f8b --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/legacy/inbound.rs @@ -0,0 +1,52 @@ +// 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 std::sync::Arc; + +use arrow_flight::FlightData; +use async_channel::Receiver; +use databend_common_base::runtime::drop_guard; +use databend_common_exception::Result; + +use crate::pipelines::executor::WatchNotify; + +/// Raw do_get response stream with its transport cancellation handle. +pub struct LegacyInbound { + notify: Arc, + receiver: Receiver>, +} + +impl LegacyInbound { + pub fn create(notify: Arc, receiver: Receiver>) -> Self { + Self { notify, receiver } + } + + pub async fn recv(&self) -> Result> { + match self.receiver.recv().await { + Err(_) => Ok(None), + Ok(result) => result.map(Some), + } + } + + pub fn close(&self) { + self.receiver.close(); + self.notify.notify_waiters(); + } +} + +impl Drop for LegacyInbound { + fn drop(&mut self) { + drop_guard(move || self.close()) + } +} diff --git a/src/query/service/src/servers/flight/v1/transport/legacy/mod.rs b/src/query/service/src/servers/flight/v1/transport/legacy/mod.rs new file mode 100644 index 00000000000..0568a306f3f --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/legacy/mod.rs @@ -0,0 +1,27 @@ +// 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. + +mod buffer; +mod inbound; +mod outbound; +mod ping_pong; + +pub use buffer::ExchangeBufferConfig; +pub use buffer::ExchangeSinkBuffer; +pub use inbound::LegacyInbound; +pub use outbound::LegacyOutbound; +pub use ping_pong::PingPongCallback; +pub use ping_pong::PingPongExchange; +pub use ping_pong::PingPongExchangeInner; +pub use ping_pong::PingPongResponse; diff --git a/src/query/service/src/servers/flight/v1/transport/legacy/outbound.rs b/src/query/service/src/servers/flight/v1/transport/legacy/outbound.rs new file mode 100644 index 00000000000..b6c7dd9c646 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/legacy/outbound.rs @@ -0,0 +1,71 @@ +// 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 async_channel::Sender; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use log::warn; +use tonic::Status; + +use crate::servers::flight::v1::packets::DataPacket; +use crate::servers::flight::v1::transport::OutboundStream; +use crate::servers::flight::v1::transport::OutboundStreamRef; +use crate::servers::flight::v1::transport::StreamSendOutcome; + +/// Existing do_get sender expressed through the logical stream interface. +pub struct LegacyOutbound { + sender: Sender>, +} + +impl LegacyOutbound { + pub fn create(sender: Sender>) -> OutboundStreamRef { + std::sync::Arc::new(Self { sender }) + } +} + +#[async_trait::async_trait] +impl OutboundStream for LegacyOutbound { + async fn send(&self, _lane: usize, data: FlightData) -> Result { + if self.sender.send(Ok(data)).await.is_err() { + return Ok(StreamSendOutcome::ConsumerClosed); + } + Ok(StreamSendOutcome::Accepted) + } + + async fn finish(&self) -> Result<()> { + self.sender.close(); + Ok(()) + } + + async fn fail(&self, cause: ErrorCode) { + match FlightData::try_from(DataPacket::ErrorCode(cause)) { + Ok(data) => { + let _ = self.sender.send(Ok(data)).await; + } + Err(error) => { + warn!("cannot encode legacy Flight failure packet: {}", error); + } + } + self.sender.close(); + } + + fn abort(&self) { + self.sender.close(); + } + + fn is_closed(&self) -> bool { + self.sender.is_closed() + } +} diff --git a/src/query/service/src/servers/flight/v1/network/outbound_transport.rs b/src/query/service/src/servers/flight/v1/transport/legacy/ping_pong.rs similarity index 100% rename from src/query/service/src/servers/flight/v1/network/outbound_transport.rs rename to src/query/service/src/servers/flight/v1/transport/legacy/ping_pong.rs diff --git a/src/query/service/src/servers/flight/v1/transport/mod.rs b/src/query/service/src/servers/flight/v1/transport/mod.rs new file mode 100644 index 00000000000..c6cf9d36de5 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/mod.rs @@ -0,0 +1,26 @@ +// 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. + +pub mod batch; +pub mod legacy; +pub mod reliable; +mod stream; + +pub use stream::DeliveryOutcome; +pub use stream::InboundDelivery; +pub use stream::OutboundStream; +pub use stream::OutboundStreamRef; +pub use stream::StreamSendOutcome; +pub(crate) use stream::frame_lane; +pub(crate) use stream::take_lane; diff --git a/src/query/service/src/servers/flight/v1/transport/reliable/inbound.rs b/src/query/service/src/servers/flight/v1/transport/reliable/inbound.rs new file mode 100644 index 00000000000..7f52bf41d33 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/reliable/inbound.rs @@ -0,0 +1,419 @@ +// 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 std::sync::Arc; +use std::time::Duration; + +use arrow_flight::FlightData; +use databend_common_base::base::WatchNotify; +use databend_common_base::runtime::Runtime; +use databend_common_exception::ErrorCode; +use log::info; +use log::warn; +use parking_lot::Mutex; +use tokio_stream::StreamExt; +use tonic::Status; +use tonic::Streaming; + +use super::DoExchangeRequest; +use super::DoExchangeResponse; +use crate::servers::flight::v1::transport::DeliveryOutcome; +use crate::servers::flight::v1::transport::InboundDelivery; +use crate::servers::flight::v1::transport::batch; +use crate::servers::flight::v1::transport::take_lane; + +pub struct ReliableInboundSource { + delivery: Arc, + next_sequence: tokio::sync::Mutex, + lifecycle: Mutex, + /// Grace period after the last physical attachment disconnects. Expiry without a replacement + /// fails the logical source and releases its inbound delivery. + reconnect_lease: Duration, + source_label: String, + terminal_notified: WatchNotify, +} + +#[derive(Clone)] +enum InboundTerminal { + Completed, + SenderFailed(ErrorCode), + ReceiverFailed(ErrorCode), +} + +struct InboundLifecycle { + terminal: Option, + attachments: usize, + generation: u64, + receiver_monitor_started: bool, +} + +pub struct ReliableInboundConnection { + source: Arc, + runtime: Arc, + disconnect_error: Option, +} + +impl ReliableInboundSource { + pub fn new( + delivery: Arc, + reconnect_lease: Duration, + source_label: String, + ) -> Self { + Self { + delivery, + next_sequence: tokio::sync::Mutex::new(0), + lifecycle: Mutex::new(InboundLifecycle { + terminal: None, + attachments: 0, + generation: 0, + receiver_monitor_started: false, + }), + reconnect_lease, + source_label, + terminal_notified: WatchNotify::new(), + } + } + + pub fn connect( + self: &Arc, + runtime: Arc, + disconnect_error: ErrorCode, + ) -> ReliableInboundConnection { + let (disconnect_error, consumer_closed) = { + let mut lifecycle = self.lifecycle.lock(); + if lifecycle.terminal.is_some() { + (None, None) + } else { + let reconnect = lifecycle.generation != 0; + lifecycle.attachments += 1; + lifecycle.generation += 1; + let consumer_closed = if lifecycle.receiver_monitor_started { + None + } else { + lifecycle.receiver_monitor_started = true; + self.delivery.consumer_closed() + }; + if reconnect { + warn!( + "do_exchange receiver accepted replacement connection: {}", + self.source_label + ); + } + (Some(disconnect_error), consumer_closed) + } + }; + + if let Some(consumer_closed) = consumer_closed { + let source = self.clone(); + runtime.spawn(async move { + tokio::select! { + _ = consumer_closed => { + source.terminate(InboundTerminal::Completed); + } + _ = source.terminal_notified.notified() => {} + } + }); + } + + ReliableInboundConnection { + source: self.clone(), + runtime, + disconnect_error, + } + } + + async fn add_data( + &self, + sequence: u64, + data: FlightData, + ) -> Result { + let mut next_sequence = self.next_sequence.lock().await; + if let Some(response) = self.terminal_response() { + return Ok(response); + } + if sequence < *next_sequence { + info!( + "do_exchange receiver ignored replayed data: {}, sequence={}, expected={}", + self.source_label, sequence, *next_sequence + ); + return Ok(DoExchangeResponse::ack(sequence)); + } + if sequence > *next_sequence { + return Err(ErrorCode::Internal(format!( + "out-of-order do_exchange packet: expected {}, got {}", + *next_sequence, sequence + ))); + } + + let accepted = self.deliver(data).await; + *next_sequence += 1; + if let Some(response) = self.terminal_response() { + return Ok(response); + } + + match accepted? { + DeliveryOutcome::Accepted => Ok(DoExchangeResponse::ack(sequence)), + DeliveryOutcome::ConsumerClosed => { + // No downstream consumer can accept more data, so tell the sender to stop. + Ok(self.terminate(InboundTerminal::Completed).response()) + } + } + } + + async fn deliver(&self, data: FlightData) -> Result { + if !batch::is_batch(&data) { + let (lane, data) = take_lane(data)?; + return self.delivery.deliver(lane, data).await; + } + + for item in batch::split(data) { + let (lane, item) = take_lane(item)?; + if self.delivery.deliver(lane, item).await? == DeliveryOutcome::ConsumerClosed { + return Ok(DeliveryOutcome::ConsumerClosed); + } + } + Ok(DeliveryOutcome::Accepted) + } + + async fn finish(&self) -> Result { + let _next_sequence = self.next_sequence.lock().await; + if let Some(response) = self.terminal_response() { + return Ok(response); + } + Ok(self.terminate(InboundTerminal::Completed).response()) + } + + async fn sender_fail(&self, cause: ErrorCode) -> DoExchangeResponse { + let _next_sequence = self.next_sequence.lock().await; + self.terminate(InboundTerminal::SenderFailed(cause)) + .response() + } + + /// Terminates the logical source because its local consumers can no longer accept data. + /// This releases its destinations and wakes active connections with a `FAIL` response. + pub fn fail(&self, cause: ErrorCode) { + self.terminate(InboundTerminal::ReceiverFailed(cause)); + } + + fn terminate(&self, requested: InboundTerminal) -> InboundTerminal { + let (terminal, installed) = { + let mut lifecycle = self.lifecycle.lock(); + match &lifecycle.terminal { + Some(terminal) => (terminal.clone(), false), + None => { + lifecycle.terminal = Some(requested.clone()); + (requested, true) + } + } + }; + if installed { + self.terminal_notified.notify_waiters(); + self.release(&terminal); + } + terminal + } + + fn detach(self: &Arc, runtime: &Arc, cause: ErrorCode) { + let lease = { + let mut lifecycle = self.lifecycle.lock(); + if lifecycle.terminal.is_some() { + return; + } + lifecycle.attachments -= 1; + if lifecycle.attachments != 0 { + return; + } + + lifecycle.generation += 1; + let generation = lifecycle.generation; + if self.reconnect_lease.is_zero() { + lifecycle.terminal = Some(InboundTerminal::ReceiverFailed(cause.clone())); + None + } else { + Some((generation, self.reconnect_lease)) + } + }; + + let Some((generation, reconnect_lease)) = lease else { + self.release(&InboundTerminal::ReceiverFailed(cause)); + return; + }; + + info!( + "do_exchange receiver waiting for replacement connection: {}, lease={:?}, generation={}", + self.source_label, reconnect_lease, generation + ); + let source = Arc::downgrade(self); + runtime.spawn(async move { + tokio::time::sleep(reconnect_lease).await; + if let Some(source) = source.upgrade() { + source.expire_lease(generation, cause); + } + }); + } + + fn expire_lease(&self, generation: u64, cause: ErrorCode) { + let failed = { + let mut lifecycle = self.lifecycle.lock(); + if lifecycle.terminal.is_some() + || lifecycle.attachments != 0 + || lifecycle.generation != generation + { + false + } else { + lifecycle.terminal = Some(InboundTerminal::ReceiverFailed(cause.clone())); + true + } + }; + if failed { + self.release(&InboundTerminal::ReceiverFailed(cause)); + } + } + + fn release(&self, terminal: &InboundTerminal) { + if let Some(cause) = terminal.cause() { + warn!( + "do_exchange logical source failed: {}, error={}", + self.source_label, cause + ); + } + self.delivery.terminate(terminal.cause().cloned()); + } + + fn terminal_response(&self) -> Option { + self.lifecycle + .lock() + .terminal + .as_ref() + .map(InboundTerminal::response) + } +} + +impl InboundTerminal { + fn response(&self) -> DoExchangeResponse { + match self { + Self::Completed => DoExchangeResponse::receiver_closed(), + Self::SenderFailed(_) => DoExchangeResponse::receiver_closed(), + Self::ReceiverFailed(cause) => DoExchangeResponse::fail(cause.clone()), + } + } + + fn cause(&self) -> Option<&ErrorCode> { + match self { + Self::Completed => None, + Self::SenderFailed(cause) | Self::ReceiverFailed(cause) => Some(cause), + } + } +} + +impl ReliableInboundConnection { + fn disconnect(&mut self) { + if let Some(cause) = self.disconnect_error.take() { + self.source.detach(&self.runtime, cause); + } + } + + pub(crate) async fn handle_request( + &self, + request: DoExchangeRequest, + ) -> Result { + match request { + DoExchangeRequest::Data { sequence, payload } => { + self.source.add_data(sequence, payload).await + } + DoExchangeRequest::Finish => self.source.finish().await, + DoExchangeRequest::SenderFail(cause) => Ok(self.source.sender_fail(cause).await), + } + } + + fn fail(&self, cause: ErrorCode) -> DoExchangeResponse { + self.source + .terminate(InboundTerminal::ReceiverFailed(cause)) + .response() + } + + pub async fn serve( + self, + mut stream: Streaming, + tx: async_channel::Sender>, + ) { + if let Some(response) = self.source.terminal_response() { + info!( + "do_exchange receiver serving terminal response to a late connection: {}", + self.source.source_label + ); + let _ = tx.send(Ok(response.encode())).await; + return; + } + + loop { + let result = tokio::select! { + result = stream.next() => result, + _ = self.source.terminal_notified.notified() => { + if let Some(response) = self.source.terminal_response() { + let _ = tx.send(Ok(response.encode())).await; + return; + } + continue; + } + }; + let Some(result) = result else { + return; + }; + let flight_data = match result { + Ok(flight_data) => flight_data, + Err(status) => { + info!( + "do_exchange receiver request stream failed: {}, status={}", + self.source.source_label, status + ); + return; + } + }; + let request = match DoExchangeRequest::decode(flight_data) { + Ok(request) => request, + Err(cause) => { + let response = self.fail(cause); + let _ = tx.send(Ok(response.encode())).await; + return; + } + }; + let response = match self.handle_request(request).await { + Ok(response) => response, + Err(cause) => self.fail(cause), + }; + let terminal = matches!( + response, + DoExchangeResponse::ReceiverClosed | DoExchangeResponse::Fail(_) + ); + if tx.send(Ok(response.encode())).await.is_err() || terminal { + return; + } + } + } +} + +impl Drop for ReliableInboundConnection { + fn drop(&mut self) { + self.disconnect(); + } +} + +impl Drop for ReliableInboundSource { + fn drop(&mut self) { + self.terminate(InboundTerminal::ReceiverFailed(ErrorCode::AbortedQuery( + "do_exchange logical source dropped before reaching a terminal state", + ))); + } +} diff --git a/src/query/service/src/servers/flight/v1/transport/reliable/mod.rs b/src/query/service/src/servers/flight/v1/transport/reliable/mod.rs new file mode 100644 index 00000000000..2daac901cdc --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/reliable/mod.rs @@ -0,0 +1,29 @@ +// 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. + +mod inbound; +mod outbound; +mod protocol; +mod reconnect; + +pub use inbound::ReliableInboundConnection; +pub use inbound::ReliableInboundSource; +pub use outbound::DoExchangeConnector; +pub use outbound::DoExchangeTransport; +pub use outbound::PendingReliableOutbound; +pub use outbound::ReliableOutbound; +pub(crate) use protocol::DoExchangeRequest; +pub(crate) use protocol::DoExchangeResponse; +pub(crate) use reconnect::FlightConnectionAttempts; +pub use reconnect::FlightReconnectPolicy; diff --git a/src/query/service/src/servers/flight/v1/transport/reliable/outbound.rs b/src/query/service/src/servers/flight/v1/transport/reliable/outbound.rs new file mode 100644 index 00000000000..09af544592b --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/reliable/outbound.rs @@ -0,0 +1,724 @@ +// 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 std::collections::VecDeque; +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use arrow_flight::FlightData; +use async_channel::Receiver; +use async_channel::Sender; +use async_channel::TrySendError; +use databend_common_base::base::WatchNotify; +use databend_common_base::runtime::Runtime; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use futures::StreamExt; +use futures::stream::BoxStream; +use log::warn; +use parking_lot::Mutex; +use tokio::sync::OwnedSemaphorePermit; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; +use tonic::Status; + +use super::DoExchangeRequest; +use super::DoExchangeResponse; +use super::FlightConnectionAttempts; +use super::FlightReconnectPolicy; +use crate::servers::flight::FlightOperation; +use crate::servers::flight::add_flight_error_context; +use crate::servers::flight::v1::transport::OutboundStream; +use crate::servers::flight::v1::transport::StreamSendOutcome; +use crate::servers::flight::v1::transport::batch; +use crate::servers::flight::v1::transport::frame_lane; + +type FlightDataStream = BoxStream<'static, std::result::Result>; + +pub struct DoExchangeTransport { + pub send_tx: Sender, + pub response_stream: FlightDataStream, +} + +pub type DoExchangeConnector = Arc< + dyn Fn() -> Pin> + Send>> + Send + Sync, +>; + +/// A queued DATA packet holds its permit until it becomes the next in-flight request. +struct PendingPacket { + data: FlightData, + _permit: OwnedSemaphorePermit, +} + +/// The sole request awaiting a response, retained for reconnect replay. +struct InFlightPacket { + encoded: FlightData, + expected: ExpectedResponse, + // A replacement transport consumes this budget; only logical progress installs a fresh one. + reconnect_attempts: FlightConnectionAttempts, +} + +enum ExpectedResponse { + Ack(u64), + ReceiverClosed, +} + +#[derive(Clone)] +enum OutboundTerminal { + ReceiverClosed, + Failed(ErrorCode), +} + +impl OutboundTerminal { + fn into_send_result(self) -> Result { + match self { + Self::ReceiverClosed => Ok(StreamSendOutcome::ConsumerClosed), + Self::Failed(cause) => Err(cause), + } + } +} + +struct CompletionState { + terminal: Mutex>, + notified: WatchNotify, +} + +impl CompletionState { + fn new() -> Self { + Self { + terminal: Mutex::new(None), + notified: WatchNotify::new(), + } + } + + fn current(&self) -> Option> { + self.terminal + .lock() + .clone() + .map(OutboundTerminal::into_send_result) + } + + fn complete(&self, terminal: OutboundTerminal) { + let mut current = self.terminal.lock(); + if current.is_some() { + return; + } + *current = Some(terminal); + drop(current); + self.notified.notify_waiters(); + } + + async fn wait(&self) -> Result { + if let Some(result) = self.current() { + return result; + } + self.notified.notified().await; + self.current() + .expect("do_exchange completion must publish a terminal state") + } +} + +enum OutboundCommand { + Data { + channel: usize, + data: FlightData, + permit: OwnedSemaphorePermit, + }, + Finish, + SenderFail(ErrorCode), +} + +struct OutboundDriver { + physical: PhysicalConnection, + commands: Receiver, + cancellation: CancellationToken, + logical: LogicalConnection, +} + +struct PhysicalConnection { + transport: Option, + connector: DoExchangeConnector, + reconnect: FlightReconnectPolicy, + local_node_id: String, + remote_node_id: String, +} + +struct LogicalConnection { + close: Option, + in_flight: Option, + next_sequence: u64, + pending: Vec>, + max_batch_bytes: Option, +} + +enum CloseIntent { + /// Producer input ended normally; drain accepted DATA before FINISH. + Finish, + /// Producer failed; discard unsent DATA and preserve this error during terminal cleanup. + SenderFail(ErrorCode), +} + +pub struct ReliableOutbound { + num_threads: usize, + commands: Sender, + slots: Arc, + cancellation: CancellationToken, + completion: Arc, +} + +pub struct PendingReliableOutbound { + num_threads: usize, + physical: PhysicalConnection, +} + +impl PendingReliableOutbound { + pub async fn connect( + num_threads: usize, + connector: DoExchangeConnector, + reconnect: FlightReconnectPolicy, + local_node_id: String, + remote_node_id: String, + ) -> Result { + Ok(Self { + num_threads, + physical: PhysicalConnection::open(connector, reconnect, local_node_id, remote_node_id) + .await?, + }) + } + pub fn start( + self, + slots: Arc, + max_batch_bytes: Option, + runtime: &Runtime, + ) -> ReliableOutbound { + let (commands_tx, commands_rx) = async_channel::unbounded(); + let cancellation = CancellationToken::new(); + let completion = Arc::new(CompletionState::new()); + let driver = OutboundDriver { + physical: self.physical, + commands: commands_rx, + cancellation: cancellation.clone(), + logical: LogicalConnection { + close: None, + in_flight: None, + next_sequence: 0, + pending: (0..self.num_threads).map(|_| VecDeque::new()).collect(), + max_batch_bytes, + }, + }; + let task_completion = completion.clone(); + runtime.spawn(async move { + task_completion.complete(driver.run().await); + }); + + ReliableOutbound { + num_threads: self.num_threads, + commands: commands_tx, + slots, + cancellation, + completion, + } + } +} + +#[async_trait::async_trait] +impl OutboundStream for ReliableOutbound { + async fn send(&self, lane: usize, data: FlightData) -> Result { + debug_assert!(lane < self.num_threads, "too many channels"); + let data = frame_lane(lane, data)?; + if let Some(result) = self.completion.current() { + return result; + } + let permit = tokio::select! { + permit = self.slots.clone().acquire_owned() => permit.unwrap(), + result = self.completion.wait() => return result, + }; + if let Some(result) = self.completion.current() { + return result; + } + if self + .commands + .send(OutboundCommand::Data { + channel: lane, + data, + permit, + }) + .await + .is_err() + { + return self.completion.wait().await; + } + Ok(StreamSendOutcome::Accepted) + } + + /// Drains all accepted DATA before sending FINISH, then waits for ReceiverClosed. + async fn finish(&self) -> Result<()> { + if self.completion.current().is_none() { + let _ = self.commands.send(OutboundCommand::Finish).await; + } + self.completion.wait().await.map(|_| ()) + } + + /// Keeps the transport alive for the bounded failure handshake. This is cleanup and must not + /// delay returning the producer's original error to the query. + async fn fail(&self, cause: ErrorCode) { + if self.completion.current().is_none() { + let _ = self.commands.send(OutboundCommand::SenderFail(cause)).await; + } + let _ = self.completion.wait().await; + } + + fn abort(&self) { + self.cancellation.cancel(); + } + + fn is_closed(&self) -> bool { + self.completion.current().is_some() + } +} + +impl Drop for ReliableOutbound { + fn drop(&mut self) { + OutboundStream::abort(self); + } +} + +impl OutboundDriver { + async fn run(mut self) -> OutboundTerminal { + let terminal = self.drive().await; + self.commands.close(); + match &self.logical.close { + Some(CloseIntent::SenderFail(cause)) => OutboundTerminal::Failed(cause.clone()), + _ => terminal, + } + } + + async fn drive(&mut self) -> OutboundTerminal { + loop { + enum Event { + Cancelled, + Command(std::result::Result), + Response(std::result::Result), + } + + let event = tokio::select! { + _ = self.cancellation.cancelled() => Event::Cancelled, + command = self.commands.recv() => Event::Command(command), + response = self.physical.response_stream().next() => Event::Response( + response.unwrap_or_else(|| Err(Status::unavailable( + "do_exchange response stream ended before a terminal packet", + ))) + ), + }; + + match event { + Event::Cancelled | Event::Command(Err(_)) => { + return self.failed(ErrorCode::AbortedQuery("do_exchange was cancelled")); + } + Event::Command(Ok(command)) => { + if let Err(cause) = self.handle_command(command) { + return self.failed(cause); + } + } + Event::Response(Ok(data)) => match DoExchangeResponse::decode(data) { + Ok(DoExchangeResponse::ReceiverClosed) => { + return OutboundTerminal::ReceiverClosed; + } + Ok(DoExchangeResponse::Fail(cause)) => { + return self.failed(cause); + } + Ok(DoExchangeResponse::Ack { sequence }) => { + if let Err(cause) = self.acknowledge(sequence) { + return self.failed(cause); + } + } + Err(cause) => { + return self.failed(cause); + } + }, + Event::Response(Err(status)) => { + if let Err(cause) = self.reconnect_transport(status).await { + return OutboundTerminal::Failed(cause); + } + } + } + } + } + + fn handle_command(&mut self, command: OutboundCommand) -> Result<()> { + match command { + OutboundCommand::Data { + channel, + data, + permit, + } => { + match &self.logical.close { + Some(CloseIntent::SenderFail(_)) => return Ok(()), + Some(CloseIntent::Finish) => { + return Err(ErrorCode::Internal( + "cannot send data after do_exchange producer completion", + )); + } + None => {} + } + if self.logical.in_flight.is_none() { + self.install_data(data); + } else { + self.logical.pending[channel].push_back(PendingPacket { + data, + _permit: permit, + }); + } + } + OutboundCommand::Finish => { + if self.logical.close.is_none() { + self.logical.close = Some(CloseIntent::Finish); + } + self.send_next(); + } + OutboundCommand::SenderFail(cause) => { + if !matches!(&self.logical.close, Some(CloseIntent::SenderFail(_))) { + self.logical.close = Some(CloseIntent::SenderFail(cause)); + } + for channel in &mut self.logical.pending { + channel.clear(); + } + self.send_next(); + } + } + Ok(()) + } + + fn acknowledge(&mut self, sequence: u64) -> Result<()> { + match self.logical.in_flight.take() { + Some(InFlightPacket { + expected: ExpectedResponse::Ack(expected), + .. + }) if sequence == expected => { + self.send_next(); + Ok(()) + } + Some(InFlightPacket { + expected: ExpectedResponse::Ack(expected), + .. + }) => Err(ErrorCode::Internal(format!( + "received do_exchange ACK for sequence {}, expected {}", + sequence, expected + ))), + Some(InFlightPacket { + expected: ExpectedResponse::ReceiverClosed, + .. + }) => Err(ErrorCode::Internal( + "received do_exchange ACK while waiting for ReceiverClosed", + )), + None => Err(ErrorCode::Internal( + "received do_exchange ACK without an in-flight request", + )), + } + } + + fn send_next(&mut self) { + if self.logical.in_flight.is_some() { + return; + } + if let Some(CloseIntent::SenderFail(cause)) = &self.logical.close { + self.install_sender_fail(cause.clone()); + return; + } + if let Some(data) = pop_pending(&mut self.logical.pending, self.logical.max_batch_bytes) { + self.install_data(data); + return; + } + if matches!(&self.logical.close, Some(CloseIntent::Finish)) { + self.install_finish(); + } + } + + fn install_data(&mut self, data: FlightData) { + let sequence = self.logical.next_sequence; + self.logical.next_sequence += 1; + self.install_request( + DoExchangeRequest::data(sequence, data), + ExpectedResponse::Ack(sequence), + ); + } + + fn install_finish(&mut self) { + self.install_request( + DoExchangeRequest::finish(), + ExpectedResponse::ReceiverClosed, + ); + } + + fn install_sender_fail(&mut self, cause: ErrorCode) { + self.install_request( + DoExchangeRequest::sender_fail(cause), + ExpectedResponse::ReceiverClosed, + ); + } + + fn install_request(&mut self, request: DoExchangeRequest, expected: ExpectedResponse) { + let encoded = request.encode(); + self.physical.send(&encoded); + self.logical.in_flight = Some(InFlightPacket { + encoded, + expected, + reconnect_attempts: self.physical.reconnect.reconnect_attempts(), + }); + } + + async fn reconnect_transport(&mut self, status: Status) -> Result<()> { + let (replay, attempts) = match &self.logical.in_flight { + Some(packet) => (Some(packet.encoded.clone()), packet.reconnect_attempts), + None => (None, self.physical.reconnect.reconnect_attempts()), + }; + let attempts_used = self + .physical + .reconnect(status, replay, attempts, &self.cancellation) + .await?; + if let Some(packet) = &mut self.logical.in_flight { + packet.reconnect_attempts = packet.reconnect_attempts.consume(attempts_used); + } + Ok(()) + } + + fn failed(&self, cause: ErrorCode) -> OutboundTerminal { + let cause = self.physical.contextualize(cause); + self.physical.warn_failure(&cause); + OutboundTerminal::Failed(cause) + } +} + +impl PhysicalConnection { + async fn open( + connector: DoExchangeConnector, + reconnect: FlightReconnectPolicy, + local_node_id: String, + remote_node_id: String, + ) -> Result { + let mut connection = Self { + transport: None, + connector, + reconnect, + local_node_id, + remote_node_id, + }; + let attempts = reconnect.initial_attempts(); + let (transport, _) = connection.establish(None, attempts).await?; + connection.transport = Some(transport); + Ok(connection) + } + + fn response_stream(&mut self) -> &mut FlightDataStream { + &mut self + .transport + .as_mut() + .expect("driver must reconnect before polling a transport") + .response_stream + } + + fn send(&self, encoded: &FlightData) { + let send_tx = &self + .transport + .as_ref() + .expect("sending requires a physical transport") + .send_tx; + match send_tx.try_send(encoded.clone()) { + Ok(()) | Err(TrySendError::Closed(_)) => {} + Err(TrySendError::Full(_)) => { + unreachable!("stop-and-wait request channel unexpectedly full") + } + } + } + + async fn reconnect( + &mut self, + status: Status, + replay: Option, + attempts: FlightConnectionAttempts, + cancellation: &CancellationToken, + ) -> Result { + self.transport = None; + if !is_retryable_status(&status) || attempts.is_empty() { + let cause = status_to_error(status, &self.local_node_id, &self.remote_node_id); + self.warn_failure(&cause); + return Err(cause); + } + + let cancelled = self.contextualize(ErrorCode::AbortedQuery("do_exchange was cancelled")); + let establish = self.establish(replay, attempts); + let (transport, attempts_used) = tokio::select! { + _ = cancellation.cancelled() => { + self.warn_failure(&cancelled); + return Err(cancelled); + } + result = establish => result, + }?; + warn!( + "do_exchange sender reconnected: client={}, service={}, attempts={}, initial_status={}", + self.local_node_id, self.remote_node_id, attempts_used, status + ); + self.transport = Some(transport); + Ok(attempts_used) + } + + fn establish( + &self, + replay: Option, + attempts: FlightConnectionAttempts, + ) -> impl Future> + Send + 'static { + let connector = self.connector.clone(); + let reconnect = self.reconnect; + let local_node_id = self.local_node_id.clone(); + let remote_node_id = self.remote_node_id.clone(); + + async move { + let attempts = attempts.remaining(); + for attempt in 0..attempts { + if attempt > 0 { + tokio::time::sleep(reconnect.retry_interval).await; + } + + let failure = match tokio::time::timeout(reconnect.timeout, (connector)()).await { + Ok(Ok(transport)) => { + let Some(encoded) = &replay else { + return Ok((transport, attempt + 1)); + }; + match transport.send_tx.try_send(encoded.clone()) { + Ok(()) => return Ok((transport, attempt + 1)), + Err(TrySendError::Full(_)) => unreachable!( + "new do_exchange request channel cannot be full before replay" + ), + Err(TrySendError::Closed(_)) => status_to_error( + Status::unavailable( + "new do_exchange transport closed before replay", + ), + &local_node_id, + &remote_node_id, + ), + } + } + Ok(Err(cause)) if cause.code() == ErrorCode::CANNOT_CONNECT_NODE => cause, + Ok(Err(cause)) => { + warn!( + "do_exchange connection attempt failed: client={}, service={}, attempt={}/{}, error={}", + local_node_id, + remote_node_id, + attempt + 1, + attempts, + cause + ); + return Err(cause); + } + Err(_) => status_to_error( + Status::deadline_exceeded(format!( + "connection attempt {}/{} exceeded its {:?} deadline", + attempt + 1, + attempts, + reconnect.timeout + )), + &local_node_id, + &remote_node_id, + ), + }; + + warn!( + "do_exchange connection attempt failed: client={}, service={}, attempt={}/{}, error={}", + local_node_id, + remote_node_id, + attempt + 1, + attempts, + failure + ); + + if attempt + 1 == attempts { + return Err(failure.add_message_back(format!( + "do_exchange connection exhausted after {} attempts", + attempts + ))); + } + } + + unreachable!("establish returns from its final attempt") + } + } + + fn contextualize(&self, cause: ErrorCode) -> ErrorCode { + add_flight_error_context( + cause, + FlightOperation::DoExchange, + &self.local_node_id, + &self.remote_node_id, + ) + } + + fn warn_failure(&self, cause: &ErrorCode) { + warn!( + "do_exchange sender failed: client={}, service={}, error={}", + self.local_node_id, self.remote_node_id, cause + ); + } +} + +fn status_to_error(status: Status, local_node_id: &str, remote_node_id: &str) -> ErrorCode { + let cause = if status.code() == tonic::Code::Aborted { + ErrorCode::AbortedQuery(status.message()) + } else { + status.into() + }; + add_flight_error_context( + cause, + FlightOperation::DoExchange, + local_node_id, + remote_node_id, + ) +} + +fn pop_pending( + pending: &mut [VecDeque], + max_batch_bytes: Option, +) -> Option { + let channel = pending.iter_mut().max_by_key(|channel| channel.len())?; + let first = channel.pop_front()?; + let Some(max_batch_bytes) = max_batch_bytes else { + return Some(first.data); + }; + + let mut total = first.data.data_body.len(); + let mut items = vec![first.data]; + while total < max_batch_bytes { + let Some(next) = channel.pop_front() else { + break; + }; + total += next.data.data_body.len(); + items.push(next.data); + } + if items.len() == 1 { + return items.pop(); + } + Some(batch::merge(items)) +} + +fn is_retryable_status(status: &Status) -> bool { + status.details().is_empty() + && matches!( + status.code(), + tonic::Code::Cancelled + | tonic::Code::Unknown + | tonic::Code::DeadlineExceeded + | tonic::Code::Internal + | tonic::Code::Unavailable + ) +} diff --git a/src/query/service/src/servers/flight/v1/transport/reliable/protocol.rs b/src/query/service/src/servers/flight/v1/transport/reliable/protocol.rs new file mode 100644 index 00000000000..946c3184126 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/reliable/protocol.rs @@ -0,0 +1,219 @@ +// 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 std::mem::size_of; + +use arrow_flight::FlightData; +use bytes::BufMut; +use bytes::BytesMut; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; + +// do_exchange is a logical stop-and-wait stream carried by replaceable physical Flight streams. +// Each logical stream has at most one unacknowledged request. A DATA response is sent only after +// the payload has crossed the receiver's deduplication boundary: it has been accepted by the +// logical inbound queue and the expected sequence has advanced. The sender can therefore replay +// the in-flight request without delivering its payload twice: +// +// sender receiver +// | | +// |-------- DATA(sequence=N) ---------->| accept N and advance expected sequence +// |<------------- ACK(N) ---------------| +// | | +// |-------- DATA(sequence=N+1) -------->| accept N+1 +// | physical stream lost | ACK is lost +// | | +// | open replacement stream | +// |-------- DATA(sequence=N+1) -------->| duplicate: do not deliver again +// |<------------- ACK(N+1) -------------| +// | | +// |------------- FINISH --------------->| complete the logical receiver +// |<-------- RECEIVER_CLOSED ------------| +// | | +// +// Either endpoint may also terminate the logical stream with an error. The outbound client sends +// SENDER_FAIL when its producer fails; the inbound service returns FAIL when its consumer fails. +// Physical EOF or transport failure is never logical completion. Until a response is received, +// the sender retains the encoded request so it can be replayed after reconnect. + +const DATA_KIND: u8 = 1; +const ACK_KIND: u8 = 2; +const FINISH_KIND: u8 = 3; +const RECEIVER_CLOSED_KIND: u8 = 4; +const FAIL_KIND: u8 = 5; +const SENDER_FAIL_KIND: u8 = 6; +const HEADER_LEN: usize = 1 + size_of::(); + +pub(crate) enum DoExchangeRequest { + Data { sequence: u64, payload: FlightData }, + Finish, + SenderFail(ErrorCode), +} + +impl DoExchangeRequest { + pub(crate) fn data(sequence: u64, payload: FlightData) -> Self { + Self::Data { sequence, payload } + } + + pub(crate) fn finish() -> Self { + Self::Finish + } + + pub(crate) fn sender_fail(cause: ErrorCode) -> Self { + Self::SenderFail(cause) + } + + pub(crate) fn encode(&self) -> FlightData { + match self { + Self::Data { sequence, payload } => { + let mut payload = payload.clone(); + let mut metadata = BytesMut::with_capacity(HEADER_LEN + payload.app_metadata.len()); + encode_header(&mut metadata, DATA_KIND, *sequence); + metadata.extend_from_slice(&payload.app_metadata); + payload.app_metadata = metadata.freeze(); + payload + } + Self::Finish => encode_control_packet(FINISH_KIND, 0), + Self::SenderFail(cause) => encode_error_packet(SENDER_FAIL_KIND, cause.clone()), + } + } + + pub(crate) fn decode(mut data: FlightData) -> Result { + let (kind, sequence) = decode_header(&data)?; + match kind { + DATA_KIND => { + data.app_metadata = data.app_metadata.slice(HEADER_LEN..); + Ok(Self::Data { + sequence, + payload: data, + }) + } + FINISH_KIND => { + validate_control_packet(&data)?; + Ok(Self::Finish) + } + SENDER_FAIL_KIND => { + data.app_metadata = data.app_metadata.slice(HEADER_LEN..); + Ok(Self::SenderFail(ErrorCode::try_from(data)?)) + } + ACK_KIND | RECEIVER_CLOSED_KIND | FAIL_KIND => Err(ErrorCode::Internal( + "received a do_exchange response packet on the request stream", + )), + _ => Err(unknown_packet_kind(kind)), + } + } +} + +pub(crate) enum DoExchangeResponse { + Ack { sequence: u64 }, + ReceiverClosed, + Fail(ErrorCode), +} + +impl DoExchangeResponse { + pub(crate) fn ack(sequence: u64) -> Self { + Self::Ack { sequence } + } + + pub(crate) fn receiver_closed() -> Self { + Self::ReceiverClosed + } + + pub(crate) fn fail(cause: ErrorCode) -> Self { + Self::Fail(cause) + } + + pub(crate) fn encode(self) -> FlightData { + match self { + Self::Ack { sequence } => encode_control_packet(ACK_KIND, sequence), + Self::ReceiverClosed => encode_control_packet(RECEIVER_CLOSED_KIND, 0), + Self::Fail(cause) => encode_error_packet(FAIL_KIND, cause), + } + } + + pub(crate) fn decode(mut data: FlightData) -> Result { + let (kind, sequence) = decode_header(&data)?; + match kind { + ACK_KIND => { + validate_control_packet(&data)?; + Ok(Self::Ack { sequence }) + } + RECEIVER_CLOSED_KIND => { + validate_control_packet(&data)?; + Ok(Self::ReceiverClosed) + } + FAIL_KIND => { + data.app_metadata = data.app_metadata.slice(HEADER_LEN..); + Ok(Self::Fail(ErrorCode::try_from(data)?)) + } + DATA_KIND | FINISH_KIND | SENDER_FAIL_KIND => Err(ErrorCode::Internal( + "received a do_exchange request packet on the response stream", + )), + _ => Err(unknown_packet_kind(kind)), + } + } +} + +fn encode_header(metadata: &mut BytesMut, kind: u8, sequence: u64) { + metadata.put_u8(kind); + metadata.put_u64_le(sequence); +} + +fn encode_control_packet(kind: u8, sequence: u64) -> FlightData { + let mut metadata = BytesMut::with_capacity(HEADER_LEN); + encode_header(&mut metadata, kind, sequence); + FlightData { + app_metadata: metadata.freeze(), + ..Default::default() + } +} + +fn encode_error_packet(kind: u8, cause: ErrorCode) -> FlightData { + let mut data = FlightData::from(cause); + let mut metadata = BytesMut::with_capacity(HEADER_LEN + data.app_metadata.len()); + encode_header(&mut metadata, kind, 0); + metadata.extend_from_slice(&data.app_metadata); + data.app_metadata = metadata.freeze(); + data +} + +fn decode_header(data: &FlightData) -> Result<(u8, u64)> { + if data.app_metadata.len() < HEADER_LEN { + return Err(ErrorCode::Internal( + "do_exchange packet has an incomplete header", + )); + } + + let mut sequence_bytes = [0; size_of::()]; + sequence_bytes.copy_from_slice(&data.app_metadata[1..HEADER_LEN]); + Ok((data.app_metadata[0], u64::from_le_bytes(sequence_bytes))) +} + +fn validate_control_packet(data: &FlightData) -> Result<()> { + if data.app_metadata.len() == HEADER_LEN + && data.flight_descriptor.is_none() + && data.data_header.is_empty() + && data.data_body.is_empty() + { + return Ok(()); + } + + Err(ErrorCode::Internal( + "do_exchange control packet contains a data payload", + )) +} + +fn unknown_packet_kind(kind: u8) -> ErrorCode { + ErrorCode::Internal(format!("unknown do_exchange packet kind {}", kind)) +} diff --git a/src/query/service/src/servers/flight/v1/transport/reliable/reconnect.rs b/src/query/service/src/servers/flight/v1/transport/reliable/reconnect.rs new file mode 100644 index 00000000000..3f46865c72c --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/reliable/reconnect.rs @@ -0,0 +1,109 @@ +// 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 std::time::Duration; + +use databend_common_exception::Result; +use databend_common_settings::Settings; + +const RECEIVER_LEASE_MARGIN: Duration = Duration::from_secs(5); + +#[derive(Clone, Copy)] +pub struct FlightReconnectPolicy { + retry_times: u64, + pub(crate) retry_interval: Duration, + pub(crate) timeout: Duration, +} + +#[derive(Clone, Copy)] +pub(crate) struct FlightConnectionAttempts { + remaining: u64, +} + +impl FlightConnectionAttempts { + pub(crate) fn remaining(self) -> u64 { + self.remaining + } + + pub(crate) fn is_empty(self) -> bool { + self.remaining == 0 + } + + pub(crate) fn consume(mut self, used: u64) -> Self { + self.remaining = self + .remaining + .checked_sub(used) + .expect("connection attempts used must not exceed the available budget"); + self + } + + fn max_elapsed(self, timeout: Duration, retry_interval: Duration) -> Duration { + let attempts = self.remaining.min(u32::MAX as u64) as u32; + let intervals = attempts.saturating_sub(1); + timeout + .saturating_mul(attempts) + .saturating_add(retry_interval.saturating_mul(intervals)) + } +} + +impl FlightReconnectPolicy { + pub fn new(retry_times: u64, retry_interval: Duration, timeout: Duration) -> Self { + Self { + retry_times, + retry_interval, + timeout, + } + } + + /// Returns the reconnect policy for New Flight, or `None` when the query keeps the + /// existing Flight path. Production query setup is the only caller. + pub fn from_settings(settings: &Settings) -> Result> { + if !settings.get_enable_experiment_new_flight()? { + return Ok(None); + } + + Ok(Some(Self::new( + settings.get_flight_max_retry_times()?, + Duration::from_secs(settings.get_flight_retry_interval()?), + Duration::from_secs(settings.get_flight_client_timeout()?), + ))) + } + + pub fn receiver_lease_secs(self) -> u64 { + self.receiver_lease().as_secs() + } + + pub(crate) fn initial_attempts(self) -> FlightConnectionAttempts { + FlightConnectionAttempts { + remaining: self.retry_times.saturating_add(1), + } + } + + pub(crate) fn reconnect_attempts(self) -> FlightConnectionAttempts { + FlightConnectionAttempts { + remaining: self.retry_times, + } + } + + pub fn receiver_lease(self) -> Duration { + let attempts = self.reconnect_attempts(); + if attempts.is_empty() { + return Duration::ZERO; + } + + attempts + .max_elapsed(self.timeout, self.retry_interval) + .saturating_add(std::cmp::max(self.retry_interval, RECEIVER_LEASE_MARGIN)) + } +} diff --git a/src/query/service/src/servers/flight/v1/transport/stream.rs b/src/query/service/src/servers/flight/v1/transport/stream.rs new file mode 100644 index 00000000000..008c6986341 --- /dev/null +++ b/src/query/service/src/servers/flight/v1/transport/stream.rs @@ -0,0 +1,89 @@ +// 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 std::sync::Arc; + +use arrow_flight::FlightData; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use futures::future::BoxFuture; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StreamSendOutcome { + Accepted, + ConsumerClosed, +} + +#[async_trait::async_trait] +pub trait OutboundStream: Send + Sync { + async fn send(&self, lane: usize, data: FlightData) -> Result; + + /// Completes a logical stream normally. Reliable implementations wait for the peer's terminal + /// response; legacy implementations may complete as soon as their local channel is closed. + async fn finish(&self) -> Result<()>; + + /// Best-effort failure handshake. Cleanup must not replace the producer's original error. + async fn fail(&self, cause: ErrorCode); + + /// Cancels without completing the logical stream. + fn abort(&self); + + fn is_closed(&self) -> bool; +} + +pub type OutboundStreamRef = Arc; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum DeliveryOutcome { + Accepted, + ConsumerClosed, +} + +/// Execution-side destination for one logical inbound stream. +/// +/// Transport implementations own framing, sequencing, replay, and terminal handshakes. Delivery +/// adapters own payload routing, backpressure, and consumer lifecycle. Every `deliver` call +/// contains one original logical payload; transport batches are removed before crossing this seam. +#[async_trait::async_trait] +pub trait InboundDelivery: Send + Sync { + async fn deliver(&self, lane: usize, data: FlightData) -> Result; + + fn is_closed(&self) -> bool; + + /// Returns a sticky consumer-close notification when the delivery can detect idle closure. + fn consumer_closed(&self) -> Option>; + + /// Releases the delivery after the logical stream reaches its first terminal state. + fn terminate(&self, cause: Option); +} + +pub(crate) fn frame_lane(lane: usize, mut data: FlightData) -> Result { + let lane = u16::try_from(lane) + .map_err(|_| ErrorCode::Internal(format!("Flight stream lane {lane} exceeds u16")))?; + let mut metadata = lane.to_le_bytes().to_vec(); + metadata.extend_from_slice(&data.app_metadata); + data.app_metadata = metadata.into(); + Ok(data) +} + +pub(crate) fn take_lane(mut data: FlightData) -> Result<(usize, FlightData)> { + if data.app_metadata.len() < 2 { + return Err(ErrorCode::BadBytes( + "Flight stream payload is missing its lane", + )); + } + let lane = u16::from_le_bytes([data.app_metadata[0], data.app_metadata[1]]) as usize; + data.app_metadata = data.app_metadata.slice(2..); + Ok((lane, data)) +} diff --git a/src/query/service/tests/it/servers/flight/mod.rs b/src/query/service/tests/it/servers/flight/mod.rs index e40099a15aa..01b04afac19 100644 --- a/src/query/service/tests/it/servers/flight/mod.rs +++ b/src/query/service/tests/it/servers/flight/mod.rs @@ -13,3 +13,4 @@ // limitations under the License. mod flight_service; +mod reliable_transport; diff --git a/src/query/service/tests/it/servers/flight/reliable_transport.rs b/src/query/service/tests/it/servers/flight/reliable_transport.rs new file mode 100644 index 00000000000..c762cfbef5d --- /dev/null +++ b/src/query/service/tests/it/servers/flight/reliable_transport.rs @@ -0,0 +1,525 @@ +// 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 std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use arrow_flight::Action; +use arrow_flight::ActionType; +use arrow_flight::Criteria; +use arrow_flight::Empty; +use arrow_flight::FlightData; +use arrow_flight::FlightDescriptor; +use arrow_flight::FlightEndpoint; +use arrow_flight::FlightInfo; +use arrow_flight::HandshakeRequest; +use arrow_flight::HandshakeResponse; +use arrow_flight::PollInfo; +use arrow_flight::PutResult; +use arrow_flight::SchemaResult; +use arrow_flight::Ticket; +use arrow_flight::flight_service_client::FlightServiceClient; +use arrow_flight::flight_service_server::FlightService; +use arrow_flight::flight_service_server::FlightServiceServer; +use async_channel::Receiver; +use databend_common_base::runtime::Runtime; +use databend_common_exception::ErrorCode; +use databend_common_exception::Result; +use databend_common_grpc::ConnectionFactory; +use databend_query::servers::flight::v1::transport::DeliveryOutcome; +use databend_query::servers::flight::v1::transport::InboundDelivery; +use databend_query::servers::flight::v1::transport::OutboundStream; +use databend_query::servers::flight::v1::transport::StreamSendOutcome; +use databend_query::servers::flight::v1::transport::reliable::DoExchangeConnector; +use databend_query::servers::flight::v1::transport::reliable::DoExchangeTransport; +use databend_query::servers::flight::v1::transport::reliable::FlightReconnectPolicy; +use databend_query::servers::flight::v1::transport::reliable::PendingReliableOutbound; +use databend_query::servers::flight::v1::transport::reliable::ReliableInboundSource; +use databend_query::servers::flight::v1::transport::reliable::ReliableOutbound; +use futures::Stream; +use futures::StreamExt; +use parking_lot::Mutex; +use socket2::SockRef; +use tokio::io::copy_bidirectional; +use tokio::net::TcpListener; +use tokio::net::TcpStream; +use tokio::sync::Semaphore; +use tokio::sync::oneshot; +use tokio_stream::wrappers::TcpListenerStream; +use tokio_util::sync::CancellationToken; +use tonic::Request; +use tonic::Response; +use tonic::Status; +use tonic::Streaming; +use tonic::transport::Server; + +struct TestDelivery { + sender: async_channel::Sender, + terminal: Mutex>>, + consumer_closed: CancellationToken, +} + +impl TestDelivery { + fn create() -> (Arc, Receiver) { + let (sender, receiver) = async_channel::unbounded(); + ( + Arc::new(Self { + sender, + terminal: Mutex::new(None), + consumer_closed: CancellationToken::new(), + }), + receiver, + ) + } + + fn close_consumer(&self) { + self.consumer_closed.cancel(); + } + + fn terminal_error(&self) -> Option { + self.terminal.lock().clone().flatten() + } +} + +#[async_trait::async_trait] +impl InboundDelivery for TestDelivery { + async fn deliver(&self, _lane: usize, data: FlightData) -> Result { + if self.consumer_closed.is_cancelled() { + return Ok(DeliveryOutcome::ConsumerClosed); + } + self.sender + .send(data) + .await + .map(|_| DeliveryOutcome::Accepted) + .map_err(|_| ErrorCode::AbortedQuery("test delivery closed")) + } + + fn is_closed(&self) -> bool { + self.consumer_closed.is_cancelled() || self.sender.is_closed() + } + + fn consumer_closed(&self) -> Option> { + let closed = self.consumer_closed.clone(); + Some(Box::pin(async move { closed.cancelled().await })) + } + + fn terminate(&self, cause: Option) { + let mut terminal = self.terminal.lock(); + if terminal.is_none() { + *terminal = Some(cause); + self.sender.close(); + } + } +} + +type FlightStream = Pin> + Send + 'static>>; + +#[derive(Clone)] +struct ReliableFlightService { + source: Arc, + runtime: Arc, +} + +#[tonic::async_trait] +impl FlightService for ReliableFlightService { + type HandshakeStream = FlightStream; + + async fn handshake( + &self, + _: Request>, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("handshake")) + } + + type ListFlightsStream = FlightStream; + + async fn list_flights( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("list_flights")) + } + + async fn get_flight_info( + &self, + _: Request, + ) -> std::result::Result, Status> { + Ok(Response::new( + FlightInfo::new().with_endpoint(FlightEndpoint::new()), + )) + } + + async fn poll_flight_info( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("poll_flight_info")) + } + + async fn get_schema( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("get_schema")) + } + + type DoGetStream = FlightStream; + + async fn do_get( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("do_get")) + } + + type DoPutStream = FlightStream; + + async fn do_put( + &self, + _: Request>, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("do_put")) + } + + type DoExchangeStream = FlightStream; + + async fn do_exchange( + &self, + request: Request>, + ) -> std::result::Result, Status> { + let connection = self.source.connect( + self.runtime.clone(), + ErrorCode::CannotConnectNode("TCP attachment disconnected"), + ); + let (tx, rx) = async_channel::bounded(1); + let stream = request.into_inner(); + databend_common_base::runtime::spawn(async move { + connection.serve(stream, tx).await; + }); + Ok(Response::new(Box::pin(rx))) + } + + type DoActionStream = FlightStream; + + async fn do_action( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("do_action")) + } + + type ListActionsStream = FlightStream; + + async fn list_actions( + &self, + _: Request, + ) -> std::result::Result, Status> { + Err(Status::unimplemented("list_actions")) + } +} + +#[derive(Debug)] +enum ProxyFault { + Reset, + ResetAndRejectReconnects, +} + +struct ProxyConnection { + fault: oneshot::Sender, +} + +struct TransportHarness { + _runtime: Arc, + delivery: Arc, + receiver: Receiver, + source: Arc, + outbound: Arc, + connections: Receiver, + server_shutdown: Option>, + proxy_task: tokio::task::JoinHandle<()>, + server_task: tokio::task::JoinHandle<()>, +} + +impl TransportHarness { + async fn create(retry_times: u64, slots: usize) -> Self { + let runtime = Arc::new(Runtime::with_worker_threads(1, None).unwrap()); + let (delivery, receiver) = TestDelivery::create(); + let source = Arc::new(ReliableInboundSource::new( + delivery.clone(), + Duration::from_secs(5), + "test receiver".to_string(), + )); + + let backend_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let backend_addr = backend_listener.local_addr().unwrap(); + let (server_shutdown, server_shutdown_rx) = oneshot::channel(); + let service = ReliableFlightService { + source: source.clone(), + runtime: runtime.clone(), + }; + let server_task = databend_common_base::runtime::spawn(async move { + Server::builder() + .add_service(FlightServiceServer::new(service)) + .serve_with_incoming_shutdown( + TcpListenerStream::new(backend_listener), + async move { + let _ = server_shutdown_rx.await; + }, + ) + .await + .unwrap(); + }); + + let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_addr = proxy_listener.local_addr().unwrap(); + let reject_reconnects = Arc::new(AtomicBool::new(false)); + let (connection_tx, connections) = async_channel::unbounded(); + let proxy_task = { + let reject_reconnects = reject_reconnects.clone(); + databend_common_base::runtime::spawn(async move { + while let Ok((mut client, _)) = proxy_listener.accept().await { + if reject_reconnects.load(Ordering::SeqCst) { + let _ = SockRef::from(&client).set_linger(Some(Duration::ZERO)); + continue; + } + let mut backend = TcpStream::connect(backend_addr).await.unwrap(); + let (fault, fault_rx) = oneshot::channel(); + if connection_tx.send(ProxyConnection { fault }).await.is_err() { + return; + } + let reject_reconnects = reject_reconnects.clone(); + databend_common_base::runtime::spawn(async move { + tokio::select! { + _ = copy_bidirectional(&mut client, &mut backend) => {} + command = fault_rx => { + if let Ok(command) = command { + if matches!(command, ProxyFault::ResetAndRejectReconnects) { + reject_reconnects.store(true, Ordering::SeqCst); + } + let _ = SockRef::from(&client) + .set_linger(Some(Duration::ZERO)); + let _ = SockRef::from(&backend) + .set_linger(Some(Duration::ZERO)); + } + } + } + }); + } + }) + }; + + let connector: DoExchangeConnector = Arc::new(move || { + Box::pin(async move { + let channel = ConnectionFactory::create_rpc_channel( + proxy_addr, + Some(Duration::from_millis(500)), + None, + None, + ) + .await + .map_err(ErrorCode::from)?; + let mut client = FlightServiceClient::new(channel); + let (send_tx, send_rx) = async_channel::bounded(1); + let response = client + .do_exchange(Request::new(send_rx)) + .await + .map_err(ErrorCode::from)?; + Ok(DoExchangeTransport { + send_tx, + response_stream: response.into_inner().boxed(), + }) + }) + }); + let reconnect = + FlightReconnectPolicy::new(retry_times, Duration::ZERO, Duration::from_millis(500)); + let pending = PendingReliableOutbound::connect( + 1, + connector, + reconnect, + "test sender".to_string(), + "test receiver".to_string(), + ) + .await + .unwrap(); + let outbound = Arc::new(pending.start(Arc::new(Semaphore::new(slots)), None, &runtime)); + + Self { + _runtime: runtime, + delivery, + receiver, + source, + outbound, + connections, + server_shutdown: Some(server_shutdown), + proxy_task, + server_task, + } + } + + fn payload(value: u8) -> FlightData { + FlightData { + data_body: vec![value].into(), + ..Default::default() + } + } + + async fn next_connection(&self) -> ProxyConnection { + tokio::time::timeout(Duration::from_secs(5), self.connections.recv()) + .await + .expect("timed out waiting for a TCP connection") + .expect("TCP proxy stopped") + } +} + +impl Drop for TransportHarness { + fn drop(&mut self) { + if let Some(shutdown) = self.server_shutdown.take() { + let _ = shutdown.send(()); + } + self.proxy_task.abort(); + self.server_task.abort(); + } +} + +// Scenario: DATA reaches the consumer and normal producer completion closes both endpoints. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn normal_transfer_completes_both_endpoints() { + let harness = TransportHarness::create(3, 8).await; + + assert_eq!( + harness + .outbound + .send(0, TransportHarness::payload(1)) + .await + .unwrap(), + StreamSendOutcome::Accepted + ); + harness.outbound.finish().await.unwrap(); + + assert_eq!( + harness.receiver.recv().await.unwrap(), + TransportHarness::payload(1) + ); + assert!(harness.receiver.recv().await.is_err()); + assert!(harness.delivery.terminal_error().is_none()); +} + +// Scenario: An idle downstream consumer can close the logical stream without more DATA. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn idle_consumer_close_reaches_the_sender() { + let harness = TransportHarness::create(3, 8).await; + + harness.delivery.close_consumer(); + tokio::time::timeout(Duration::from_secs(1), harness.outbound.finish()) + .await + .expect("consumer close must wake the sender") + .unwrap(); + + assert!(harness.outbound.is_closed()); + assert!(harness.delivery.terminal_error().is_none()); +} + +// Scenario: An idle receiver pipeline failure reaches the sender with its original cause. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn idle_receiver_failure_reaches_the_sender() { + let harness = TransportHarness::create(3, 8).await; + let cause = ErrorCode::AbortedQuery("receiver pipeline failed"); + + harness.source.fail(cause.clone()); + let error = tokio::time::timeout(Duration::from_secs(1), harness.outbound.finish()) + .await + .expect("receiver failure must wake the sender") + .unwrap_err(); + + assert_eq!(error.code(), cause.code()); + assert!(error.message().contains(&cause.message())); +} + +// Scenario: A producer failure terminates the receiver with the producer's original error. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn producer_failure_reaches_the_receiver() { + let harness = TransportHarness::create(3, 8).await; + let cause = ErrorCode::AbortedQuery("producer serialization failed"); + + harness.outbound.fail(cause.clone()).await; + + let error = harness + .delivery + .terminal_error() + .expect("producer failure must terminate the receiver"); + assert_eq!(error.code(), cause.code()); + assert_eq!(error.message(), cause.message()); +} + +// Scenario: Cancellation releases a producer blocked by transport backpressure. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn cancellation_releases_a_backpressured_sender() { + let harness = TransportHarness::create(3, 0).await; + let mut send = Box::pin(harness.outbound.send(0, TransportHarness::payload(2))); + + assert!(futures::poll!(&mut send).is_pending()); + harness.outbound.abort(); + let error = tokio::time::timeout(Duration::from_secs(1), send) + .await + .expect("cancellation must release the sender") + .unwrap_err(); + + assert_eq!(error.code(), ErrorCode::ABORTED_QUERY); +} + +// Scenario: A real TCP reset after DATA is accepted reconnects and delivers it exactly once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn tcp_reset_during_transfer_recovers_without_duplicate_data() { + let harness = TransportHarness::create(3, 8).await; + let first = harness.next_connection().await; + + assert_eq!( + harness + .outbound + .send(0, TransportHarness::payload(3)) + .await + .unwrap(), + StreamSendOutcome::Accepted + ); + first.fault.send(ProxyFault::Reset).unwrap(); + let _replacement = harness.next_connection().await; + + harness.outbound.finish().await.unwrap(); + assert_eq!( + harness.receiver.recv().await.unwrap(), + TransportHarness::payload(3) + ); + assert!(harness.receiver.recv().await.is_err()); +} + +// Scenario: Reconnect failure stops after the configured retry budget instead of hanging. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn reconnect_budget_exhaustion_returns_an_error() { + let harness = TransportHarness::create(3, 8).await; + let first = harness.next_connection().await; + + first + .fault + .send(ProxyFault::ResetAndRejectReconnects) + .unwrap(); + let error = tokio::time::timeout(Duration::from_secs(5), harness.outbound.finish()) + .await + .expect("reconnect exhaustion must not hang") + .unwrap_err(); + + assert!(!error.message().is_empty()); +} diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs index 358ae5cc87d..6789648c5c5 100644 --- a/src/query/settings/src/settings_default.rs +++ b/src/query/settings/src/settings_default.rs @@ -1514,6 +1514,13 @@ impl DefaultSettings { scope: SettingScope::Both, range: Some(SettingRange::Numeric(0..=10)), }), + ("enable_experiment_new_flight", DefaultSettingValue { + value: UserSettingValue::UInt64(0), + desc: "Enable the experimental New Flight data path.", + mode: SettingMode::Both, + scope: SettingScope::Both, + range: Some(SettingRange::Numeric(0..=1)), + }), ("network_policy", DefaultSettingValue { value: UserSettingValue::String("".to_owned()), desc: "Network policy for all users in the tenant", diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs index 6b7013155f1..48662517499 100644 --- a/src/query/settings/src/settings_getter_setter.rs +++ b/src/query/settings/src/settings_getter_setter.rs @@ -1090,6 +1090,10 @@ impl Settings { self.try_get_u64("flight_connection_retry_interval") } + pub fn get_enable_experiment_new_flight(&self) -> Result { + Ok(self.try_get_u64("enable_experiment_new_flight")? != 0) + } + pub fn get_hash_shuffle_rows_threshold(&self) -> Result { Ok(self.try_get_u64("hash_shuffle_rows_threshold")? as usize) } diff --git a/tests/query-flight-reconnect/test_tpch_reconnect.py b/tests/query-flight-reconnect/test_tpch_reconnect.py new file mode 100755 index 00000000000..450321abb8a --- /dev/null +++ b/tests/query-flight-reconnect/test_tpch_reconnect.py @@ -0,0 +1,636 @@ +#!/usr/bin/env python3 + +"""Run the complete TPC-H SF1 suite across three nodes while partitioning a +randomly selected live Flight TCP link. Both the injection time and partition +duration vary; packet counters and reconnect logs keep the assertions stable. +""" + +import argparse +import random +import subprocess +import tempfile +import threading +import time +from dataclasses import dataclass +from datetime import UTC +from datetime import datetime +from pathlib import Path +from typing import Any + +import mysql.connector + + +MYSQL_PORT = 3307 +FLIGHT_PORTS = frozenset({9091, 9092, 9093}) +IPTABLES_COMMENT_PREFIX = "databend-tpch-flight-reconnect" +RESET_IPTABLES_COMMENT = "databend-tpch-flight-reconnect-reset" +PARTITION_IPTABLES_COMMENT = "databend-tpch-flight-reconnect-partition" +POLL_INTERVAL_SECONDS = 0.1 +QUERY_DISCOVERY_TIMEOUT_SECONDS = 30 +FAULT_MATCH_TIMEOUT_SECONDS = 5 +RECONNECT_TIMEOUT_SECONDS = 30 +SUITE_TIMEOUT_SECONDS = 300 +MAX_WORKLOAD_ROUNDS = 5 +MAX_FAULT_ATTEMPTS = 60 +FAULT_DELAY_RANGE_SECONDS = (0.2, 3.0) +PARTITION_DURATION_RANGE_SECONDS = (0.5, 2.5) +RECONNECT_LOG_MARKERS = ( + "do_exchange connection attempt failed", + "do_exchange sender reconnected", +) +RECONNECT_CONFIRMED_MARKER = "do_exchange sender reconnected" + + +def chaos_log(message: str) -> None: + timestamp = datetime.now(UTC).isoformat(timespec="milliseconds") + print(f"{timestamp} TPCH_FLIGHT_CHAOS {message}", flush=True) + + +@dataclass(frozen=True, order=True) +class FlightLink: + source_host: str + source_port: int + destination_host: str + destination_port: int + + def __str__(self) -> str: + return ( + f"{self.source_host}:{self.source_port}->" + f"{self.destination_host}:{self.destination_port}" + ) + + +@dataclass(frozen=True) +class ActiveQuery: + query_id: str + sql: str + + +class OperationLog: + def __init__(self, path: Path): + path.parent.mkdir(parents=True, exist_ok=True) + self._file = path.open("w", encoding="utf-8", buffering=1) + self._lock = threading.Lock() + + def record(self, message: str) -> None: + timestamp = datetime.now(UTC).isoformat(timespec="milliseconds") + line = f"{timestamp} {message}" + with self._lock: + chaos_log(message) + self._file.write(f"{line}\n") + + def close(self) -> None: + self._file.close() + + +class QueryLogWatcher: + def __init__(self, repo_dir: Path): + self._repo_dir = repo_dir + self._offsets: dict[Path, int] = {} + self._partials: dict[Path, str] = {} + self._appended: list[tuple[Path, str]] = [] + self.checkpoint() + + def _log_paths(self) -> set[Path]: + paths = set((self._repo_dir / ".databend").glob("query-*.out")) + for node in range(1, 4): + log_dir = self._repo_dir / f".databend/logs_{node}" + if log_dir.exists(): + paths.update(path for path in log_dir.rglob("*") if path.is_file()) + return paths + + def _read_appended(self) -> None: + for path in self._log_paths(): + try: + size = path.stat().st_size + offset = self._offsets.get(path, size) + if size < offset: + offset = 0 + self._partials.pop(path, None) + if size > offset: + with path.open("rb") as log_file: + log_file.seek(offset) + chunk = log_file.read().decode(errors="replace") + text = self._partials.pop(path, "") + chunk + lines = text.splitlines(keepends=True) + for line in lines: + if line.endswith(("\n", "\r")): + self._appended.append( + (path, line.rstrip("\r\n").replace("\0", "")) + ) + else: + self._partials[path] = line + self._offsets[path] = size + except FileNotFoundError: + continue + + def checkpoint(self) -> None: + self._read_appended() + self._appended = [] + + def _print_raw_log(self, path: Path, line: str) -> None: + try: + source = path.relative_to(self._repo_dir) + except ValueError: + source = path + print(f"TPCH_FLIGHT_RAW_LOG source={source}", flush=True) + print(line, flush=True) + + def wait_for_reconnect(self, stop: threading.Event) -> bool: + deadline = time.monotonic() + RECONNECT_TIMEOUT_SECONDS + inspected = 0 + while time.monotonic() < deadline and not stop.is_set(): + self._read_appended() + reconnect_confirmed = False + for path, line in self._appended[inspected:]: + if any(marker in line for marker in RECONNECT_LOG_MARKERS): + self._print_raw_log(path, line) + if RECONNECT_CONFIRMED_MARKER in line: + reconnect_confirmed = True + inspected = len(self._appended) + if reconnect_confirmed: + return True + stop.wait(POLL_INTERVAL_SECONDS) + return False + + +class FlightPartitionFault: + def __init__(self) -> None: + self._link: FlightLink | None = None + + @staticmethod + def _iptables( + *args: str, + check: bool = True, + ) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["sudo", "-n", "iptables", "-w", "5", *args], + check=check, + capture_output=True, + text=True, + ) + + @staticmethod + def _reset_rule(operation: str, link: FlightLink) -> list[str]: + return [ + operation, + "OUTPUT", + "-p", + "tcp", + "-s", + link.source_host, + "--sport", + str(link.source_port), + "-d", + link.destination_host, + "--dport", + str(link.destination_port), + "-m", + "comment", + "--comment", + RESET_IPTABLES_COMMENT, + "-j", + "REJECT", + "--reject-with", + "tcp-reset", + ] + + @staticmethod + def _partition_rule(operation: str, link: FlightLink) -> list[str]: + return [ + operation, + "OUTPUT", + "-p", + "tcp", + "-s", + link.source_host, + "-d", + link.destination_host, + "--dport", + str(link.destination_port), + "-m", + "comment", + "--comment", + PARTITION_IPTABLES_COMMENT, + "-j", + "DROP", + ] + + def apply(self, link: FlightLink) -> None: + self.clear() + self._link = link + try: + # Install the broad partition first, then put the exact reset rule + # ahead of it. The live connection receives RST while replacement + # connections remain blackholed until the partition is removed. + self._iptables(*self._partition_rule("-I", link)) + self._iptables(*self._reset_rule("-I", link)) + except BaseException: + self.clear() + raise + + def _clear_rule(self, rule: list[str]) -> None: + while self._iptables(*rule, check=False).returncode == 0: + pass + + def clear_reset(self) -> None: + if self._link is not None: + self._clear_rule(self._reset_rule("-D", self._link)) + + def clear_partition(self) -> None: + if self._link is not None: + self._clear_rule(self._partition_rule("-D", self._link)) + + def clear(self) -> None: + if self._link is not None: + self.clear_reset() + self.clear_partition() + self._link = None + + # Remove stale rules left by an interrupted previous run, but only when + # they carry this test's unique comment prefix. + while True: + result = self._iptables("-L", "OUTPUT", "--line-numbers", "-n", check=False) + matching_line = next( + ( + line.split()[0] + for line in result.stdout.splitlines() + if IPTABLES_COMMENT_PREFIX in line + ), + None, + ) + if matching_line is None: + return + self._iptables("-D", "OUTPUT", matching_line, check=False) + + def matched_reset_packets(self) -> int: + result = self._iptables("-L", "OUTPUT", "-v", "-n", "-x") + return sum( + int(line.split()[0]) + for line in result.stdout.splitlines() + if RESET_IPTABLES_COMMENT in line + ) + + +def split_endpoint(endpoint: str) -> tuple[str, int]: + host, port = endpoint.rsplit(":", 1) + return host.strip("[]"), int(port) + + +def flight_links() -> set[FlightLink]: + output = subprocess.run( + ["ss", "-Hnt", "state", "established"], + check=True, + capture_output=True, + text=True, + ).stdout + links = set() + for line in output.splitlines(): + fields = line.split() + if len(fields) < 4: + continue + try: + source_host, source_port = split_endpoint(fields[-2]) + destination_host, destination_port = split_endpoint(fields[-1]) + except ValueError: + continue + if destination_port not in FLIGHT_PORTS or source_port in FLIGHT_PORTS: + continue + links.add( + FlightLink( + source_host, + source_port, + destination_host, + destination_port, + ) + ) + return links + + +def connect_mysql() -> Any: + deadline = time.monotonic() + QUERY_DISCOVERY_TIMEOUT_SECONDS + last_error: Exception | None = None + while time.monotonic() < deadline: + try: + return mysql.connector.connect( + host="127.0.0.1", + user="root", + passwd="root", + port=MYSQL_PORT, + connection_timeout=5, + autocommit=True, + ) + except Exception as error: + last_error = error + time.sleep(0.5) + raise AssertionError(f"could not connect to Databend: {last_error}") from last_error + + +def active_workload_query(cursor: Any) -> ActiveQuery | None: + cursor.execute( + "SELECT current_query_id, extra_info FROM system.processes " + "WHERE current_query_id != '' AND extra_info != '' " + "AND extra_info NOT LIKE '%system.processes%' LIMIT 1" + ) + row = cursor.fetchone() + if row is None: + return None + return ActiveQuery(str(row[0]), " ".join(str(row[1]).split())[:240]) + + +def wait_for_fault_match( + fault: FlightPartitionFault, + stop: threading.Event, +) -> int: + deadline = time.monotonic() + FAULT_MATCH_TIMEOUT_SECONDS + while time.monotonic() < deadline and not stop.is_set(): + matched = fault.matched_reset_packets() + if matched > 0: + return matched + stop.wait(POLL_INTERVAL_SECONDS) + return 0 + + +class RandomLinkFaultInjector: + def __init__( + self, + cursor: Any, + fault: FlightPartitionFault, + log_watcher: QueryLogWatcher, + operations: OperationLog, + stop: threading.Event, + ): + self._cursor = cursor + self._fault = fault + self._log_watcher = log_watcher + self._operations = operations + self._stop = stop + self._random = random.SystemRandom() + self.attempted = 0 + self.confirmed_reconnects = 0 + self.error: Exception | None = None + + def _wait_for_candidate(self) -> bool: + while not self._stop.is_set(): + query = active_workload_query(self._cursor) + if query is not None and flight_links(): + return True + self._stop.wait(POLL_INTERVAL_SECONDS) + return False + + def run(self) -> None: + try: + while not self._stop.is_set() and self.attempted < MAX_FAULT_ATTEMPTS: + if not self._wait_for_candidate(): + return + + delay = self._random.uniform(*FAULT_DELAY_RANGE_SECONDS) + if self._stop.wait(delay): + return + + query = active_workload_query(self._cursor) + before = flight_links() + current_links = sorted(before) + if query is None or not current_links: + continue + + selected = self._random.choice(current_links) + partition_duration = self._random.uniform( + *PARTITION_DURATION_RANGE_SECONDS + ) + self.attempted += 1 + self._operations.record( + f"attempt={self.attempted} random_delay={delay:.3f}s " + f"partition_duration={partition_duration:.3f}s " + f"query_id={query.query_id} sql={query.sql!r} " + f"selected_link={selected} candidates={len(current_links)}" + ) + self._log_watcher.checkpoint() + + matched = 0 + partition_interrupted = False + try: + self._fault.apply(selected) + matched = wait_for_fault_match(self._fault, self._stop) + if matched == 0: + self._operations.record( + f"attempt={self.attempted} result=no_packet_matched " + f"selected_link={selected}" + ) + continue + + # Once the selected live connection has received RST, keep + # the broader DROP rule active so reconnect attempts observe + # a real, randomly timed one-way network partition. + self._fault.clear_reset() + self._operations.record( + f"attempt={self.attempted} fault=tcp_reset_and_partition " + f"matched_packets={matched} selected_link={selected} " + f"partition_active=1" + ) + partition_started = time.monotonic() + partition_interrupted = self._stop.wait(partition_duration) + actual_duration = time.monotonic() - partition_started + self._fault.clear_partition() + self._operations.record( + f"attempt={self.attempted} partition=healed " + f"planned_duration={partition_duration:.3f}s " + f"actual_duration={actual_duration:.3f}s" + ) + finally: + self._fault.clear() + + if partition_interrupted: + return + + reconnect_logged = self._log_watcher.wait_for_reconnect(self._stop) + if not reconnect_logged: + self._operations.record( + f"attempt={self.attempted} result=reconnect_not_confirmed " + "reconnect_log=0" + ) + continue + + current = flight_links() + replacement = next( + ( + link + for link in current - before + if link.destination_host == selected.destination_host + and link.destination_port == selected.destination_port + ), + None, + ) + self.confirmed_reconnects += 1 + self._operations.record( + f"attempt={self.attempted} result=reconnected " + f"old_link={selected} new_link={replacement or 'not_observed'} " + f"reconnect_log=1 confirmed_reconnects={self.confirmed_reconnects}" + ) + except Exception as error: + self.error = error + self._operations.record( + f"injector=failed error={type(error).__name__}: {error}" + ) + self._stop.set() + finally: + self._fault.clear() + + +def write_new_flight_suite(source: Path, destination: Path) -> None: + prefix = """statement ok +set enable_experiment_new_flight = 1; + +statement ok +set flight_connection_max_retry_times = 10; + +statement ok +set flight_connection_retry_interval = 1; + +statement ok +set group_by_shuffle_mode = 'before_partial'; + +""" + destination.write_text( + prefix + source.read_text(encoding="utf-8"), encoding="utf-8" + ) + + +def forward_output(process: subprocess.Popen[str]) -> None: + assert process.stdout is not None + for line in process.stdout: + print(line, end="", flush=True) + + +def run_tpch_round(command: list[str], round_number: int) -> None: + chaos_log(f"workload_round={round_number} starting command={' '.join(command)}") + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + output_thread = threading.Thread( + target=forward_output, + args=(process,), + name="tpch-output", + daemon=True, + ) + output_thread.start() + try: + return_code = process.wait(timeout=SUITE_TIMEOUT_SECONDS) + except subprocess.TimeoutExpired: + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + raise AssertionError( + f"TPC-H round {round_number} exceeded {SUITE_TIMEOUT_SECONDS} seconds" + ) + finally: + output_thread.join(timeout=10) + + if return_code != 0: + raise AssertionError( + f"TPC-H sqllogictest round {round_number} exited with status {return_code}" + ) + chaos_log(f"workload_round={round_number} passed") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--sqllogictests", required=True, type=Path) + parser.add_argument("--tpch-suite", required=True, type=Path) + parser.add_argument("--operation-log", required=True, type=Path) + parser.add_argument("--repo-dir", required=True, type=Path) + args = parser.parse_args() + + operations = OperationLog(args.operation_log) + fault = FlightPartitionFault() + fault.clear() + stop = threading.Event() + connection = connect_mysql() + cursor = connection.cursor() + injector = RandomLinkFaultInjector( + cursor, + fault, + QueryLogWatcher(args.repo_dir), + operations, + stop, + ) + injector_thread = threading.Thread( + target=injector.run, + name="tpch-flight-fault-injector", + daemon=True, + ) + injector_started = False + + try: + operations.record( + "suite=starting randomness=system max_fault_attempts={} max_rounds={} " + "fault_delay_range={} partition_duration_range={}".format( + MAX_FAULT_ATTEMPTS, + MAX_WORKLOAD_ROUNDS, + FAULT_DELAY_RANGE_SECONDS, + PARTITION_DURATION_RANGE_SECONDS, + ) + ) + with tempfile.TemporaryDirectory(prefix="databend-tpch-reconnect-") as temp_dir: + suite = Path(temp_dir) / "queries.test" + write_new_flight_suite(args.tpch_suite, suite) + command = [ + str(args.sqllogictests), + "--handlers", + "mysql", + "--run", + str(suite), + "--parallel", + "1", + ] + injector_thread.start() + injector_started = True + completed_rounds = 0 + for round_number in range(1, MAX_WORKLOAD_ROUNDS + 1): + run_tpch_round(command, round_number) + completed_rounds = round_number + if injector.error is not None: + message = "fault injector failed during round {}: {}".format( + round_number, injector.error + ) + raise AssertionError(message) from injector.error + + stop.set() + injector_thread.join(timeout=RECONNECT_TIMEOUT_SECONDS) + if injector_thread.is_alive(): + raise AssertionError("fault injector did not stop") + if injector.error is not None: + raise AssertionError( + f"fault injector failed: {injector.error}" + ) from injector.error + if injector.confirmed_reconnects == 0: + raise AssertionError( + f"no reconnect was confirmed after {completed_rounds} complete TPC-H " + f"rounds and {injector.attempted} random link selections" + ) + operations.record( + f"suite=passed rounds={completed_rounds} attempts={injector.attempted} " + f"confirmed_reconnects={injector.confirmed_reconnects}" + ) + except BaseException as error: + operations.record(f"suite=failed error={type(error).__name__}: {error}") + raise + finally: + stop.set() + fault.clear() + if injector_started: + injector_thread.join(timeout=5) + cursor.close() + connection.close() + operations.close() + + +if __name__ == "__main__": + main()