feat: add health-aware gRPC client routing (#8684)

* feat: add gRPC client health routing

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: harden gRPC client health routing

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: defer gRPC client health checks until first use

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>

Cargo.lock

@@ -2068,6 +2068,7 @@ dependencies = [
  "common-meta",
  "common-query",
  "common-recordbatch",
+ "common-runtime",
  "common-telemetry",
  "datatypes",
  "enum_dispatch",

src/client/Cargo.toml

@@ -24,6 +24,7 @@ common-macro.workspace = true
 common-meta.workspace = true
 common-query.workspace = true
 common-recordbatch.workspace = true
+common-runtime.workspace = true
 common-telemetry.workspace = true
 datatypes.workspace = true
 enum_dispatch = "0.3"

src/client/src/client.rs

@@ -13,6 +13,8 @@
 // limitations under the License.
 
 use std::sync::Arc;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::time::Duration;
 
 use api::v1::HealthCheckRequest;
 use api::v1::flow::flow_client::FlowClient as PbFlowClient;
@@ -31,6 +33,27 @@ use tonic::transport::Channel;
 use crate::load_balance::{LoadBalance, Loadbalancer};
 use crate::{Result, error};
 
+const DEFAULT_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(30);
+const DEFAULT_HEALTH_CHECK_TIMEOUT: Duration = Duration::from_secs(1);
+
+/// Options for a gRPC client.
+#[derive(Clone, Debug)]
+pub struct ClientOptions {
+    /// Interval for refreshing peer health. `Duration::ZERO` disables background health checks.
+    pub health_check_interval: Duration,
+    /// Timeout for checking the health of a peer.
+    pub health_check_timeout: Duration,
+}
+
+impl Default for ClientOptions {
+    fn default() -> Self {
+        Self {
+            health_check_interval: DEFAULT_HEALTH_CHECK_INTERVAL,
+            health_check_timeout: DEFAULT_HEALTH_CHECK_TIMEOUT,
+        }
+    }
+}
+
 pub struct FlightClient {
     addr: String,
     client: FlightServiceClient<Channel>,
@@ -51,30 +74,140 @@ pub struct Client {
     inner: Arc<Inner>,
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug)]
 struct Inner {
     channel_manager: ChannelManager,
-    peers: Arc<RwLock<Vec<String>>>,
+    peers: RwLock<Peers>,
     load_balance: Loadbalancer,
+    health_check_interval: Duration,
+    health_check_timeout: Duration,
+    health_check_started: AtomicBool,
+}
+
+impl Default for Inner {
+    fn default() -> Self {
+        Self::with_manager_and_peers(ChannelManager::new(), Vec::new(), ClientOptions::default())
+    }
+}
+
+#[derive(Debug, Default, PartialEq, Eq)]
+struct PeerStates {
+    active: Vec<usize>,
+    inactive: Vec<usize>,
+}
+
+#[derive(Debug, Default)]
+struct Peers {
+    addresses: Vec<String>,
+    states: PeerStates,
+    generation: u64,
 }
 
 impl Inner {
-    fn with_manager(channel_manager: ChannelManager) -> Self {
+    fn with_manager_and_peers(
+        channel_manager: ChannelManager,
+        peers: Vec<String>,
+        options: ClientOptions,
+    ) -> Self {
+        let peer_count = peers.len();
         Self {
             channel_manager,
-            ..Default::default()
+            peers: RwLock::new(Peers {
+                addresses: peers,
+                states: PeerStates {
+                    active: (0..peer_count).collect(),
+                    inactive: Vec::new(),
+                },
+                generation: 0,
+            }),
+            load_balance: Loadbalancer::default(),
+            health_check_interval: options.health_check_interval,
+            health_check_timeout: options.health_check_timeout,
+            health_check_started: AtomicBool::new(false),
         }
     }
 
-    fn set_peers(&self, peers: Vec<String>) {
-        let mut guard = self.peers.write();
-        *guard = peers;
+    fn set_peers(&self, addresses: Vec<String>) {
+        let peer_count = addresses.len();
+        let mut peers = self.peers.write();
+        peers.addresses = addresses;
+        peers.states = PeerStates {
+            active: (0..peer_count).collect(),
+            inactive: Vec::new(),
+        };
+        peers.generation = peers.generation.wrapping_add(1);
+    }
+
+    fn peer_count(&self) -> usize {
+        self.peers.read().addresses.len()
     }
 
     fn get_peer(&self) -> Option<String> {
-        let guard = self.peers.read();
-        self.load_balance.get_peer(&guard).cloned()
+        let peers = self.peers.read();
+        let index = self
+            .load_balance
+            .get_index(&peers.states.active)
+            .or_else(|| self.load_balance.get_index(&peers.states.inactive))?;
+        Some(peers.addresses[*index].clone())
     }
+
+    async fn refresh_peer_states(&self) {
+        let (generation, peers) = {
+            let peers = self.peers.read();
+            let addresses = peers
+                .states
+                .active
+                .iter()
+                .chain(&peers.states.inactive)
+                .map(|&index| (index, peers.addresses[index].clone()))
+                .collect::<Vec<_>>();
+            (peers.generation, addresses)
+        };
+        let health_checks = peers.into_iter().map(|(index, addr)| async move {
+            let is_active = self.check_peer_health(&addr).await;
+            (index, is_active)
+        });
+        let results = futures::future::join_all(health_checks).await;
+
+        let (active, inactive) = results.into_iter().fold(
+            (Vec::new(), Vec::new()),
+            |(mut active, mut inactive), (index, is_active)| {
+                if is_active {
+                    active.push(index);
+                } else {
+                    inactive.push(index);
+                }
+                (active, inactive)
+            },
+        );
+
+        let mut peers = self.peers.write();
+        if peers.generation == generation {
+            peers.states = PeerStates { active, inactive };
+        }
+    }
+
+    async fn check_peer_health(&self, addr: &str) -> bool {
+        let Ok(channel) = self.channel_manager.get(addr) else {
+            return false;
+        };
+        let mut client = HealthCheckClient::new(channel);
+        tokio::time::timeout(
+            self.health_check_timeout,
+            client.health_check(HealthCheckRequest {}),
+        )
+        .await
+        .is_ok_and(|result| result.is_ok())
+    }
+}
+
+fn random_initial_delay(max_delay: Duration) -> Duration {
+    let max_nanos = max_delay.as_nanos().min(u64::MAX as u128) as u64;
+    if max_nanos == 0 {
+        return Duration::ZERO;
+    }
+
+    Duration::from_nanos(rand::random_range(0..max_nanos))
 }
 
 impl Client {
@@ -87,10 +220,32 @@ impl Client {
         U: AsRef<str>,
         A: AsRef<[U]>,
     {
-        Self::with_manager_and_urls(ChannelManager::new(), urls)
+        Self::with_urls_and_options(urls, ClientOptions::default())
+    }
+
+    /// Creates a client with URLs and custom options.
+    pub fn with_urls_and_options<U, A>(urls: A, options: ClientOptions) -> Self
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        Self::with_manager_and_urls_and_options(ChannelManager::new(), urls, options)
     }
 
     pub fn with_tls_and_urls<U, A>(urls: A, client_tls: ClientTlsOption) -> Result<Self>
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
+        Self::with_tls_and_urls_and_options(urls, client_tls, ClientOptions::default())
+    }
+
+    /// Creates a client with TLS URLs and custom options.
+    pub fn with_tls_and_urls_and_options<U, A>(
+        urls: A,
+        client_tls: ClientTlsOption,
+        options: ClientOptions,
+    ) -> Result<Self>
     where
         U: AsRef<str>,
         A: AsRef<[U]>,
@@ -99,7 +254,11 @@ impl Client {
         let tls_config =
             load_client_tls_config(Some(client_tls)).context(error::CreateTlsChannelSnafu)?;
         let channel_manager = ChannelManager::with_config(channel_config, tls_config);
-        Ok(Self::with_manager_and_urls(channel_manager, urls))
+        Ok(Self::with_manager_and_urls_and_options(
+            channel_manager,
+            urls,
+            options,
+        ))
     }
 
     pub fn with_manager_and_urls<U, A>(channel_manager: ChannelManager, urls: A) -> Self
@@ -107,15 +266,30 @@ impl Client {
         U: AsRef<str>,
         A: AsRef<[U]>,
     {
-        let inner = Inner::with_manager(channel_manager);
+        Self::with_manager_and_urls_and_options(channel_manager, urls, ClientOptions::default())
+    }
+
+    /// Creates a client with a channel manager, URLs, and custom options.
+    pub fn with_manager_and_urls_and_options<U, A>(
+        channel_manager: ChannelManager,
+        urls: A,
+        options: ClientOptions,
+    ) -> Self
+    where
+        U: AsRef<str>,
+        A: AsRef<[U]>,
+    {
         let urls: Vec<String> = urls
             .as_ref()
             .iter()
             .map(|peer| peer.as_ref().to_string())
             .collect();
-        inner.set_peers(urls);
         Self {
-            inner: Arc::new(inner),
+            inner: Arc::new(Inner::with_manager_and_peers(
+                channel_manager,
+                urls,
+                options,
+            )),
         }
     }
 
@@ -124,16 +298,49 @@ impl Client {
         U: AsRef<str>,
         A: AsRef<[U]>,
     {
-        let urls: Vec<String> = urls
+        let urls = urls
             .as_ref()
             .iter()
             .map(|peer| peer.as_ref().to_string())
             .collect();
-
         self.inner.set_peers(urls);
     }
 
+    fn trigger_health_check(&self) {
+        if self.inner.health_check_interval.is_zero() || self.inner.peer_count() <= 1 {
+            return;
+        }
+
+        if self
+            .inner
+            .health_check_started
+            .swap(true, Ordering::Relaxed)
+        {
+            return;
+        }
+
+        let inner = Arc::downgrade(&self.inner);
+        let health_check_interval = self.inner.health_check_interval;
+        common_runtime::spawn_global(async move {
+            tokio::time::sleep(random_initial_delay(health_check_interval)).await;
+            let mut interval = tokio::time::interval(health_check_interval);
+            interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
+
+            loop {
+                interval.tick().await;
+                let Some(inner) = inner.upgrade() else {
+                    return;
+                };
+                if inner.peer_count() > 1 {
+                    inner.refresh_peer_states().await;
+                }
+            }
+        });
+    }
+
     pub fn find_channel(&self) -> Result<(String, Channel)> {
+        self.trigger_health_check();
+
         let addr = self
             .inner
             .get_peer()
@@ -220,15 +427,124 @@ impl Client {
         let _ = client.health_check(HealthCheckRequest {}).await?;
         Ok(())
     }
+
+    /// Returns peer addresses grouped by active and inactive state for tests.
+    #[cfg(feature = "testing")]
+    pub fn peer_addresses_by_state(&self) -> (Vec<String>, Vec<String>) {
+        let peers = self.inner.peers.read();
+        let addresses = |indices: &[usize]| {
+            indices
+                .iter()
+                .map(|&index| peers.addresses[index].clone())
+                .collect()
+        };
+        (
+            addresses(&peers.states.active),
+            addresses(&peers.states.inactive),
+        )
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use std::collections::HashSet;
+    use std::sync::Arc;
+    use std::sync::atomic::Ordering;
+    use std::time::Duration;
 
-    use super::Inner;
+    use api::v1::health_check_server::{HealthCheck, HealthCheckServer};
+    use api::v1::{HealthCheckRequest, HealthCheckResponse};
+    use common_grpc::channel_manager::ChannelManager;
+    use tokio::net::TcpListener;
+    use tokio::sync::Notify;
+    use tokio::task::JoinHandle;
+    use tokio::time::{interval, timeout};
+    use tokio_stream::wrappers::TcpListenerStream;
+    use tonic::{Request, Response, Status};
+
+    use super::{Client, ClientOptions, Inner, PeerStates};
     use crate::load_balance::Loadbalancer;
 
+    const HEALTH_REFRESH_INTERVAL: Duration = Duration::from_millis(10);
+    const STATE_REFRESH_TIMEOUT: Duration = Duration::from_secs(1);
+
+    struct HealthyHealthCheck;
+
+    #[tonic::async_trait]
+    impl HealthCheck for HealthyHealthCheck {
+        async fn health_check(
+            &self,
+            _request: Request<HealthCheckRequest>,
+        ) -> Result<Response<HealthCheckResponse>, Status> {
+            Ok(Response::new(HealthCheckResponse {}))
+        }
+    }
+
+    struct UnhealthyHealthCheck;
+
+    #[tonic::async_trait]
+    impl HealthCheck for UnhealthyHealthCheck {
+        async fn health_check(
+            &self,
+            _request: Request<HealthCheckRequest>,
+        ) -> Result<Response<HealthCheckResponse>, Status> {
+            Err(Status::unavailable("peer is unavailable"))
+        }
+    }
+
+    struct PendingHealthCheck {
+        started: Option<Arc<Notify>>,
+    }
+
+    #[tonic::async_trait]
+    impl HealthCheck for PendingHealthCheck {
+        async fn health_check(
+            &self,
+            _request: Request<HealthCheckRequest>,
+        ) -> Result<Response<HealthCheckResponse>, Status> {
+            if let Some(started) = &self.started {
+                started.notify_one();
+            }
+            std::future::pending().await
+        }
+    }
+
+    async fn start_health_check_server<T>(handler: T) -> (String, JoinHandle<()>)
+    where
+        T: HealthCheck + Send + Sync + 'static,
+    {
+        let listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .expect("bind health check server");
+        let addr = listener
+            .local_addr()
+            .expect("read health check server address")
+            .to_string();
+        let server = tokio::spawn(async move {
+            tonic::transport::Server::builder()
+                .add_service(HealthCheckServer::new(handler))
+                .serve_with_incoming(TcpListenerStream::new(listener))
+                .await
+                .expect("serve health check server");
+        });
+
+        (addr, server)
+    }
+
+    async fn wait_for_peer_states(client: &Client, expected: PeerStates) {
+        let mut poll = interval(HEALTH_REFRESH_INTERVAL);
+        timeout(STATE_REFRESH_TIMEOUT, async {
+            loop {
+                poll.tick().await;
+                if client.inner.peers.read().states == expected {
+                    return;
+                }
+            }
+        })
+        .await
+        .expect("health refresh did not reach expected peer states");
+    }
+
     fn mock_peers() -> Vec<String> {
         vec![
             "127.0.0.1:3001".to_string(),
@@ -248,11 +564,166 @@ mod tests {
         assert!(inner.get_peer().is_none());
 
         let peers = mock_peers();
-        inner.set_peers(peers.clone());
-        let all: HashSet<String> = peers.into_iter().collect();
+        let all: HashSet<String> = peers.iter().cloned().collect();
+        let inner =
+            Inner::with_manager_and_peers(ChannelManager::new(), peers, ClientOptions::default());
 
         for _ in 0..20 {
             assert!(all.contains(&inner.get_peer().unwrap()));
         }
     }
+
+    #[test]
+    fn test_inner_prefers_active_peer() {
+        let peers = mock_peers();
+        let inner = Inner::with_manager_and_peers(
+            ChannelManager::new(),
+            peers.clone(),
+            ClientOptions::default(),
+        );
+        inner.peers.write().states = PeerStates {
+            active: vec![0],
+            inactive: vec![1, 2],
+        };
+
+        assert_eq!(Some(peers[0].clone()), inner.get_peer());
+    }
+
+    #[test]
+    fn test_zero_health_check_interval_disables_background_task() {
+        let client = Client::with_urls_and_options(
+            mock_peers(),
+            ClientOptions {
+                health_check_interval: Duration::ZERO,
+                ..Default::default()
+            },
+        );
+
+        assert!(!client.inner.health_check_started.load(Ordering::Relaxed));
+        let peers = client.inner.peers.read();
+        assert_eq!(mock_peers(), peers.addresses);
+        assert_eq!(vec![0, 1, 2], peers.states.active);
+        assert!(peers.states.inactive.is_empty());
+    }
+
+    #[test]
+    fn test_multi_peer_constructor_defers_background_task() {
+        let client = Client::with_urls(mock_peers());
+
+        assert!(!client.inner.health_check_started.load(Ordering::Relaxed));
+    }
+
+    #[tokio::test]
+    async fn test_single_peer_does_not_start_background_task() {
+        let client = Client::with_urls(["127.0.0.1:3001"]);
+
+        client.find_channel().unwrap();
+
+        assert!(!client.inner.health_check_started.load(Ordering::Relaxed));
+    }
+
+    #[test]
+    fn test_start_initializes_new_client_without_starting_background_task() {
+        let client = Client::new();
+        let peers = mock_peers();
+
+        client.start(peers.clone());
+
+        assert!(peers.contains(&client.inner.get_peer().unwrap()));
+        assert!(!client.inner.health_check_started.load(Ordering::Relaxed));
+    }
+
+    #[tokio::test]
+    async fn test_health_refresh_marks_unhealthy_peer_inactive_and_selects_healthy_peer() {
+        // Arrange: one peer responds to health checks and the other rejects them.
+        let (healthy_addr, healthy_server) = start_health_check_server(HealthyHealthCheck).await;
+        let (unhealthy_addr, unhealthy_server) =
+            start_health_check_server(UnhealthyHealthCheck).await;
+        let client = Client::with_urls_and_options(
+            [healthy_addr.clone(), unhealthy_addr],
+            ClientOptions {
+                health_check_interval: HEALTH_REFRESH_INTERVAL,
+                ..Default::default()
+            },
+        );
+        assert!(!client.inner.health_check_started.load(Ordering::Relaxed));
+
+        // Act: trigger lazy health checks, then poll until the background refresh completes.
+        client.find_channel().unwrap();
+        assert!(client.inner.health_check_started.load(Ordering::Relaxed));
+        wait_for_peer_states(
+            &client,
+            PeerStates {
+                active: vec![0],
+                inactive: vec![1],
+            },
+        )
+        .await;
+
+        // Assert: an inactive peer does not prevent selection of its active peer.
+        assert_eq!(Some(healthy_addr), client.inner.get_peer());
+
+        healthy_server.abort();
+        unhealthy_server.abort();
+    }
+
+    #[tokio::test]
+    async fn test_health_refresh_times_out_pending_peer() {
+        let (healthy_addr, healthy_server) = start_health_check_server(HealthyHealthCheck).await;
+        let (pending_addr, pending_server) =
+            start_health_check_server(PendingHealthCheck { started: None }).await;
+        let client = Client::with_urls_and_options(
+            [healthy_addr, pending_addr],
+            ClientOptions {
+                health_check_interval: HEALTH_REFRESH_INTERVAL,
+                health_check_timeout: Duration::from_millis(20),
+            },
+        );
+
+        client.find_channel().unwrap();
+        wait_for_peer_states(
+            &client,
+            PeerStates {
+                active: vec![0],
+                inactive: vec![1],
+            },
+        )
+        .await;
+
+        healthy_server.abort();
+        pending_server.abort();
+    }
+
+    #[tokio::test]
+    async fn test_peer_update_ignores_in_flight_health_result() {
+        let started = Arc::new(Notify::new());
+        let (pending_addr, pending_server) = start_health_check_server(PendingHealthCheck {
+            started: Some(started.clone()),
+        })
+        .await;
+        let inner = Arc::new(Inner::with_manager_and_peers(
+            ChannelManager::new(),
+            vec![pending_addr],
+            ClientOptions {
+                health_check_timeout: Duration::from_millis(20),
+                ..Default::default()
+            },
+        ));
+        let refresh_inner = inner.clone();
+        let refresh = tokio::spawn(async move {
+            refresh_inner.refresh_peer_states().await;
+        });
+        timeout(STATE_REFRESH_TIMEOUT, started.notified())
+            .await
+            .expect("pending health check did not start");
+
+        inner.set_peers(vec!["127.0.0.1:3001".to_string()]);
+        refresh.await.unwrap();
+
+        let peers = inner.peers.read();
+        assert_eq!(vec!["127.0.0.1:3001"], peers.addresses);
+        assert_eq!(vec![0], peers.states.active);
+        assert!(peers.states.inactive.is_empty());
+        pending_server.abort();
+    }
 }

src/client/src/lib.rs

@@ -32,7 +32,7 @@ pub use common_query::{Output, OutputData, OutputMeta};
 pub use common_recordbatch::{RecordBatches, SendableRecordBatchStream};
 use snafu::OptionExt;
 
-pub use self::client::Client;
+pub use self::client::{Client, ClientOptions};
 pub use self::database::{Database, OutputMetrics, OutputWithMetrics};
 pub use self::error::{Error, Result};
 use crate::error::{IllegalDatabaseResponseSnafu, ServerSnafu};

src/client/src/load_balance.rs

@@ -17,7 +17,7 @@ use rand::seq::IndexedRandom;
 
 #[enum_dispatch]
 pub trait LoadBalance {
-    fn get_peer<'a>(&self, peers: &'a [String]) -> Option<&'a String>;
+    fn get_index<'a>(&self, candidates: &'a [usize]) -> Option<&'a usize>;
 }
 
 #[enum_dispatch(LoadBalance)]
@@ -36,8 +36,8 @@ impl Default for Loadbalancer {
 pub struct Random;
 
 impl LoadBalance for Random {
-    fn get_peer<'a>(&self, peers: &'a [String]) -> Option<&'a String> {
-        peers.choose(&mut rand::rng())
+    fn get_index<'a>(&self, candidates: &'a [usize]) -> Option<&'a usize> {
+        candidates.choose(&mut rand::rng())
     }
 }
 
@@ -49,18 +49,13 @@ mod tests {
 
     #[test]
     fn test_random_lb() {
-        let peers = vec![
-            "127.0.0.1:3001".to_string(),
-            "127.0.0.1:3002".to_string(),
-            "127.0.0.1:3003".to_string(),
-            "127.0.0.1:3004".to_string(),
-        ];
-        let all: HashSet<String> = peers.clone().into_iter().collect();
+        let candidates = vec![0, 1, 2, 3];
+        let all: HashSet<usize> = candidates.iter().copied().collect();
 
         let random = Random;
         for _ in 0..100 {
-            let peer = random.get_peer(&peers).unwrap();
-            assert!(all.contains(peer));
+            let index = random.get_index(&candidates).unwrap();
+            assert!(all.contains(index));
         }
     }
 }

src/datanode/src/service.rs

@@ -97,3 +97,56 @@ impl<'a> DatanodeServiceBuilder<'a> {
             .region_server_handler(Arc::new(region_server.clone()))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use api::v1::HealthCheckRequest;
+    use api::v1::health_check_client::HealthCheckClient;
+    use servers::grpc::GRPC_SERVER;
+
+    use super::DatanodeServiceBuilder;
+    use crate::config::DatanodeOptions;
+    use crate::tests::mock_region_server;
+
+    #[tokio::test]
+    async fn test_default_grpc_server_health_check_is_reachable() {
+        // Arrange
+        let opts = DatanodeOptions {
+            grpc: servers::grpc::GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
+            ..Default::default()
+        };
+        let region_server = mock_region_server();
+        let mut services = DatanodeServiceBuilder::new(&opts)
+            .with_default_grpc_server(&region_server)
+            .build()
+            .unwrap();
+
+        // Act
+        services.start_all().await.unwrap();
+        let addr = services.addr(GRPC_SERVER).unwrap();
+        let health_check = HealthCheckClient::connect(format!("http://{addr}"))
+            .await
+            .unwrap()
+            .health_check(HealthCheckRequest {})
+            .await;
+        services.shutdown_all().await.unwrap();
+
+        // Assert
+        assert!(health_check.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_service_builder_without_grpc_server_does_not_expose_grpc_address() {
+        // Arrange
+        let opts = DatanodeOptions::default();
+        let mut services = DatanodeServiceBuilder::new(&opts).build().unwrap();
+
+        // Act
+        services.start_all().await.unwrap();
+        let addr = services.addr(GRPC_SERVER);
+        services.shutdown_all().await.unwrap();
+
+        // Assert
+        assert!(addr.is_none());
+    }
+}

src/flow/src/server.rs

@@ -723,6 +723,8 @@ mod tests {
     use std::sync::Arc;
     use std::time::Duration;
 
+    use api::v1::HealthCheckRequest;
+    use api::v1::health_check_client::HealthCheckClient;
     use api::v1::meta::Role;
     use catalog::memory::new_memory_catalog_manager;
     use common_base::Plugins;
@@ -731,6 +733,7 @@ mod tests {
     use common_meta::kv_backend::memory::MemoryKvBackend;
     use meta_client::client::MetaClient;
     use query::options::QueryOptions;
+    use servers::grpc::GRPC_SERVER;
 
     use super::*;
     use crate::adapter::flownode_impl::FlowDualEngine;
@@ -834,4 +837,29 @@ mod tests {
                 .is_none()
         );
     }
+
+    #[tokio::test]
+    async fn test_service_builder_registers_reachable_health_check() {
+        // Arrange: compose the production gRPC service with an ephemeral local listener.
+        let (flownode_server, _report_sender) = new_test_flownode_server().await;
+        let mut opts = FlownodeOptions::default();
+        opts.grpc.bind_addr = "127.0.0.1:0".to_string();
+        let mut services = FlownodeServiceBuilder::new(&opts)
+            .with_default_grpc_server(&flownode_server)
+            .build()
+            .unwrap();
+        services.start_all().await.unwrap();
+        let addr = services.addr(GRPC_SERVER).unwrap();
+
+        // Act: call the shared health handler through the registered production server.
+        let mut client = HealthCheckClient::connect(format!("http://{addr}"))
+            .await
+            .unwrap();
+        let result = client.health_check(HealthCheckRequest {}).await;
+
+        services.shutdown_all().await.unwrap();
+
+        // Assert: the service composition exposes a healthy endpoint.
+        assert!(result.is_ok());
+    }
 }

src/frontend/src/server.rs

@@ -447,7 +447,14 @@ fn parse_addr(addr: &str) -> Result<SocketAddr> {
 mod tests {
     use std::time::Duration;
 
+    use api::v1::HealthCheckRequest;
+    use api::v1::health_check_client::HealthCheckClient;
+    use api::v1::meta::Role;
+    use meta_client::client::MetaClientBuilder;
+    use servers::grpc::GRPC_SERVER;
+
     use super::*;
+    use crate::instance::builder::FrontendBuilder;
 
     #[test]
     fn test_effective_http_timeout_for_pending_rows() {
@@ -538,4 +545,52 @@ mod tests {
             );
         }
     }
+
+    #[tokio::test]
+    async fn test_services_builder_health_check_is_reachable() {
+        // Arrange
+        let options = FrontendOptions {
+            http: HttpOptions {
+                addr: "127.0.0.1:0".to_string(),
+                ..Default::default()
+            },
+            grpc: GrpcOptions::default().with_bind_addr("127.0.0.1:0"),
+            mysql: crate::service_config::MysqlOptions {
+                enable: false,
+                ..Default::default()
+            },
+            postgres: crate::service_config::PostgresOptions {
+                enable: false,
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+        let meta_client = Arc::new(
+            MetaClientBuilder::new(0, Role::Frontend)
+                .enable_procedure()
+                .build(),
+        );
+        let instance = Arc::new(
+            FrontendBuilder::new_test(&options, meta_client)
+                .try_build()
+                .await
+                .unwrap(),
+        );
+        let mut services = Services::new(options, instance, Default::default())
+            .build()
+            .unwrap();
+
+        // Act
+        services.start_all().await.unwrap();
+        let addr = services.addr(GRPC_SERVER).unwrap();
+        let health_check = HealthCheckClient::connect(format!("http://{addr}"))
+            .await
+            .unwrap()
+            .health_check(HealthCheckRequest {})
+            .await;
+        services.shutdown_all().await.unwrap();
+
+        // Assert
+        assert!(health_check.is_ok());
+    }
 }

tests-integration/Cargo.toml

@@ -26,7 +26,7 @@ auth.workspace = true
 axum.workspace = true
 base64.workspace = true
 cache.workspace = true
-catalog.workspace = true
+catalog = { workspace = true, features = ["testing"] }
 chrono.workspace = true
 clap.workspace = true
 client = { workspace = true, features = ["testing"] }

tests-integration/src/cluster.rs

@@ -19,6 +19,7 @@ use std::ops::RangeInclusive;
 use std::sync::Arc;
 use std::time::Duration;
 
+use api::v1::health_check_server::HealthCheckServer;
 use api::v1::region::region_server::RegionServer;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use cache::{
@@ -65,9 +66,9 @@ use mito2::gc::GcConfig;
 use mito2::region::MitoRegionRef;
 use object_store::config::ObjectStoreConfig;
 use rand::Rng;
-use servers::grpc::GrpcOptions;
 use servers::grpc::flight::FlightCraftWrapper;
 use servers::grpc::region_server::RegionServerRequestHandler;
+use servers::grpc::{GrpcOptions, HealthCheckHandler};
 use servers::server::ServerHandlers;
 use store_api::storage::RegionId;
 use tempfile::TempDir;
@@ -77,7 +78,7 @@ use tower::service_fn;
 use uuid::Uuid;
 
 use crate::test_util::{
-    self, FileDirGuard, PEER_PLACEHOLDER_ADDR, StorageType, TestGuard, create_datanode_opts,
+    self, FileDirGuard, StorageType, TestGuard, create_datanode_opts,
     create_tmp_dir_and_datanode_opts,
 };
 
@@ -303,13 +304,17 @@ impl GreptimeDbClusterBuilder {
             .build_datanodes_with_options(&metasrv, &datanode_options)
             .await;
 
-        build_datanode_clients(datanode_clients.clone(), &datanode_instances, datanodes).await;
+        build_datanode_clients(datanode_clients.clone(), &datanode_instances).await;
 
         self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes)
             .await;
 
         let mut frontend = self
-            .build_frontend(metasrv.clone(), datanode_clients, start_frontend_servers)
+            .build_frontend(
+                metasrv.clone(),
+                datanode_clients.clone(),
+                start_frontend_servers,
+            )
             .await;
 
         frontend.start().await.unwrap();
@@ -433,7 +438,7 @@ impl GreptimeDbClusterBuilder {
                 .build(),
         );
 
-        let mut builder = DatanodeBuilder::new(opts, Plugins::default(), meta_backend);
+        let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend);
         builder
             .with_cache_registry(layered_cache_registry)
             .with_meta_client(meta_client);
@@ -578,11 +583,8 @@ impl GreptimeDbClusterBuilder {
 async fn build_datanode_clients(
     clients: Arc<NodeClients>,
     instances: &HashMap<DatanodeId, Datanode>,
-    datanodes: usize,
 ) {
-    for i in 0..datanodes {
-        let datanode_id = i as u64 + 1;
-        let instance = instances.get(&datanode_id).unwrap();
+    for (&datanode_id, instance) in instances {
         let (addr, client) = create_datanode_client(instance).await;
         clients
             .insert_client(Peer::new(datanode_id, addr), client)
@@ -600,7 +602,6 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
         .unwrap();
 
     let flight_handler = FlightCraftWrapper(datanode.region_server());
-
     let region_server_handler =
         RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
 
@@ -620,34 +621,298 @@ async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
                     .send_compressed(CompressionEncoding::Gzip)
                     .send_compressed(CompressionEncoding::Zstd),
             )
+            .add_service(HealthCheckServer::new(HealthCheckHandler))
             .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
             .await
     });
 
-    // Move client to an option so we can _move_ the inner value
-    // on the first attempt to connect. All other attempts will fail.
     let mut client = Some(client);
-    // `PEER_PLACEHOLDER_ADDR` is just a placeholder, does not actually connect to it.
-    let addr = PEER_PLACEHOLDER_ADDR;
+    let addr = test_util::PEER_PLACEHOLDER_ADDR;
     let channel_manager = ChannelManager::new();
-    let _ = channel_manager
+    channel_manager
         .reset_with_connector(
             addr,
             service_fn(move |_| {
                 let client = client.take();
 
                 async move {
-                    if let Some(client) = client {
-                        Ok(TokioIo::new(client))
-                    } else {
-                        Err(std::io::Error::other("Client already taken"))
-                    }
+                    client
+                        .map(TokioIo::new)
+                        .ok_or_else(|| std::io::Error::other("Client already taken"))
                 }
             }),
         )
         .unwrap();
+
     (
         addr.to_string(),
-        Client::with_manager_and_urls(channel_manager, vec![addr]),
+        Client::with_manager_and_urls(channel_manager, [addr]),
     )
 }
+
+#[cfg(test)]
+mod tests {
+    use api::v1::flow::FlowRequest;
+    use api::v1::region::{
+        ListMetadataRequest, RegionRequest, RegionRequestHeader, region_request,
+    };
+    use catalog::memory::new_memory_catalog_manager;
+    use client::Client;
+    use common_error::ext::ErrorExt;
+    use common_error::status_code::StatusCode;
+    use common_meta::key::TableMetadataManager;
+    use common_meta::key::flow::FlowMetadataManager;
+    use common_meta::kv_backend::memory::MemoryKvBackend;
+    use common_meta::node_manager::{DatanodeManager, FlownodeManager};
+    use common_meta::peer::Peer;
+    use flow::{FlownodeBuilder, FlownodeOptions, FlownodeServiceBuilder, FrontendClient};
+
+    use super::*;
+
+    const HEALTH_REFRESH_INTERVAL: Duration = Duration::from_millis(10);
+    const HEALTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
+
+    async fn wait_for_active_route(client: &Client, active_addr: &str) {
+        client.find_channel().unwrap();
+        tokio::time::timeout(HEALTH_REFRESH_TIMEOUT, async {
+            let mut interval = tokio::time::interval(HEALTH_REFRESH_INTERVAL);
+            loop {
+                interval.tick().await;
+                let (active, inactive) = client.peer_addresses_by_state();
+                if active.len() == 1 && active[0] == active_addr && inactive.len() == 1 {
+                    return;
+                }
+            }
+        })
+        .await
+        .expect("background health routing did not select the active peer");
+    }
+
+    async fn create_region_requester_health_client(datanode: &Datanode) -> (Peer, Peer, Client) {
+        const ACTIVE_ADDR: &str = "127.0.0.1:3001";
+        const INACTIVE_ADDR: &str = "127.0.0.1:3002";
+
+        let (active_client, server) = tokio::io::duplex(1024);
+        let runtime = RuntimeBuilder::default()
+            .worker_threads(2)
+            .thread_name("grpc-handlers")
+            .build()
+            .unwrap();
+        let region_server_handler =
+            RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
+
+        tokio::spawn(async move {
+            Server::builder()
+                .add_service(RegionServer::new(region_server_handler))
+                .add_service(HealthCheckServer::new(HealthCheckHandler))
+                .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
+                .await
+        });
+
+        let channel_manager = ChannelManager::new();
+        let mut active_client = Some(active_client);
+        channel_manager
+            .reset_with_connector(
+                ACTIVE_ADDR,
+                service_fn(move |_| {
+                    let client = active_client.take();
+
+                    async move {
+                        client
+                            .map(TokioIo::new)
+                            .ok_or_else(|| std::io::Error::other("Active client already taken"))
+                    }
+                }),
+            )
+            .unwrap();
+        channel_manager
+            .reset_with_connector(
+                INACTIVE_ADDR,
+                service_fn(|_| async {
+                    Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
+                        "Inactive synthetic peer",
+                    ))
+                }),
+            )
+            .unwrap();
+
+        let client = Client::with_manager_and_urls_and_options(
+            channel_manager,
+            [INACTIVE_ADDR, ACTIVE_ADDR],
+            client::ClientOptions {
+                health_check_interval: HEALTH_REFRESH_INTERVAL,
+                ..Default::default()
+            },
+        );
+        (
+            Peer::new(1, ACTIVE_ADDR),
+            Peer::new(2, INACTIVE_ADDR),
+            client,
+        )
+    }
+
+    async fn create_flow_requester_health_client(
+        flownode: &flow::FlownodeInstance,
+    ) -> (Peer, Peer, Client) {
+        const ACTIVE_ADDR: &str = "127.0.0.1:3003";
+        const INACTIVE_ADDR: &str = "127.0.0.1:3004";
+
+        let (active_client, server) = tokio::io::duplex(1024);
+        let flow_service = flownode.flownode_server().create_flow_service();
+
+        tokio::spawn(async move {
+            Server::builder()
+                .add_service(flow_service)
+                .add_service(HealthCheckServer::new(HealthCheckHandler))
+                .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
+                .await
+        });
+
+        let channel_manager = ChannelManager::new();
+        let mut active_client = Some(active_client);
+        channel_manager
+            .reset_with_connector(
+                ACTIVE_ADDR,
+                service_fn(move |_| {
+                    let client = active_client.take();
+
+                    async move {
+                        client
+                            .map(TokioIo::new)
+                            .ok_or_else(|| std::io::Error::other("Active client already taken"))
+                    }
+                }),
+            )
+            .unwrap();
+        channel_manager
+            .reset_with_connector(
+                INACTIVE_ADDR,
+                service_fn(|_| async {
+                    Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
+                        "Inactive synthetic peer",
+                    ))
+                }),
+            )
+            .unwrap();
+
+        let client = Client::with_manager_and_urls_and_options(
+            channel_manager,
+            [INACTIVE_ADDR, ACTIVE_ADDR],
+            client::ClientOptions {
+                health_check_interval: HEALTH_REFRESH_INTERVAL,
+                ..Default::default()
+            },
+        );
+        (
+            Peer::new(1, ACTIVE_ADDR),
+            Peer::new(2, INACTIVE_ADDR),
+            client,
+        )
+    }
+
+    async fn start_flownode(addr: String) -> flow::FlownodeInstance {
+        let kv_backend = Arc::new(MemoryKvBackend::new());
+        let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
+        table_meta.init().await.unwrap();
+        let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend));
+        let catalog_manager = new_memory_catalog_manager().unwrap();
+        let (frontend_client, _handler) =
+            FrontendClient::from_empty_grpc_handler(Default::default());
+        let mut opts = FlownodeOptions::default();
+        opts.grpc.bind_addr = addr.clone();
+        opts.grpc.server_addr = addr;
+
+        let mut flownode = FlownodeBuilder::new(
+            opts.clone(),
+            Plugins::default(),
+            table_meta,
+            catalog_manager,
+            flow_meta,
+            Arc::new(frontend_client),
+        )
+        .build()
+        .await
+        .unwrap();
+        let services = FlownodeServiceBuilder::new(&opts)
+            .with_default_grpc_server(flownode.flownode_server())
+            .build()
+            .unwrap();
+        flownode.setup_services(services);
+        flownode.start().await.unwrap();
+        flownode
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_region_requester_health_check_and_list_metadata() {
+        // Arrange: configure active and inactive synthetic peers over the duplex harness.
+        let cluster = GreptimeDbClusterBuilder::new("region_requester_health_check")
+            .await
+            .with_datanodes(1)
+            .build(false)
+            .await;
+        let (active_peer, inactive_peer, client) =
+            create_region_requester_health_client(cluster.datanode_instances.get(&1).unwrap())
+                .await;
+        let clients = NodeClients::default();
+        clients
+            .insert_client(active_peer.clone(), client.clone())
+            .await;
+        clients
+            .insert_client(inactive_peer.clone(), client.clone())
+            .await;
+
+        // Act: wait for background health routing, then send a real Region RPC.
+        wait_for_active_route(&client, &active_peer.addr).await;
+        let requester = clients.datanode(&active_peer).await;
+        let response = requester
+            .handle(RegionRequest {
+                header: Some(RegionRequestHeader::default()),
+                body: Some(region_request::Body::ListMetadata(ListMetadataRequest {
+                    region_ids: vec![],
+                })),
+            })
+            .await
+            .unwrap();
+
+        // Assert: the inactive peer was rejected and the RegionRequester reached the active service.
+        assert_eq!(
+            (vec![active_peer.addr], vec![inactive_peer.addr]),
+            client.peer_addresses_by_state()
+        );
+        assert_eq!(b"[]", response.metadata.as_slice());
+    }
+
+    #[tokio::test(flavor = "multi_thread")]
+    async fn test_flow_requester_health_check_and_handle_request() {
+        // Arrange: start a flownode and expose its production flow service through a duplex harness.
+        let mut flownode = start_flownode("127.0.0.1:0".to_string()).await;
+        let (active_peer, inactive_peer, client) =
+            create_flow_requester_health_client(&flownode).await;
+        let clients = NodeClients::default();
+        clients
+            .insert_client(active_peer.clone(), client.clone())
+            .await;
+        clients
+            .insert_client(inactive_peer.clone(), client.clone())
+            .await;
+
+        // Act: wait for background health routing, then send a real Flow RPC.
+        wait_for_active_route(&client, &active_peer.addr).await;
+        let response = clients
+            .flownode(&active_peer)
+            .await
+            .handle(FlowRequest::default())
+            .await;
+
+        // Assert: the inactive peer was rejected and the production Flow service handled the RPC.
+        assert_eq!(
+            (vec![active_peer.addr], vec![inactive_peer.addr]),
+            client.peer_addresses_by_state()
+        );
+        assert_eq!(
+            StatusCode::InvalidArguments,
+            response.unwrap_err().status_code()
+        );
+        flownode.shutdown().await.unwrap();
+    }
+}

tests-integration/src/grpc/flight.rs

@@ -17,6 +17,7 @@ mod test {
     use std::collections::HashMap;
     use std::net::SocketAddr;
     use std::sync::Arc;
+    use std::time::Duration;
 
     use api::v1::auth_header::AuthScheme;
     use api::v1::query_request::Query;
@@ -130,13 +131,14 @@ mod test {
     }
 
     #[tokio::test(flavor = "multi_thread")]
-    async fn test_distributed_flight_do_put() {
+    async fn test_distributed_frontend_database_and_flight_health() {
         common_telemetry::init_default_ut_logging();
 
-        let db = GreptimeDbClusterBuilder::new("test_distributed_flight_do_put")
-            .await
-            .build(false)
-            .await;
+        let db =
+            GreptimeDbClusterBuilder::new("test_distributed_frontend_database_and_flight_health")
+                .await
+                .build(false)
+                .await;
 
         let runtime = common_runtime::global_runtime().clone();
         let greptime_request_handler = GreptimeRequestHandler::new(
@@ -146,6 +148,7 @@ mod test {
             FlightCompression::default(),
         );
         let mut grpc_server = GrpcServerBuilder::new(GrpcServerConfig::default(), runtime)
+            .database_handler(greptime_request_handler.clone())
             .flight_handler(Arc::new(greptime_request_handler))
             .build();
         grpc_server
@@ -155,6 +158,7 @@ mod test {
         let addr = grpc_server.bind_addr().unwrap().to_string();
 
         let client = Client::with_urls(vec![addr]);
+        wait_for_client_health(&client).await;
         let mut client = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, client);
         client.set_auth(AuthScheme::Basic(Basic {
             username: "greptime_user".to_string(),
@@ -282,6 +286,19 @@ mod test {
         );
     }
 
+    async fn wait_for_client_health(client: &Client) {
+        tokio::time::timeout(Duration::from_secs(10), async {
+            loop {
+                if client.health_check().await.is_ok() {
+                    return;
+                }
+                tokio::time::sleep(Duration::from_millis(10)).await;
+            }
+        })
+        .await
+        .expect("cluster frontend did not become healthy within 10 seconds");
+    }
+
     #[tokio::test(flavor = "multi_thread")]
     async fn test_distributed_flight_snapshot_seqs_rejects_stale_sst_fence() {
         common_telemetry::init_default_ut_logging();

tests-integration/tests/grpc.rs

@@ -789,6 +789,9 @@ pub async fn test_health_check(store_type: StorageType) {
     let grpc_client = Client::with_urls(vec![addr]);
     grpc_client.health_check().await.unwrap();
 
+    let db = Database::new(DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME, grpc_client);
+    assert!(db.sql("SHOW TABLES").await.is_ok());
+
     let _ = fe_grpc_server.shutdown().await;
 }