feat(frontend): add heartbeat between frontend and compute

proto/task_service.proto

@@ -68,11 +68,16 @@ message ExecuteRequest {
   plan_common.ExprContext expr_context = 5;
 }
 
+message HeartbeatRequest {}
+
+message HeartbeatResponse {}
+
 service TaskService {
   rpc CreateTask(CreateTaskRequest) returns (stream TaskInfoResponse);
   // Cancel an already-died (self execution-failure, previous aborted, completed) task will still succeed.
   rpc CancelTask(CancelTaskRequest) returns (CancelTaskResponse);
   rpc Execute(ExecuteRequest) returns (stream GetDataResponse);
+  rpc Heartbeat(HeartbeatRequest) returns (HeartbeatResponse);
 }
 
 message GetDataRequest {

src/batch/src/rpc/service/task_service.rs

@@ -20,7 +20,7 @@ use risingwave_pb::batch_plan::TaskOutputId;
 use risingwave_pb::task_service::task_service_server::TaskService;
 use risingwave_pb::task_service::{
     CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, ExecuteRequest, GetDataResponse,
-    TaskInfoResponse,
+    HeartbeatRequest, HeartbeatResponse, TaskInfoResponse,
 };
 use thiserror_ext::AsReport;
 use tokio_stream::wrappers::ReceiverStream;
@@ -122,6 +122,14 @@ impl TaskService for BatchServiceImpl {
         let mgr = self.mgr.clone();
         BatchServiceImpl::get_execute_stream(env, mgr, req).await
     }
+
+    #[cfg_attr(coverage, coverage(off))]
+    async fn heartbeat(
+        &self,
+        _req: Request<HeartbeatRequest>,
+    ) -> Result<Response<HeartbeatResponse>, Status> {
+        Ok(Response::new(HeartbeatResponse {}))
+    }
 }
 
 impl BatchServiceImpl {

src/frontend/src/session.rs

@@ -286,7 +286,7 @@ impl FrontendEnv {
             Arc::new(ComputeClientPool::new(config.server.connection_pool_size));
         let query_manager = QueryManager::new(
             worker_node_manager.clone(),
-            compute_client_pool,
+            compute_client_pool.clone(),
             catalog_reader.clone(),
             Arc::new(GLOBAL_DISTRIBUTED_QUERY_METRICS.clone()),
             batch_config.distributed_query_limit,
@@ -393,6 +393,52 @@ impl FrontendEnv {
         });
         join_handles.push(join_handle);
 
+        // Heartbeat between compute nodes (worker nodes) and frontend
+        let heartbeat_worker_node_manager = worker_node_manager.clone();
+        let join_handle = tokio::spawn(async move {
+            let duration = Duration::from_secs(std::cmp::max(
+                batch_config.mask_worker_temporary_secs as u64,
+                1,
+            ));
+
+            let mut check_heartbeat_interval =
+                tokio::time::interval(core::time::Duration::from_secs(2));
+            check_heartbeat_interval
+                .set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+            check_heartbeat_interval.reset();
+            let heartbeat_fail_retry_num = 3;
+            loop {
+                check_heartbeat_interval.tick().await;
+                for worker in heartbeat_worker_node_manager
+                    .list_worker_nodes()
+                    .into_iter()
+                    .filter(|w| w.property.as_ref().map_or(false, |p| p.is_serving))
+                {
+                    let addr: HostAddr = worker.get_host().unwrap().into();
+                    for i in 0..heartbeat_fail_retry_num {
+                        let success = if let Ok(client) =
+                            compute_client_pool.get_by_addr(addr.clone()).await
+                        {
+                            client.heartbeat().await.is_ok()
+                        } else {
+                            false
+                        };
+
+                        if success {
+                            break;
+                        } else {
+                            compute_client_pool.invalidate(&addr).await;
+                            if i == heartbeat_fail_retry_num - 1 {
+                                info!("Worker node {} is not reachable, mask it", worker.id);
+                                heartbeat_worker_node_manager.mask_worker_node(worker.id, duration);
+                            }
+                        }
+                    }
+                }
+            }
+        });
+        join_handles.push(join_handle);
+
         let total_memory_bytes = resource_util::memory::system_memory_available_bytes();
         let heap_profiler =
             HeapProfiler::new(total_memory_bytes, config.server.heap_profiling.clone());

src/rpc_client/src/compute_client.rs

@@ -37,8 +37,8 @@ use risingwave_pb::task_service::exchange_service_client::ExchangeServiceClient;
 use risingwave_pb::task_service::task_service_client::TaskServiceClient;
 use risingwave_pb::task_service::{
     permits, CancelTaskRequest, CancelTaskResponse, CreateTaskRequest, ExecuteRequest,
-    GetDataRequest, GetDataResponse, GetStreamRequest, GetStreamResponse, PbPermits,
-    TaskInfoResponse,
+    GetDataRequest, GetDataResponse, GetStreamRequest, GetStreamResponse, HeartbeatRequest,
+    HeartbeatResponse, PbPermits, TaskInfoResponse,
 };
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::UnboundedReceiverStream;
@@ -189,6 +189,15 @@ impl ComputeClient {
             .into_inner())
     }
 
+    pub async fn heartbeat(&self) -> Result<HeartbeatResponse> {
+        Ok(self
+            .task_client
+            .to_owned()
+            .heartbeat(HeartbeatRequest {})
+            .await?
+            .into_inner())
+    }
+
     pub async fn stack_trace(&self) -> Result<StackTraceResponse> {
         Ok(self
             .monitor_client

src/rpc_client/src/lib.rs

@@ -129,6 +129,11 @@ where
             .unwrap()
             .clone())
     }
+
+    /// Invalidates the clients for the given addr.
+    pub async fn invalidate(&self, addr: &HostAddr) {
+        self.clients.invalidate(addr).await;
+    }
 }
 
 /// `ExtraInfoSource` is used by heartbeat worker to pull extra info that needs to be piggybacked.