Merge remote-tracking branch 'up/main' into range-shuffle

Signed-off-by: coldWater <forsaken628@gmail.com>

Author: forsaken628

src/query/service/src/pipelines/executor/processor_async_task.rs

@@ -32,7 +32,6 @@ use databend_common_pipeline_core::processors::ProcessorPtr;
 use futures_util::future::BoxFuture;
 use futures_util::future::Either;
 use futures_util::FutureExt;
-use log::error;
 use log::warn;
 use petgraph::prelude::NodeIndex;
 
@@ -125,11 +124,9 @@ impl ProcessorAsyncTask {
         let processor_id = unsafe { processor.id() };
         let processor_name = unsafe { processor.name() };
         let queue_clone = queue.clone();
-        let graph_clone = graph.clone();
         let inner = async move {
             let start = Instant::now();
             let mut inner = inner.boxed();
-            let mut log_graph = false;
 
             loop {
                 let interval = Box::pin(sleep(Duration::from_secs(5)));
@@ -138,29 +135,10 @@ impl ProcessorAsyncTask {
                         inner = right;
                         let elapsed = start.elapsed();
                         let active_workers = queue_clone.active_workers();
-                        match elapsed >= Duration::from_secs(200)
-                            && active_workers == 0
-                            && !log_graph
-                        {
-                            false => {
-                                warn!(
-                                    "Very slow processor async task, query_id:{:?}, processor id: {:?}, name: {:?}, elapsed: {:?}, active sync workers: {:?}",
-                                    query_id, processor_id, processor_name, elapsed, active_workers
-                                );
-                            }
-                            true => {
-                                log_graph = true;
-                                error!(
-                                    "Very slow processor async task, query_id:{:?}, processor id: {:?}, name: {:?}, elapsed: {:?}, active sync workers: {:?}, {}",
-                                    query_id,
-                                    processor_id,
-                                    processor_name,
-                                    elapsed,
-                                    active_workers,
-                                    graph_clone.format_graph_nodes(false)
-                                );
-                            }
-                        };
+                        warn!(
+                            "Very slow processor async task, query_id:{:?}, processor id: {:?}, name: {:?}, elapsed: {:?}, active sync workers: {:?}",
+                            query_id, processor_id, processor_name, elapsed, active_workers
+                        );
                     }
                     Either::Right((res, _)) => {
                         return res;

src/query/service/src/servers/admin/admin_service.rs

@@ -79,6 +79,10 @@ impl AdminService {
                 "v1/queries/:query_id/profiling",
                 get(super::v1::query_profiling::query_profiling_handler),
             )
+            .at(
+                "v1/queries/:query_id/graph_dump",
+                get(super::v1::query_dump::running_query_dump),
+            )
             .at("/debug/home", get(debug_home_handler))
             .at("/debug/pprof/profile", get(debug_pprof_handler))
             .at("/debug/async_tasks/dump", get(debug_dump_stack));

src/query/service/src/servers/admin/v1/mod.rs

@@ -16,6 +16,7 @@ pub mod cluster;
 pub mod config;
 pub mod instance_status;
 pub mod processes;
+pub mod query_dump;
 pub mod query_profiling;
 pub mod settings;
 pub mod stream_status;

src/query/service/src/servers/admin/v1/query_dump.rs

@@ -0,0 +1,72 @@
+// 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 databend_common_config::GlobalConfig;
+use databend_common_exception::Result;
+use http::StatusCode;
+use poem::web::Json;
+use poem::web::Path;
+use poem::IntoResponse;
+
+use crate::clusters::ClusterDiscovery;
+use crate::clusters::ClusterHelper;
+use crate::clusters::FlightParams;
+use crate::servers::flight::v1::actions::GET_RUNNING_QUERY_DUMP;
+
+#[poem::handler]
+#[async_backtrace::framed]
+pub async fn running_query_dump(Path(query_id): Path<String>) -> poem::Result<impl IntoResponse> {
+    #[derive(serde::Serialize)]
+    struct QueryRunningGraphDump {
+        query_id: String,
+        graph_dump: HashMap<String, String>,
+    }
+
+    let graph_dump = match get_running_query_dump(&query_id).await {
+        Ok(graph_dump) => graph_dump,
+        Err(cause) => {
+            return Err(poem::Error::from_string(
+                format!("Failed to fetch executor dump. cause: {cause}"),
+                StatusCode::INTERNAL_SERVER_ERROR,
+            ))
+        }
+    };
+
+    Ok(Json(QueryRunningGraphDump {
+        graph_dump,
+        query_id: query_id.clone(),
+    }))
+}
+
+async fn get_running_query_dump(query_id: &str) -> Result<HashMap<String, String>> {
+    let config = GlobalConfig::instance();
+    let cluster = ClusterDiscovery::instance().discover(&config).await?;
+
+    let mut message = HashMap::with_capacity(cluster.nodes.len());
+
+    for node_info in &cluster.nodes {
+        message.insert(node_info.id.clone(), query_id.to_owned());
+    }
+
+    let flight_params = FlightParams {
+        timeout: 60,
+        retry_times: 3,
+        retry_interval: 3,
+    };
+    cluster
+        .do_action::<_, String>(GET_RUNNING_QUERY_DUMP, message, flight_params)
+        .await
+}

src/query/service/src/servers/flight/v1/actions/flight_actions.rs

@@ -26,6 +26,7 @@ use serde::Deserialize;
 use serde::Serialize;
 
 use crate::servers::flight::v1::actions::get_profile::get_profile;
+use crate::servers::flight::v1::actions::get_running_query_dump::get_running_query_dump;
 use crate::servers::flight::v1::actions::init_query_env::init_query_env;
 use crate::servers::flight::v1::actions::init_query_env::INIT_QUERY_ENV;
 use crate::servers::flight::v1::actions::init_query_fragments::init_query_fragments;
@@ -37,6 +38,7 @@ use crate::servers::flight::v1::actions::system_action::system_action;
 use crate::servers::flight::v1::actions::truncate_table::truncate_table;
 use crate::servers::flight::v1::actions::truncate_table::TRUNCATE_TABLE;
 use crate::servers::flight::v1::actions::GET_PROFILE;
+use crate::servers::flight::v1::actions::GET_RUNNING_QUERY_DUMP;
 use crate::servers::flight::v1::actions::INIT_QUERY_FRAGMENTS;
 use crate::servers::flight::v1::actions::KILL_QUERY;
 use crate::servers::flight::v1::actions::START_PREPARED_QUERY;
@@ -135,4 +137,5 @@ pub fn flight_actions() -> FlightActions {
         .action(SET_PRIORITY, set_priority)
         .action(SYSTEM_ACTION, system_action)
         .action(GET_PROFILE, get_profile)
+        .action(GET_RUNNING_QUERY_DUMP, get_running_query_dump)
 }

src/query/service/src/servers/flight/v1/actions/get_running_query_dump.rs

@@ -0,0 +1,31 @@
+// 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 databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+
+use crate::servers::flight::v1::exchange::DataExchangeManager;
+use crate::sessions::SessionManager;
+
+pub static GET_RUNNING_QUERY_DUMP: &str = "/actions/get_running_query_dump";
+
+pub async fn get_running_query_dump(query_id: String) -> Result<String> {
+    match SessionManager::instance().get_running_graph_dump(&query_id) {
+        Ok(running_graph_dump) => Ok(running_graph_dump),
+        Err(cause) => match cause.code() == ErrorCode::UNKNOWN_QUERY {
+            false => Err(cause),
+            true => DataExchangeManager::instance().get_running_query_graph_dump(&query_id),
+        },
+    }
+}

src/query/service/src/servers/flight/v1/actions/mod.rs

@@ -14,6 +14,7 @@
 
 mod flight_actions;
 mod get_profile;
+mod get_running_query_dump;
 mod init_query_env;
 mod init_query_fragments;
 mod kill_query;
@@ -30,6 +31,7 @@ use databend_common_settings::Settings;
 pub use flight_actions::flight_actions;
 pub use flight_actions::FlightActions;
 pub use get_profile::GET_PROFILE;
+pub use get_running_query_dump::GET_RUNNING_QUERY_DUMP;
 pub use init_query_env::INIT_QUERY_ENV;
 pub use init_query_fragments::init_query_fragments;
 pub use init_query_fragments::INIT_QUERY_FRAGMENTS;

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

@@ -93,6 +93,27 @@ impl DataExchangeManager {
         GlobalInstance::get()
     }
 
+    pub fn get_running_query_graph_dump(&self, query_id: &str) -> Result<String> {
+        let running_executor = {
+            let queries_coordinator_guard = self.queries_coordinator.lock();
+            let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };
+            let Some(coordinator) = queries_coordinator.get(query_id) else {
+                return Ok(format!("Unknown query {}", query_id));
+            };
+
+            let Some(info) = &coordinator.info else {
+                return Ok(format!("Unknown running query {}", query_id));
+            };
+
+            info.query_executor.clone()
+        };
+
+        Ok(match running_executor {
+            None => format!("Unknown running query {}", query_id),
+            Some(executor) => executor.get_inner().format_graph_nodes(),
+        })
+    }
+
     pub fn get_query_ctx(&self, query_id: &str) -> Result<Arc<QueryContext>> {
         let queries_coordinator_guard = self.queries_coordinator.lock();
         let queries_coordinator = unsafe { &mut *queries_coordinator_guard.deref().get() };

src/query/service/src/sessions/query_ctx_shared.rs

@@ -678,6 +678,13 @@ impl QueryContextShared {
         }
     }
 
+    pub fn get_executor_graph_dump(&self) -> String {
+        match self.executor.read().upgrade() {
+            None => String::new(),
+            Some(executor) => executor.format_graph_nodes(),
+        }
+    }
+
     pub fn get_query_profiles(&self) -> Vec<PlanProfile> {
         if let Some(executor) = self.executor.read().upgrade() {
             self.add_query_profiles(&executor.fetch_profiling(false));

src/query/service/src/sessions/session_mgr.rs

@@ -373,6 +373,27 @@ impl SessionManager {
         )))
     }
 
+    pub fn get_running_graph_dump(&self, query_id: &str) -> Result<String> {
+        for weak_ptr in self.active_sessions_snapshot() {
+            let Some(arc_session) = weak_ptr.upgrade() else {
+                continue;
+            };
+
+            let session_ctx = arc_session.session_ctx.as_ref();
+
+            if let Some(context_shared) = session_ctx.get_query_context_shared() {
+                if query_id == *context_shared.init_query_id.as_ref().read() {
+                    return Ok(context_shared.get_executor_graph_dump());
+                }
+            }
+        }
+
+        Err(ErrorCode::UnknownQuery(format!(
+            "Unknown query {}",
+            query_id
+        )))
+    }
+
     fn active_sessions_snapshot(&self) -> Vec<Weak<Session>> {
         // Here the situation is the same of method `graceful_shutdown`:
         //