Skip to main content

tests_integration/
cluster.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use 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    /// List all SST files from all datanodes.
101    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    /// List all SST files from the manifests of all datanodes.
122    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            // Each retry requires a new isolation namespace.
191            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                // Due to large network delay during cross data-center.
276                // We only make max_retry_times and retry_delay large than the default in tests.
277                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        // Build the servers for the frontend.
516        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        // Choose a random unused port between [14000, 24000] for local test to avoid conflicts.
535        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    // Choose a random unused port between [start, end].
562    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        // Arrange: configure active and inactive synthetic peers over the duplex harness.
848        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        // Act: wait for background health routing, then send a real Region RPC.
865        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: the inactive peer was rejected and the RegionRequester reached the active service.
878        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        // Arrange: start a flownode and expose its production flow service through a duplex harness.
888        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        // Act: wait for background health routing, then send a real Flow RPC.
900        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: the inactive peer was rejected and the production Flow service handled the RPC.
908        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}