Merge branch 'main' into main

Author: mergify[bot]

src/common/base/src/base/mod.rs

@@ -18,6 +18,7 @@ mod profiling;
 mod progress;
 mod runtime;
 mod runtime_tracker;
+mod select;
 mod shutdown_signal;
 mod singleton_instance;
 mod stop_handle;
@@ -37,6 +38,8 @@ pub use runtime::Runtime;
 pub use runtime::TrySpawn;
 pub use runtime_tracker::RuntimeTracker;
 pub use runtime_tracker::ThreadTracker;
+pub use select::select3;
+pub use select::Select3Output;
 pub use shutdown_signal::signal_stream;
 pub use shutdown_signal::DummySignalStream;
 pub use shutdown_signal::SignalStream;

src/common/base/src/base/select.rs

@@ -0,0 +1,71 @@
+// Copyright 2022 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::future::Future;
+use std::pin::Pin;
+use std::task::Context;
+use std::task::Poll;
+
+use futures::FutureExt;
+
+pub fn select3<A, B, C>(fut1: A, fut2: B, fut3: C) -> Select3<A, B, C>
+where
+    A: Future + Unpin,
+    B: Future + Unpin,
+    C: Future + Unpin,
+{
+    Select3 {
+        inner: Some((fut1, fut2, fut3)),
+    }
+}
+
+pub struct Select3<A, B, C> {
+    inner: Option<(A, B, C)>,
+}
+
+impl<A: Unpin, B: Unpin, C: Unpin> Unpin for Select3<A, B, C> {}
+
+impl<A, B, C> Future for Select3<A, B, C>
+where
+    A: Future + Unpin,
+    B: Future + Unpin,
+    C: Future + Unpin,
+{
+    type Output = Select3Output<(A::Output, B, C), (B::Output, A, C), (C::Output, A, B)>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let (mut a, mut b, mut c) = self.inner.take().expect("cannot poll Select twice");
+
+        if let Poll::Ready(val) = a.poll_unpin(cx) {
+            return Poll::Ready(Select3Output::Left((val, b, c)));
+        }
+
+        if let Poll::Ready(val) = b.poll_unpin(cx) {
+            return Poll::Ready(Select3Output::Middle((val, a, c)));
+        }
+
+        if let Poll::Ready(val) = c.poll_unpin(cx) {
+            return Poll::Ready(Select3Output::Right((val, a, b)));
+        }
+
+        self.inner = Some((a, b, c));
+        Poll::Pending
+    }
+}
+
+pub enum Select3Output<A, B, C> {
+    Left(A),
+    Middle(B),
+    Right(C),
+}

src/query/service/src/api/mod.rs

@@ -30,6 +30,7 @@ pub use rpc::FragmentPayload;
 pub use rpc::FragmentPlanPacket;
 pub use rpc::InitNodesChannelPacket;
 pub use rpc::MergeExchange;
+pub use rpc::PrecommitBlock;
 pub use rpc::QueryFragmentsPlanPacket;
 pub use rpc::ServerFlightExchange;
 pub use rpc::ShuffleDataExchange;

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

@@ -40,7 +40,6 @@ use crate::api::rpc::flight_client::FlightExchange;
 use crate::api::rpc::flight_scatter_broadcast::BroadcastFlightScatter;
 use crate::api::rpc::flight_scatter_hash::HashFlightScatter;
 use crate::api::rpc::flight_scatter_hash_v2::HashFlightScatterV2;
-use crate::api::rpc::packets::DataPacket;
 use crate::api::rpc::Packet;
 use crate::api::DataExchange;
 use crate::api::FlightClient;
@@ -294,20 +293,14 @@ impl DataExchangeManager {
                     let query_id = ctx.get_id();
                     let mut statistics_receiver = statistics_receiver.lock();
 
-                    if let Some(error) = may_error {
-                        // TODO: try get recv error if network error.
-                        statistics_receiver.shutdown();
-                        ctx.get_exchange_manager().on_finished_query(&query_id);
-
-                        if error.code() == common_exception::ABORT_QUERY {
-                            statistics_receiver.take_receive_error()?;
-                        }
+                    statistics_receiver.shutdown();
+                    ctx.get_exchange_manager().on_finished_query(&query_id);
+                    statistics_receiver.wait_shutdown()?;
 
-                        return Err(error.clone());
+                    match may_error {
+                        None => Ok(()),
+                        Some(error_code) => Err(error_code.clone()),
                     }
-
-                    ctx.get_exchange_manager().on_finished_query(&query_id);
-                    statistics_receiver.wait_shutdown()
                 });
 
                 Ok(build_res)
@@ -584,39 +577,12 @@ impl QueryCoordinator {
         }
 
         let ctx = query_ctx.clone();
-        let request_server_exchange = request_server_exchanges[0].clone();
         let mut statistics_sender =
-            StatisticsSender::create(&query_id, ctx, request_server_exchange);
+            StatisticsSender::create(&query_id, ctx, request_server_exchanges.remove(0));
         statistics_sender.start();
 
         Thread::named_spawn(Some(String::from("Distributed-Executor")), move || {
-            if let Err(cause) = executor.execute() {
-                let request_server_exchange = request_server_exchanges.remove(0);
-
-                futures::executor::block_on(async move {
-                    statistics_sender.shutdown();
-
-                    if let Err(_cause) = request_server_exchange
-                        .send(DataPacket::ErrorCode(cause))
-                        .await
-                    {
-                        tracing::warn!(
-                            "Cannot send message to request server when executor failure."
-                        );
-                    };
-
-                    // Wait request server close channel.
-                    // request_server_exchange.close_output();
-                    // while let Ok(Some(_data)) = request_server_exchange.recv().await {}
-                });
-
-                return;
-            }
-
-            // TODO: destroy query resource when panic?
-            // Destroy coordinator_server_exchange if executor is not failure.
-            statistics_sender.shutdown();
-            drop(request_server_exchanges);
+            statistics_sender.shutdown(executor.execute().err());
             query_ctx
                 .get_exchange_manager()
                 .on_finished_query(&query_id);

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

@@ -35,8 +35,6 @@ use crate::api::rpc::exchange::exchange_params::SerializeParams;
 use crate::api::rpc::exchange::exchange_params::ShuffleExchangeParams;
 use crate::api::rpc::exchange::exchange_transform_source::ExchangeSourceTransform;
 use crate::api::rpc::flight_client::FlightExchange;
-use crate::api::rpc::packets::PrecommitBlock;
-use crate::api::rpc::packets::ProgressInfo;
 use crate::api::DataPacket;
 use crate::api::FragmentData;
 use crate::clusters::ClusterHelper;
@@ -256,9 +254,9 @@ impl Processor for ExchangeTransform {
         if let Some(remote_data) = self.remote_data.take() {
             return match remote_data {
                 DataPacket::ErrorCode(v) => self.on_recv_error(v),
-                DataPacket::Progress(v) => self.on_recv_progress(v),
+                DataPacket::ProgressAndPrecommit { .. } => unreachable!(),
+                DataPacket::FetchProgressAndPrecommit => unreachable!(),
                 DataPacket::FragmentData(v) => self.on_recv_data(v),
-                DataPacket::PrecommitBlock(v) => self.on_recv_precommit(v),
             };
         }
 
@@ -331,12 +329,4 @@ impl ExchangeTransform {
 
         Ok(())
     }
-
-    fn on_recv_progress(&mut self, _progress: ProgressInfo) -> Result<()> {
-        unimplemented!()
-    }
-
-    fn on_recv_precommit(&mut self, _fragment_data: PrecommitBlock) -> Result<()> {
-        unimplemented!()
-    }
 }

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

@@ -25,8 +25,6 @@ use common_exception::Result;
 use crate::api::rpc::exchange::exchange_params::ExchangeParams;
 use crate::api::rpc::exchange::exchange_params::MergeExchangeParams;
 use crate::api::rpc::flight_client::FlightExchange;
-use crate::api::rpc::packets::PrecommitBlock;
-use crate::api::rpc::packets::ProgressInfo;
 use crate::api::DataPacket;
 use crate::api::FragmentData;
 use crate::pipelines::processors::port::InputPort;
@@ -129,9 +127,9 @@ impl Processor for ExchangeSourceTransform {
         if let Some(remote_data) = self.remote_flight_data.take() {
             return match remote_data {
                 DataPacket::ErrorCode(v) => self.on_recv_error(v),
-                DataPacket::Progress(v) => self.on_recv_progress(v),
                 DataPacket::FragmentData(v) => self.on_recv_data(v),
-                DataPacket::PrecommitBlock(v) => self.on_recv_precommit(v),
+                DataPacket::FetchProgressAndPrecommit => unreachable!(),
+                DataPacket::ProgressAndPrecommit { .. } => unreachable!(),
             };
         }
 
@@ -178,12 +176,4 @@ impl ExchangeSourceTransform {
 
         Ok(())
     }
-
-    fn on_recv_progress(&mut self, _progress: ProgressInfo) -> Result<()> {
-        unimplemented!()
-    }
-
-    fn on_recv_precommit(&mut self, _fragment_data: PrecommitBlock) -> Result<()> {
-        unimplemented!()
-    }
 }