1use std::collections::{BTreeSet, HashMap};
16use std::env;
17use std::net::TcpListener;
18use std::ops::RangeInclusive;
19use std::sync::Arc;
20use std::time::Duration;
21
22use api::v1::health_check_server::HealthCheckServer;
23use api::v1::region::region_server::RegionServer;
24use arrow_flight::flight_service_server::FlightServiceServer;
25use cache::{
26 build_datanode_cache_registry, build_fundamental_cache_registry,
27 with_default_composite_cache_registry,
28};
29use catalog::information_extension::DistributedInformationExtension;
30use catalog::kvbackend::{
31 CachedKvBackendBuilder, KvBackendCatalogManagerBuilder, new_read_only_meta_kv_backend,
32};
33use catalog::process_manager::ProcessManager;
34use client::Client;
35use client::client_manager::NodeClients;
36use cmd::frontend::create_heartbeat_task;
37use common_base::Plugins;
38use common_grpc::channel_manager::{ChannelConfig, ChannelManager};
39use common_meta::DatanodeId;
40use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder};
41use common_meta::kv_backend::KvBackendRef;
42use common_meta::kv_backend::chroot::ChrootKvBackend;
43use common_meta::kv_backend::etcd::EtcdStore;
44use common_meta::kv_backend::memory::MemoryKvBackend;
45use common_meta::peer::Peer;
46use common_runtime::Builder as RuntimeBuilder;
47use common_runtime::runtime::BuilderBuild;
48use common_test_util::temp_dir::create_temp_dir;
49use common_time::util::DefaultSystemTimer;
50use common_wal::config::{DatanodeWalConfig, MetasrvWalConfig};
51use datanode::config::DatanodeOptions;
52use datanode::datanode::{Datanode, DatanodeBuilder, ProcedureConfig};
53use frontend::frontend::{Frontend, FrontendOptions};
54use frontend::instance::Instance as FeInstance;
55use frontend::instance::builder::FrontendBuilder;
56use frontend::server::Services;
57use futures::TryStreamExt;
58use hyper_util::rt::TokioIo;
59use meta_client::client::MetaClientBuilder;
60use meta_srv::cluster::MetaPeerClientRef;
61use meta_srv::discovery;
62use meta_srv::gc::GcSchedulerOptions;
63use meta_srv::metasrv::{Metasrv, MetasrvOptions, SelectorRef};
64use meta_srv::mocks::MockInfo;
65use mito2::gc::GcConfig;
66use mito2::region::MitoRegionRef;
67use object_store::config::ObjectStoreConfig;
68use rand::Rng;
69use servers::grpc::flight::FlightCraftWrapper;
70use servers::grpc::region_server::RegionServerRequestHandler;
71use servers::grpc::{GrpcOptions, HealthCheckHandler};
72use servers::server::ServerHandlers;
73use store_api::storage::RegionId;
74use tempfile::TempDir;
75use tonic::codec::CompressionEncoding;
76use tonic::transport::Server;
77use tower::service_fn;
78use uuid::Uuid;
79
80use crate::test_util::{
81 self, FileDirGuard, StorageType, TestGuard, create_datanode_opts,
82 create_tmp_dir_and_datanode_opts,
83};
84
85pub struct GreptimeDbCluster {
86 pub guards: Vec<TestGuard>,
87 pub datanode_options: Vec<DatanodeOptions>,
88
89 pub datanode_instances: HashMap<DatanodeId, Datanode>,
90 pub kv_backend: KvBackendRef,
91 pub metasrv: Arc<Metasrv>,
92 pub frontend: Arc<Frontend>,
93}
94
95impl GreptimeDbCluster {
96 pub fn fe_instance(&self) -> &Arc<FeInstance> {
97 &self.frontend.instance
98 }
99
100 pub async fn list_sst_files_from_all_datanodes(&self) -> BTreeSet<String> {
102 let mut sst_files = BTreeSet::new();
103
104 for datanode in self.datanode_instances.values() {
105 let region_server = datanode.region_server();
106 let mito = region_server.mito_engine().unwrap();
107 let all_files = mito
108 .all_ssts_from_storage()
109 .try_collect::<Vec<_>>()
110 .await
111 .unwrap()
112 .into_iter()
113 .map(|e| e.file_path)
114 .collect::<Vec<_>>();
115 sst_files.extend(all_files);
116 }
117
118 sst_files
119 }
120
121 pub async fn list_sst_files_from_manifests(&self) -> BTreeSet<String> {
123 let mut sst_files = BTreeSet::new();
124
125 for datanode in self.datanode_instances.values() {
126 let region_server = datanode.region_server();
127 let mito = region_server.mito_engine().unwrap();
128 let all_files = mito
129 .all_ssts_from_manifest()
130 .await
131 .into_iter()
132 .flat_map(|e| {
133 if let Some(index_file_path) = e.index_file_path {
134 vec![e.file_path, index_file_path]
135 } else {
136 vec![e.file_path]
137 }
138 })
139 .collect::<BTreeSet<_>>();
140
141 sst_files.extend(all_files);
142 }
143
144 sst_files
145 }
146
147 pub async fn list_all_regions(&self) -> HashMap<RegionId, MitoRegionRef> {
148 let mut regions = HashMap::new();
149
150 for datanode in self.datanode_instances.values() {
151 let region_server = datanode.region_server();
152 let mito = region_server.mito_engine().unwrap();
153 for region in mito.regions() {
154 regions.insert(region.region_id(), region);
155 }
156 }
157
158 regions
159 }
160}
161
162pub struct GreptimeDbClusterBuilder {
163 cluster_name: String,
164 kv_backend: KvBackendRef,
165 store_config: Option<ObjectStoreConfig>,
166 store_providers: Option<Vec<StorageType>>,
167 datanodes: Option<u32>,
168 datanode_wal_config: DatanodeWalConfig,
169 metasrv_wal_config: MetasrvWalConfig,
170 datanode_gc_config: GcConfig,
171 metasrv_gc_config: GcSchedulerOptions,
172 shared_home_dir: Option<Arc<TempDir>>,
173 meta_selector: Option<SelectorRef>,
174}
175
176impl GreptimeDbClusterBuilder {
177 pub async fn new(cluster_name: &str) -> Self {
178 let endpoints = env::var("GT_ETCD_ENDPOINTS").unwrap_or_default();
179
180 let kv_backend: KvBackendRef = if endpoints.is_empty() {
181 Arc::new(MemoryKvBackend::new())
182 } else {
183 let endpoints = endpoints
184 .split(',')
185 .map(|s| s.to_string())
186 .collect::<Vec<String>>();
187 let backend = EtcdStore::with_endpoints(endpoints, 128)
188 .await
189 .expect("malformed endpoints");
190 let chroot = format!("{}{}", cluster_name, Uuid::new_v4());
192 Arc::new(ChrootKvBackend::new(chroot.into(), backend))
193 };
194
195 Self {
196 cluster_name: cluster_name.to_string(),
197 kv_backend,
198 store_config: None,
199 store_providers: None,
200 datanodes: None,
201 datanode_wal_config: DatanodeWalConfig::default(),
202 metasrv_wal_config: MetasrvWalConfig::default(),
203 datanode_gc_config: GcConfig::default(),
204 metasrv_gc_config: GcSchedulerOptions::default(),
205 shared_home_dir: None,
206 meta_selector: None,
207 }
208 }
209
210 #[must_use]
211 pub fn with_store_config(mut self, store_config: ObjectStoreConfig) -> Self {
212 self.store_config = Some(store_config);
213 self
214 }
215
216 #[must_use]
217 pub fn with_store_providers(mut self, store_providers: Vec<StorageType>) -> Self {
218 self.store_providers = Some(store_providers);
219 self
220 }
221
222 #[must_use]
223 pub fn with_datanodes(mut self, datanodes: u32) -> Self {
224 self.datanodes = Some(datanodes);
225 self
226 }
227
228 #[must_use]
229 pub fn with_datanode_wal_config(mut self, datanode_wal_config: DatanodeWalConfig) -> Self {
230 self.datanode_wal_config = datanode_wal_config;
231 self
232 }
233
234 #[must_use]
235 pub fn with_metasrv_wal_config(mut self, metasrv_wal_config: MetasrvWalConfig) -> Self {
236 self.metasrv_wal_config = metasrv_wal_config;
237 self
238 }
239
240 #[must_use]
241 pub fn with_datanode_gc_config(mut self, datanode_gc_config: GcConfig) -> Self {
242 self.datanode_gc_config = datanode_gc_config;
243 self
244 }
245
246 pub fn with_metasrv_gc_config(mut self, metasrv_gc_config: GcSchedulerOptions) -> Self {
247 self.metasrv_gc_config = metasrv_gc_config;
248 self
249 }
250
251 #[must_use]
252 pub fn with_shared_home_dir(mut self, shared_home_dir: Arc<TempDir>) -> Self {
253 self.shared_home_dir = Some(shared_home_dir);
254 self
255 }
256
257 #[must_use]
258 pub fn with_meta_selector(mut self, selector: SelectorRef) -> Self {
259 self.meta_selector = Some(selector);
260 self
261 }
262
263 pub async fn build_with(
264 &self,
265 datanode_options: Vec<DatanodeOptions>,
266 start_frontend_servers: bool,
267 guards: Vec<TestGuard>,
268 ) -> GreptimeDbCluster {
269 let datanodes = datanode_options.len();
270 let channel_config = ChannelConfig::new().timeout(Some(Duration::from_secs(20)));
271 let datanode_clients = Arc::new(NodeClients::new(channel_config));
272
273 let opt = MetasrvOptions {
274 procedure: ProcedureConfig {
275 max_retry_times: 5,
278 retry_delay: Duration::from_secs(1),
279 max_metadata_value_size: None,
280 max_running_procedures: 128,
281 },
282 wal: self.metasrv_wal_config.clone(),
283 grpc: GrpcOptions {
284 server_addr: "127.0.0.1:3002".to_string(),
285 ..Default::default()
286 },
287 gc: self.metasrv_gc_config.clone(),
288 ..Default::default()
289 };
290
291 test_util::prepare_another_catalog_and_schema_with_kv_backend(self.kv_backend.clone())
292 .await;
293
294 let metasrv = meta_srv::mocks::mock(
295 opt,
296 self.kv_backend.clone(),
297 self.meta_selector.clone(),
298 Some(datanode_clients.clone()),
299 None,
300 )
301 .await;
302
303 let datanode_instances = self
304 .build_datanodes_with_options(&metasrv, &datanode_options)
305 .await;
306
307 build_datanode_clients(datanode_clients.clone(), &datanode_instances).await;
308
309 self.wait_datanodes_alive(metasrv.metasrv.meta_peer_client(), datanodes)
310 .await;
311
312 let mut frontend = self
313 .build_frontend(
314 metasrv.clone(),
315 datanode_clients.clone(),
316 start_frontend_servers,
317 )
318 .await;
319
320 frontend.start().await.unwrap();
321
322 GreptimeDbCluster {
323 datanode_options,
324 guards,
325 datanode_instances,
326 kv_backend: self.kv_backend.clone(),
327 metasrv: metasrv.metasrv,
328 frontend: Arc::new(frontend),
329 }
330 }
331
332 pub async fn build(&self, start_frontend_servers: bool) -> GreptimeDbCluster {
333 let datanodes = self.datanodes.unwrap_or(4);
334 let (datanode_options, guards) = self.build_datanode_options_and_guards(datanodes).await;
335 self.build_with(datanode_options, start_frontend_servers, guards)
336 .await
337 }
338
339 async fn build_datanode_options_and_guards(
340 &self,
341 datanodes: u32,
342 ) -> (Vec<DatanodeOptions>, Vec<TestGuard>) {
343 let mut options = Vec::with_capacity(datanodes as usize);
344 let mut guards = Vec::with_capacity(datanodes as usize);
345
346 for i in 0..datanodes {
347 let datanode_id = i as u64 + 1;
348 let mut opts = if let Some(store_config) = &self.store_config {
349 let home_dir = if let Some(home_dir) = &self.shared_home_dir {
350 home_dir.path().to_str().unwrap().to_string()
351 } else {
352 let home_tmp_dir = create_temp_dir(&format!("gt_home_{}", &self.cluster_name));
353 let home_dir = home_tmp_dir.path().to_str().unwrap().to_string();
354 guards.push(TestGuard {
355 home_guard: FileDirGuard::new(home_tmp_dir),
356 storage_guards: Vec::new(),
357 });
358
359 home_dir
360 };
361
362 create_datanode_opts(
363 store_config.clone(),
364 vec![],
365 home_dir,
366 self.datanode_wal_config.clone(),
367 self.datanode_gc_config.clone(),
368 )
369 } else {
370 let (opts, guard) = create_tmp_dir_and_datanode_opts(
371 StorageType::File,
372 self.store_providers.clone().unwrap_or_default(),
373 &format!("{}-dn-{}", self.cluster_name, datanode_id),
374 self.datanode_wal_config.clone(),
375 self.datanode_gc_config.clone(),
376 );
377 guards.push(guard);
378
379 opts
380 };
381 opts.node_id = Some(datanode_id);
382
383 options.push(opts);
384 }
385 (options, guards)
386 }
387
388 async fn build_datanodes_with_options(
389 &self,
390 metasrv: &MockInfo,
391 options: &[DatanodeOptions],
392 ) -> HashMap<DatanodeId, Datanode> {
393 let mut instances = HashMap::with_capacity(options.len());
394
395 for opts in options {
396 let datanode = self.create_datanode(opts.clone(), metasrv.clone()).await;
397 instances.insert(opts.node_id.unwrap(), datanode);
398 }
399
400 instances
401 }
402
403 async fn wait_datanodes_alive(
404 &self,
405 meta_peer_client: &MetaPeerClientRef,
406 expected_datanodes: usize,
407 ) {
408 for _ in 0..100 {
409 let alive_datanodes = discovery::utils::alive_datanode_infos(
410 &DefaultSystemTimer,
411 meta_peer_client.as_ref(),
412 Duration::from_secs(u64::MAX),
413 None,
414 )
415 .await
416 .unwrap()
417 .len();
418 if alive_datanodes == expected_datanodes {
419 return;
420 }
421 tokio::time::sleep(Duration::from_micros(100)).await
422 }
423 panic!("Some Datanodes are not alive in 10 seconds!")
424 }
425
426 async fn create_datanode(&self, opts: DatanodeOptions, metasrv: MockInfo) -> Datanode {
427 let mut meta_client = MetaClientBuilder::datanode_default_options(opts.node_id.unwrap())
428 .channel_manager(metasrv.channel_manager)
429 .build();
430 meta_client.start(&[&metasrv.server_addr]).await.unwrap();
431 let meta_client = Arc::new(meta_client);
432
433 let meta_backend = new_read_only_meta_kv_backend(meta_client.clone());
434
435 let layered_cache_registry = Arc::new(
436 LayeredCacheRegistryBuilder::default()
437 .add_cache_registry(build_datanode_cache_registry(meta_backend.clone()))
438 .build(),
439 );
440
441 let mut builder = DatanodeBuilder::new(opts.clone(), Plugins::default(), meta_backend);
442 builder
443 .with_cache_registry(layered_cache_registry)
444 .with_meta_client(meta_client);
445 let mut datanode = builder.build().await.unwrap();
446
447 datanode.start_heartbeat().await.unwrap();
448
449 datanode
450 }
451
452 async fn build_frontend(
453 &self,
454 metasrv: MockInfo,
455 datanode_clients: Arc<NodeClients>,
456 start_frontend_servers: bool,
457 ) -> Frontend {
458 let mut meta_client = MetaClientBuilder::frontend_default_options()
459 .channel_manager(metasrv.channel_manager)
460 .enable_access_cluster_info()
461 .build();
462 meta_client.start(&[&metasrv.server_addr]).await.unwrap();
463 let meta_client = Arc::new(meta_client);
464
465 let readonly_meta_backend = new_read_only_meta_kv_backend(meta_client.clone());
466 let cached_meta_backend =
467 Arc::new(CachedKvBackendBuilder::new(readonly_meta_backend.clone()).build());
468
469 let layered_cache_builder = LayeredCacheRegistryBuilder::default().add_cache_registry(
470 CacheRegistryBuilder::default()
471 .add_cache(cached_meta_backend.clone())
472 .build(),
473 );
474 let fundamental_cache_registry =
475 build_fundamental_cache_registry(readonly_meta_backend.clone());
476 let cache_registry = Arc::new(
477 with_default_composite_cache_registry(
478 layered_cache_builder.add_cache_registry(fundamental_cache_registry),
479 )
480 .unwrap()
481 .build(),
482 );
483
484 let information_extension = Arc::new(DistributedInformationExtension::new(
485 meta_client.clone(),
486 datanode_clients.clone(),
487 ));
488 let catalog_manager = KvBackendCatalogManagerBuilder::new(
489 information_extension,
490 cached_meta_backend.clone(),
491 cache_registry.clone(),
492 )
493 .build();
494
495 let fe_opts = self.build_frontend_options();
496
497 let instance = FrontendBuilder::new(
498 fe_opts.clone(),
499 cached_meta_backend.clone(),
500 cache_registry.clone(),
501 catalog_manager,
502 datanode_clients,
503 meta_client.clone(),
504 Arc::new(ProcessManager::new(fe_opts.grpc.server_addr.clone(), None)),
505 )
506 .with_local_cache_invalidator(cache_registry)
507 .try_build()
508 .await
509 .unwrap();
510
511 let heartbeat_task = create_heartbeat_task(&fe_opts, meta_client, &instance);
512
513 let instance = Arc::new(instance);
514
515 let servers = if start_frontend_servers {
517 Services::new(fe_opts, instance.clone(), Plugins::default())
518 .build()
519 .unwrap()
520 } else {
521 ServerHandlers::default()
522 };
523
524 Frontend {
525 instance,
526 servers,
527 heartbeat_task: Some(heartbeat_task),
528 }
529 }
530
531 fn build_frontend_options(&self) -> FrontendOptions {
532 let mut fe_opts = FrontendOptions::default();
533
534 let port_range = 14000..=24000;
536 let max_attempts = 10;
537 let localhost = "127.0.0.1";
538 let construct_addr = |port: u16| format!("{}:{}", localhost, port);
539
540 fe_opts.http.addr = construct_addr(self.choose_random_unused_port(
541 port_range.clone(),
542 max_attempts,
543 localhost,
544 ));
545
546 let grpc_port = self.choose_random_unused_port(port_range.clone(), max_attempts, localhost);
547 fe_opts.grpc.bind_addr = construct_addr(grpc_port);
548 fe_opts.grpc.server_addr = construct_addr(grpc_port);
549
550 fe_opts.mysql.addr = construct_addr(self.choose_random_unused_port(
551 port_range.clone(),
552 max_attempts,
553 localhost,
554 ));
555 fe_opts.postgres.addr =
556 construct_addr(self.choose_random_unused_port(port_range, max_attempts, localhost));
557
558 fe_opts
559 }
560
561 fn choose_random_unused_port(
563 &self,
564 port_range: RangeInclusive<u16>,
565 max_attempts: u16,
566 addr: &str,
567 ) -> u16 {
568 let mut rng = rand::rng();
569
570 let mut attempts = 0;
571 while attempts < max_attempts {
572 let port = rng.random_range(port_range.clone());
573 if TcpListener::bind(format!("{}:{}", addr, port)).is_ok() {
574 return port;
575 }
576 attempts += 1;
577 }
578
579 panic!("No unused port found");
580 }
581}
582
583async fn build_datanode_clients(
584 clients: Arc<NodeClients>,
585 instances: &HashMap<DatanodeId, Datanode>,
586) {
587 for (&datanode_id, instance) in instances {
588 let (addr, client) = create_datanode_client(instance).await;
589 clients
590 .insert_client(Peer::new(datanode_id, addr), client)
591 .await;
592 }
593}
594
595async fn create_datanode_client(datanode: &Datanode) -> (String, Client) {
596 let (client, server) = tokio::io::duplex(1024);
597
598 let runtime = RuntimeBuilder::default()
599 .worker_threads(2)
600 .thread_name("grpc-handlers")
601 .build()
602 .unwrap();
603
604 let flight_handler = FlightCraftWrapper(datanode.region_server());
605 let region_server_handler =
606 RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
607
608 let _handle = tokio::spawn(async move {
609 Server::builder()
610 .add_service(
611 FlightServiceServer::new(flight_handler)
612 .accept_compressed(CompressionEncoding::Gzip)
613 .accept_compressed(CompressionEncoding::Zstd)
614 .send_compressed(CompressionEncoding::Gzip)
615 .send_compressed(CompressionEncoding::Zstd),
616 )
617 .add_service(
618 RegionServer::new(region_server_handler)
619 .accept_compressed(CompressionEncoding::Gzip)
620 .accept_compressed(CompressionEncoding::Zstd)
621 .send_compressed(CompressionEncoding::Gzip)
622 .send_compressed(CompressionEncoding::Zstd),
623 )
624 .add_service(HealthCheckServer::new(HealthCheckHandler))
625 .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
626 .await
627 });
628
629 let mut client = Some(client);
630 let addr = test_util::PEER_PLACEHOLDER_ADDR;
631 let channel_manager = ChannelManager::new();
632 channel_manager
633 .reset_with_connector(
634 addr,
635 service_fn(move |_| {
636 let client = client.take();
637
638 async move {
639 client
640 .map(TokioIo::new)
641 .ok_or_else(|| std::io::Error::other("Client already taken"))
642 }
643 }),
644 )
645 .unwrap();
646
647 (
648 addr.to_string(),
649 Client::with_manager_and_urls(channel_manager, [addr]),
650 )
651}
652
653#[cfg(test)]
654mod tests {
655 use api::v1::flow::FlowRequest;
656 use api::v1::region::{
657 ListMetadataRequest, RegionRequest, RegionRequestHeader, region_request,
658 };
659 use catalog::memory::new_memory_catalog_manager;
660 use client::Client;
661 use common_error::ext::ErrorExt;
662 use common_error::status_code::StatusCode;
663 use common_meta::key::TableMetadataManager;
664 use common_meta::key::flow::FlowMetadataManager;
665 use common_meta::kv_backend::memory::MemoryKvBackend;
666 use common_meta::node_manager::{DatanodeManager, FlownodeManager};
667 use common_meta::peer::Peer;
668 use flow::{FlownodeBuilder, FlownodeOptions, FlownodeServiceBuilder, FrontendClient};
669
670 use super::*;
671
672 const HEALTH_REFRESH_INTERVAL: Duration = Duration::from_millis(10);
673 const HEALTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
674
675 async fn wait_for_active_route(client: &Client, active_addr: &str) {
676 client.find_channel().unwrap();
677 tokio::time::timeout(HEALTH_REFRESH_TIMEOUT, async {
678 let mut interval = tokio::time::interval(HEALTH_REFRESH_INTERVAL);
679 loop {
680 interval.tick().await;
681 let (active, inactive) = client.peer_addresses_by_state();
682 if active.len() == 1 && active[0] == active_addr && inactive.len() == 1 {
683 return;
684 }
685 }
686 })
687 .await
688 .expect("background health routing did not select the active peer");
689 }
690
691 async fn create_region_requester_health_client(datanode: &Datanode) -> (Peer, Peer, Client) {
692 const ACTIVE_ADDR: &str = "127.0.0.1:3001";
693 const INACTIVE_ADDR: &str = "127.0.0.1:3002";
694
695 let (active_client, server) = tokio::io::duplex(1024);
696 let runtime = RuntimeBuilder::default()
697 .worker_threads(2)
698 .thread_name("grpc-handlers")
699 .build()
700 .unwrap();
701 let region_server_handler =
702 RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
703
704 tokio::spawn(async move {
705 Server::builder()
706 .add_service(RegionServer::new(region_server_handler))
707 .add_service(HealthCheckServer::new(HealthCheckHandler))
708 .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
709 .await
710 });
711
712 let channel_manager = ChannelManager::new();
713 let mut active_client = Some(active_client);
714 channel_manager
715 .reset_with_connector(
716 ACTIVE_ADDR,
717 service_fn(move |_| {
718 let client = active_client.take();
719
720 async move {
721 client
722 .map(TokioIo::new)
723 .ok_or_else(|| std::io::Error::other("Active client already taken"))
724 }
725 }),
726 )
727 .unwrap();
728 channel_manager
729 .reset_with_connector(
730 INACTIVE_ADDR,
731 service_fn(|_| async {
732 Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
733 "Inactive synthetic peer",
734 ))
735 }),
736 )
737 .unwrap();
738
739 let client = Client::with_manager_and_urls_and_options(
740 channel_manager,
741 [INACTIVE_ADDR, ACTIVE_ADDR],
742 client::ClientOptions {
743 health_check_interval: HEALTH_REFRESH_INTERVAL,
744 ..Default::default()
745 },
746 );
747 (
748 Peer::new(1, ACTIVE_ADDR),
749 Peer::new(2, INACTIVE_ADDR),
750 client,
751 )
752 }
753
754 async fn create_flow_requester_health_client(
755 flownode: &flow::FlownodeInstance,
756 ) -> (Peer, Peer, Client) {
757 const ACTIVE_ADDR: &str = "127.0.0.1:3003";
758 const INACTIVE_ADDR: &str = "127.0.0.1:3004";
759
760 let (active_client, server) = tokio::io::duplex(1024);
761 let flow_service = flownode.flownode_server().create_flow_service();
762
763 tokio::spawn(async move {
764 Server::builder()
765 .add_service(flow_service)
766 .add_service(HealthCheckServer::new(HealthCheckHandler))
767 .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
768 .await
769 });
770
771 let channel_manager = ChannelManager::new();
772 let mut active_client = Some(active_client);
773 channel_manager
774 .reset_with_connector(
775 ACTIVE_ADDR,
776 service_fn(move |_| {
777 let client = active_client.take();
778
779 async move {
780 client
781 .map(TokioIo::new)
782 .ok_or_else(|| std::io::Error::other("Active client already taken"))
783 }
784 }),
785 )
786 .unwrap();
787 channel_manager
788 .reset_with_connector(
789 INACTIVE_ADDR,
790 service_fn(|_| async {
791 Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
792 "Inactive synthetic peer",
793 ))
794 }),
795 )
796 .unwrap();
797
798 let client = Client::with_manager_and_urls_and_options(
799 channel_manager,
800 [INACTIVE_ADDR, ACTIVE_ADDR],
801 client::ClientOptions {
802 health_check_interval: HEALTH_REFRESH_INTERVAL,
803 ..Default::default()
804 },
805 );
806 (
807 Peer::new(1, ACTIVE_ADDR),
808 Peer::new(2, INACTIVE_ADDR),
809 client,
810 )
811 }
812
813 async fn start_flownode(addr: String) -> flow::FlownodeInstance {
814 let kv_backend = Arc::new(MemoryKvBackend::new());
815 let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
816 table_meta.init().await.unwrap();
817 let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend));
818 let catalog_manager = new_memory_catalog_manager().unwrap();
819 let (frontend_client, _handler) =
820 FrontendClient::from_empty_grpc_handler(Default::default());
821 let mut opts = FlownodeOptions::default();
822 opts.grpc.bind_addr = addr.clone();
823 opts.grpc.server_addr = addr;
824
825 let mut flownode = FlownodeBuilder::new(
826 opts.clone(),
827 Plugins::default(),
828 table_meta,
829 catalog_manager,
830 flow_meta,
831 Arc::new(frontend_client),
832 )
833 .build()
834 .await
835 .unwrap();
836 let services = FlownodeServiceBuilder::new(&opts)
837 .with_default_grpc_server(flownode.flownode_server())
838 .build()
839 .unwrap();
840 flownode.setup_services(services);
841 flownode.start().await.unwrap();
842 flownode
843 }
844
845 #[tokio::test(flavor = "multi_thread")]
846 async fn test_region_requester_health_check_and_list_metadata() {
847 let cluster = GreptimeDbClusterBuilder::new("region_requester_health_check")
849 .await
850 .with_datanodes(1)
851 .build(false)
852 .await;
853 let (active_peer, inactive_peer, client) =
854 create_region_requester_health_client(cluster.datanode_instances.get(&1).unwrap())
855 .await;
856 let clients = NodeClients::default();
857 clients
858 .insert_client(active_peer.clone(), client.clone())
859 .await;
860 clients
861 .insert_client(inactive_peer.clone(), client.clone())
862 .await;
863
864 wait_for_active_route(&client, &active_peer.addr).await;
866 let requester = clients.datanode(&active_peer).await;
867 let response = requester
868 .handle(RegionRequest {
869 header: Some(RegionRequestHeader::default()),
870 body: Some(region_request::Body::ListMetadata(ListMetadataRequest {
871 region_ids: vec![],
872 })),
873 })
874 .await
875 .unwrap();
876
877 assert_eq!(
879 (vec![active_peer.addr], vec![inactive_peer.addr]),
880 client.peer_addresses_by_state()
881 );
882 assert_eq!(b"[]", response.metadata.as_slice());
883 }
884
885 #[tokio::test(flavor = "multi_thread")]
886 async fn test_flow_requester_health_check_and_handle_request() {
887 let mut flownode = start_flownode("127.0.0.1:0".to_string()).await;
889 let (active_peer, inactive_peer, client) =
890 create_flow_requester_health_client(&flownode).await;
891 let clients = NodeClients::default();
892 clients
893 .insert_client(active_peer.clone(), client.clone())
894 .await;
895 clients
896 .insert_client(inactive_peer.clone(), client.clone())
897 .await;
898
899 wait_for_active_route(&client, &active_peer.addr).await;
901 let response = clients
902 .flownode(&active_peer)
903 .await
904 .handle(FlowRequest::default())
905 .await;
906
907 assert_eq!(
909 (vec![active_peer.addr], vec![inactive_peer.addr]),
910 client.peer_addresses_by_state()
911 );
912 assert_eq!(
913 StatusCode::InvalidArguments,
914 response.unwrap_err().status_code()
915 );
916 flownode.shutdown().await.unwrap();
917 }
918}