fix(query): fix the 'not found' error caused by heartbeat loss on the local (#18512)

Author: zhang2014

src/query/config/src/config.rs

@@ -2031,6 +2031,7 @@ impl TryInto<InnerQueryConfig> for QueryConfig {
                 .collect(),
             resources_management: self.resources_management,
             enable_queries_executor: self.enable_queries_executor,
+            check_connection_before_schedule: true,
         })
     }
 }

src/query/config/src/inner.rs

@@ -254,6 +254,7 @@ pub struct QueryConfig {
     pub resources_management: Option<ResourcesManagementConfig>,
 
     pub enable_queries_executor: bool,
+    pub check_connection_before_schedule: bool,
 }
 
 impl Default for QueryConfig {
@@ -343,6 +344,7 @@ impl Default for QueryConfig {
             settings: HashMap::new(),
             resources_management: None,
             enable_queries_executor: false,
+            check_connection_before_schedule: true,
         }
     }
 }

src/query/service/src/clusters/cluster.rs

@@ -294,9 +294,14 @@ impl ClusterDiscovery {
                 Err(cause.add_message_back("(while cluster api get_nodes)."))
             }
             Ok(cluster_nodes) => {
+                let mut has_local_node = false;
                 let mut res = Vec::with_capacity(cluster_nodes.len());
                 for node in &cluster_nodes {
-                    if node.id != self.local_id {
+                    if node.id == self.local_id {
+                        has_local_node = true;
+                    }
+
+                    if config.query.check_connection_before_schedule && node.id != self.local_id {
                         let start_at = Instant::now();
                         if let Err(cause) = create_client(config, &node.flight_address).await {
                             warn!(
@@ -313,6 +318,26 @@ impl ClusterDiscovery {
                     res.push(Arc::new(node.clone()));
                 }
 
+                // When this node loses heartbeat with the meta node but receives an SQL request from the client, we should attempt to use this node.
+                if !has_local_node && !res.is_empty() {
+                    let mut local_node_info = NodeInfo::create(
+                        config.query.node_id.clone(),
+                        config.query.node_secret.clone(),
+                        format!(
+                            "{}:{}",
+                            config.query.http_handler_host, config.query.http_handler_port
+                        ),
+                        config.query.flight_api_address.clone(),
+                        config.query.discovery_address.clone(),
+                        String::new(),
+                        String::new(),
+                    );
+
+                    local_node_info.cluster_id = res[0].cluster_id.clone();
+                    local_node_info.warehouse_id = res[0].warehouse_id.clone();
+                    res.push(Arc::new(local_node_info));
+                }
+
                 metrics_gauge_discovered_nodes(
                     &self.local_id,
                     &self.cluster_id,

src/query/service/tests/it/clusters.rs

@@ -86,6 +86,38 @@ async fn test_remove_invalid_nodes() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+async fn test_lost_local_cluster_discovery() -> Result<()> {
+    let _guard = TestFixture::setup().await?;
+
+    let mut config_1 = ConfigBuilder::create()
+        .query_flight_address("10.0.0.1")
+        .build();
+    let mut config_2 = ConfigBuilder::create()
+        .query_flight_address("10.0.0.2")
+        .build();
+
+    config_1.query.check_connection_before_schedule = false;
+    config_2.query.check_connection_before_schedule = false;
+
+    let metastore = ClusterDiscovery::create_meta_client(&config_1).await?;
+    let cluster_discovery_1 = ClusterDiscovery::try_create(&config_1, metastore.clone()).await?;
+    let cluster_discovery_2 = ClusterDiscovery::try_create(&config_2, metastore.clone()).await?;
+
+    cluster_discovery_2.register_to_metastore(&config_2).await?;
+
+    let cluster_1 = cluster_discovery_1.discover(&config_1).await?;
+    assert_eq!(cluster_1.get_nodes().len(), 2);
+    assert_eq!(cluster_1.is_local(&cluster_1.nodes[0]), false);
+    assert_eq!(cluster_1.is_local(&cluster_1.nodes[1]), true);
+
+    let cluster_2 = cluster_discovery_2.discover(&config_2).await?;
+    assert_eq!(cluster_2.get_nodes().len(), 1);
+    assert_eq!(cluster_2.is_local(&cluster_2.nodes[0]), true);
+
+    Ok(())
+}
+
 // TODO:(Winter) need kvapi::KVApi for cluster multiple nodes test
 // #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
 // async fn test_multiple_cluster_discovery() -> Result<()> {