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_query_and_control_managers(channel_manager.clone(), channel_manager, [addr]),
688 )
689}
690
691#[cfg(test)]
693#[allow(deprecated)]
694mod tests {
695 use api::v1::flow::FlowRequest;
696 use api::v1::region::{
697 ListMetadataRequest, RegionRequest, RegionRequestHeader, region_request,
698 };
699 use catalog::memory::new_memory_catalog_manager;
700 use client::Client;
701 use common_error::ext::ErrorExt;
702 use common_error::status_code::StatusCode;
703 use common_meta::key::TableMetadataManager;
704 use common_meta::key::flow::FlowMetadataManager;
705 use common_meta::kv_backend::memory::MemoryKvBackend;
706 use common_meta::node_manager::{DatanodeManager, FlownodeManager};
707 use common_meta::peer::Peer;
708 use flow::{FlownodeBuilder, FlownodeOptions, FlownodeServiceBuilder, FrontendClient};
709
710 use super::*;
711
712 const HEALTH_REFRESH_INTERVAL: Duration = Duration::from_millis(10);
713 const HEALTH_REFRESH_TIMEOUT: Duration = Duration::from_secs(5);
714
715 async fn wait_for_active_route(client: &Client, active_addr: &str) {
716 client.find_channel().unwrap();
717 tokio::time::timeout(HEALTH_REFRESH_TIMEOUT, async {
718 let mut interval = tokio::time::interval(HEALTH_REFRESH_INTERVAL);
719 loop {
720 interval.tick().await;
721 let (active, inactive) = client.peer_addresses_by_state();
722 if active.len() == 1 && active[0] == active_addr && inactive.len() == 1 {
723 return;
724 }
725 }
726 })
727 .await
728 .expect("background health routing did not select the active peer");
729 }
730
731 async fn create_region_requester_health_client(datanode: &Datanode) -> (Peer, Peer, Client) {
732 const ACTIVE_ADDR: &str = "127.0.0.1:3001";
733 const INACTIVE_ADDR: &str = "127.0.0.1:3002";
734
735 let (active_client, server) = tokio::io::duplex(1024);
736 let runtime = RuntimeBuilder::default()
737 .worker_threads(2)
738 .thread_name("grpc-handlers")
739 .build()
740 .unwrap();
741 let region_server_handler =
742 RegionServerRequestHandler::new(Arc::new(datanode.region_server()), runtime);
743
744 tokio::spawn(async move {
745 Server::builder()
746 .add_service(RegionServer::new(region_server_handler))
747 .add_service(HealthCheckServer::new(HealthCheckHandler))
748 .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
749 .await
750 });
751
752 let channel_manager = ChannelManager::new();
753 let mut active_client = Some(active_client);
754 channel_manager
755 .reset_with_connector(
756 ACTIVE_ADDR,
757 service_fn(move |_| {
758 let client = active_client.take();
759
760 async move {
761 client
762 .map(TokioIo::new)
763 .ok_or_else(|| std::io::Error::other("Active client already taken"))
764 }
765 }),
766 )
767 .unwrap();
768 channel_manager
769 .reset_with_connector(
770 INACTIVE_ADDR,
771 service_fn(|_| async {
772 Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
773 "Inactive synthetic peer",
774 ))
775 }),
776 )
777 .unwrap();
778
779 let client = Client::with_manager_and_urls_and_options(
780 channel_manager,
781 [INACTIVE_ADDR, ACTIVE_ADDR],
782 client::ClientOptions {
783 health_check_interval: HEALTH_REFRESH_INTERVAL,
784 ..Default::default()
785 },
786 );
787 (
788 Peer::new(1, ACTIVE_ADDR),
789 Peer::new(2, INACTIVE_ADDR),
790 client,
791 )
792 }
793
794 async fn create_flow_requester_health_client(
795 flownode: &flow::FlownodeInstance,
796 ) -> (Peer, Peer, Client) {
797 const ACTIVE_ADDR: &str = "127.0.0.1:3003";
798 const INACTIVE_ADDR: &str = "127.0.0.1:3004";
799
800 let (active_client, server) = tokio::io::duplex(1024);
801 let flow_service = flownode.flownode_server().create_flow_service();
802
803 tokio::spawn(async move {
804 Server::builder()
805 .add_service(flow_service)
806 .add_service(HealthCheckServer::new(HealthCheckHandler))
807 .serve_with_incoming(futures::stream::iter(vec![Ok::<_, std::io::Error>(server)]))
808 .await
809 });
810
811 let channel_manager = ChannelManager::new();
812 let mut active_client = Some(active_client);
813 channel_manager
814 .reset_with_connector(
815 ACTIVE_ADDR,
816 service_fn(move |_| {
817 let client = active_client.take();
818
819 async move {
820 client
821 .map(TokioIo::new)
822 .ok_or_else(|| std::io::Error::other("Active client already taken"))
823 }
824 }),
825 )
826 .unwrap();
827 channel_manager
828 .reset_with_connector(
829 INACTIVE_ADDR,
830 service_fn(|_| async {
831 Err::<TokioIo<tokio::io::DuplexStream>, _>(std::io::Error::other(
832 "Inactive synthetic peer",
833 ))
834 }),
835 )
836 .unwrap();
837
838 let client = Client::with_manager_and_urls_and_options(
839 channel_manager,
840 [INACTIVE_ADDR, ACTIVE_ADDR],
841 client::ClientOptions {
842 health_check_interval: HEALTH_REFRESH_INTERVAL,
843 ..Default::default()
844 },
845 );
846 (
847 Peer::new(1, ACTIVE_ADDR),
848 Peer::new(2, INACTIVE_ADDR),
849 client,
850 )
851 }
852
853 async fn start_flownode(addr: String) -> flow::FlownodeInstance {
854 let kv_backend = Arc::new(MemoryKvBackend::new());
855 let table_meta = Arc::new(TableMetadataManager::new(kv_backend.clone()));
856 table_meta.init().await.unwrap();
857 let flow_meta = Arc::new(FlowMetadataManager::new(kv_backend));
858 let catalog_manager = new_memory_catalog_manager().unwrap();
859 let (frontend_client, _handler) =
860 FrontendClient::from_empty_grpc_handler(Default::default());
861 let mut opts = FlownodeOptions::default();
862 opts.grpc.bind_addr = addr.clone();
863 opts.grpc.server_addr = addr;
864
865 let mut flownode = FlownodeBuilder::new(
866 opts.clone(),
867 Plugins::default(),
868 table_meta,
869 catalog_manager,
870 flow_meta,
871 Arc::new(frontend_client),
872 )
873 .build()
874 .await
875 .unwrap();
876 let services = FlownodeServiceBuilder::new(&opts)
877 .with_default_grpc_server(flownode.flownode_server())
878 .build()
879 .unwrap();
880 flownode.setup_services(services);
881 flownode.start().await.unwrap();
882 flownode
883 }
884
885 #[tokio::test(flavor = "multi_thread")]
886 async fn test_region_requester_health_check_and_list_metadata() {
887 let cluster = GreptimeDbClusterBuilder::new("region_requester_health_check")
889 .await
890 .with_datanodes(1)
891 .build(false)
892 .await;
893 let (active_peer, inactive_peer, client) =
894 create_region_requester_health_client(cluster.datanode_instances.get(&1).unwrap())
895 .await;
896 let clients = NodeClients::default();
897 clients
898 .insert_client(active_peer.clone(), client.clone())
899 .await;
900 clients
901 .insert_client(inactive_peer.clone(), client.clone())
902 .await;
903
904 wait_for_active_route(&client, &active_peer.addr).await;
906 let requester = clients.datanode(&active_peer).await;
907 let response = requester
908 .handle(RegionRequest {
909 header: Some(RegionRequestHeader::default()),
910 body: Some(region_request::Body::ListMetadata(ListMetadataRequest {
911 region_ids: vec![],
912 })),
913 })
914 .await
915 .unwrap();
916
917 assert_eq!(
919 (vec![active_peer.addr], vec![inactive_peer.addr]),
920 client.peer_addresses_by_state()
921 );
922 assert_eq!(b"[]", response.metadata.as_slice());
923 }
924
925 #[tokio::test(flavor = "multi_thread")]
926 async fn test_flow_requester_health_check_and_handle_request() {
927 let mut flownode = start_flownode("127.0.0.1:0".to_string()).await;
929 let (active_peer, inactive_peer, client) =
930 create_flow_requester_health_client(&flownode).await;
931 let clients = NodeClients::default();
932 clients
933 .insert_client(active_peer.clone(), client.clone())
934 .await;
935 clients
936 .insert_client(inactive_peer.clone(), client.clone())
937 .await;
938
939 wait_for_active_route(&client, &active_peer.addr).await;
941 let response = clients
942 .flownode(&active_peer)
943 .await
944 .handle(FlowRequest::default())
945 .await;
946
947 assert_eq!(
949 (vec![active_peer.addr], vec![inactive_peer.addr]),
950 client.peer_addresses_by_state()
951 );
952 assert_eq!(
953 StatusCode::InvalidArguments,
954 response.unwrap_err().status_code()
955 );
956 flownode.shutdown().await.unwrap();
957 }
958}