datanode/

datanode.rs

// Copyright 2023 Greptime Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Datanode implementation.

use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;
use std::time::{Duration, Instant};

use catalog::information_schema::NoopInformationExtension;
use catalog::kvbackend::KvBackendCatalogManagerBuilder;
use client::client_manager::NodeClients;
use common_base::Plugins;
use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE};
use common_datasource::object_store::LocalFileAccess;
use common_error::ext::BoxedError;
use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
use common_grpc::channel_manager::ChannelConfig;
use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef};
use common_meta::cache_invalidator::CacheInvalidatorRef;
use common_meta::datanode::TopicStatsReporter;
use common_meta::key::runtime_switch::RuntimeSwitchManager;
use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef};
use common_meta::kv_backend::KvBackendRef;
use common_meta::node_manager::NodeManagerRef;
pub use common_procedure::options::ProcedureConfig;
use common_query::prelude::set_default_prefix;
use common_stat::ResourceStatImpl;
use common_telemetry::{error, info, warn};
use common_wal::config::DatanodeWalConfig;
use common_wal::config::kafka::DatanodeKafkaConfig;
use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
use common_wal::config::raft_engine::RaftEngineConfig;
use file_engine::engine::FileRegionEngine;
use log_store::ObjectStoreLogStore;
use log_store::kafka::log_store::KafkaLogStore;
use log_store::noop::log_store::NoopLogStore;
use log_store::raft_engine::log_store::RaftEngineLogStore;
use meta_client::MetaClientRef;
use metric_engine::engine::MetricEngine;
use mito2::config::MitoConfig;
use mito2::engine::{MitoEngine, MitoEngineBuilder};
use mito2::region::opener::PartitionExprFetcherRef;
use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef};
use object_store::ObjectStore;
use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
use object_store::util::normalize_dir;
use query::QueryEngineFactory;
use query::dummy_catalog::TableProviderFactoryRef;
use query::region_query::RegionQueryHandlerFactoryRef;
use servers::server::ServerHandlers;
use snafu::{OptionExt, ResultExt, ensure};
use store_api::logstore::LogStore;
use store_api::path_utils::WAL_DIR;
use store_api::region_engine::{
    RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState,
};
use tokio::fs;
use tokio::sync::Notify;

use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
use crate::error::{
    self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu,
    DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu,
    InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu,
    ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu,
    ShutdownServerSnafu, StartServerSnafu,
};
use crate::event_listener::{
    NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver,
    new_region_server_event_channel,
};
use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
use crate::heartbeat::HeartbeatTask;
use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
use crate::region_query::DatanodeRegionQueryHandler;
use crate::region_server::{DummyTableProviderFactory, RegionServer};
use crate::store;
use crate::utils::{RegionOpenRequests, build_region_open_requests};

/// Datanode service.
pub struct Datanode {
    services: ServerHandlers,
    heartbeat_task: Option<HeartbeatTask>,
    region_event_receiver: Option<RegionServerEventReceiver>,
    region_server: RegionServer,
    greptimedb_telemetry_task: Arc<GreptimeDBTelemetryTask>,
    leases_notifier: Option<Arc<Notify>>,
    plugins: Plugins,
    object_store_log_store: Option<Arc<ObjectStoreLogStore>>,
}

impl Datanode {
    pub async fn start(&mut self) -> Result<()> {
        info!("Starting datanode instance...");

        self.start_heartbeat().await?;
        self.wait_coordinated().await;

        self.start_telemetry();

        self.services.start_all().await.context(StartServerSnafu)
    }

    pub fn server_handlers(&self) -> &ServerHandlers {
        &self.services
    }

    pub fn start_telemetry(&self) {
        if let Err(e) = self.greptimedb_telemetry_task.start() {
            warn!(e; "Failed to start telemetry task!");
        }
    }

    pub async fn start_heartbeat(&mut self) -> Result<()> {
        if let Some(task) = &self.heartbeat_task {
            // Safety: The event_receiver must exist.
            let receiver = self.region_event_receiver.take().unwrap();

            task.start(receiver, self.leases_notifier.clone()).await?;
        }
        Ok(())
    }

    /// If `leases_notifier` exists, it waits until leases have been obtained in all regions.
    pub async fn wait_coordinated(&mut self) {
        if let Some(notifier) = self.leases_notifier.take() {
            notifier.notified().await;
        }
    }

    pub fn setup_services(&mut self, services: ServerHandlers) {
        self.services = services;
    }

    /// Shuts down every component in order. A failed step does not skip the
    /// later ones; the first error is returned and the rest are logged.
    pub async fn shutdown(&mut self) -> Result<()> {
        let mut first_error = None;
        record_shutdown_error(
            &mut first_error,
            self.services
                .shutdown_all()
                .await
                .context(ShutdownServerSnafu),
        );

        let _ = self.greptimedb_telemetry_task.stop().await;
        if let Some(heartbeat_task) = &self.heartbeat_task {
            record_shutdown_error(
                &mut first_error,
                heartbeat_task
                    .close()
                    .map_err(BoxedError::new)
                    .context(ShutdownInstanceSnafu),
            );
        }
        // The object store WAL stops before the regions close: its stop
        // uploads the backlog or drops it, and answers every wait for its
        // entries, so a flush that waits for the WAL, queued or run by the
        // close of a region, cannot hold up the shutdown.
        if let Some(log_store) = &self.object_store_log_store {
            record_shutdown_error(
                &mut first_error,
                log_store
                    .stop()
                    .await
                    .map_err(BoxedError::new)
                    .context(ShutdownInstanceSnafu),
            );
        }
        record_shutdown_error(&mut first_error, self.region_server.stop().await);
        first_error.map_or(Ok(()), Err)
    }

    pub fn region_server(&self) -> RegionServer {
        self.region_server.clone()
    }

    /// Returns the object store log store when the datanode runs on the object
    /// store WAL.
    #[cfg(any(test, feature = "testing"))]
    pub fn object_store_log_store(&self) -> Option<Arc<ObjectStoreLogStore>> {
        self.object_store_log_store.clone()
    }

    pub fn plugins(&self) -> Plugins {
        self.plugins.clone()
    }
}

pub struct DatanodeBuilder {
    opts: DatanodeOptions,
    table_provider_factory: Option<TableProviderFactoryRef>,
    plugins: Plugins,
    meta_client: Option<MetaClientRef>,
    kv_backend: KvBackendRef,
    cache_registry: Option<Arc<LayeredCacheRegistry>>,
    topic_stats_reporter: Option<Box<dyn TopicStatsReporter>>,
    open_regions_writable_override: Option<bool>,
    local_file_access: LocalFileAccess,
    #[cfg(feature = "enterprise")]
    extension_range_provider_factory: Option<mito2::extension::BoxedExtensionRangeProviderFactory>,
    /// Set as soon as the object store log store exists so that a failed build can stop it.
    object_store_log_store: Option<Arc<ObjectStoreLogStore>>,
}

impl DatanodeBuilder {
    pub fn new(opts: DatanodeOptions, plugins: Plugins, kv_backend: KvBackendRef) -> Self {
        Self {
            opts,
            table_provider_factory: None,
            plugins,
            meta_client: None,
            kv_backend,
            cache_registry: None,
            open_regions_writable_override: None,
            local_file_access: LocalFileAccess::Disabled,
            #[cfg(feature = "enterprise")]
            extension_range_provider_factory: None,
            topic_stats_reporter: None,
            object_store_log_store: None,
        }
    }

    pub fn options(&self) -> &DatanodeOptions {
        &self.opts
    }

    pub fn with_meta_client(&mut self, client: MetaClientRef) -> &mut Self {
        self.meta_client = Some(client);
        self
    }

    pub fn with_cache_registry(&mut self, registry: Arc<LayeredCacheRegistry>) -> &mut Self {
        self.cache_registry = Some(registry);
        self
    }

    pub fn with_local_file_access(&mut self, local_file_access: LocalFileAccess) -> &mut Self {
        self.local_file_access = local_file_access;
        self
    }

    pub fn kv_backend(&self) -> &KvBackendRef {
        &self.kv_backend
    }

    pub fn meta_client(&self) -> Option<&MetaClientRef> {
        self.meta_client.as_ref()
    }

    pub fn cache_registry(&self) -> Option<&Arc<LayeredCacheRegistry>> {
        self.cache_registry.as_ref()
    }

    pub fn set_plugins(&mut self, plugins: Plugins) {
        self.plugins = plugins;
    }

    pub fn with_table_provider_factory(&mut self, factory: TableProviderFactoryRef) -> &mut Self {
        self.table_provider_factory = Some(factory);
        self
    }

    /// Overrides whether regions opened during datanode startup should become writable.
    ///
    /// When unset, the builder uses its default writable policy for reopened regions
    /// (writable only when no metasrv client is configured).
    ///
    /// Warning: setting this to `true` on a metasrv-controlled datanode (one built
    /// with `with_meta_client`) will promote regions to Leader before heartbeat and
    /// lease coordination begin, bypassing the metasrv safety contract and creating a
    /// potential split-brain window during startup.
    pub fn with_open_regions_writable_override(&mut self, writable: bool) -> &mut Self {
        self.open_regions_writable_override = Some(writable);
        self
    }

    #[cfg(feature = "enterprise")]
    pub fn with_extension_range_provider(
        &mut self,
        extension_range_provider_factory: mito2::extension::BoxedExtensionRangeProviderFactory,
    ) -> &mut Self {
        self.extension_range_provider_factory = Some(extension_range_provider_factory);
        self
    }

    pub async fn build(mut self) -> Result<Datanode> {
        self.try_build().await
    }

    /// Builds the datanode and stops the object store log store if a later step fails.
    async fn try_build(&mut self) -> Result<Datanode> {
        let result = self.build_inner().await;
        if result.is_err()
            && let Some(log_store) = &self.object_store_log_store
            && let Err(err) = log_store.stop().await
        {
            warn!(err; "Failed to stop the object store log store after a failed build");
        }
        result
    }

    async fn build_inner(&mut self) -> Result<Datanode> {
        validate_region_engine_config(&self.opts.region_engine)?;
        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;
        set_default_prefix(self.opts.default_column_prefix.as_deref())
            .map_err(BoxedError::new)
            .context(BuildDatanodeSnafu)?;

        let meta_client = self.meta_client.take();

        // If metasrv client is provided, we will use it to control the region server.
        // Otherwise the region server is self-controlled, meaning no heartbeat and immediately
        // writable upon open.
        let controlled_by_metasrv = meta_client.is_some();
        if let DatanodeWalConfig::ObjectStore(config) = &self.opts.wal {
            ensure!(!controlled_by_metasrv, ObjectStoreWalNotStandaloneSnafu);
            validate_object_store_wal_config(config)?;
        }

        // build and initialize region server
        let (region_event_listener, region_event_receiver) = if controlled_by_metasrv {
            let (tx, rx) = new_region_server_event_channel();
            (Box::new(tx) as _, Some(rx))
        } else {
            (Box::new(NoopRegionServerEventListener) as _, None)
        };

        let cache_registry = self.cache_registry.take().context(MissingCacheSnafu)?;
        let schema_cache: SchemaCacheRef = cache_registry.get().context(MissingCacheSnafu)?;
        let table_id_schema_cache: TableSchemaCacheRef =
            cache_registry.get().context(MissingCacheSnafu)?;

        let schema_metadata_manager = Arc::new(SchemaMetadataManager::new(
            table_id_schema_cache,
            schema_cache,
        ));

        let gc_enabled = self.opts.region_engine.iter().any(|engine| {
            if let RegionEngineConfig::Mito(config) = engine {
                config.gc.enable
            } else {
                false
            }
        });

        let file_ref_manager = Arc::new(FileReferenceManager::with_gc_enabled(
            Some(node_id),
            gc_enabled,
        ));
        let region_server = self
            .new_region_server(
                schema_metadata_manager,
                region_event_listener,
                file_ref_manager,
                cache_registry.clone(),
            )
            .await?;

        // TODO(weny): Considering introducing a readonly kv_backend trait.
        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
        let is_recovery_mode = runtime_switch_manager
            .recovery_mode()
            .await
            .context(GetMetadataSnafu)?;

        let region_open_requests =
            build_region_open_requests(node_id, self.kv_backend.clone()).await?;
        let open_with_writable = self
            .open_regions_writable_override
            .unwrap_or(!controlled_by_metasrv);
        let open_all_regions = open_all_regions(
            region_server.clone(),
            region_open_requests,
            open_with_writable,
            self.opts.init_regions_parallelism,
            // Ignore nonexistent regions in recovery mode.
            is_recovery_mode,
        );

        if self.opts.init_regions_in_background {
            // Opens regions in background.
            common_runtime::spawn_global(async move {
                if let Err(err) = open_all_regions.await {
                    error!(err; "Failed to open regions during the startup.");
                }
            });
        } else {
            open_all_regions.await?;
        }

        let heartbeat_task = if let Some(meta_client) = meta_client {
            let task = self
                .create_heartbeat_task(&region_server, meta_client, cache_registry)
                .await?;
            Some(task)
        } else {
            None
        };

        let is_standalone = heartbeat_task.is_none();
        let greptimedb_telemetry_task = get_greptimedb_telemetry_task(
            Some(self.opts.storage.data_home.clone()),
            is_standalone && self.opts.enable_telemetry,
        )
        .await;

        let leases_notifier = if self.opts.require_lease_before_startup && !is_standalone {
            Some(Arc::new(Notify::new()))
        } else {
            None
        };

        Ok(Datanode {
            services: ServerHandlers::default(),
            heartbeat_task,
            region_server,
            greptimedb_telemetry_task,
            region_event_receiver,
            leases_notifier,
            plugins: self.plugins.clone(),
            object_store_log_store: self.object_store_log_store.take(),
        })
    }

    async fn create_heartbeat_task(
        &self,
        region_server: &RegionServer,
        meta_client: MetaClientRef,
        cache_invalidator: CacheInvalidatorRef,
    ) -> Result<HeartbeatTask> {
        let stat = {
            let mut stat = ResourceStatImpl::default();
            stat.start_collect_cpu_usage();
            Arc::new(stat)
        };

        HeartbeatTask::try_new(
            &self.opts,
            region_server.clone(),
            meta_client,
            self.kv_backend.clone(),
            cache_invalidator,
            self.plugins.clone(),
            stat,
        )
        .await
    }

    /// Builds [ObjectStoreManager] from [StorageConfig].
    pub async fn build_object_store_manager(cfg: &StorageConfig) -> Result<ObjectStoreManagerRef> {
        let object_store = store::new_object_store(cfg.store.clone(), &cfg.data_home).await?;
        let default_name = cfg.store.config_name();
        let mut object_store_manager = ObjectStoreManager::new(default_name, object_store);
        for store in &cfg.providers {
            object_store_manager.add(
                store.config_name(),
                store::new_object_store(store.clone(), &cfg.data_home).await?,
            );
        }
        Ok(Arc::new(object_store_manager))
    }

    #[cfg(test)]
    /// Open all regions belong to this datanode.
    async fn initialize_region_server(
        &self,
        region_server: &RegionServer,
        open_with_writable: bool,
    ) -> Result<()> {
        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;

        // TODO(weny): Considering introducing a readonly kv_backend trait.
        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
        let is_recovery_mode = runtime_switch_manager
            .recovery_mode()
            .await
            .context(GetMetadataSnafu)?;
        let region_open_requests =
            build_region_open_requests(node_id, self.kv_backend.clone()).await?;

        open_all_regions(
            region_server.clone(),
            region_open_requests,
            open_with_writable,
            self.opts.init_regions_parallelism,
            is_recovery_mode,
        )
        .await
    }

    async fn new_region_server(
        &mut self,
        schema_metadata_manager: SchemaMetadataManagerRef,
        event_listener: RegionServerEventListenerRef,
        file_ref_manager: FileReferenceManagerRef,
        cache_registry: Arc<LayeredCacheRegistry>,
    ) -> Result<RegionServer> {
        let opts: &DatanodeOptions = &self.opts;

        // Keep `with_dist_planner` false: datanodes execute received MergeScan plans.
        let catalog_manager = KvBackendCatalogManagerBuilder::new(
            Arc::new(NoopInformationExtension),
            self.kv_backend.clone(),
            cache_registry.clone(),
        )
        .build();
        let partition_manager = catalog_manager.partition_manager();

        // MergeScan region requests stream until completion, so they have no client timeout.
        let mut channel_config = ChannelConfig {
            timeout: None,
            ..Default::default()
        };
        if opts.grpc.flight_compression.transport_compression() {
            channel_config.accept_compression = true;
            channel_config.send_compression = true;
        }
        let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config));
        let region_query_handler =
            if let Some(factory) = self.plugins.get::<RegionQueryHandlerFactoryRef>() {
                factory.build(partition_manager.clone(), node_manager)
            } else {
                DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager)
            };

        let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
            // query engine in datanode only executes plan with resolved table source.
            catalog_manager,
            Some(partition_manager),
            Some(region_query_handler),
            None,
            None,
            None,
            false,
            self.plugins.clone(),
            opts.query.clone(),
        )
        .context(DataFusionSnafu)?;
        let query_engine = query_engine_factory.query_engine();

        let table_provider_factory = self
            .table_provider_factory
            .clone()
            .unwrap_or_else(|| Arc::new(DummyTableProviderFactory));
        let mut region_server = RegionServer::with_table_provider(
            query_engine,
            common_runtime::global_runtime(),
            event_listener,
            table_provider_factory,
            opts.max_concurrent_queries,
            opts.concurrent_query_limiter_timeout,
            opts.grpc.flight_compression,
        );
        region_server.install_remote_dyn_filter_receiver_injector(&self.plugins);

        let object_store_manager = Self::build_object_store_manager(&opts.storage).await?;
        let engines = self
            .build_store_engines(
                object_store_manager,
                schema_metadata_manager,
                file_ref_manager,
                self.plugins.clone(),
            )
            .await?;
        for engine in engines {
            region_server.register_engine(engine);
        }
        if let Some(topic_stats_reporter) = self.topic_stats_reporter.take() {
            region_server.set_topic_stats_reporter(topic_stats_reporter);
        }

        Ok(region_server)
    }

    // internal utils

    /// Builds [RegionEngineRef] from `store_engine` section in `opts`
    async fn build_store_engines(
        &mut self,
        object_store_manager: ObjectStoreManagerRef,
        schema_metadata_manager: SchemaMetadataManagerRef,
        file_ref_manager: FileReferenceManagerRef,
        plugins: Plugins,
    ) -> Result<Vec<RegionEngineRef>> {
        let mut metric_engine_config = metric_engine::config::EngineConfig::default();
        let mut mito_engine_config = MitoConfig::default();
        let mut file_engine_config = file_engine::config::EngineConfig::default();

        for engine in &self.opts.region_engine {
            match engine {
                RegionEngineConfig::Mito(config) => {
                    mito_engine_config = config.clone();
                }
                RegionEngineConfig::File(config) => {
                    file_engine_config = config.clone();
                }
                RegionEngineConfig::Metric(metric_config) => {
                    metric_engine_config = metric_config.clone();
                }
            }
        }

        // Build a fetcher to backfill partition_expr on open.
        let fetcher = Arc::new(MetaPartitionExprFetcher::new(self.kv_backend.clone()));
        let mito_engine = self
            .build_mito_engine(
                object_store_manager.clone(),
                mito_engine_config,
                schema_metadata_manager.clone(),
                file_ref_manager.clone(),
                fetcher.clone(),
                plugins.clone(),
            )
            .await?;

        let metric_engine = MetricEngine::try_new(mito_engine.clone(), metric_engine_config)
            .context(BuildMetricEngineSnafu)?;

        let file_engine = FileRegionEngine::new(
            file_engine_config,
            object_store_manager.default_object_store().clone(), // TODO: implement custom storage for file engine
            self.local_file_access.clone(),
        );

        Ok(vec![
            Arc::new(mito_engine) as _,
            Arc::new(metric_engine) as _,
            Arc::new(file_engine) as _,
        ])
    }

    /// Builds [MitoEngine] according to options.
    async fn build_mito_engine(
        &mut self,
        object_store_manager: ObjectStoreManagerRef,
        mut config: MitoConfig,
        schema_metadata_manager: SchemaMetadataManagerRef,
        file_ref_manager: FileReferenceManagerRef,
        partition_expr_fetcher: PartitionExprFetcherRef,
        plugins: Plugins,
    ) -> Result<MitoEngine> {
        let opts = &self.opts;
        if opts.storage.is_object_storage() {
            // Enable the write cache when setting object storage
            config.enable_write_cache = true;
            info!("Configured 'enable_write_cache=true' for mito engine.");
        }

        let mito_engine = match &opts.wal {
            DatanodeWalConfig::RaftEngine(raft_engine_config) => {
                let log_store =
                    Self::build_raft_engine_log_store(&opts.storage.data_home, raft_engine_config)
                        .await?;

                let builder = MitoEngineBuilder::new(
                    &opts.storage.data_home,
                    config,
                    log_store,
                    object_store_manager,
                    schema_metadata_manager,
                    file_ref_manager,
                    partition_expr_fetcher.clone(),
                    plugins,
                );

                #[cfg(feature = "enterprise")]
                let builder = builder.with_extension_range_provider_factory(
                    self.extension_range_provider_factory.take(),
                );

                builder.try_build().await.context(BuildMitoEngineSnafu)?
            }
            DatanodeWalConfig::Kafka(kafka_config) => {
                let log_store = Self::build_kafka_log_store(kafka_config).await?;
                self.topic_stats_reporter = Some(log_store.topic_stats_reporter());
                let builder = MitoEngineBuilder::new(
                    &opts.storage.data_home,
                    config,
                    log_store,
                    object_store_manager,
                    schema_metadata_manager,
                    file_ref_manager,
                    partition_expr_fetcher,
                    plugins,
                );

                #[cfg(feature = "enterprise")]
                let builder = builder.with_extension_range_provider_factory(
                    self.extension_range_provider_factory.take(),
                );

                builder.try_build().await.context(BuildMitoEngineSnafu)?
            }
            DatanodeWalConfig::Noop => {
                let log_store = Arc::new(NoopLogStore);

                let builder = MitoEngineBuilder::new(
                    &opts.storage.data_home,
                    config,
                    log_store,
                    object_store_manager,
                    schema_metadata_manager,
                    file_ref_manager,
                    partition_expr_fetcher.clone(),
                    plugins,
                );

                #[cfg(feature = "enterprise")]
                let builder = builder.with_extension_range_provider_factory(
                    self.extension_range_provider_factory.take(),
                );

                builder.try_build().await.context(BuildMitoEngineSnafu)?
            }
            DatanodeWalConfig::ObjectStore(object_store_config) => {
                let object_store = wal_object_store(
                    &object_store_manager,
                    &opts.storage,
                    &object_store_config.storage_provider,
                )?;
                let node_id = opts.node_id.context(MissingNodeIdSnafu)?;
                let log_store =
                    Self::build_object_store_log_store(object_store, object_store_config, node_id)
                        .await?;
                self.object_store_log_store = Some(log_store.clone());

                let builder = MitoEngineBuilder::new(
                    &opts.storage.data_home,
                    config,
                    log_store,
                    object_store_manager,
                    schema_metadata_manager,
                    file_ref_manager,
                    partition_expr_fetcher,
                    plugins,
                );

                #[cfg(feature = "enterprise")]
                let builder = builder.with_extension_range_provider_factory(
                    self.extension_range_provider_factory.take(),
                );

                builder.try_build().await.context(BuildMitoEngineSnafu)?
            }
        };
        Ok(mito_engine)
    }

    /// Builds [ObjectStoreLogStore] under the node prefix of `node_id` in the
    /// standalone generation, the prefix the regions persist in their WAL options.
    async fn build_object_store_log_store(
        object_store: ObjectStore,
        config: &ObjectStoreWalConfig,
        node_id: u64,
    ) -> Result<Arc<ObjectStoreLogStore>> {
        info!(
            "Creating object store logstore with config: {:?}, node id: {}",
            config, node_id
        );
        ObjectStoreLogStore::try_new(object_store, config, node_id, STANDALONE_GENERATION)
            .await
            .map_err(Box::new)
            .context(OpenLogStoreSnafu)
    }

    /// Builds [RaftEngineLogStore].
    async fn build_raft_engine_log_store(
        data_home: &str,
        config: &RaftEngineConfig,
    ) -> Result<Arc<RaftEngineLogStore>> {
        let data_home = normalize_dir(data_home);
        let wal_dir = match &config.dir {
            Some(dir) => dir.clone(),
            None => format!("{}{WAL_DIR}", data_home),
        };

        // create WAL directory
        fs::create_dir_all(Path::new(&wal_dir))
            .await
            .context(CreateDirSnafu { dir: &wal_dir })?;
        info!(
            "Creating raft-engine logstore with config: {:?} and storage path: {}",
            config, &wal_dir
        );
        let logstore = RaftEngineLogStore::try_new(wal_dir, config)
            .await
            .map_err(Box::new)
            .context(OpenLogStoreSnafu)?;

        Ok(Arc::new(logstore))
    }

    /// Builds [`KafkaLogStore`].
    async fn build_kafka_log_store(config: &DatanodeKafkaConfig) -> Result<Arc<KafkaLogStore>> {
        KafkaLogStore::try_new(config)
            .await
            .map_err(Box::new)
            .context(OpenLogStoreSnafu)
            .map(Arc::new)
    }
}

/// Keeps the first shutdown error and logs the later ones.
pub(crate) fn record_shutdown_error(first_error: &mut Option<error::Error>, result: Result<()>) {
    if let Err(err) = result {
        if first_error.is_none() {
            *first_error = Some(err);
        } else {
            warn!(err; "Ignored a later shutdown error");
        }
    }
}

/// Returns the object store that holds the WAL objects: the default one for an
/// empty `storage_provider`, otherwise the configured provider with that name.
fn wal_object_store(
    object_store_manager: &ObjectStoreManager,
    storage: &StorageConfig,
    storage_provider: &str,
) -> Result<ObjectStore> {
    if storage_provider.is_empty() {
        return Ok(object_store_manager.default_object_store().clone());
    }
    object_store_manager
        .find(storage_provider)
        .cloned()
        .with_context(|| {
            let configured = std::iter::once(&storage.store)
                .chain(storage.providers.iter())
                .map(|store| store.config_name())
                .collect::<Vec<_>>();
            InvalidObjectStoreWalConfigSnafu {
                field: "storage_provider",
                value: storage_provider,
                reason: format!("must name a configured storage provider, one of {configured:?}"),
            }
        })
}

/// Rejects repeated engine types before their configs can silently overwrite each other.
fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> {
    let mut engine_indices = HashMap::new();
    for (index, config) in configs.iter().enumerate() {
        let engine = match config {
            RegionEngineConfig::Mito(_) => MITO_ENGINE,
            RegionEngineConfig::File(_) => FILE_ENGINE,
            RegionEngineConfig::Metric(_) => METRIC_ENGINE,
        };
        if let Some(first_index) = engine_indices.insert(engine, index) {
            return DuplicateRegionEngineConfigSnafu {
                engine,
                first_index,
                duplicate_index: index,
            }
            .fail();
        }
    }
    Ok(())
}

/// Rejects an object store WAL config the log store cannot run on.
fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> {
    let prefix = config.prefix.trim();
    ensure!(
        !prefix.is_empty(),
        InvalidObjectStoreWalConfigSnafu {
            field: "prefix",
            value: &config.prefix,
            reason: "must not be empty",
        }
    );
    ensure!(
        !prefix.starts_with('/'),
        InvalidObjectStoreWalConfigSnafu {
            field: "prefix",
            value: &config.prefix,
            reason: "must be a relative path",
        }
    );
    ensure!(
        !prefix
            .split('/')
            .any(|component| component.is_empty() || component == ".."),
        InvalidObjectStoreWalConfigSnafu {
            field: "prefix",
            value: &config.prefix,
            reason: "must not contain empty or `..` components",
        }
    );
    ensure!(
        config.flush_interval >= Duration::from_millis(10),
        InvalidObjectStoreWalConfigSnafu {
            field: "flush_interval",
            value: format!("{:?}", config.flush_interval),
            reason: "must be at least 10ms",
        }
    );
    ensure!(
        config.max_batch_bytes.as_bytes() > 0,
        InvalidObjectStoreWalConfigSnafu {
            field: "max_batch_bytes",
            value: config.max_batch_bytes.to_string(),
            reason: "must be greater than 0",
        }
    );
    ensure!(
        config.max_unpersisted_bytes.as_bytes() > 0,
        InvalidObjectStoreWalConfigSnafu {
            field: "max_unpersisted_bytes",
            value: config.max_unpersisted_bytes.to_string(),
            reason: "must be greater than 0",
        }
    );
    ensure!(
        config.max_unpersisted_age > Duration::ZERO,
        InvalidObjectStoreWalConfigSnafu {
            field: "max_unpersisted_age",
            value: format!("{:?}", config.max_unpersisted_age),
            reason: "must be greater than 0",
        }
    );
    Ok(())
}

/// Open all regions belong to this datanode.
async fn open_all_regions(
    region_server: RegionServer,
    region_open_requests: RegionOpenRequests,
    open_with_writable: bool,
    init_regions_parallelism: usize,
    ignore_nonexistent_region: bool,
) -> Result<()> {
    let RegionOpenRequests {
        leader_regions,
        #[cfg(feature = "enterprise")]
        follower_regions,
    } = region_open_requests;

    let leader_region_num = leader_regions.len();
    info!("going to open {} region(s)", leader_region_num);
    let now = Instant::now();
    let open_regions = region_server
        .handle_batch_open_requests(
            init_regions_parallelism,
            leader_regions,
            ignore_nonexistent_region,
        )
        .await?;
    info!(
        "Opened {} regions in {:?}",
        open_regions.len(),
        now.elapsed()
    );
    if !ignore_nonexistent_region {
        ensure!(
            open_regions.len() == leader_region_num,
            error::UnexpectedSnafu {
                violated: format!(
                    "Expected to open {} of regions, only {} of regions has opened",
                    leader_region_num,
                    open_regions.len()
                )
            }
        );
    } else if open_regions.len() != leader_region_num {
        warn!(
            "ignore nonexistent region, expected to open {} of regions, only {} of regions has opened",
            leader_region_num,
            open_regions.len()
        );
    }

    for region_id in open_regions {
        if open_with_writable {
            let res = region_server.set_region_role(region_id, RegionRole::Leader);
            match res {
                Ok(_) => {
                    // Finalize leadership: persist backfilled metadata.
                    if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server
                        .set_region_role_state_gracefully(
                            region_id,
                            SettableRegionRoleState::Leader,
                        )
                        .await?
                    {
                        error!(err; "failed to convert region {region_id} to leader");
                    }
                }
                Err(e) => {
                    error!(e; "failed to convert region {region_id} to leader");
                }
            }
        }
    }

    #[cfg(feature = "enterprise")]
    if !follower_regions.is_empty() {
        use tokio::time::Instant;

        let follower_region_num = follower_regions.len();
        info!("going to open {} follower region(s)", follower_region_num);

        let now = Instant::now();
        let open_regions = region_server
            .handle_batch_open_requests(
                init_regions_parallelism,
                follower_regions,
                ignore_nonexistent_region,
            )
            .await?;
        info!(
            "Opened {} follower regions in {:?}",
            open_regions.len(),
            now.elapsed()
        );

        if !ignore_nonexistent_region {
            ensure!(
                open_regions.len() == follower_region_num,
                error::UnexpectedSnafu {
                    violated: format!(
                        "Expected to open {} of follower regions, only {} of regions has opened",
                        follower_region_num,
                        open_regions.len()
                    )
                }
            );
        } else if open_regions.len() != follower_region_num {
            warn!(
                "ignore nonexistent region, expected to open {} of follower regions, only {} of regions has opened",
                follower_region_num,
                open_regions.len()
            );
        }
    }

    info!("all regions are opened");

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::assert_matches;
    use std::collections::{BTreeMap, HashMap};
    use std::io::Write;
    use std::ops::Range;
    use std::path::Path;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::Duration;

    use api::v1::{ColumnSchema, Rows};
    use cache::build_datanode_layered_cache_registry;
    use common_base::Plugins;
    use common_base::readable_size::ReadableSize;
    use common_config::Configurable;
    use common_error::ext::ErrorExt;
    use common_error::status_code::StatusCode;
    use common_meta::key::RegionRoleSet;
    use common_meta::key::datanode_table::DatanodeTableManager;
    use common_meta::kv_backend::KvBackendRef;
    use common_meta::kv_backend::memory::MemoryKvBackend;
    use common_meta::node_manager::NodeManagerRef;
    use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
    use common_wal::config::DatanodeWalConfig;
    use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig, STANDALONE_GENERATION};
    use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions};
    use log_store::ObjectStoreLogStore;
    use log_store::error::Error as LogStoreError;
    use meta_client::client::MetaClientBuilder;
    use mito2::engine::MITO_ENGINE_NAME;
    use mito2::test_util::{CreateRequestBuilder, build_rows, rows_schema};
    use object_store::ObjectStore;
    use object_store::manager::ObjectStoreManager;
    use partition::manager::PartitionRuleManagerRef;
    use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef};
    use store_api::logstore::LogStore;
    use store_api::logstore::provider::Provider;
    use store_api::region_request::{PathType, RegionOpenRequest, RegionPutRequest, RegionRequest};
    use store_api::storage::RegionId;

    use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
    use crate::datanode::{
        Datanode, DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config,
        wal_object_store,
    };
    use crate::error::{self, Error};
    use crate::region_query::DatanodeRegionQueryHandler;
    use crate::region_server::RegionServer;
    use crate::tests::{MockRegionEngine, mock_region_server};

    const WAIT: Duration = Duration::from_secs(30);

    async fn setup_table_datanode(kv: &KvBackendRef) {
        let mgr = DatanodeTableManager::new(kv.clone());
        let txn = mgr
            .build_create_txn(
                1028,
                MITO_ENGINE_NAME,
                "foo/bar/weny",
                HashMap::from([("foo".to_string(), "bar".to_string())]),
                HashMap::default(),
                BTreeMap::from([(0, RegionRoleSet::new(vec![0, 1, 2], vec![]))]),
            )
            .unwrap();

        let r = kv.txn(txn).await.unwrap();
        assert!(r.succeeded);
    }

    #[tokio::test]
    async fn test_initialize_region_server() {
        common_telemetry::init_default_ut_logging();
        let mut mock_region_server = mock_region_server();
        let (mock_region, mut mock_region_handler) = MockRegionEngine::new(MITO_ENGINE_NAME);

        mock_region_server.register_engine(mock_region.clone());

        let kv_backend = Arc::new(MemoryKvBackend::new());
        let layered_cache_registry =
            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));

        let mut builder = DatanodeBuilder::new(
            DatanodeOptions {
                node_id: Some(0),
                ..Default::default()
            },
            Plugins::default(),
            kv_backend.clone(),
        );
        builder.with_cache_registry(layered_cache_registry);
        setup_table_datanode(&(kv_backend as _)).await;

        builder
            .initialize_region_server(&mock_region_server, false)
            .await
            .unwrap();

        for i in 0..3 {
            let (region_id, req) = mock_region_handler.recv().await.unwrap();
            assert_eq!(region_id, RegionId::new(1028, i));
            if let RegionRequest::Open(req) = req {
                assert_eq!(
                    req.options,
                    HashMap::from([("foo".to_string(), "bar".to_string())])
                )
            } else {
                unreachable!()
            }
        }

        assert_matches!(
            mock_region_handler.try_recv(),
            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
        );
    }

    fn datanode_builder(
        data_home: &str,
        wal: DatanodeWalConfig,
        kv_backend: KvBackendRef,
    ) -> DatanodeBuilder {
        let layered_cache_registry =
            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
        let mut opts = DatanodeOptions {
            node_id: Some(0),
            wal,
            enable_telemetry: false,
            ..Default::default()
        };
        opts.storage.data_home = data_home.to_string();

        let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend);
        builder.with_cache_registry(layered_cache_registry);
        builder
    }

    fn object_store_wal_builder(data_home: &str, config: ObjectStoreWalConfig) -> DatanodeBuilder {
        datanode_builder(
            data_home,
            DatanodeWalConfig::ObjectStore(config),
            Arc::new(MemoryKvBackend::new()),
        )
    }

    async fn build_err(builder: DatanodeBuilder) -> Error {
        match builder.build().await {
            Ok(_) => panic!("build must fail"),
            Err(err) => err,
        }
    }

    fn is_empty_dir(dir: &Path) -> bool {
        std::fs::read_dir(dir).unwrap().next().is_none()
    }

    /// Returns true if the store rejects even an empty append, which it does only once stopped.
    async fn is_stopped(log_store: &ObjectStoreLogStore) -> bool {
        matches!(
            log_store.append_batch(vec![]).await,
            Err(LogStoreError::ObjectStoreWalStopped { .. })
        )
    }

    #[test]
    fn test_validate_unique_region_engine_config() {
        let mut opts = DatanodeOptions::default();
        validate_region_engine_config(&opts.region_engine).unwrap();

        opts.region_engine
            .push(RegionEngineConfig::Metric(Default::default()));
        validate_region_engine_config(&opts.region_engine).unwrap();

        // Omitted engines still use their defaults during engine construction.
        validate_region_engine_config(&[]).unwrap();
    }

    #[tokio::test]
    async fn test_build_rejects_duplicate_region_engine_config() {
        for (engine, first_index) in [("mito", 0), ("file", 1), ("metric", 2)] {
            let mut file = create_named_temp_file();
            write!(
                file,
                r#"
                [[region_engine]]
                [region_engine.mito]
                global_write_buffer_size = "4GiB"
                min_compaction_interval = "600s"

                [[region_engine]]
                [region_engine.file]

                [[region_engine]]
                [region_engine.metric]

                [[region_engine]]
                [region_engine.{engine}]
                "#
            )
            .unwrap();
            let opts = DatanodeOptions::load_layered_options(
                Some(file.path().to_str().unwrap()),
                "DATANODE_DUPLICATE_REGION_ENGINE_UT",
            )
            .unwrap();
            let data_home = create_temp_dir("duplicate-region-engine-config");
            let mut builder = datanode_builder(
                data_home.path().to_str().unwrap(),
                DatanodeWalConfig::default(),
                Arc::new(MemoryKvBackend::new()),
            );
            builder.opts.region_engine = opts.region_engine;
            let err = build_err(builder).await;

            let Error::DuplicateRegionEngineConfig {
                engine: actual_engine,
                first_index: actual_first_index,
                duplicate_index,
                ..
            } = &err
            else {
                panic!("unexpected error for {engine}: {err:?}");
            };
            assert_eq!(*actual_engine, engine);
            assert_eq!(*actual_first_index, first_index);
            assert_eq!(*duplicate_index, 3);
            assert_eq!(err.status_code(), StatusCode::InvalidArguments);
            // Rejected before the builder creates any storage or log store.
            assert!(is_empty_dir(data_home.path()));
        }
    }

    #[tokio::test]
    async fn test_build_rejects_invalid_object_store_wal_config() {
        let cases = [
            (
                ObjectStoreWalConfig {
                    prefix: String::new(),
                    ..Default::default()
                },
                "prefix",
            ),
            (
                ObjectStoreWalConfig {
                    prefix: "/wal".to_string(),
                    ..Default::default()
                },
                "prefix",
            ),
            (
                ObjectStoreWalConfig {
                    prefix: "wal/../other".to_string(),
                    ..Default::default()
                },
                "prefix",
            ),
            (
                ObjectStoreWalConfig {
                    prefix: "wal//objects".to_string(),
                    ..Default::default()
                },
                "prefix",
            ),
            (
                ObjectStoreWalConfig {
                    prefix: "wal/".to_string(),
                    ..Default::default()
                },
                "prefix",
            ),
            (
                ObjectStoreWalConfig {
                    flush_interval: Duration::from_millis(9),
                    ..Default::default()
                },
                "flush_interval",
            ),
            (
                ObjectStoreWalConfig {
                    max_batch_bytes: ReadableSize(0),
                    ..Default::default()
                },
                "max_batch_bytes",
            ),
            (
                ObjectStoreWalConfig {
                    max_unpersisted_bytes: ReadableSize(0),
                    ..Default::default()
                },
                "max_unpersisted_bytes",
            ),
            (
                ObjectStoreWalConfig {
                    max_unpersisted_age: Duration::ZERO,
                    ..Default::default()
                },
                "max_unpersisted_age",
            ),
        ];

        for (config, expected_field) in cases {
            let expected_value = match expected_field {
                "prefix" => config.prefix.clone(),
                "flush_interval" => format!("{:?}", config.flush_interval),
                "max_unpersisted_bytes" => config.max_unpersisted_bytes.to_string(),
                "max_unpersisted_age" => format!("{:?}", config.max_unpersisted_age),
                _ => config.max_batch_bytes.to_string(),
            };
            let data_home = create_temp_dir("object-store-wal-invalid-config");
            let builder = object_store_wal_builder(data_home.path().to_str().unwrap(), config);
            let err = build_err(builder).await;

            let Error::InvalidObjectStoreWalConfig { field, value, .. } = &err else {
                panic!("unexpected error for {expected_field}: {err:?}");
            };
            assert_eq!(*field, expected_field);
            assert_eq!(*value, expected_value);
            assert_eq!(StatusCode::InvalidArguments, err.status_code());
            // The builder stops before it creates any storage or log store.
            assert!(is_empty_dir(data_home.path()));
        }
    }

    #[tokio::test]
    async fn test_build_rejects_object_store_wal_with_meta_client() {
        let data_home = create_temp_dir("object-store-wal-meta-client");
        let meta_client = Arc::new(MetaClientBuilder::datanode_default_options(0).build());

        let mut builder = object_store_wal_builder(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig::default(),
        );
        builder.with_meta_client(meta_client);
        let err = build_err(builder).await;

        assert_matches!(err, Error::ObjectStoreWalNotStandalone { .. });
        assert_eq!(StatusCode::InvalidArguments, err.status_code());
        // The builder stops before it creates any storage or log store.
        assert!(is_empty_dir(data_home.path()));
    }

    #[tokio::test]
    async fn test_build_rejects_unknown_wal_storage_provider() {
        let data_home = create_temp_dir("object-store-wal-unknown-provider");
        let builder = object_store_wal_builder(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig {
                storage_provider: "s3".to_string(),
                ..Default::default()
            },
        );
        let err = build_err(builder).await;

        let Error::InvalidObjectStoreWalConfig {
            field,
            value,
            reason,
            ..
        } = &err
        else {
            panic!("unexpected error: {err:?}");
        };
        assert_eq!(*field, "storage_provider");
        assert_eq!(value, "s3");
        assert!(reason.contains(r#"["File"]"#), "{reason}");
        assert_eq!(StatusCode::InvalidArguments, err.status_code());
        // No log store is created under the prefix.
        assert!(!data_home.path().join("wal").exists());
    }

    #[test]
    fn test_wal_object_store_selects_provider() {
        fn fs_object_store(dir: &Path) -> ObjectStore {
            let builder = object_store::services::Fs::default().root(dir.to_str().unwrap());
            ObjectStore::new(builder).unwrap()
        }

        let default_dir = create_temp_dir("wal-object-store-default");
        let named_dir = create_temp_dir("wal-object-store-named");
        let mut manager = ObjectStoreManager::new("File", fs_object_store(default_dir.path()));
        manager.add("Wal-Store", fs_object_store(named_dir.path()));
        let storage = StorageConfig::default();

        let root = |store: ObjectStore| store.info().root().clone();
        assert_eq!(
            root(wal_object_store(&manager, &storage, "").unwrap()),
            root(manager.default_object_store().clone())
        );
        assert_eq!(
            root(wal_object_store(&manager, &storage, "wal-store").unwrap()),
            root(manager.find("Wal-Store").unwrap().clone())
        );
        assert_matches!(
            wal_object_store(&manager, &storage, "missing").unwrap_err(),
            Error::InvalidObjectStoreWalConfig {
                field: "storage_provider",
                ..
            }
        );
    }

    struct CountingRegionQueryHandlerFactory(Arc<AtomicUsize>);

    impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory {
        fn build(
            &self,
            partition_manager: PartitionRuleManagerRef,
            node_manager: NodeManagerRef,
        ) -> query::region_query::RegionQueryHandlerRef {
            self.0.fetch_add(1, Ordering::SeqCst);
            DatanodeRegionQueryHandler::arc(partition_manager, node_manager)
        }
    }

    #[tokio::test]
    async fn test_build_uses_installed_region_query_handler_factory() {
        let data_home = create_temp_dir("region-query-handler-factory");
        let builds = Arc::new(AtomicUsize::new(0));
        let plugins = Plugins::new();
        plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone()))
            as RegionQueryHandlerFactoryRef);
        let mut builder = datanode_builder(
            data_home.path().to_str().unwrap(),
            DatanodeWalConfig::default(),
            Arc::new(MemoryKvBackend::new()),
        );
        builder.set_plugins(plugins);

        let mut datanode = builder.build().await.unwrap();
        assert_eq!(builds.load(Ordering::SeqCst), 1);
        datanode.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_build_and_shutdown_with_raft_engine_wal() {
        let data_home = create_temp_dir("raft-engine-wal");
        let builder = datanode_builder(
            data_home.path().to_str().unwrap(),
            DatanodeWalConfig::default(),
            Arc::new(MemoryKvBackend::new()),
        );

        let mut datanode = builder.build().await.unwrap();
        assert!(datanode.object_store_log_store.is_none());
        assert!(data_home.path().join("wal").exists());
        datanode.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_build_and_shutdown_with_object_store_wal() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal");
        let builder = object_store_wal_builder(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig::default(),
        );

        let mut datanode = builder.build().await.unwrap();
        let log_store = datanode.object_store_log_store.clone().unwrap();
        assert!(!is_stopped(&log_store).await);
        // The store runs under the node prefix derived from the configured root.
        assert!(format!("{log_store:?}").contains(r#"prefix: "wal/datanodes/0/epochs/0""#));

        datanode.shutdown().await.unwrap();
        assert!(is_stopped(&log_store).await);
    }

    #[tokio::test]
    async fn test_build_with_named_wal_storage_provider() {
        let data_home = create_temp_dir("object-store-wal-named-provider");
        // The default file store is registered under its provider name.
        let builder = object_store_wal_builder(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig {
                storage_provider: "file".to_string(),
                ..Default::default()
            },
        );

        let mut datanode = builder.build().await.unwrap();
        assert!(datanode.object_store_log_store.is_some());
        datanode.shutdown().await.unwrap();
    }

    #[tokio::test]
    async fn test_shutdown_stops_every_engine_and_the_log_store_after_engine_failures() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal-shutdown-failure");
        let builder = object_store_wal_builder(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig::default(),
        );
        let mut datanode = builder.build().await.unwrap();
        let log_store = datanode.object_store_log_store.clone().unwrap();
        // Both engines fail to stop, so whichever the region server stops first fails.
        let stop_calls = Arc::new(AtomicUsize::new(0));
        for name in ["failing-a", "failing-b"] {
            let stop_calls = stop_calls.clone();
            let (engine, _) = MockRegionEngine::with_custom_apply_fn(name, move |engine| {
                engine.handle_stop_mock_fn = Some(Box::new(move || {
                    stop_calls.fetch_add(1, Ordering::Relaxed);
                    error::UnexpectedSnafu {
                        violated: "stop failed",
                    }
                    .fail()
                }));
            });
            datanode.region_server().register_engine(engine);
        }

        let err = datanode.shutdown().await.unwrap_err();

        assert_matches!(err, Error::StopRegionEngine { .. });
        assert_eq!(2, stop_calls.load(Ordering::Relaxed));
        assert!(is_stopped(&log_store).await);
    }

    #[tokio::test]
    async fn test_failed_build_stops_object_store_log_store() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal-failed-build");
        let kv_backend = Arc::new(MemoryKvBackend::new());
        // Regions registered in the metadata but missing on disk fail the region open
        // that runs after the log store and the engines are built.
        setup_table_datanode(&(kv_backend.clone() as _)).await;
        let mut builder = datanode_builder(
            data_home.path().to_str().unwrap(),
            DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()),
            kv_backend,
        );

        assert!(builder.try_build().await.is_err());

        let log_store = builder.object_store_log_store.as_ref().unwrap();
        assert!(is_stopped(log_store).await);
    }

    const REGION: RegionId = RegionId::new(1024, 1);

    fn enqueued_wal_config(flush_interval: Duration) -> ObjectStoreWalConfig {
        ObjectStoreWalConfig {
            ack_mode: AckMode::Enqueued,
            flush_interval,
            ..Default::default()
        }
    }

    /// The options that put a region on the WAL of a datanode with `config`.
    fn region_wal_options(config: &ObjectStoreWalConfig) -> HashMap<String, String> {
        let prefix = config.node_prefix(0, STANDALONE_GENERATION);
        let wal_options =
            serde_json::to_string(&WalOptions::ObjectStore(ObjectStoreWalOptions::new(prefix)))
                .unwrap();
        HashMap::from([(WAL_OPTIONS_KEY.to_string(), wal_options)])
    }

    async fn put_rows(
        region_server: &RegionServer,
        schema: &[ColumnSchema],
        rows: Range<usize>,
        skip_wal: bool,
    ) {
        if rows.is_empty() {
            return;
        }
        region_server
            .handle_request(
                REGION,
                RegionRequest::Put(RegionPutRequest {
                    rows: Rows {
                        schema: schema.to_vec(),
                        rows: build_rows(rows.start, rows.end),
                    },
                    skip_wal,
                    hint: None,
                    partition_expr_version: None,
                }),
            )
            .await
            .unwrap();
    }

    /// Builds a datanode on an `enqueued` object store WAL with
    /// `flush_interval` and creates [`REGION`] on it with `durable_rows`
    /// written through the WAL and uploaded. Then makes every later create
    /// of the WAL fail with a transient error if `uploads_fail`, and writes
    /// `rows` through the WAL and `skipped_rows` past it. Returns the datanode
    /// and its log store.
    async fn datanode_with_object_store_wal(
        data_home: &str,
        flush_interval: Duration,
        durable_rows: usize,
        uploads_fail: bool,
        rows: usize,
        skipped_rows: usize,
    ) -> (Datanode, Arc<ObjectStoreLogStore>) {
        let config = enqueued_wal_config(flush_interval);
        let prefix = config.node_prefix(0, STANDALONE_GENERATION);
        let mut request = CreateRequestBuilder::new().build();
        request.options.extend(region_wal_options(&config));
        let datanode = object_store_wal_builder(data_home, config)
            .build()
            .await
            .unwrap();
        let log_store = datanode.object_store_log_store.clone().unwrap();
        let region_server = datanode.region_server();
        let schema = rows_schema(&request);
        region_server
            .handle_request(REGION, RegionRequest::Create(request))
            .await
            .unwrap();

        put_rows(&region_server, &schema, 0..durable_rows, false).await;
        if durable_rows > 0 {
            let provider = Provider::object_store_provider(REGION, prefix);
            tokio::time::timeout(WAIT, async {
                while log_store.latest_entry_id(&provider).unwrap() == 0 {
                    tokio::time::sleep(Duration::from_millis(10)).await;
                }
            })
            .await
            .expect("the rows must become durable");
        }
        if uploads_fail {
            log_store.fail_creates();
        }
        let skipped = durable_rows + rows;
        put_rows(&region_server, &schema, durable_rows..skipped, false).await;
        put_rows(
            &region_server,
            &schema,
            skipped..skipped + skipped_rows,
            true,
        )
        .await;
        (datanode, log_store)
    }

    /// Restarts a datanode on the files of `data_home`, opens [`REGION`] and
    /// returns the number of rows it holds after replaying the WAL.
    async fn rows_after_restart(data_home: &str) -> u64 {
        let config = enqueued_wal_config(ObjectStoreWalConfig::default().flush_interval);
        let options = region_wal_options(&config);
        let mut datanode = object_store_wal_builder(data_home, config)
            .build()
            .await
            .unwrap();
        let region_server = datanode.region_server();
        region_server
            .handle_request(
                REGION,
                RegionRequest::Open(RegionOpenRequest {
                    engine: MITO_ENGINE_NAME.to_string(),
                    table_dir: CreateRequestBuilder::new().build().table_dir,
                    path_type: PathType::Bare,
                    options,
                    skip_wal_replay: false,
                    checkpoint: None,
                    requirements: Default::default(),
                }),
            )
            .await
            .unwrap();
        let rows = region_server.region_statistic(REGION).unwrap().num_rows;
        datanode.shutdown().await.unwrap();
        rows
    }

    /// Returns the Parquet files under `dir` and its subdirectories.
    fn sst_files(dir: &Path) -> Vec<String> {
        let mut files = Vec::new();
        for entry in std::fs::read_dir(dir).unwrap() {
            let path = entry.unwrap().path();
            if path.is_dir() {
                files.extend(sst_files(&path));
            } else if path.extension().is_some_and(|ext| ext == "parquet") {
                files.push(path.display().to_string());
            }
        }
        files
    }

    /// A flush that waits for WAL uploads that keep failing does not hold up
    /// the shutdown: the WAL stops first and drops its backlog, which fails
    /// the flush with the stopped error before it writes its manifest edit,
    /// and the shutdown reports the dropped backlog. After a restart the rows
    /// that were durable before the outage replay once.
    #[tokio::test]
    async fn test_shutdown_while_a_flush_waits_for_failing_wal_uploads() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal-shutdown-flush");
        let data_home = data_home.path().to_str().unwrap();
        let (mut datanode, log_store) = datanode_with_object_store_wal(
            data_home,
            ObjectStoreWalConfig::default().flush_interval,
            3,
            true,
            2,
            0,
        )
        .await;
        let flush = {
            let region_server = datanode.region_server();
            tokio::spawn(async move {
                region_server
                    .handle_request(REGION, RegionRequest::Flush(Default::default()))
                    .await
            })
        };
        tokio::time::timeout(WAIT, log_store.wait_for_durability_waits(1))
            .await
            .unwrap()
            .unwrap();

        let err = tokio::time::timeout(WAIT, datanode.shutdown())
            .await
            .expect("the shutdown must finish")
            .unwrap_err();

        assert_matches!(err, Error::ShutdownInstance { .. });
        assert!(
            format!("{err:?}").contains("injected create failure"),
            "{err:?}"
        );
        let flush_err = flush.await.unwrap().unwrap_err();
        assert!(
            format!("{flush_err:?}").contains("Object store WAL log store is stopped"),
            "{flush_err:?}"
        );
        assert_eq!(Vec::<String>::new(), sst_files(Path::new(data_home)));
        assert!(is_stopped(&log_store).await);
        drop(datanode);

        assert_eq!(3, rows_after_restart(data_home).await);
    }

    /// A region that holds writes that skipped the WAL is flushed when it
    /// closes. When the WAL's uploads keep failing, the WAL stops first and
    /// drops its backlog, so that flush fails at once instead of waiting: the
    /// close fails, the writes that skipped the WAL are lost, and the shutdown
    /// returns the dropped backlog.
    #[tokio::test]
    async fn test_shutdown_closes_a_region_after_the_wal_dropped_its_backlog() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal-shutdown-closing-flush");
        let (mut datanode, log_store) = datanode_with_object_store_wal(
            data_home.path().to_str().unwrap(),
            ObjectStoreWalConfig::default().flush_interval,
            0,
            true,
            3,
            2,
        )
        .await;

        let err = tokio::time::timeout(WAIT, datanode.shutdown())
            .await
            .expect("the shutdown must finish")
            .unwrap_err();

        assert_matches!(err, Error::ShutdownInstance { .. });
        assert!(
            format!("{err:?}").contains("injected create failure"),
            "{err:?}"
        );
        assert_eq!(Vec::<String>::new(), sst_files(data_home.path()));
        assert!(is_stopped(&log_store).await);
    }

    /// With a healthy WAL whose seal interval is longer than any test, the
    /// close-time flush of a region that holds writes that skipped the WAL
    /// commits: the WAL's stop seals and uploads its open batch before the
    /// region closes, so the flush finds its entries durable, and every row
    /// survives a restart.
    #[tokio::test]
    async fn test_shutdown_flushes_a_region_while_the_wal_is_healthy() {
        common_telemetry::init_default_ut_logging();
        let data_home = create_temp_dir("object-store-wal-shutdown-healthy-flush");
        let data_home = data_home.path().to_str().unwrap();
        let (mut datanode, log_store) =
            datanode_with_object_store_wal(data_home, Duration::from_secs(3600), 0, false, 3, 2)
                .await;

        tokio::time::timeout(WAIT, datanode.shutdown())
            .await
            .expect("the shutdown must finish")
            .unwrap();
        assert_eq!(1, sst_files(Path::new(data_home)).len());
        assert!(is_stopped(&log_store).await);
        drop(datanode);

        assert_eq!(5, rows_after_restart(data_home).await);
    }

    #[test]
    fn test_node_prefix_passes_validation() {
        let config = ObjectStoreWalConfig::default();
        let derived = ObjectStoreWalConfig {
            prefix: config.node_prefix(0, STANDALONE_GENERATION),
            ..config
        };
        validate_object_store_wal_config(&derived).unwrap();
    }
}