diff --git a/Cargo.lock b/Cargo.lock index 8f7fc61ae7..ccfb23c51a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2068,6 +2068,7 @@ dependencies = [ "common-meta", "common-query", "common-recordbatch", + "common-runtime", "common-telemetry", "datatypes", "enum_dispatch", diff --git a/src/client/Cargo.toml b/src/client/Cargo.toml index 963656e17b..5b479b1fbf 100644 --- a/src/client/Cargo.toml +++ b/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" diff --git a/src/client/src/client.rs b/src/client/src/client.rs index 39cb5c30aa..8b6671cb5c 100644 --- a/src/client/src/client.rs +++ b/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, @@ -51,30 +74,140 @@ pub struct Client { inner: Arc, } -#[derive(Debug, Default)] +#[derive(Debug)] struct Inner { channel_manager: ChannelManager, - peers: Arc>>, + peers: RwLock, 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, + inactive: Vec, +} + +#[derive(Debug, Default)] +struct Peers { + addresses: Vec, + states: PeerStates, + generation: u64, } impl Inner { - fn with_manager(channel_manager: ChannelManager) -> Self { + fn with_manager_and_peers( + channel_manager: ChannelManager, + peers: Vec, + 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) { - let mut guard = self.peers.write(); - *guard = peers; + fn set_peers(&self, addresses: Vec) { + 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 { - 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::>(); + (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, 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(urls: A, options: ClientOptions) -> Self + where + U: AsRef, + A: AsRef<[U]>, + { + Self::with_manager_and_urls_and_options(ChannelManager::new(), urls, options) } pub fn with_tls_and_urls(urls: A, client_tls: ClientTlsOption) -> Result + where + U: AsRef, + 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( + urls: A, + client_tls: ClientTlsOption, + options: ClientOptions, + ) -> Result where U: AsRef, 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(channel_manager: ChannelManager, urls: A) -> Self @@ -107,15 +266,30 @@ impl Client { U: AsRef, 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( + channel_manager: ChannelManager, + urls: A, + options: ClientOptions, + ) -> Self + where + U: AsRef, + A: AsRef<[U]>, + { let urls: Vec = 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, A: AsRef<[U]>, { - let urls: Vec = 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, Vec) { + 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, + ) -> Result, Status> { + Ok(Response::new(HealthCheckResponse {})) + } + } + + struct UnhealthyHealthCheck; + + #[tonic::async_trait] + impl HealthCheck for UnhealthyHealthCheck { + async fn health_check( + &self, + _request: Request, + ) -> Result, Status> { + Err(Status::unavailable("peer is unavailable")) + } + } + + struct PendingHealthCheck { + started: Option>, + } + + #[tonic::async_trait] + impl HealthCheck for PendingHealthCheck { + async fn health_check( + &self, + _request: Request, + ) -> Result, Status> { + if let Some(started) = &self.started { + started.notify_one(); + } + std::future::pending().await + } + } + + async fn start_health_check_server(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 { 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 = peers.into_iter().collect(); + let all: HashSet = 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(); + } } diff --git a/src/client/src/lib.rs b/src/client/src/lib.rs index 10f97a2b75..4b9d543e35 100644 --- a/src/client/src/lib.rs +++ b/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}; diff --git a/src/client/src/load_balance.rs b/src/client/src/load_balance.rs index f3dc33910d..9c90e4dc8f 100644 --- a/src/client/src/load_balance.rs +++ b/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 = peers.clone().into_iter().collect(); + let candidates = vec![0, 1, 2, 3]; + let all: HashSet = 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)); } } } diff --git a/src/datanode/src/service.rs b/src/datanode/src/service.rs index 04256527e0..dec233350e 100644 --- a/src/datanode/src/service.rs +++ b/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(®ion_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()); + } +} diff --git a/src/flow/src/server.rs b/src/flow/src/server.rs index c9a24a4880..9bda516f24 100644 --- a/src/flow/src/server.rs +++ b/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()); + } } diff --git a/src/frontend/src/server.rs b/src/frontend/src/server.rs index 7153208f6b..8fdebd8683 100644 --- a/src/frontend/src/server.rs +++ b/src/frontend/src/server.rs @@ -447,7 +447,14 @@ fn parse_addr(addr: &str) -> Result { 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()); + } } diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index c0622eed35..1456693704 100644 --- a/tests-integration/Cargo.toml +++ b/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"] } diff --git a/tests-integration/src/cluster.rs b/tests-integration/src/cluster.rs index 3021f684fe..d0cd2e584c 100644 --- a/tests-integration/src/cluster.rs +++ b/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, instances: &HashMap, - 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::, _>(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::, _>(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(); + } +} diff --git a/tests-integration/src/grpc/flight.rs b/tests-integration/src/grpc/flight.rs index 068cb98142..520635518a 100644 --- a/tests-integration/src/grpc/flight.rs +++ b/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(); diff --git a/tests-integration/tests/grpc.rs b/tests-integration/tests/grpc.rs index ea003b6ff3..792a2ce965 100644 --- a/tests-integration/tests/grpc.rs +++ b/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; }