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