Skip to main content

datanode/
datanode.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
15//! Datanode implementation.
16
17use std::collections::HashMap;
18use std::path::Path;
19use std::sync::Arc;
20use std::time::{Duration, Instant};
21
22use common_base::Plugins;
23use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE};
24use common_datasource::object_store::LocalFileAccess;
25use common_error::ext::BoxedError;
26use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
27use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef};
28use common_meta::cache_invalidator::CacheInvalidatorRef;
29use common_meta::datanode::TopicStatsReporter;
30use common_meta::key::runtime_switch::RuntimeSwitchManager;
31use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef};
32use common_meta::kv_backend::KvBackendRef;
33pub use common_procedure::options::ProcedureConfig;
34use common_query::prelude::set_default_prefix;
35use common_stat::ResourceStatImpl;
36use common_telemetry::{error, info, warn};
37use common_wal::config::DatanodeWalConfig;
38use common_wal::config::kafka::DatanodeKafkaConfig;
39use common_wal::config::object_store::ObjectStoreWalConfig;
40use common_wal::config::raft_engine::RaftEngineConfig;
41use file_engine::engine::FileRegionEngine;
42use log_store::kafka::log_store::KafkaLogStore;
43use log_store::kafka::{GlobalIndexCollector, default_index_file};
44use log_store::noop::log_store::NoopLogStore;
45use log_store::raft_engine::log_store::RaftEngineLogStore;
46use meta_client::MetaClientRef;
47use metric_engine::engine::MetricEngine;
48use mito2::config::MitoConfig;
49use mito2::engine::{MitoEngine, MitoEngineBuilder};
50use mito2::region::opener::PartitionExprFetcherRef;
51use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef};
52use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
53use object_store::util::normalize_dir;
54use query::QueryEngineFactory;
55use query::dummy_catalog::{DummyCatalogManager, TableProviderFactoryRef};
56use servers::server::ServerHandlers;
57use snafu::{OptionExt, ResultExt, ensure};
58use store_api::path_utils::WAL_DIR;
59use store_api::region_engine::{
60    RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState,
61};
62use tokio::fs;
63use tokio::sync::Notify;
64
65use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
66use crate::error::{
67    self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu,
68    DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu,
69    InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu,
70    ObjectStoreWalNotSupportedSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu,
71    ShutdownServerSnafu, StartServerSnafu,
72};
73use crate::event_listener::{
74    NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver,
75    new_region_server_event_channel,
76};
77use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
78use crate::heartbeat::HeartbeatTask;
79use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
80use crate::region_server::{DummyTableProviderFactory, RegionServer};
81use crate::store::{self, new_object_store_without_cache};
82use crate::utils::{RegionOpenRequests, build_region_open_requests};
83
84/// Datanode service.
85pub struct Datanode {
86    services: ServerHandlers,
87    heartbeat_task: Option<HeartbeatTask>,
88    region_event_receiver: Option<RegionServerEventReceiver>,
89    region_server: RegionServer,
90    greptimedb_telemetry_task: Arc<GreptimeDBTelemetryTask>,
91    leases_notifier: Option<Arc<Notify>>,
92    plugins: Plugins,
93}
94
95impl Datanode {
96    pub async fn start(&mut self) -> Result<()> {
97        info!("Starting datanode instance...");
98
99        self.start_heartbeat().await?;
100        self.wait_coordinated().await;
101
102        self.start_telemetry();
103
104        self.services.start_all().await.context(StartServerSnafu)
105    }
106
107    pub fn server_handlers(&self) -> &ServerHandlers {
108        &self.services
109    }
110
111    pub fn start_telemetry(&self) {
112        if let Err(e) = self.greptimedb_telemetry_task.start() {
113            warn!(e; "Failed to start telemetry task!");
114        }
115    }
116
117    pub async fn start_heartbeat(&mut self) -> Result<()> {
118        if let Some(task) = &self.heartbeat_task {
119            // Safety: The event_receiver must exist.
120            let receiver = self.region_event_receiver.take().unwrap();
121
122            task.start(receiver, self.leases_notifier.clone()).await?;
123        }
124        Ok(())
125    }
126
127    /// If `leases_notifier` exists, it waits until leases have been obtained in all regions.
128    pub async fn wait_coordinated(&mut self) {
129        if let Some(notifier) = self.leases_notifier.take() {
130            notifier.notified().await;
131        }
132    }
133
134    pub fn setup_services(&mut self, services: ServerHandlers) {
135        self.services = services;
136    }
137
138    pub async fn shutdown(&mut self) -> Result<()> {
139        self.services
140            .shutdown_all()
141            .await
142            .context(ShutdownServerSnafu)?;
143
144        let _ = self.greptimedb_telemetry_task.stop().await;
145        if let Some(heartbeat_task) = &self.heartbeat_task {
146            heartbeat_task
147                .close()
148                .map_err(BoxedError::new)
149                .context(ShutdownInstanceSnafu)?;
150        }
151        self.region_server.stop().await?;
152        Ok(())
153    }
154
155    pub fn region_server(&self) -> RegionServer {
156        self.region_server.clone()
157    }
158
159    pub fn plugins(&self) -> Plugins {
160        self.plugins.clone()
161    }
162}
163
164pub struct DatanodeBuilder {
165    opts: DatanodeOptions,
166    table_provider_factory: Option<TableProviderFactoryRef>,
167    plugins: Plugins,
168    meta_client: Option<MetaClientRef>,
169    kv_backend: KvBackendRef,
170    cache_registry: Option<Arc<LayeredCacheRegistry>>,
171    topic_stats_reporter: Option<Box<dyn TopicStatsReporter>>,
172    open_regions_writable_override: Option<bool>,
173    local_file_access: LocalFileAccess,
174    #[cfg(feature = "enterprise")]
175    extension_range_provider_factory: Option<mito2::extension::BoxedExtensionRangeProviderFactory>,
176}
177
178impl DatanodeBuilder {
179    pub fn new(opts: DatanodeOptions, plugins: Plugins, kv_backend: KvBackendRef) -> Self {
180        Self {
181            opts,
182            table_provider_factory: None,
183            plugins,
184            meta_client: None,
185            kv_backend,
186            cache_registry: None,
187            open_regions_writable_override: None,
188            local_file_access: LocalFileAccess::Disabled,
189            #[cfg(feature = "enterprise")]
190            extension_range_provider_factory: None,
191            topic_stats_reporter: None,
192        }
193    }
194
195    pub fn options(&self) -> &DatanodeOptions {
196        &self.opts
197    }
198
199    pub fn with_meta_client(&mut self, client: MetaClientRef) -> &mut Self {
200        self.meta_client = Some(client);
201        self
202    }
203
204    pub fn with_cache_registry(&mut self, registry: Arc<LayeredCacheRegistry>) -> &mut Self {
205        self.cache_registry = Some(registry);
206        self
207    }
208
209    pub fn with_local_file_access(&mut self, local_file_access: LocalFileAccess) -> &mut Self {
210        self.local_file_access = local_file_access;
211        self
212    }
213
214    pub fn kv_backend(&self) -> &KvBackendRef {
215        &self.kv_backend
216    }
217
218    pub fn meta_client(&self) -> Option<&MetaClientRef> {
219        self.meta_client.as_ref()
220    }
221
222    pub fn cache_registry(&self) -> Option<&Arc<LayeredCacheRegistry>> {
223        self.cache_registry.as_ref()
224    }
225
226    pub fn set_plugins(&mut self, plugins: Plugins) {
227        self.plugins = plugins;
228    }
229
230    pub fn with_table_provider_factory(&mut self, factory: TableProviderFactoryRef) -> &mut Self {
231        self.table_provider_factory = Some(factory);
232        self
233    }
234
235    /// Overrides whether regions opened during datanode startup should become writable.
236    ///
237    /// When unset, the builder uses its default writable policy for reopened regions
238    /// (writable only when no metasrv client is configured).
239    ///
240    /// Warning: setting this to `true` on a metasrv-controlled datanode (one built
241    /// with `with_meta_client`) will promote regions to Leader before heartbeat and
242    /// lease coordination begin, bypassing the metasrv safety contract and creating a
243    /// potential split-brain window during startup.
244    pub fn with_open_regions_writable_override(&mut self, writable: bool) -> &mut Self {
245        self.open_regions_writable_override = Some(writable);
246        self
247    }
248
249    #[cfg(feature = "enterprise")]
250    pub fn with_extension_range_provider(
251        &mut self,
252        extension_range_provider_factory: mito2::extension::BoxedExtensionRangeProviderFactory,
253    ) -> &mut Self {
254        self.extension_range_provider_factory = Some(extension_range_provider_factory);
255        self
256    }
257
258    pub async fn build(mut self) -> Result<Datanode> {
259        validate_region_engine_config(&self.opts.region_engine)?;
260        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;
261        set_default_prefix(self.opts.default_column_prefix.as_deref())
262            .map_err(BoxedError::new)
263            .context(BuildDatanodeSnafu)?;
264
265        let meta_client = self.meta_client.take();
266
267        // If metasrv client is provided, we will use it to control the region server.
268        // Otherwise the region server is self-controlled, meaning no heartbeat and immediately
269        // writable upon open.
270        let controlled_by_metasrv = meta_client.is_some();
271        if let DatanodeWalConfig::ObjectStore(config) = &self.opts.wal {
272            validate_object_store_wal_config(config)?;
273        }
274
275        // build and initialize region server
276        let (region_event_listener, region_event_receiver) = if controlled_by_metasrv {
277            let (tx, rx) = new_region_server_event_channel();
278            (Box::new(tx) as _, Some(rx))
279        } else {
280            (Box::new(NoopRegionServerEventListener) as _, None)
281        };
282
283        let cache_registry = self.cache_registry.take().context(MissingCacheSnafu)?;
284        let schema_cache: SchemaCacheRef = cache_registry.get().context(MissingCacheSnafu)?;
285        let table_id_schema_cache: TableSchemaCacheRef =
286            cache_registry.get().context(MissingCacheSnafu)?;
287
288        let schema_metadata_manager = Arc::new(SchemaMetadataManager::new(
289            table_id_schema_cache,
290            schema_cache,
291        ));
292
293        let gc_enabled = self.opts.region_engine.iter().any(|engine| {
294            if let RegionEngineConfig::Mito(config) = engine {
295                config.gc.enable
296            } else {
297                false
298            }
299        });
300
301        let file_ref_manager = Arc::new(FileReferenceManager::with_gc_enabled(
302            Some(node_id),
303            gc_enabled,
304        ));
305        let region_server = self
306            .new_region_server(
307                schema_metadata_manager,
308                region_event_listener,
309                file_ref_manager,
310            )
311            .await?;
312
313        // TODO(weny): Considering introducing a readonly kv_backend trait.
314        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
315        let is_recovery_mode = runtime_switch_manager
316            .recovery_mode()
317            .await
318            .context(GetMetadataSnafu)?;
319
320        let region_open_requests =
321            build_region_open_requests(node_id, self.kv_backend.clone()).await?;
322        let open_with_writable = self
323            .open_regions_writable_override
324            .unwrap_or(!controlled_by_metasrv);
325        let open_all_regions = open_all_regions(
326            region_server.clone(),
327            region_open_requests,
328            open_with_writable,
329            self.opts.init_regions_parallelism,
330            // Ignore nonexistent regions in recovery mode.
331            is_recovery_mode,
332        );
333
334        if self.opts.init_regions_in_background {
335            // Opens regions in background.
336            common_runtime::spawn_global(async move {
337                if let Err(err) = open_all_regions.await {
338                    error!(err; "Failed to open regions during the startup.");
339                }
340            });
341        } else {
342            open_all_regions.await?;
343        }
344
345        let heartbeat_task = if let Some(meta_client) = meta_client {
346            let task = self
347                .create_heartbeat_task(&region_server, meta_client, cache_registry)
348                .await?;
349            Some(task)
350        } else {
351            None
352        };
353
354        let is_standalone = heartbeat_task.is_none();
355        let greptimedb_telemetry_task = get_greptimedb_telemetry_task(
356            Some(self.opts.storage.data_home.clone()),
357            is_standalone && self.opts.enable_telemetry,
358        )
359        .await;
360
361        let leases_notifier = if self.opts.require_lease_before_startup && !is_standalone {
362            Some(Arc::new(Notify::new()))
363        } else {
364            None
365        };
366
367        Ok(Datanode {
368            services: ServerHandlers::default(),
369            heartbeat_task,
370            region_server,
371            greptimedb_telemetry_task,
372            region_event_receiver,
373            leases_notifier,
374            plugins: self.plugins.clone(),
375        })
376    }
377
378    async fn create_heartbeat_task(
379        &self,
380        region_server: &RegionServer,
381        meta_client: MetaClientRef,
382        cache_invalidator: CacheInvalidatorRef,
383    ) -> Result<HeartbeatTask> {
384        let stat = {
385            let mut stat = ResourceStatImpl::default();
386            stat.start_collect_cpu_usage();
387            Arc::new(stat)
388        };
389
390        HeartbeatTask::try_new(
391            &self.opts,
392            region_server.clone(),
393            meta_client,
394            self.kv_backend.clone(),
395            cache_invalidator,
396            self.plugins.clone(),
397            stat,
398        )
399        .await
400    }
401
402    /// Builds [ObjectStoreManager] from [StorageConfig].
403    pub async fn build_object_store_manager(cfg: &StorageConfig) -> Result<ObjectStoreManagerRef> {
404        let object_store = store::new_object_store(cfg.store.clone(), &cfg.data_home).await?;
405        let default_name = cfg.store.config_name();
406        let mut object_store_manager = ObjectStoreManager::new(default_name, object_store);
407        for store in &cfg.providers {
408            object_store_manager.add(
409                store.config_name(),
410                store::new_object_store(store.clone(), &cfg.data_home).await?,
411            );
412        }
413        Ok(Arc::new(object_store_manager))
414    }
415
416    #[cfg(test)]
417    /// Open all regions belong to this datanode.
418    async fn initialize_region_server(
419        &self,
420        region_server: &RegionServer,
421        open_with_writable: bool,
422    ) -> Result<()> {
423        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;
424
425        // TODO(weny): Considering introducing a readonly kv_backend trait.
426        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
427        let is_recovery_mode = runtime_switch_manager
428            .recovery_mode()
429            .await
430            .context(GetMetadataSnafu)?;
431        let region_open_requests =
432            build_region_open_requests(node_id, self.kv_backend.clone()).await?;
433
434        open_all_regions(
435            region_server.clone(),
436            region_open_requests,
437            open_with_writable,
438            self.opts.init_regions_parallelism,
439            is_recovery_mode,
440        )
441        .await
442    }
443
444    async fn new_region_server(
445        &mut self,
446        schema_metadata_manager: SchemaMetadataManagerRef,
447        event_listener: RegionServerEventListenerRef,
448        file_ref_manager: FileReferenceManagerRef,
449    ) -> Result<RegionServer> {
450        let opts: &DatanodeOptions = &self.opts;
451
452        let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
453            // query engine in datanode only executes plan with resolved table source.
454            DummyCatalogManager::arc(),
455            None,
456            None,
457            None,
458            None,
459            None,
460            false,
461            self.plugins.clone(),
462            opts.query.clone(),
463        )
464        .context(DataFusionSnafu)?;
465        let query_engine = query_engine_factory.query_engine();
466
467        let table_provider_factory = self
468            .table_provider_factory
469            .clone()
470            .unwrap_or_else(|| Arc::new(DummyTableProviderFactory));
471        let mut region_server = RegionServer::with_table_provider(
472            query_engine,
473            common_runtime::global_runtime(),
474            event_listener,
475            table_provider_factory,
476            opts.max_concurrent_queries,
477            opts.concurrent_query_limiter_timeout,
478            opts.grpc.flight_compression,
479        );
480        region_server.install_remote_dyn_filter_receiver_injector(&self.plugins);
481
482        let object_store_manager = Self::build_object_store_manager(&opts.storage).await?;
483        let engines = self
484            .build_store_engines(
485                object_store_manager,
486                schema_metadata_manager,
487                file_ref_manager,
488                self.plugins.clone(),
489            )
490            .await?;
491        for engine in engines {
492            region_server.register_engine(engine);
493        }
494        if let Some(topic_stats_reporter) = self.topic_stats_reporter.take() {
495            region_server.set_topic_stats_reporter(topic_stats_reporter);
496        }
497
498        Ok(region_server)
499    }
500
501    // internal utils
502
503    /// Builds [RegionEngineRef] from `store_engine` section in `opts`
504    async fn build_store_engines(
505        &mut self,
506        object_store_manager: ObjectStoreManagerRef,
507        schema_metadata_manager: SchemaMetadataManagerRef,
508        file_ref_manager: FileReferenceManagerRef,
509        plugins: Plugins,
510    ) -> Result<Vec<RegionEngineRef>> {
511        let mut metric_engine_config = metric_engine::config::EngineConfig::default();
512        let mut mito_engine_config = MitoConfig::default();
513        let mut file_engine_config = file_engine::config::EngineConfig::default();
514
515        for engine in &self.opts.region_engine {
516            match engine {
517                RegionEngineConfig::Mito(config) => {
518                    mito_engine_config = config.clone();
519                }
520                RegionEngineConfig::File(config) => {
521                    file_engine_config = config.clone();
522                }
523                RegionEngineConfig::Metric(metric_config) => {
524                    metric_engine_config = metric_config.clone();
525                }
526            }
527        }
528
529        // Build a fetcher to backfill partition_expr on open.
530        let fetcher = Arc::new(MetaPartitionExprFetcher::new(self.kv_backend.clone()));
531        let mito_engine = self
532            .build_mito_engine(
533                object_store_manager.clone(),
534                mito_engine_config,
535                schema_metadata_manager.clone(),
536                file_ref_manager.clone(),
537                fetcher.clone(),
538                plugins.clone(),
539            )
540            .await?;
541
542        let metric_engine = MetricEngine::try_new(mito_engine.clone(), metric_engine_config)
543            .context(BuildMetricEngineSnafu)?;
544
545        let file_engine = FileRegionEngine::new(
546            file_engine_config,
547            object_store_manager.default_object_store().clone(), // TODO: implement custom storage for file engine
548            self.local_file_access.clone(),
549        );
550
551        Ok(vec![
552            Arc::new(mito_engine) as _,
553            Arc::new(metric_engine) as _,
554            Arc::new(file_engine) as _,
555        ])
556    }
557
558    /// Builds [MitoEngine] according to options.
559    async fn build_mito_engine(
560        &mut self,
561        object_store_manager: ObjectStoreManagerRef,
562        mut config: MitoConfig,
563        schema_metadata_manager: SchemaMetadataManagerRef,
564        file_ref_manager: FileReferenceManagerRef,
565        partition_expr_fetcher: PartitionExprFetcherRef,
566        plugins: Plugins,
567    ) -> Result<MitoEngine> {
568        let opts = &self.opts;
569        if opts.storage.is_object_storage() {
570            // Enable the write cache when setting object storage
571            config.enable_write_cache = true;
572            info!("Configured 'enable_write_cache=true' for mito engine.");
573        }
574
575        let mito_engine = match &opts.wal {
576            DatanodeWalConfig::RaftEngine(raft_engine_config) => {
577                let log_store =
578                    Self::build_raft_engine_log_store(&opts.storage.data_home, raft_engine_config)
579                        .await?;
580
581                let builder = MitoEngineBuilder::new(
582                    &opts.storage.data_home,
583                    config,
584                    log_store,
585                    object_store_manager,
586                    schema_metadata_manager,
587                    file_ref_manager,
588                    partition_expr_fetcher.clone(),
589                    plugins,
590                );
591
592                #[cfg(feature = "enterprise")]
593                let builder = builder.with_extension_range_provider_factory(
594                    self.extension_range_provider_factory.take(),
595                );
596
597                builder.try_build().await.context(BuildMitoEngineSnafu)?
598            }
599            DatanodeWalConfig::Kafka(kafka_config) => {
600                if kafka_config.create_index && opts.node_id.is_none() {
601                    warn!("The WAL index creation only available in distributed mode.")
602                }
603                let global_index_collector = if kafka_config.create_index
604                    && let Some(node_id) = opts.node_id
605                {
606                    let operator = new_object_store_without_cache(
607                        &opts.storage.store,
608                        &opts.storage.data_home,
609                    )
610                    .await?;
611                    let path = default_index_file(node_id);
612                    Some(Self::build_global_index_collector(
613                        kafka_config.dump_index_interval,
614                        operator,
615                        path,
616                    ))
617                } else {
618                    None
619                };
620
621                let log_store =
622                    Self::build_kafka_log_store(kafka_config, global_index_collector).await?;
623                self.topic_stats_reporter = Some(log_store.topic_stats_reporter());
624                let builder = MitoEngineBuilder::new(
625                    &opts.storage.data_home,
626                    config,
627                    log_store,
628                    object_store_manager,
629                    schema_metadata_manager,
630                    file_ref_manager,
631                    partition_expr_fetcher,
632                    plugins,
633                );
634
635                #[cfg(feature = "enterprise")]
636                let builder = builder.with_extension_range_provider_factory(
637                    self.extension_range_provider_factory.take(),
638                );
639
640                builder.try_build().await.context(BuildMitoEngineSnafu)?
641            }
642            DatanodeWalConfig::Noop => {
643                let log_store = Arc::new(NoopLogStore);
644
645                let builder = MitoEngineBuilder::new(
646                    &opts.storage.data_home,
647                    config,
648                    log_store,
649                    object_store_manager,
650                    schema_metadata_manager,
651                    file_ref_manager,
652                    partition_expr_fetcher.clone(),
653                    plugins,
654                );
655
656                #[cfg(feature = "enterprise")]
657                let builder = builder.with_extension_range_provider_factory(
658                    self.extension_range_provider_factory.take(),
659                );
660
661                builder.try_build().await.context(BuildMitoEngineSnafu)?
662            }
663            DatanodeWalConfig::ObjectStore(_) => return ObjectStoreWalNotSupportedSnafu.fail(),
664        };
665        Ok(mito_engine)
666    }
667
668    /// Builds [RaftEngineLogStore].
669    async fn build_raft_engine_log_store(
670        data_home: &str,
671        config: &RaftEngineConfig,
672    ) -> Result<Arc<RaftEngineLogStore>> {
673        let data_home = normalize_dir(data_home);
674        let wal_dir = match &config.dir {
675            Some(dir) => dir.clone(),
676            None => format!("{}{WAL_DIR}", data_home),
677        };
678
679        // create WAL directory
680        fs::create_dir_all(Path::new(&wal_dir))
681            .await
682            .context(CreateDirSnafu { dir: &wal_dir })?;
683        info!(
684            "Creating raft-engine logstore with config: {:?} and storage path: {}",
685            config, &wal_dir
686        );
687        let logstore = RaftEngineLogStore::try_new(wal_dir, config)
688            .await
689            .map_err(Box::new)
690            .context(OpenLogStoreSnafu)?;
691
692        Ok(Arc::new(logstore))
693    }
694
695    /// Builds [`KafkaLogStore`].
696    async fn build_kafka_log_store(
697        config: &DatanodeKafkaConfig,
698        global_index_collector: Option<GlobalIndexCollector>,
699    ) -> Result<Arc<KafkaLogStore>> {
700        KafkaLogStore::try_new(config, global_index_collector)
701            .await
702            .map_err(Box::new)
703            .context(OpenLogStoreSnafu)
704            .map(Arc::new)
705    }
706
707    /// Builds [`GlobalIndexCollector`]
708    fn build_global_index_collector(
709        dump_index_interval: Duration,
710        operator: object_store::ObjectStore,
711        path: String,
712    ) -> GlobalIndexCollector {
713        GlobalIndexCollector::new(dump_index_interval, operator, path)
714    }
715}
716
717/// Rejects repeated engine types before their configs can silently overwrite each other.
718fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> {
719    let mut engine_indices = HashMap::new();
720    for (index, config) in configs.iter().enumerate() {
721        let engine = match config {
722            RegionEngineConfig::Mito(_) => MITO_ENGINE,
723            RegionEngineConfig::File(_) => FILE_ENGINE,
724            RegionEngineConfig::Metric(_) => METRIC_ENGINE,
725        };
726        if let Some(first_index) = engine_indices.insert(engine, index) {
727            return DuplicateRegionEngineConfigSnafu {
728                engine,
729                first_index,
730                duplicate_index: index,
731            }
732            .fail();
733        }
734    }
735    Ok(())
736}
737
738/// Rejects an object store WAL config the log store cannot run on.
739fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> {
740    let prefix = config.prefix.trim();
741    ensure!(
742        !prefix.is_empty(),
743        InvalidObjectStoreWalConfigSnafu {
744            field: "prefix",
745            value: &config.prefix,
746            reason: "must not be empty",
747        }
748    );
749    ensure!(
750        !prefix.starts_with('/'),
751        InvalidObjectStoreWalConfigSnafu {
752            field: "prefix",
753            value: &config.prefix,
754            reason: "must be a relative path",
755        }
756    );
757    ensure!(
758        !prefix
759            .split('/')
760            .any(|component| component.is_empty() || component == ".."),
761        InvalidObjectStoreWalConfigSnafu {
762            field: "prefix",
763            value: &config.prefix,
764            reason: "must not contain empty or `..` components",
765        }
766    );
767    ensure!(
768        config.flush_interval >= Duration::from_millis(10),
769        InvalidObjectStoreWalConfigSnafu {
770            field: "flush_interval",
771            value: format!("{:?}", config.flush_interval),
772            reason: "must be at least 10ms",
773        }
774    );
775    ensure!(
776        config.max_batch_bytes.as_bytes() > 0,
777        InvalidObjectStoreWalConfigSnafu {
778            field: "max_batch_bytes",
779            value: config.max_batch_bytes.to_string(),
780            reason: "must be greater than 0",
781        }
782    );
783    ensure!(
784        config.max_unpersisted_bytes.as_bytes() > 0,
785        InvalidObjectStoreWalConfigSnafu {
786            field: "max_unpersisted_bytes",
787            value: config.max_unpersisted_bytes.to_string(),
788            reason: "must be greater than 0",
789        }
790    );
791    ensure!(
792        config.max_unpersisted_age > Duration::ZERO,
793        InvalidObjectStoreWalConfigSnafu {
794            field: "max_unpersisted_age",
795            value: format!("{:?}", config.max_unpersisted_age),
796            reason: "must be greater than 0",
797        }
798    );
799    Ok(())
800}
801
802/// Open all regions belong to this datanode.
803async fn open_all_regions(
804    region_server: RegionServer,
805    region_open_requests: RegionOpenRequests,
806    open_with_writable: bool,
807    init_regions_parallelism: usize,
808    ignore_nonexistent_region: bool,
809) -> Result<()> {
810    let RegionOpenRequests {
811        leader_regions,
812        #[cfg(feature = "enterprise")]
813        follower_regions,
814    } = region_open_requests;
815
816    let leader_region_num = leader_regions.len();
817    info!("going to open {} region(s)", leader_region_num);
818    let now = Instant::now();
819    let open_regions = region_server
820        .handle_batch_open_requests(
821            init_regions_parallelism,
822            leader_regions,
823            ignore_nonexistent_region,
824        )
825        .await?;
826    info!(
827        "Opened {} regions in {:?}",
828        open_regions.len(),
829        now.elapsed()
830    );
831    if !ignore_nonexistent_region {
832        ensure!(
833            open_regions.len() == leader_region_num,
834            error::UnexpectedSnafu {
835                violated: format!(
836                    "Expected to open {} of regions, only {} of regions has opened",
837                    leader_region_num,
838                    open_regions.len()
839                )
840            }
841        );
842    } else if open_regions.len() != leader_region_num {
843        warn!(
844            "ignore nonexistent region, expected to open {} of regions, only {} of regions has opened",
845            leader_region_num,
846            open_regions.len()
847        );
848    }
849
850    for region_id in open_regions {
851        if open_with_writable {
852            let res = region_server.set_region_role(region_id, RegionRole::Leader);
853            match res {
854                Ok(_) => {
855                    // Finalize leadership: persist backfilled metadata.
856                    if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server
857                        .set_region_role_state_gracefully(
858                            region_id,
859                            SettableRegionRoleState::Leader,
860                        )
861                        .await?
862                    {
863                        error!(err; "failed to convert region {region_id} to leader");
864                    }
865                }
866                Err(e) => {
867                    error!(e; "failed to convert region {region_id} to leader");
868                }
869            }
870        }
871    }
872
873    #[cfg(feature = "enterprise")]
874    if !follower_regions.is_empty() {
875        use tokio::time::Instant;
876
877        let follower_region_num = follower_regions.len();
878        info!("going to open {} follower region(s)", follower_region_num);
879
880        let now = Instant::now();
881        let open_regions = region_server
882            .handle_batch_open_requests(
883                init_regions_parallelism,
884                follower_regions,
885                ignore_nonexistent_region,
886            )
887            .await?;
888        info!(
889            "Opened {} follower regions in {:?}",
890            open_regions.len(),
891            now.elapsed()
892        );
893
894        if !ignore_nonexistent_region {
895            ensure!(
896                open_regions.len() == follower_region_num,
897                error::UnexpectedSnafu {
898                    violated: format!(
899                        "Expected to open {} of follower regions, only {} of regions has opened",
900                        follower_region_num,
901                        open_regions.len()
902                    )
903                }
904            );
905        } else if open_regions.len() != follower_region_num {
906            warn!(
907                "ignore nonexistent region, expected to open {} of follower regions, only {} of regions has opened",
908                follower_region_num,
909                open_regions.len()
910            );
911        }
912    }
913
914    info!("all regions are opened");
915
916    Ok(())
917}
918
919#[cfg(test)]
920mod tests {
921    use std::assert_matches;
922    use std::collections::{BTreeMap, HashMap};
923    use std::io::Write;
924    use std::path::Path;
925    use std::sync::Arc;
926    use std::time::Duration;
927
928    use cache::build_datanode_cache_registry;
929    use common_base::Plugins;
930    use common_base::readable_size::ReadableSize;
931    use common_config::Configurable;
932    use common_error::ext::ErrorExt;
933    use common_error::status_code::StatusCode;
934    use common_meta::cache::LayeredCacheRegistryBuilder;
935    use common_meta::key::RegionRoleSet;
936    use common_meta::key::datanode_table::DatanodeTableManager;
937    use common_meta::kv_backend::KvBackendRef;
938    use common_meta::kv_backend::memory::MemoryKvBackend;
939    use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
940    use common_wal::config::DatanodeWalConfig;
941    use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
942    use mito2::engine::MITO_ENGINE_NAME;
943    use store_api::region_request::RegionRequest;
944    use store_api::storage::RegionId;
945
946    use crate::config::{DatanodeOptions, RegionEngineConfig};
947    use crate::datanode::{
948        DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config,
949    };
950    use crate::error::Error;
951    use crate::tests::{MockRegionEngine, mock_region_server};
952
953    async fn setup_table_datanode(kv: &KvBackendRef) {
954        let mgr = DatanodeTableManager::new(kv.clone());
955        let txn = mgr
956            .build_create_txn(
957                1028,
958                MITO_ENGINE_NAME,
959                "foo/bar/weny",
960                HashMap::from([("foo".to_string(), "bar".to_string())]),
961                HashMap::default(),
962                BTreeMap::from([(0, RegionRoleSet::new(vec![0, 1, 2], vec![]))]),
963            )
964            .unwrap();
965
966        let r = kv.txn(txn).await.unwrap();
967        assert!(r.succeeded);
968    }
969
970    #[tokio::test]
971    async fn test_initialize_region_server() {
972        common_telemetry::init_default_ut_logging();
973        let mut mock_region_server = mock_region_server();
974        let (mock_region, mut mock_region_handler) = MockRegionEngine::new(MITO_ENGINE_NAME);
975
976        mock_region_server.register_engine(mock_region.clone());
977
978        let kv_backend = Arc::new(MemoryKvBackend::new());
979        let layered_cache_registry = Arc::new(
980            LayeredCacheRegistryBuilder::default()
981                .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
982                .build(),
983        );
984
985        let mut builder = DatanodeBuilder::new(
986            DatanodeOptions {
987                node_id: Some(0),
988                ..Default::default()
989            },
990            Plugins::default(),
991            kv_backend.clone(),
992        );
993        builder.with_cache_registry(layered_cache_registry);
994        setup_table_datanode(&(kv_backend as _)).await;
995
996        builder
997            .initialize_region_server(&mock_region_server, false)
998            .await
999            .unwrap();
1000
1001        for i in 0..3 {
1002            let (region_id, req) = mock_region_handler.recv().await.unwrap();
1003            assert_eq!(region_id, RegionId::new(1028, i));
1004            if let RegionRequest::Open(req) = req {
1005                assert_eq!(
1006                    req.options,
1007                    HashMap::from([("foo".to_string(), "bar".to_string())])
1008                )
1009            } else {
1010                unreachable!()
1011            }
1012        }
1013
1014        assert_matches!(
1015            mock_region_handler.try_recv(),
1016            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1017        );
1018    }
1019
1020    fn datanode_builder(
1021        data_home: &str,
1022        wal: DatanodeWalConfig,
1023        kv_backend: KvBackendRef,
1024    ) -> DatanodeBuilder {
1025        let layered_cache_registry = Arc::new(
1026            LayeredCacheRegistryBuilder::default()
1027                .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
1028                .build(),
1029        );
1030        let mut opts = DatanodeOptions {
1031            node_id: Some(0),
1032            wal,
1033            enable_telemetry: false,
1034            ..Default::default()
1035        };
1036        opts.storage.data_home = data_home.to_string();
1037
1038        let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend);
1039        builder.with_cache_registry(layered_cache_registry);
1040        builder
1041    }
1042
1043    fn object_store_wal_builder(data_home: &str, config: ObjectStoreWalConfig) -> DatanodeBuilder {
1044        datanode_builder(
1045            data_home,
1046            DatanodeWalConfig::ObjectStore(config),
1047            Arc::new(MemoryKvBackend::new()),
1048        )
1049    }
1050
1051    async fn build_err(builder: DatanodeBuilder) -> Error {
1052        match builder.build().await {
1053            Ok(_) => panic!("build must fail"),
1054            Err(err) => err,
1055        }
1056    }
1057
1058    fn is_empty_dir(dir: &Path) -> bool {
1059        std::fs::read_dir(dir).unwrap().next().is_none()
1060    }
1061
1062    #[test]
1063    fn test_validate_unique_region_engine_config() {
1064        let mut opts = DatanodeOptions::default();
1065        validate_region_engine_config(&opts.region_engine).unwrap();
1066
1067        opts.region_engine
1068            .push(RegionEngineConfig::Metric(Default::default()));
1069        validate_region_engine_config(&opts.region_engine).unwrap();
1070
1071        // Omitted engines still use their defaults during engine construction.
1072        validate_region_engine_config(&[]).unwrap();
1073    }
1074
1075    #[tokio::test]
1076    async fn test_build_rejects_duplicate_region_engine_config() {
1077        for (engine, first_index) in [("mito", 0), ("file", 1), ("metric", 2)] {
1078            let mut file = create_named_temp_file();
1079            write!(
1080                file,
1081                r#"
1082                [[region_engine]]
1083                [region_engine.mito]
1084                global_write_buffer_size = "4GiB"
1085                min_compaction_interval = "600s"
1086
1087                [[region_engine]]
1088                [region_engine.file]
1089
1090                [[region_engine]]
1091                [region_engine.metric]
1092
1093                [[region_engine]]
1094                [region_engine.{engine}]
1095                "#
1096            )
1097            .unwrap();
1098            let opts = DatanodeOptions::load_layered_options(
1099                Some(file.path().to_str().unwrap()),
1100                "DATANODE_DUPLICATE_REGION_ENGINE_UT",
1101            )
1102            .unwrap();
1103            let data_home = create_temp_dir("duplicate-region-engine-config");
1104            let mut builder = datanode_builder(
1105                data_home.path().to_str().unwrap(),
1106                DatanodeWalConfig::default(),
1107                Arc::new(MemoryKvBackend::new()),
1108            );
1109            builder.opts.region_engine = opts.region_engine;
1110            let err = build_err(builder).await;
1111
1112            let Error::DuplicateRegionEngineConfig {
1113                engine: actual_engine,
1114                first_index: actual_first_index,
1115                duplicate_index,
1116                ..
1117            } = &err
1118            else {
1119                panic!("unexpected error for {engine}: {err:?}");
1120            };
1121            assert_eq!(*actual_engine, engine);
1122            assert_eq!(*actual_first_index, first_index);
1123            assert_eq!(*duplicate_index, 3);
1124            assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1125            // Rejected before the builder creates any storage or log store.
1126            assert!(is_empty_dir(data_home.path()));
1127        }
1128    }
1129
1130    #[tokio::test]
1131    async fn test_build_rejects_invalid_object_store_wal_config() {
1132        let cases = [
1133            (
1134                ObjectStoreWalConfig {
1135                    prefix: String::new(),
1136                    ..Default::default()
1137                },
1138                "prefix",
1139            ),
1140            (
1141                ObjectStoreWalConfig {
1142                    prefix: "/wal".to_string(),
1143                    ..Default::default()
1144                },
1145                "prefix",
1146            ),
1147            (
1148                ObjectStoreWalConfig {
1149                    prefix: "wal/../other".to_string(),
1150                    ..Default::default()
1151                },
1152                "prefix",
1153            ),
1154            (
1155                ObjectStoreWalConfig {
1156                    prefix: "wal//objects".to_string(),
1157                    ..Default::default()
1158                },
1159                "prefix",
1160            ),
1161            (
1162                ObjectStoreWalConfig {
1163                    prefix: "wal/".to_string(),
1164                    ..Default::default()
1165                },
1166                "prefix",
1167            ),
1168            (
1169                ObjectStoreWalConfig {
1170                    flush_interval: Duration::from_millis(9),
1171                    ..Default::default()
1172                },
1173                "flush_interval",
1174            ),
1175            (
1176                ObjectStoreWalConfig {
1177                    max_batch_bytes: ReadableSize(0),
1178                    ..Default::default()
1179                },
1180                "max_batch_bytes",
1181            ),
1182            (
1183                ObjectStoreWalConfig {
1184                    max_unpersisted_bytes: ReadableSize(0),
1185                    ..Default::default()
1186                },
1187                "max_unpersisted_bytes",
1188            ),
1189            (
1190                ObjectStoreWalConfig {
1191                    max_unpersisted_age: Duration::ZERO,
1192                    ..Default::default()
1193                },
1194                "max_unpersisted_age",
1195            ),
1196        ];
1197
1198        for (config, expected_field) in cases {
1199            let expected_value = match expected_field {
1200                "prefix" => config.prefix.clone(),
1201                "flush_interval" => format!("{:?}", config.flush_interval),
1202                "max_unpersisted_bytes" => config.max_unpersisted_bytes.to_string(),
1203                "max_unpersisted_age" => format!("{:?}", config.max_unpersisted_age),
1204                _ => config.max_batch_bytes.to_string(),
1205            };
1206            let data_home = create_temp_dir("object-store-wal-invalid-config");
1207            let builder = object_store_wal_builder(data_home.path().to_str().unwrap(), config);
1208            let err = build_err(builder).await;
1209
1210            let Error::InvalidObjectStoreWalConfig { field, value, .. } = &err else {
1211                panic!("unexpected error for {expected_field}: {err:?}");
1212            };
1213            assert_eq!(*field, expected_field);
1214            assert_eq!(*value, expected_value);
1215            assert_eq!(StatusCode::InvalidArguments, err.status_code());
1216            // The builder stops before it creates any storage or log store.
1217            assert!(is_empty_dir(data_home.path()));
1218        }
1219    }
1220
1221    #[tokio::test]
1222    async fn test_build_rejects_object_store_wal_as_not_supported() {
1223        let data_home = create_temp_dir("object-store-wal-not-supported");
1224        let builder = object_store_wal_builder(
1225            data_home.path().to_str().unwrap(),
1226            ObjectStoreWalConfig::default(),
1227        );
1228        let err = build_err(builder).await;
1229
1230        assert_matches!(err, Error::ObjectStoreWalNotSupported { .. });
1231        assert_eq!(StatusCode::Unsupported, err.status_code());
1232        // No log store is created under the prefix.
1233        assert!(!data_home.path().join("wal").exists());
1234    }
1235
1236    #[test]
1237    fn test_node_prefix_passes_validation() {
1238        let config = ObjectStoreWalConfig::default();
1239        let derived = ObjectStoreWalConfig {
1240            prefix: config.node_prefix(0, STANDALONE_GENERATION),
1241            ..config
1242        };
1243        validate_object_store_wal_config(&derived).unwrap();
1244    }
1245}