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_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    /// List all SST files from all datanodes.
103    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    /// List all SST files from the manifests of all datanodes.
124    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            // Each retry requires a new isolation namespace.
196            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    /// Sets whether the test frontend automatically creates tables on write.
269    #[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    /// Configure the frontend COPY sandbox for filesystem integration tests.
288    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                // Due to large network delay during cross data-center.
306                // We only make max_retry_times and retry_delay large than the default in tests.
307                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        // Build the servers for the frontend.
548        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        // Choose a random unused port between [14000, 24000] for local test to avoid conflicts.
571        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    // Choose a random unused port between [start, end].
598    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        // The mock connector pool lives on this single manager, so both lanes
686        // intentionally share it (same semantics as the legacy constructor).
687        Client::with_query_and_control_managers(channel_manager.clone(), channel_manager, [addr]),
688    )
689}
690
691// Mock connectors are registered on a shared channel manager.
692#[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        // Arrange: configure active and inactive synthetic peers over the duplex harness.
888        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        // Act: wait for background health routing, then send a real Region RPC.
905        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: the inactive peer was rejected and the RegionRequester reached the active service.
918        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        // Arrange: start a flownode and expose its production flow service through a duplex harness.
928        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        // Act: wait for background health routing, then send a real Flow RPC.
940        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: the inactive peer was rejected and the production Flow service handled the RPC.
948        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}