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