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 catalog::information_schema::NoopInformationExtension;
23use catalog::kvbackend::KvBackendCatalogManagerBuilder;
24use client::client_manager::NodeClients;
25use common_base::Plugins;
26use common_catalog::consts::{FILE_ENGINE, METRIC_ENGINE, MITO_ENGINE};
27use common_datasource::object_store::LocalFileAccess;
28use common_error::ext::BoxedError;
29use common_greptimedb_telemetry::GreptimeDBTelemetryTask;
30use common_grpc::channel_manager::ChannelConfig;
31use common_meta::cache::{LayeredCacheRegistry, SchemaCacheRef, TableSchemaCacheRef};
32use common_meta::cache_invalidator::CacheInvalidatorRef;
33use common_meta::datanode::TopicStatsReporter;
34use common_meta::key::runtime_switch::RuntimeSwitchManager;
35use common_meta::key::{SchemaMetadataManager, SchemaMetadataManagerRef};
36use common_meta::kv_backend::KvBackendRef;
37use common_meta::node_manager::NodeManagerRef;
38pub use common_procedure::options::ProcedureConfig;
39use common_query::prelude::set_default_prefix;
40use common_stat::ResourceStatImpl;
41use common_telemetry::{error, info, warn};
42use common_wal::config::DatanodeWalConfig;
43use common_wal::config::kafka::DatanodeKafkaConfig;
44use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
45use common_wal::config::raft_engine::RaftEngineConfig;
46use file_engine::engine::FileRegionEngine;
47use futures::future::join_all;
48use log_store::ObjectStoreLogStore;
49use log_store::kafka::log_store::KafkaLogStore;
50use log_store::noop::log_store::NoopLogStore;
51use log_store::raft_engine::log_store::RaftEngineLogStore;
52use meta_client::MetaClientRef;
53use metric_engine::engine::MetricEngine;
54use mito2::config::MitoConfig;
55use mito2::engine::{MitoEngine, MitoEngineBuilder};
56use mito2::region::opener::PartitionExprFetcherRef;
57use mito2::sst::file_ref::{FileReferenceManager, FileReferenceManagerRef};
58use object_store::ObjectStore;
59use object_store::manager::{ObjectStoreManager, ObjectStoreManagerRef};
60use object_store::util::normalize_dir;
61use query::QueryEngineFactory;
62use query::dummy_catalog::TableProviderFactoryRef;
63use query::region_query::RegionQueryHandlerFactoryRef;
64use relational_engine::engine::RelationalEngine;
65use servers::server::ServerHandlers;
66use snafu::{OptionExt, ResultExt, ensure};
67use store_api::logstore::LogStore;
68use store_api::path_utils::WAL_DIR;
69use store_api::region_engine::{
70    RegionEngineRef, RegionRole, SetRegionRoleStateResponse, SettableRegionRoleState,
71};
72use tokio::fs;
73use tokio::sync::Notify;
74
75use crate::config::{DatanodeOptions, RegionEngineConfig, StorageConfig};
76use crate::error::{
77    self, BuildDatanodeSnafu, BuildMetricEngineSnafu, BuildMitoEngineSnafu, CreateDirSnafu,
78    DataFusionSnafu, DuplicateRegionEngineConfigSnafu, GetMetadataSnafu,
79    InvalidObjectStoreWalConfigSnafu, MissingCacheSnafu, MissingNodeIdSnafu,
80    ObjectStoreWalNotStandaloneSnafu, OpenLogStoreSnafu, Result, ShutdownInstanceSnafu,
81    ShutdownServerSnafu, StartServerSnafu,
82};
83use crate::event_listener::{
84    NoopRegionServerEventListener, RegionServerEventListenerRef, RegionServerEventReceiver,
85    new_region_server_event_channel,
86};
87use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
88use crate::heartbeat::HeartbeatTask;
89use crate::partition_expr_fetcher::MetaPartitionExprFetcher;
90use crate::region_query::DatanodeRegionQueryHandler;
91use crate::region_server::{DummyTableProviderFactory, RegionServer};
92use crate::store;
93use crate::utils::{RegionOpenRequests, build_region_open_requests};
94
95/// Datanode service.
96pub struct Datanode {
97    services: ServerHandlers,
98    heartbeat_task: Option<HeartbeatTask>,
99    region_event_receiver: Option<RegionServerEventReceiver>,
100    region_server: RegionServer,
101    greptimedb_telemetry_task: Arc<GreptimeDBTelemetryTask>,
102    leases_notifier: Option<Arc<Notify>>,
103    plugins: Plugins,
104    object_store_log_store: Option<Arc<ObjectStoreLogStore>>,
105    /// The raft-engine WAL, if any. Stopped after the region server, see [Datanode::shutdown].
106    raft_engine_log_store: Option<Arc<RaftEngineLogStore>>,
107}
108
109impl Datanode {
110    pub async fn start(&mut self) -> Result<()> {
111        info!("Starting datanode instance...");
112
113        self.start_heartbeat().await?;
114        self.wait_coordinated().await;
115
116        self.start_telemetry();
117
118        self.services.start_all().await.context(StartServerSnafu)
119    }
120
121    pub fn server_handlers(&self) -> &ServerHandlers {
122        &self.services
123    }
124
125    pub fn start_telemetry(&self) {
126        if let Err(e) = self.greptimedb_telemetry_task.start() {
127            warn!(e; "Failed to start telemetry task!");
128        }
129    }
130
131    pub async fn start_heartbeat(&mut self) -> Result<()> {
132        if let Some(task) = &self.heartbeat_task {
133            // Safety: The event_receiver must exist.
134            let receiver = self.region_event_receiver.take().unwrap();
135
136            task.start(receiver, self.leases_notifier.clone()).await?;
137        }
138        Ok(())
139    }
140
141    /// If `leases_notifier` exists, it waits until leases have been obtained in all regions.
142    pub async fn wait_coordinated(&mut self) {
143        if let Some(notifier) = self.leases_notifier.take() {
144            notifier.notified().await;
145        }
146    }
147
148    pub fn setup_services(&mut self, services: ServerHandlers) {
149        self.services = services;
150    }
151
152    /// Shuts down every component in order. A failed step does not skip the
153    /// later ones; the first error is returned and the rest are logged.
154    pub async fn shutdown(&mut self) -> Result<()> {
155        let mut first_error = None;
156        record_shutdown_error(
157            &mut first_error,
158            self.services
159                .shutdown_all()
160                .await
161                .context(ShutdownServerSnafu),
162        );
163
164        let _ = self.greptimedb_telemetry_task.stop().await;
165        if let Some(heartbeat_task) = &self.heartbeat_task {
166            record_shutdown_error(
167                &mut first_error,
168                heartbeat_task
169                    .close()
170                    .map_err(BoxedError::new)
171                    .context(ShutdownInstanceSnafu),
172            );
173        }
174        // The object store WAL stops before the regions close: its stop
175        // uploads the backlog or drops it, and answers every wait for its
176        // entries, so a flush that waits for the WAL, queued or run by the
177        // close of a region, cannot hold up the shutdown.
178        if let Some(log_store) = &self.object_store_log_store {
179            record_shutdown_error(
180                &mut first_error,
181                log_store
182                    .stop()
183                    .await
184                    .map_err(BoxedError::new)
185                    .context(ShutdownInstanceSnafu),
186            );
187        }
188        record_shutdown_error(&mut first_error, self.region_server.stop().await);
189        // The stopped engines no longer hold the raft-engine WAL, so it can be closed.
190        if let Some(log_store) = self.raft_engine_log_store.take() {
191            record_shutdown_error(
192                &mut first_error,
193                close_raft_engine_log_store(log_store).await,
194            );
195        }
196        first_error.map_or(Ok(()), Err)
197    }
198
199    pub fn region_server(&self) -> RegionServer {
200        self.region_server.clone()
201    }
202
203    /// Returns the object store log store when the datanode runs on the object
204    /// store WAL.
205    #[cfg(any(test, feature = "testing"))]
206    pub fn object_store_log_store(&self) -> Option<Arc<ObjectStoreLogStore>> {
207        self.object_store_log_store.clone()
208    }
209
210    pub fn plugins(&self) -> Plugins {
211        self.plugins.clone()
212    }
213}
214
215pub struct DatanodeBuilder {
216    opts: DatanodeOptions,
217    table_provider_factory: Option<TableProviderFactoryRef>,
218    plugins: Plugins,
219    meta_client: Option<MetaClientRef>,
220    kv_backend: KvBackendRef,
221    cache_registry: Option<Arc<LayeredCacheRegistry>>,
222    topic_stats_reporter: Option<Box<dyn TopicStatsReporter>>,
223    open_regions_writable_override: Option<bool>,
224    local_file_access: LocalFileAccess,
225    /// Set as soon as the raft-engine log store exists so that a failed build can close it.
226    raft_engine_log_store: Option<Arc<RaftEngineLogStore>>,
227    #[cfg(feature = "enterprise")]
228    extension_range_provider_factory: Option<mito2::extension::BoxedExtensionRangeProviderFactory>,
229    /// Set as soon as the object store log store exists so that a failed build can stop it.
230    object_store_log_store: Option<Arc<ObjectStoreLogStore>>,
231    /// Set as soon as the region server exists so that a failed build can stop its engines,
232    /// which hold the log stores.
233    region_server: Option<RegionServer>,
234}
235
236impl DatanodeBuilder {
237    pub fn new(opts: DatanodeOptions, plugins: Plugins, kv_backend: KvBackendRef) -> Self {
238        Self {
239            opts,
240            table_provider_factory: None,
241            plugins,
242            meta_client: None,
243            kv_backend,
244            cache_registry: None,
245            open_regions_writable_override: None,
246            local_file_access: LocalFileAccess::Disabled,
247            raft_engine_log_store: None,
248            region_server: None,
249            #[cfg(feature = "enterprise")]
250            extension_range_provider_factory: None,
251            topic_stats_reporter: None,
252            object_store_log_store: None,
253        }
254    }
255
256    pub fn options(&self) -> &DatanodeOptions {
257        &self.opts
258    }
259
260    pub fn with_meta_client(&mut self, client: MetaClientRef) -> &mut Self {
261        self.meta_client = Some(client);
262        self
263    }
264
265    pub fn with_cache_registry(&mut self, registry: Arc<LayeredCacheRegistry>) -> &mut Self {
266        self.cache_registry = Some(registry);
267        self
268    }
269
270    pub fn with_local_file_access(&mut self, local_file_access: LocalFileAccess) -> &mut Self {
271        self.local_file_access = local_file_access;
272        self
273    }
274
275    pub fn kv_backend(&self) -> &KvBackendRef {
276        &self.kv_backend
277    }
278
279    pub fn meta_client(&self) -> Option<&MetaClientRef> {
280        self.meta_client.as_ref()
281    }
282
283    pub fn cache_registry(&self) -> Option<&Arc<LayeredCacheRegistry>> {
284        self.cache_registry.as_ref()
285    }
286
287    pub fn set_plugins(&mut self, plugins: Plugins) {
288        self.plugins = plugins;
289    }
290
291    pub fn with_table_provider_factory(&mut self, factory: TableProviderFactoryRef) -> &mut Self {
292        self.table_provider_factory = Some(factory);
293        self
294    }
295
296    /// Overrides whether regions opened during datanode startup should become writable.
297    ///
298    /// When unset, the builder uses its default writable policy for reopened regions
299    /// (writable only when no metasrv client is configured).
300    ///
301    /// Warning: setting this to `true` on a metasrv-controlled datanode (one built
302    /// with `with_meta_client`) will promote regions to Leader before heartbeat and
303    /// lease coordination begin, bypassing the metasrv safety contract and creating a
304    /// potential split-brain window during startup.
305    pub fn with_open_regions_writable_override(&mut self, writable: bool) -> &mut Self {
306        self.open_regions_writable_override = Some(writable);
307        self
308    }
309
310    #[cfg(feature = "enterprise")]
311    pub fn with_extension_range_provider(
312        &mut self,
313        extension_range_provider_factory: mito2::extension::BoxedExtensionRangeProviderFactory,
314    ) -> &mut Self {
315        self.extension_range_provider_factory = Some(extension_range_provider_factory);
316        self
317    }
318
319    pub async fn build(mut self) -> Result<Datanode> {
320        self.try_build().await
321    }
322
323    /// Builds the datanode and, if a later step fails, stops the engines and the log stores it
324    /// opened.
325    async fn try_build(&mut self) -> Result<Datanode> {
326        let result = self.build_inner().await;
327        if result.is_err() {
328            if let Some(log_store) = &self.object_store_log_store
329                && let Err(err) = log_store.stop().await
330            {
331                warn!(err; "Failed to stop the object store log store after a failed build");
332            }
333            if let Some(region_server) = self.region_server.take()
334                && let Err(err) = region_server.stop().await
335            {
336                warn!(err; "Failed to stop the region server after a failed build");
337            }
338            if let Some(log_store) = self.raft_engine_log_store.take()
339                && let Err(err) = close_raft_engine_log_store(log_store).await
340            {
341                warn!(err; "Failed to close the raft-engine log store after a failed build");
342            }
343        }
344        result
345    }
346
347    async fn build_inner(&mut self) -> Result<Datanode> {
348        validate_region_engine_config(&self.opts.region_engine)?;
349        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;
350        set_default_prefix(self.opts.default_column_prefix.as_deref())
351            .map_err(BoxedError::new)
352            .context(BuildDatanodeSnafu)?;
353
354        let meta_client = self.meta_client.take();
355
356        // If metasrv client is provided, we will use it to control the region server.
357        // Otherwise the region server is self-controlled, meaning no heartbeat and immediately
358        // writable upon open.
359        let controlled_by_metasrv = meta_client.is_some();
360        if let DatanodeWalConfig::ObjectStore(config) = &self.opts.wal {
361            ensure!(!controlled_by_metasrv, ObjectStoreWalNotStandaloneSnafu);
362            validate_object_store_wal_config(config)?;
363        }
364
365        // build and initialize region server
366        let (region_event_listener, region_event_receiver) = if controlled_by_metasrv {
367            let (tx, rx) = new_region_server_event_channel();
368            (Box::new(tx) as _, Some(rx))
369        } else {
370            (Box::new(NoopRegionServerEventListener) as _, None)
371        };
372
373        let cache_registry = self.cache_registry.take().context(MissingCacheSnafu)?;
374        let schema_cache: SchemaCacheRef = cache_registry.get().context(MissingCacheSnafu)?;
375        let table_id_schema_cache: TableSchemaCacheRef =
376            cache_registry.get().context(MissingCacheSnafu)?;
377
378        let schema_metadata_manager = Arc::new(SchemaMetadataManager::new(
379            table_id_schema_cache,
380            schema_cache,
381        ));
382
383        let gc_enabled = self.opts.region_engine.iter().any(|engine| {
384            if let RegionEngineConfig::Mito(config) = engine {
385                config.gc.enable
386            } else {
387                false
388            }
389        });
390
391        let file_ref_manager = Arc::new(FileReferenceManager::with_gc_enabled(
392            Some(node_id),
393            gc_enabled,
394        ));
395        let region_server = self
396            .new_region_server(
397                schema_metadata_manager,
398                region_event_listener,
399                file_ref_manager,
400                cache_registry.clone(),
401            )
402            .await?;
403        self.region_server = Some(region_server.clone());
404        // A region acknowledges only writes whose WAL append completed within its lease from
405        // the metasrv.
406        if controlled_by_metasrv && let Some(mito) = region_server.mito_engine() {
407            mito.require_write_lease();
408        }
409
410        // TODO(weny): Considering introducing a readonly kv_backend trait.
411        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
412        let is_recovery_mode = runtime_switch_manager
413            .recovery_mode()
414            .await
415            .context(GetMetadataSnafu)?;
416
417        let region_open_requests =
418            build_region_open_requests(node_id, self.kv_backend.clone()).await?;
419        let open_with_writable = self
420            .open_regions_writable_override
421            .unwrap_or(!controlled_by_metasrv);
422        let initialization_guard = self.opts.init_regions_in_background.then(|| {
423            let regions = region_open_requests
424                .leader_regions
425                .iter()
426                .map(|(id, _)| *id);
427            #[cfg(feature = "enterprise")]
428            let regions = regions.chain(
429                region_open_requests
430                    .follower_regions
431                    .iter()
432                    .map(|(id, _)| *id),
433            );
434            region_server.track_region_initialization(regions)
435        });
436        let open_all_regions = open_all_regions(
437            region_server.clone(),
438            region_open_requests,
439            open_with_writable,
440            self.opts.init_regions_parallelism,
441            // Ignore nonexistent regions in recovery mode.
442            is_recovery_mode,
443        );
444
445        if self.opts.init_regions_in_background {
446            // Opens regions in background.
447            common_runtime::spawn_global(async move {
448                let _initialization_guard = initialization_guard;
449                if let Err(err) = open_all_regions.await {
450                    error!(err; "Failed to open regions during the startup.");
451                }
452            });
453        } else {
454            open_all_regions.await?;
455        }
456
457        let heartbeat_task = if let Some(meta_client) = meta_client {
458            let task = self
459                .create_heartbeat_task(&region_server, meta_client, cache_registry)
460                .await?;
461            Some(task)
462        } else {
463            None
464        };
465
466        let is_standalone = heartbeat_task.is_none();
467        let greptimedb_telemetry_task = get_greptimedb_telemetry_task(
468            Some(self.opts.storage.data_home.clone()),
469            is_standalone && self.opts.enable_telemetry,
470        )
471        .await;
472
473        let leases_notifier = if self.opts.require_lease_before_startup && !is_standalone {
474            Some(Arc::new(Notify::new()))
475        } else {
476            None
477        };
478
479        self.region_server = None;
480        Ok(Datanode {
481            services: ServerHandlers::default(),
482            heartbeat_task,
483            region_server,
484            greptimedb_telemetry_task,
485            region_event_receiver,
486            leases_notifier,
487            plugins: self.plugins.clone(),
488            object_store_log_store: self.object_store_log_store.take(),
489            raft_engine_log_store: self.raft_engine_log_store.take(),
490        })
491    }
492
493    async fn create_heartbeat_task(
494        &self,
495        region_server: &RegionServer,
496        meta_client: MetaClientRef,
497        cache_invalidator: CacheInvalidatorRef,
498    ) -> Result<HeartbeatTask> {
499        let stat = {
500            let mut stat = ResourceStatImpl::default();
501            stat.start_collect_cpu_usage();
502            Arc::new(stat)
503        };
504
505        HeartbeatTask::try_new(
506            &self.opts,
507            region_server.clone(),
508            meta_client,
509            self.kv_backend.clone(),
510            cache_invalidator,
511            self.plugins.clone(),
512            stat,
513        )
514        .await
515    }
516
517    /// Builds [ObjectStoreManager] from [StorageConfig] and probes every remote
518    /// object store it contains, unless `probe_on_startup` is disabled.
519    ///
520    /// The probe fails the build when a store rejects any of the operations the
521    /// engines rely on, so a misconfigured storage surfaces here instead of failing
522    /// on the first write.
523    pub async fn build_object_store_manager(
524        cfg: &StorageConfig,
525        node_id: u64,
526    ) -> Result<ObjectStoreManagerRef> {
527        let object_store = store::new_object_store(cfg.store.clone(), &cfg.data_home).await?;
528        let default_name = cfg.store.config_name();
529        let mut checks = Vec::with_capacity(1 + cfg.providers.len());
530        if cfg.store.is_object_storage() {
531            checks.push((default_name.to_string(), object_store.clone()));
532        }
533        let mut object_store_manager = ObjectStoreManager::new(default_name, object_store);
534        for provider in &cfg.providers {
535            let name = provider.config_name();
536            let object_store = store::new_object_store(provider.clone(), &cfg.data_home).await?;
537            if provider.is_object_storage() {
538                checks.push((name.to_string(), object_store.clone()));
539            }
540            object_store_manager.add(name, object_store);
541        }
542
543        if cfg.probe_on_startup && !checks.is_empty() {
544            let prefix = store::probe_key_prefix(node_id);
545            Self::run_probes(&checks, &prefix).await?;
546        }
547
548        Ok(Arc::new(object_store_manager))
549    }
550
551    /// Probes every store in `checks` under its own key prefix.
552    ///
553    /// Providers often share one object storage, so when there is more than one
554    /// store each probe starts after a random delay rather than all at once.
555    ///
556    /// Each check is bounded by [`store::PROBE_TIMEOUT`], so a store that accepts
557    /// connections but never answers fails startup instead of blocking it, after a
558    /// best-effort removal of its probe.
559    async fn run_probes(checks: &[(String, ObjectStore)], prefix: &str) -> Result<()> {
560        let stagger = checks.len() > 1;
561        let checks = checks
562            .iter()
563            .enumerate()
564            .map(|(index, (name, object_store))| {
565                let prefix = format!("{prefix}/{index}");
566                async move {
567                    if stagger {
568                        tokio::time::sleep(store::probe_start_delay()).await;
569                    }
570                    match tokio::time::timeout(
571                        store::PROBE_TIMEOUT,
572                        store::probe_object_store(object_store, name, &prefix),
573                    )
574                    .await
575                    {
576                        Ok(result) => result,
577                        Err(_) => {
578                            store::remove_probe_object(object_store, &prefix).await;
579                            error::ObjectStoreProbeTimeoutSnafu {
580                                name,
581                                timeout: store::PROBE_TIMEOUT,
582                            }
583                            .fail()
584                        }
585                    }
586                }
587            });
588        for result in join_all(checks).await {
589            result?;
590        }
591        Ok(())
592    }
593
594    #[cfg(test)]
595    /// Open all regions belong to this datanode.
596    async fn initialize_region_server(
597        &self,
598        region_server: &RegionServer,
599        open_with_writable: bool,
600    ) -> Result<()> {
601        let node_id = self.opts.node_id.context(MissingNodeIdSnafu)?;
602
603        // TODO(weny): Considering introducing a readonly kv_backend trait.
604        let runtime_switch_manager = RuntimeSwitchManager::new(self.kv_backend.clone());
605        let is_recovery_mode = runtime_switch_manager
606            .recovery_mode()
607            .await
608            .context(GetMetadataSnafu)?;
609        let region_open_requests =
610            build_region_open_requests(node_id, self.kv_backend.clone()).await?;
611
612        open_all_regions(
613            region_server.clone(),
614            region_open_requests,
615            open_with_writable,
616            self.opts.init_regions_parallelism,
617            is_recovery_mode,
618        )
619        .await
620    }
621
622    async fn new_region_server(
623        &mut self,
624        schema_metadata_manager: SchemaMetadataManagerRef,
625        event_listener: RegionServerEventListenerRef,
626        file_ref_manager: FileReferenceManagerRef,
627        cache_registry: Arc<LayeredCacheRegistry>,
628    ) -> Result<RegionServer> {
629        let opts: &DatanodeOptions = &self.opts;
630
631        // Keep `with_dist_planner` false: datanodes execute received MergeScan plans.
632        let catalog_manager = KvBackendCatalogManagerBuilder::new(
633            Arc::new(NoopInformationExtension),
634            self.kv_backend.clone(),
635            cache_registry.clone(),
636        )
637        .build();
638        let partition_manager = catalog_manager.partition_manager();
639
640        // MergeScan region requests stream until completion, so they have no client timeout.
641        let mut channel_config = ChannelConfig {
642            timeout: None,
643            ..Default::default()
644        };
645        if opts.grpc.flight_compression.transport_compression() {
646            channel_config.accept_compression = true;
647            channel_config.send_compression = true;
648        }
649        let node_manager: NodeManagerRef = Arc::new(NodeClients::new(channel_config));
650        let region_query_handler =
651            if let Some(factory) = self.plugins.get::<RegionQueryHandlerFactoryRef>() {
652                factory.build(partition_manager.clone(), node_manager)
653            } else {
654                DatanodeRegionQueryHandler::arc(partition_manager.clone(), node_manager)
655            };
656
657        let query_engine_factory = QueryEngineFactory::try_new_with_plugins(
658            // query engine in datanode only executes plan with resolved table source.
659            catalog_manager,
660            Some(partition_manager),
661            Some(region_query_handler),
662            None,
663            None,
664            None,
665            false,
666            self.plugins.clone(),
667            opts.query.clone(),
668        )
669        .context(DataFusionSnafu)?;
670        let query_engine = query_engine_factory.query_engine();
671
672        let table_provider_factory = self
673            .table_provider_factory
674            .clone()
675            .unwrap_or_else(|| Arc::new(DummyTableProviderFactory));
676        let mut region_server = RegionServer::with_table_provider(
677            query_engine,
678            common_runtime::global_runtime(),
679            event_listener,
680            table_provider_factory,
681            opts.max_concurrent_queries,
682            opts.concurrent_query_limiter_timeout,
683            opts.grpc.flight_compression,
684        );
685        region_server.install_remote_dyn_filter_receiver_injector(&self.plugins);
686
687        let object_store_manager = Self::build_object_store_manager(
688            &opts.storage,
689            opts.node_id.context(MissingNodeIdSnafu)?,
690        )
691        .await?;
692        let engines = self
693            .build_store_engines(
694                object_store_manager,
695                schema_metadata_manager,
696                file_ref_manager,
697                self.plugins.clone(),
698            )
699            .await?;
700        for engine in engines {
701            region_server.register_engine(engine);
702        }
703        if let Some(topic_stats_reporter) = self.topic_stats_reporter.take() {
704            region_server.set_topic_stats_reporter(topic_stats_reporter);
705        }
706
707        Ok(region_server)
708    }
709
710    // internal utils
711
712    /// Builds [RegionEngineRef] from `store_engine` section in `opts`
713    async fn build_store_engines(
714        &mut self,
715        object_store_manager: ObjectStoreManagerRef,
716        schema_metadata_manager: SchemaMetadataManagerRef,
717        file_ref_manager: FileReferenceManagerRef,
718        plugins: Plugins,
719    ) -> Result<Vec<RegionEngineRef>> {
720        let mut metric_engine_config = metric_engine::config::EngineConfig::default();
721        let mut mito_engine_config = MitoConfig::default();
722        let mut file_engine_config = file_engine::config::EngineConfig::default();
723
724        for engine in &self.opts.region_engine {
725            match engine {
726                RegionEngineConfig::Mito(config) => {
727                    mito_engine_config = config.clone();
728                }
729                RegionEngineConfig::File(config) => {
730                    file_engine_config = config.clone();
731                }
732                RegionEngineConfig::Metric(metric_config) => {
733                    metric_engine_config = metric_config.clone();
734                }
735            }
736        }
737
738        // Build a fetcher to backfill partition_expr on open.
739        let fetcher = Arc::new(MetaPartitionExprFetcher::new(self.kv_backend.clone()));
740        let mito_engine = self
741            .build_mito_engine(
742                object_store_manager.clone(),
743                mito_engine_config,
744                schema_metadata_manager.clone(),
745                file_ref_manager.clone(),
746                fetcher.clone(),
747                plugins.clone(),
748            )
749            .await?;
750
751        let metric_engine = MetricEngine::try_new(mito_engine.clone(), metric_engine_config)
752            .context(BuildMetricEngineSnafu)?;
753
754        let file_engine = FileRegionEngine::new(
755            file_engine_config,
756            object_store_manager.default_object_store().clone(), // TODO: implement custom storage for file engine
757            self.local_file_access.clone(),
758        );
759
760        let relational_engine =
761            RelationalEngine::new(object_store_manager.default_object_store().clone());
762
763        Ok(vec![
764            Arc::new(mito_engine) as _,
765            Arc::new(metric_engine) as _,
766            Arc::new(file_engine) as _,
767            Arc::new(relational_engine) as _,
768        ])
769    }
770
771    /// Builds [MitoEngine] according to options.
772    async fn build_mito_engine(
773        &mut self,
774        object_store_manager: ObjectStoreManagerRef,
775        mut config: MitoConfig,
776        schema_metadata_manager: SchemaMetadataManagerRef,
777        file_ref_manager: FileReferenceManagerRef,
778        partition_expr_fetcher: PartitionExprFetcherRef,
779        plugins: Plugins,
780    ) -> Result<MitoEngine> {
781        let opts = &self.opts;
782        if opts.storage.is_object_storage() {
783            // Enable the write cache when setting object storage
784            config.enable_write_cache = true;
785            info!("Configured 'enable_write_cache=true' for mito engine.");
786        }
787
788        let mito_engine = match &opts.wal {
789            DatanodeWalConfig::RaftEngine(raft_engine_config) => {
790                let log_store =
791                    Self::build_raft_engine_log_store(&opts.storage.data_home, raft_engine_config)
792                        .await?;
793                self.raft_engine_log_store = Some(log_store.clone());
794
795                let builder = MitoEngineBuilder::new(
796                    &opts.storage.data_home,
797                    config,
798                    log_store,
799                    object_store_manager,
800                    schema_metadata_manager,
801                    file_ref_manager,
802                    partition_expr_fetcher.clone(),
803                    plugins,
804                );
805
806                #[cfg(feature = "enterprise")]
807                let builder = builder.with_extension_range_provider_factory(
808                    self.extension_range_provider_factory.take(),
809                );
810
811                builder.try_build().await.context(BuildMitoEngineSnafu)?
812            }
813            DatanodeWalConfig::Kafka(kafka_config) => {
814                let log_store = Self::build_kafka_log_store(kafka_config).await?;
815                self.topic_stats_reporter = Some(log_store.topic_stats_reporter());
816                let builder = MitoEngineBuilder::new(
817                    &opts.storage.data_home,
818                    config,
819                    log_store,
820                    object_store_manager,
821                    schema_metadata_manager,
822                    file_ref_manager,
823                    partition_expr_fetcher,
824                    plugins,
825                );
826
827                #[cfg(feature = "enterprise")]
828                let builder = builder.with_extension_range_provider_factory(
829                    self.extension_range_provider_factory.take(),
830                );
831
832                builder.try_build().await.context(BuildMitoEngineSnafu)?
833            }
834            DatanodeWalConfig::Noop => {
835                let log_store = Arc::new(NoopLogStore);
836
837                let builder = MitoEngineBuilder::new(
838                    &opts.storage.data_home,
839                    config,
840                    log_store,
841                    object_store_manager,
842                    schema_metadata_manager,
843                    file_ref_manager,
844                    partition_expr_fetcher.clone(),
845                    plugins,
846                );
847
848                #[cfg(feature = "enterprise")]
849                let builder = builder.with_extension_range_provider_factory(
850                    self.extension_range_provider_factory.take(),
851                );
852
853                builder.try_build().await.context(BuildMitoEngineSnafu)?
854            }
855            DatanodeWalConfig::ObjectStore(object_store_config) => {
856                let object_store = wal_object_store(
857                    &object_store_manager,
858                    &opts.storage,
859                    &object_store_config.storage_provider,
860                )?;
861                let node_id = opts.node_id.context(MissingNodeIdSnafu)?;
862                let log_store =
863                    Self::build_object_store_log_store(object_store, object_store_config, node_id)
864                        .await?;
865                self.object_store_log_store = Some(log_store.clone());
866
867                let builder = MitoEngineBuilder::new(
868                    &opts.storage.data_home,
869                    config,
870                    log_store,
871                    object_store_manager,
872                    schema_metadata_manager,
873                    file_ref_manager,
874                    partition_expr_fetcher,
875                    plugins,
876                );
877
878                #[cfg(feature = "enterprise")]
879                let builder = builder.with_extension_range_provider_factory(
880                    self.extension_range_provider_factory.take(),
881                );
882
883                builder.try_build().await.context(BuildMitoEngineSnafu)?
884            }
885        };
886        Ok(mito_engine)
887    }
888
889    /// Builds [ObjectStoreLogStore] under the node prefix of `node_id` in the
890    /// standalone generation, the prefix the regions persist in their WAL options.
891    async fn build_object_store_log_store(
892        object_store: ObjectStore,
893        config: &ObjectStoreWalConfig,
894        node_id: u64,
895    ) -> Result<Arc<ObjectStoreLogStore>> {
896        info!(
897            "Creating object store logstore with config: {:?}, node id: {}",
898            config, node_id
899        );
900        ObjectStoreLogStore::try_new(object_store, config, node_id, STANDALONE_GENERATION)
901            .await
902            .map_err(Box::new)
903            .context(OpenLogStoreSnafu)
904    }
905
906    /// Builds [RaftEngineLogStore].
907    async fn build_raft_engine_log_store(
908        data_home: &str,
909        config: &RaftEngineConfig,
910    ) -> Result<Arc<RaftEngineLogStore>> {
911        let data_home = normalize_dir(data_home);
912        let wal_dir = match &config.dir {
913            Some(dir) => dir.clone(),
914            None => format!("{}{WAL_DIR}", data_home),
915        };
916
917        // create WAL directory
918        fs::create_dir_all(Path::new(&wal_dir))
919            .await
920            .context(CreateDirSnafu { dir: &wal_dir })?;
921        info!(
922            "Creating raft-engine logstore with config: {:?} and storage path: {}",
923            config, &wal_dir
924        );
925        let logstore = RaftEngineLogStore::try_new(wal_dir, config)
926            .await
927            .map_err(Box::new)
928            .context(OpenLogStoreSnafu)?;
929
930        Ok(Arc::new(logstore))
931    }
932
933    /// Builds [`KafkaLogStore`].
934    async fn build_kafka_log_store(config: &DatanodeKafkaConfig) -> Result<Arc<KafkaLogStore>> {
935        KafkaLogStore::try_new(config)
936            .await
937            .map_err(Box::new)
938            .context(OpenLogStoreSnafu)
939            .map(Arc::new)
940    }
941}
942
943/// Keeps the first shutdown error and logs the later ones.
944pub(crate) fn record_shutdown_error(first_error: &mut Option<error::Error>, result: Result<()>) {
945    if let Err(err) = result {
946        if first_error.is_none() {
947            *first_error = Some(err);
948        } else {
949            warn!(err; "Ignored a later shutdown error");
950        }
951    }
952}
953
954/// Closes the raft-engine WAL, which takes its last reference: raft-engine only closes the
955/// active log file when its engine is dropped. If something else still holds the WAL, it is
956/// stopped instead and the file stays open until that reference is dropped.
957async fn close_raft_engine_log_store(log_store: Arc<RaftEngineLogStore>) -> Result<()> {
958    let result = match Arc::try_unwrap(log_store) {
959        Ok(log_store) => log_store.close().await,
960        Err(log_store) => {
961            warn!(
962                "The raft-engine WAL is still referenced elsewhere ({} references), stopping it without closing its active log file",
963                Arc::strong_count(&log_store) - 1
964            );
965            log_store.stop().await
966        }
967    };
968    result
969        .map_err(BoxedError::new)
970        .context(ShutdownInstanceSnafu)
971}
972
973/// Returns the object store that holds the WAL objects: the default one for an
974/// empty `storage_provider`, otherwise the configured provider with that name.
975fn wal_object_store(
976    object_store_manager: &ObjectStoreManager,
977    storage: &StorageConfig,
978    storage_provider: &str,
979) -> Result<ObjectStore> {
980    if storage_provider.is_empty() {
981        return Ok(object_store_manager.default_object_store().clone());
982    }
983    object_store_manager
984        .find(storage_provider)
985        .cloned()
986        .with_context(|| {
987            let configured = std::iter::once(&storage.store)
988                .chain(storage.providers.iter())
989                .map(|store| store.config_name())
990                .collect::<Vec<_>>();
991            InvalidObjectStoreWalConfigSnafu {
992                field: "storage_provider",
993                value: storage_provider,
994                reason: format!("must name a configured storage provider, one of {configured:?}"),
995            }
996        })
997}
998
999/// Rejects repeated engine types before their configs can silently overwrite each other.
1000fn validate_region_engine_config(configs: &[RegionEngineConfig]) -> Result<()> {
1001    let mut engine_indices = HashMap::new();
1002    for (index, config) in configs.iter().enumerate() {
1003        let engine = match config {
1004            RegionEngineConfig::Mito(_) => MITO_ENGINE,
1005            RegionEngineConfig::File(_) => FILE_ENGINE,
1006            RegionEngineConfig::Metric(_) => METRIC_ENGINE,
1007        };
1008        if let Some(first_index) = engine_indices.insert(engine, index) {
1009            return DuplicateRegionEngineConfigSnafu {
1010                engine,
1011                first_index,
1012                duplicate_index: index,
1013            }
1014            .fail();
1015        }
1016    }
1017    Ok(())
1018}
1019
1020/// Rejects an object store WAL config the log store cannot run on.
1021fn validate_object_store_wal_config(config: &ObjectStoreWalConfig) -> Result<()> {
1022    let prefix = config.prefix.trim();
1023    ensure!(
1024        !prefix.is_empty(),
1025        InvalidObjectStoreWalConfigSnafu {
1026            field: "prefix",
1027            value: &config.prefix,
1028            reason: "must not be empty",
1029        }
1030    );
1031    ensure!(
1032        !prefix.starts_with('/'),
1033        InvalidObjectStoreWalConfigSnafu {
1034            field: "prefix",
1035            value: &config.prefix,
1036            reason: "must be a relative path",
1037        }
1038    );
1039    ensure!(
1040        !prefix
1041            .split('/')
1042            .any(|component| component.is_empty() || component == ".."),
1043        InvalidObjectStoreWalConfigSnafu {
1044            field: "prefix",
1045            value: &config.prefix,
1046            reason: "must not contain empty or `..` components",
1047        }
1048    );
1049    ensure!(
1050        config.flush_interval >= Duration::from_millis(10),
1051        InvalidObjectStoreWalConfigSnafu {
1052            field: "flush_interval",
1053            value: format!("{:?}", config.flush_interval),
1054            reason: "must be at least 10ms",
1055        }
1056    );
1057    ensure!(
1058        config.max_batch_bytes.as_bytes() > 0,
1059        InvalidObjectStoreWalConfigSnafu {
1060            field: "max_batch_bytes",
1061            value: config.max_batch_bytes.to_string(),
1062            reason: "must be greater than 0",
1063        }
1064    );
1065    ensure!(
1066        config.max_unpersisted_bytes.as_bytes() > 0,
1067        InvalidObjectStoreWalConfigSnafu {
1068            field: "max_unpersisted_bytes",
1069            value: config.max_unpersisted_bytes.to_string(),
1070            reason: "must be greater than 0",
1071        }
1072    );
1073    ensure!(
1074        config.max_unpersisted_age > Duration::ZERO,
1075        InvalidObjectStoreWalConfigSnafu {
1076            field: "max_unpersisted_age",
1077            value: format!("{:?}", config.max_unpersisted_age),
1078            reason: "must be greater than 0",
1079        }
1080    );
1081    Ok(())
1082}
1083
1084/// Open all regions belong to this datanode.
1085async fn open_all_regions(
1086    region_server: RegionServer,
1087    region_open_requests: RegionOpenRequests,
1088    open_with_writable: bool,
1089    init_regions_parallelism: usize,
1090    ignore_nonexistent_region: bool,
1091) -> Result<()> {
1092    let RegionOpenRequests {
1093        leader_regions,
1094        #[cfg(feature = "enterprise")]
1095        follower_regions,
1096    } = region_open_requests;
1097
1098    let leader_region_num = leader_regions.len();
1099    info!("going to open {} region(s)", leader_region_num);
1100    let now = Instant::now();
1101    let open_regions = region_server
1102        .handle_batch_open_requests(
1103            init_regions_parallelism,
1104            leader_regions,
1105            ignore_nonexistent_region,
1106        )
1107        .await?;
1108    info!(
1109        "Opened {} regions in {:?}",
1110        open_regions.len(),
1111        now.elapsed()
1112    );
1113    if !ignore_nonexistent_region {
1114        ensure!(
1115            open_regions.len() == leader_region_num,
1116            error::UnexpectedSnafu {
1117                violated: format!(
1118                    "Expected to open {} of regions, only {} of regions has opened",
1119                    leader_region_num,
1120                    open_regions.len()
1121                )
1122            }
1123        );
1124    } else if open_regions.len() != leader_region_num {
1125        warn!(
1126            "ignore nonexistent region, expected to open {} of regions, only {} of regions has opened",
1127            leader_region_num,
1128            open_regions.len()
1129        );
1130    }
1131
1132    for region_id in open_regions {
1133        if open_with_writable {
1134            let res = region_server.set_region_role(region_id, RegionRole::Leader);
1135            match res {
1136                Ok(_) => {
1137                    // Finalize leadership: persist backfilled metadata.
1138                    if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server
1139                        .set_region_role_state_gracefully(
1140                            region_id,
1141                            SettableRegionRoleState::Leader,
1142                        )
1143                        .await?
1144                    {
1145                        error!(err; "failed to convert region {region_id} to leader");
1146                    }
1147                }
1148                Err(e) => {
1149                    error!(e; "failed to convert region {region_id} to leader");
1150                }
1151            }
1152        }
1153    }
1154
1155    #[cfg(feature = "enterprise")]
1156    if !follower_regions.is_empty() {
1157        use tokio::time::Instant;
1158
1159        let follower_region_num = follower_regions.len();
1160        info!("going to open {} follower region(s)", follower_region_num);
1161
1162        let now = Instant::now();
1163        let open_regions = region_server
1164            .handle_batch_open_requests(
1165                init_regions_parallelism,
1166                follower_regions,
1167                ignore_nonexistent_region,
1168            )
1169            .await?;
1170        info!(
1171            "Opened {} follower regions in {:?}",
1172            open_regions.len(),
1173            now.elapsed()
1174        );
1175
1176        if !ignore_nonexistent_region {
1177            ensure!(
1178                open_regions.len() == follower_region_num,
1179                error::UnexpectedSnafu {
1180                    violated: format!(
1181                        "Expected to open {} of follower regions, only {} of regions has opened",
1182                        follower_region_num,
1183                        open_regions.len()
1184                    )
1185                }
1186            );
1187        } else if open_regions.len() != follower_region_num {
1188            warn!(
1189                "ignore nonexistent region, expected to open {} of follower regions, only {} of regions has opened",
1190                follower_region_num,
1191                open_regions.len()
1192            );
1193        }
1194    }
1195
1196    info!("all regions are opened");
1197
1198    Ok(())
1199}
1200
1201#[cfg(test)]
1202mod tests {
1203    use std::assert_matches;
1204    use std::collections::{BTreeMap, HashMap};
1205    use std::io::Write;
1206    use std::ops::Range;
1207    use std::path::Path;
1208    use std::sync::Arc;
1209    use std::sync::atomic::{AtomicUsize, Ordering};
1210    use std::time::Duration;
1211
1212    use api::v1::{ColumnSchema, Rows};
1213    use cache::build_datanode_layered_cache_registry;
1214    use common_base::Plugins;
1215    use common_base::readable_size::ReadableSize;
1216    use common_base::secrets::SecretString;
1217    use common_config::Configurable;
1218    use common_error::ext::ErrorExt;
1219    use common_error::status_code::StatusCode;
1220    use common_meta::cache::{SchemaCacheRef, TableSchemaCacheRef};
1221    use common_meta::key::RegionRoleSet;
1222    use common_meta::key::SchemaMetadataManager;
1223    use common_meta::key::datanode_table::DatanodeTableManager;
1224    use common_meta::kv_backend::KvBackendRef;
1225    use common_meta::kv_backend::memory::MemoryKvBackend;
1226    use common_meta::node_manager::NodeManagerRef;
1227    use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
1228    use common_wal::config::DatanodeWalConfig;
1229    use common_wal::config::object_store::{AckMode, ObjectStoreWalConfig, STANDALONE_GENERATION};
1230    use common_wal::config::raft_engine::RaftEngineConfig;
1231    use common_wal::options::{ObjectStoreWalOptions, WAL_OPTIONS_KEY, WalOptions};
1232    use log_store::ObjectStoreLogStore;
1233    use log_store::error::Error as LogStoreError;
1234    use meta_client::client::MetaClientBuilder;
1235    use mito2::engine::MITO_ENGINE_NAME;
1236    use mito2::sst::file_ref::FileReferenceManager;
1237    use mito2::test_util::{CreateRequestBuilder, build_rows, rows_schema};
1238    use object_store::ObjectStore;
1239    use object_store::config::{S3Config, S3Connection};
1240    use object_store::manager::ObjectStoreManager;
1241    use partition::manager::PartitionRuleManagerRef;
1242    use query::region_query::{RegionQueryHandlerFactory, RegionQueryHandlerFactoryRef};
1243    use store_api::logstore::LogStore;
1244    use store_api::logstore::provider::Provider;
1245    use store_api::region_request::{PathType, RegionOpenRequest, RegionPutRequest, RegionRequest};
1246    use store_api::storage::RegionId;
1247
1248    use crate::config::{DatanodeOptions, ObjectStoreConfig, RegionEngineConfig, StorageConfig};
1249    use crate::datanode::{
1250        Datanode, DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config,
1251        wal_object_store,
1252    };
1253    use crate::error::{self, Error};
1254    use crate::region_query::DatanodeRegionQueryHandler;
1255    use crate::region_server::RegionServer;
1256    use crate::tests::{MockRegionEngine, mock_region_server};
1257
1258    const WAIT: Duration = Duration::from_secs(30);
1259
1260    async fn setup_table_datanode(kv: &KvBackendRef) {
1261        let mgr = DatanodeTableManager::new(kv.clone());
1262        let txn = mgr
1263            .build_create_txn(
1264                1028,
1265                MITO_ENGINE_NAME,
1266                "foo/bar/weny",
1267                HashMap::from([("foo".to_string(), "bar".to_string())]),
1268                HashMap::default(),
1269                BTreeMap::from([(0, RegionRoleSet::new(vec![0, 1, 2], vec![]))]),
1270            )
1271            .unwrap();
1272
1273        let r = kv.txn(txn).await.unwrap();
1274        assert!(r.succeeded);
1275    }
1276
1277    #[tokio::test]
1278    async fn test_initialize_region_server() {
1279        common_telemetry::init_default_ut_logging();
1280        let mut mock_region_server = mock_region_server();
1281        let (mock_region, mut mock_region_handler) = MockRegionEngine::new(MITO_ENGINE_NAME);
1282
1283        mock_region_server.register_engine(mock_region.clone());
1284
1285        let kv_backend = Arc::new(MemoryKvBackend::new());
1286        let layered_cache_registry =
1287            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
1288
1289        let mut builder = DatanodeBuilder::new(
1290            DatanodeOptions {
1291                node_id: Some(0),
1292                ..Default::default()
1293            },
1294            Plugins::default(),
1295            kv_backend.clone(),
1296        );
1297        builder.with_cache_registry(layered_cache_registry);
1298        setup_table_datanode(&(kv_backend as _)).await;
1299
1300        builder
1301            .initialize_region_server(&mock_region_server, false)
1302            .await
1303            .unwrap();
1304
1305        for i in 0..3 {
1306            let (region_id, req) = mock_region_handler.recv().await.unwrap();
1307            assert_eq!(region_id, RegionId::new(1028, i));
1308            if let RegionRequest::Open(req) = req {
1309                assert_eq!(
1310                    req.options,
1311                    HashMap::from([("foo".to_string(), "bar".to_string())])
1312                )
1313            } else {
1314                unreachable!()
1315            }
1316        }
1317
1318        assert_matches!(
1319            mock_region_handler.try_recv(),
1320            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1321        );
1322    }
1323
1324    fn datanode_builder(
1325        data_home: &str,
1326        wal: DatanodeWalConfig,
1327        kv_backend: KvBackendRef,
1328    ) -> DatanodeBuilder {
1329        let layered_cache_registry =
1330            Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
1331        let mut opts = DatanodeOptions {
1332            node_id: Some(0),
1333            wal,
1334            enable_telemetry: false,
1335            ..Default::default()
1336        };
1337        opts.storage.data_home = data_home.to_string();
1338
1339        let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend);
1340        builder.with_cache_registry(layered_cache_registry);
1341        builder
1342    }
1343
1344    fn object_store_wal_builder(data_home: &str, config: ObjectStoreWalConfig) -> DatanodeBuilder {
1345        datanode_builder(
1346            data_home,
1347            DatanodeWalConfig::ObjectStore(config),
1348            Arc::new(MemoryKvBackend::new()),
1349        )
1350    }
1351
1352    async fn build_err(builder: DatanodeBuilder) -> Error {
1353        match builder.build().await {
1354            Ok(_) => panic!("build must fail"),
1355            Err(err) => err,
1356        }
1357    }
1358
1359    fn is_empty_dir(dir: &Path) -> bool {
1360        std::fs::read_dir(dir).unwrap().next().is_none()
1361    }
1362
1363    /// Opens the raft-engine WAL under `data_home` again, which only succeeds once the previous
1364    /// engine was released (raft-engine holds a lock on its directory while open).
1365    async fn assert_raft_engine_wal_reopens(data_home: &str) {
1366        let reopened =
1367            DatanodeBuilder::build_raft_engine_log_store(data_home, &RaftEngineConfig::default())
1368                .await
1369                .unwrap();
1370        reopened.stop().await.unwrap();
1371    }
1372
1373    /// Returns true if the store rejects even an empty append, which it does only once stopped.
1374    async fn is_stopped(log_store: &ObjectStoreLogStore) -> bool {
1375        matches!(
1376            log_store.append_batch(vec![]).await,
1377            Err(LogStoreError::ObjectStoreWalStopped { .. })
1378        )
1379    }
1380
1381    #[test]
1382    fn test_validate_unique_region_engine_config() {
1383        let mut opts = DatanodeOptions::default();
1384        validate_region_engine_config(&opts.region_engine).unwrap();
1385
1386        opts.region_engine
1387            .push(RegionEngineConfig::Metric(Default::default()));
1388        validate_region_engine_config(&opts.region_engine).unwrap();
1389
1390        // Omitted engines still use their defaults during engine construction.
1391        validate_region_engine_config(&[]).unwrap();
1392    }
1393
1394    #[tokio::test]
1395    async fn test_build_rejects_duplicate_region_engine_config() {
1396        for (engine, first_index) in [("mito", 0), ("file", 1), ("metric", 2)] {
1397            let mut file = create_named_temp_file();
1398            write!(
1399                file,
1400                r#"
1401                [[region_engine]]
1402                [region_engine.mito]
1403                global_write_buffer_size = "4GiB"
1404                min_compaction_interval = "600s"
1405
1406                [[region_engine]]
1407                [region_engine.file]
1408
1409                [[region_engine]]
1410                [region_engine.metric]
1411
1412                [[region_engine]]
1413                [region_engine.{engine}]
1414                "#
1415            )
1416            .unwrap();
1417            let opts = DatanodeOptions::load_layered_options(
1418                Some(file.path().to_str().unwrap()),
1419                "DATANODE_DUPLICATE_REGION_ENGINE_UT",
1420            )
1421            .unwrap();
1422            let data_home = create_temp_dir("duplicate-region-engine-config");
1423            let mut builder = datanode_builder(
1424                data_home.path().to_str().unwrap(),
1425                DatanodeWalConfig::default(),
1426                Arc::new(MemoryKvBackend::new()),
1427            );
1428            builder.opts.region_engine = opts.region_engine;
1429            let err = build_err(builder).await;
1430
1431            let Error::DuplicateRegionEngineConfig {
1432                engine: actual_engine,
1433                first_index: actual_first_index,
1434                duplicate_index,
1435                ..
1436            } = &err
1437            else {
1438                panic!("unexpected error for {engine}: {err:?}");
1439            };
1440            assert_eq!(*actual_engine, engine);
1441            assert_eq!(*actual_first_index, first_index);
1442            assert_eq!(*duplicate_index, 3);
1443            assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1444            // Rejected before the builder creates any storage or log store.
1445            assert!(is_empty_dir(data_home.path()));
1446        }
1447    }
1448
1449    #[tokio::test]
1450    async fn test_build_rejects_invalid_object_store_wal_config() {
1451        let cases = [
1452            (
1453                ObjectStoreWalConfig {
1454                    prefix: String::new(),
1455                    ..Default::default()
1456                },
1457                "prefix",
1458            ),
1459            (
1460                ObjectStoreWalConfig {
1461                    prefix: "/wal".to_string(),
1462                    ..Default::default()
1463                },
1464                "prefix",
1465            ),
1466            (
1467                ObjectStoreWalConfig {
1468                    prefix: "wal/../other".to_string(),
1469                    ..Default::default()
1470                },
1471                "prefix",
1472            ),
1473            (
1474                ObjectStoreWalConfig {
1475                    prefix: "wal//objects".to_string(),
1476                    ..Default::default()
1477                },
1478                "prefix",
1479            ),
1480            (
1481                ObjectStoreWalConfig {
1482                    prefix: "wal/".to_string(),
1483                    ..Default::default()
1484                },
1485                "prefix",
1486            ),
1487            (
1488                ObjectStoreWalConfig {
1489                    flush_interval: Duration::from_millis(9),
1490                    ..Default::default()
1491                },
1492                "flush_interval",
1493            ),
1494            (
1495                ObjectStoreWalConfig {
1496                    max_batch_bytes: ReadableSize(0),
1497                    ..Default::default()
1498                },
1499                "max_batch_bytes",
1500            ),
1501            (
1502                ObjectStoreWalConfig {
1503                    max_unpersisted_bytes: ReadableSize(0),
1504                    ..Default::default()
1505                },
1506                "max_unpersisted_bytes",
1507            ),
1508            (
1509                ObjectStoreWalConfig {
1510                    max_unpersisted_age: Duration::ZERO,
1511                    ..Default::default()
1512                },
1513                "max_unpersisted_age",
1514            ),
1515        ];
1516
1517        for (config, expected_field) in cases {
1518            let expected_value = match expected_field {
1519                "prefix" => config.prefix.clone(),
1520                "flush_interval" => format!("{:?}", config.flush_interval),
1521                "max_unpersisted_bytes" => config.max_unpersisted_bytes.to_string(),
1522                "max_unpersisted_age" => format!("{:?}", config.max_unpersisted_age),
1523                _ => config.max_batch_bytes.to_string(),
1524            };
1525            let data_home = create_temp_dir("object-store-wal-invalid-config");
1526            let builder = object_store_wal_builder(data_home.path().to_str().unwrap(), config);
1527            let err = build_err(builder).await;
1528
1529            let Error::InvalidObjectStoreWalConfig { field, value, .. } = &err else {
1530                panic!("unexpected error for {expected_field}: {err:?}");
1531            };
1532            assert_eq!(*field, expected_field);
1533            assert_eq!(*value, expected_value);
1534            assert_eq!(StatusCode::InvalidArguments, err.status_code());
1535            // The builder stops before it creates any storage or log store.
1536            assert!(is_empty_dir(data_home.path()));
1537        }
1538    }
1539
1540    #[tokio::test]
1541    async fn test_build_rejects_object_store_wal_with_meta_client() {
1542        let data_home = create_temp_dir("object-store-wal-meta-client");
1543        let meta_client = Arc::new(MetaClientBuilder::datanode_default_options(0).build());
1544
1545        let mut builder = object_store_wal_builder(
1546            data_home.path().to_str().unwrap(),
1547            ObjectStoreWalConfig::default(),
1548        );
1549        builder.with_meta_client(meta_client);
1550        let err = build_err(builder).await;
1551
1552        assert_matches!(err, Error::ObjectStoreWalNotStandalone { .. });
1553        assert_eq!(StatusCode::InvalidArguments, err.status_code());
1554        // The builder stops before it creates any storage or log store.
1555        assert!(is_empty_dir(data_home.path()));
1556    }
1557
1558    #[tokio::test]
1559    async fn test_build_rejects_unknown_wal_storage_provider() {
1560        let data_home = create_temp_dir("object-store-wal-unknown-provider");
1561        let builder = object_store_wal_builder(
1562            data_home.path().to_str().unwrap(),
1563            ObjectStoreWalConfig {
1564                storage_provider: "s3".to_string(),
1565                ..Default::default()
1566            },
1567        );
1568        let err = build_err(builder).await;
1569
1570        let Error::InvalidObjectStoreWalConfig {
1571            field,
1572            value,
1573            reason,
1574            ..
1575        } = &err
1576        else {
1577            panic!("unexpected error: {err:?}");
1578        };
1579        assert_eq!(*field, "storage_provider");
1580        assert_eq!(value, "s3");
1581        assert!(reason.contains(r#"["File"]"#), "{reason}");
1582        assert_eq!(StatusCode::InvalidArguments, err.status_code());
1583        // No log store is created under the prefix.
1584        assert!(!data_home.path().join("wal").exists());
1585    }
1586
1587    #[test]
1588    fn test_wal_object_store_selects_provider() {
1589        fn fs_object_store(dir: &Path) -> ObjectStore {
1590            let builder = object_store::services::Fs::default().root(dir.to_str().unwrap());
1591            ObjectStore::new(builder).unwrap()
1592        }
1593
1594        let default_dir = create_temp_dir("wal-object-store-default");
1595        let named_dir = create_temp_dir("wal-object-store-named");
1596        let mut manager = ObjectStoreManager::new("File", fs_object_store(default_dir.path()));
1597        manager.add("Wal-Store", fs_object_store(named_dir.path()));
1598        let storage = StorageConfig::default();
1599
1600        let root = |store: ObjectStore| store.info().root().clone();
1601        assert_eq!(
1602            root(wal_object_store(&manager, &storage, "").unwrap()),
1603            root(manager.default_object_store().clone())
1604        );
1605        assert_eq!(
1606            root(wal_object_store(&manager, &storage, "wal-store").unwrap()),
1607            root(manager.find("Wal-Store").unwrap().clone())
1608        );
1609        assert_matches!(
1610            wal_object_store(&manager, &storage, "missing").unwrap_err(),
1611            Error::InvalidObjectStoreWalConfig {
1612                field: "storage_provider",
1613                ..
1614            }
1615        );
1616    }
1617
1618    #[tokio::test]
1619    async fn test_object_store_probe_can_be_disabled() {
1620        // The endpoint does not accept connections, so this only succeeds when the
1621        // probe does not run.
1622        let storage = StorageConfig {
1623            store: ObjectStoreConfig::S3(S3Config {
1624                connection: S3Connection {
1625                    bucket: "test".to_string(),
1626                    endpoint: Some("http://127.0.0.1:1".to_string()),
1627                    region: Some("us-east-1".to_string()),
1628                    access_key_id: SecretString::from("test".to_string()),
1629                    secret_access_key: SecretString::from("test".to_string()),
1630                    ..Default::default()
1631                },
1632                ..Default::default()
1633            }),
1634            probe_on_startup: false,
1635            ..Default::default()
1636        };
1637
1638        DatanodeBuilder::build_object_store_manager(&storage, 1)
1639            .await
1640            .unwrap();
1641    }
1642
1643    #[tokio::test]
1644    async fn test_file_store_is_not_probed() {
1645        // The probe is enabled by default, but a File store must skip it.
1646        let data_home = create_temp_dir("object-store-probe-file");
1647        let storage = StorageConfig {
1648            data_home: data_home.path().to_str().unwrap().to_string(),
1649            ..Default::default()
1650        };
1651
1652        DatanodeBuilder::build_object_store_manager(&storage, 1)
1653            .await
1654            .unwrap();
1655
1656        // The probe would leave this directory behind if it ran.
1657        assert!(!data_home.path().join(".greptime_probe").exists());
1658    }
1659
1660    #[tokio::test]
1661    async fn test_object_store_probe_failure_aborts_build() {
1662        // The endpoint does not accept connections, so the probe write fails and the
1663        // build must fail with it.
1664        let storage = StorageConfig {
1665            store: ObjectStoreConfig::S3(S3Config {
1666                connection: S3Connection {
1667                    bucket: "test".to_string(),
1668                    endpoint: Some("http://127.0.0.1:1".to_string()),
1669                    region: Some("us-east-1".to_string()),
1670                    access_key_id: SecretString::from("test".to_string()),
1671                    secret_access_key: SecretString::from("test".to_string()),
1672                    ..Default::default()
1673                },
1674                ..Default::default()
1675            }),
1676            ..Default::default()
1677        };
1678
1679        let err = DatanodeBuilder::build_object_store_manager(&storage, 1)
1680            .await
1681            .unwrap_err();
1682        assert_matches!(
1683            err,
1684            Error::ObjectStoreProbe {
1685                name,
1686                step: "write",
1687                ..
1688            } if name == storage.store.config_name()
1689        );
1690    }
1691
1692    struct CountingRegionQueryHandlerFactory(Arc<AtomicUsize>);
1693
1694    impl RegionQueryHandlerFactory for CountingRegionQueryHandlerFactory {
1695        fn build(
1696            &self,
1697            partition_manager: PartitionRuleManagerRef,
1698            node_manager: NodeManagerRef,
1699        ) -> query::region_query::RegionQueryHandlerRef {
1700            self.0.fetch_add(1, Ordering::SeqCst);
1701            DatanodeRegionQueryHandler::arc(partition_manager, node_manager)
1702        }
1703    }
1704
1705    #[tokio::test]
1706    async fn test_build_uses_installed_region_query_handler_factory() {
1707        let data_home = create_temp_dir("region-query-handler-factory");
1708        let builds = Arc::new(AtomicUsize::new(0));
1709        let plugins = Plugins::new();
1710        plugins.insert(Arc::new(CountingRegionQueryHandlerFactory(builds.clone()))
1711            as RegionQueryHandlerFactoryRef);
1712        let mut builder = datanode_builder(
1713            data_home.path().to_str().unwrap(),
1714            DatanodeWalConfig::default(),
1715            Arc::new(MemoryKvBackend::new()),
1716        );
1717        builder.set_plugins(plugins);
1718
1719        let mut datanode = builder.build().await.unwrap();
1720        assert_eq!(builds.load(Ordering::SeqCst), 1);
1721        datanode.shutdown().await.unwrap();
1722    }
1723
1724    #[tokio::test]
1725    async fn test_build_and_shutdown_with_raft_engine_wal() {
1726        let data_home = create_temp_dir("raft-engine-wal");
1727        let builder = datanode_builder(
1728            data_home.path().to_str().unwrap(),
1729            DatanodeWalConfig::default(),
1730            Arc::new(MemoryKvBackend::new()),
1731        );
1732
1733        let mut datanode = builder.build().await.unwrap();
1734        assert!(datanode.object_store_log_store.is_none());
1735        assert!(data_home.path().join("wal").exists());
1736        datanode.shutdown().await.unwrap();
1737        assert_raft_engine_wal_reopens(data_home.path().to_str().unwrap()).await;
1738    }
1739
1740    #[tokio::test]
1741    async fn test_build_and_shutdown_with_object_store_wal() {
1742        common_telemetry::init_default_ut_logging();
1743        let data_home = create_temp_dir("object-store-wal");
1744        let builder = object_store_wal_builder(
1745            data_home.path().to_str().unwrap(),
1746            ObjectStoreWalConfig::default(),
1747        );
1748
1749        let mut datanode = builder.build().await.unwrap();
1750        let log_store = datanode.object_store_log_store.clone().unwrap();
1751        assert!(!is_stopped(&log_store).await);
1752        // The store runs under the node prefix derived from the configured root.
1753        assert!(format!("{log_store:?}").contains(r#"prefix: "wal/datanodes/0/epochs/0""#));
1754
1755        datanode.shutdown().await.unwrap();
1756        assert!(is_stopped(&log_store).await);
1757    }
1758
1759    #[tokio::test]
1760    async fn test_build_with_named_wal_storage_provider() {
1761        let data_home = create_temp_dir("object-store-wal-named-provider");
1762        // The default file store is registered under its provider name.
1763        let builder = object_store_wal_builder(
1764            data_home.path().to_str().unwrap(),
1765            ObjectStoreWalConfig {
1766                storage_provider: "file".to_string(),
1767                ..Default::default()
1768            },
1769        );
1770
1771        let mut datanode = builder.build().await.unwrap();
1772        assert!(datanode.object_store_log_store.is_some());
1773        datanode.shutdown().await.unwrap();
1774    }
1775
1776    #[tokio::test]
1777    async fn test_shutdown_stops_every_engine_and_the_log_store_after_engine_failures() {
1778        common_telemetry::init_default_ut_logging();
1779        let data_home = create_temp_dir("object-store-wal-shutdown-failure");
1780        let builder = object_store_wal_builder(
1781            data_home.path().to_str().unwrap(),
1782            ObjectStoreWalConfig::default(),
1783        );
1784        let mut datanode = builder.build().await.unwrap();
1785        let log_store = datanode.object_store_log_store.clone().unwrap();
1786        // Both engines fail to stop, so whichever the region server stops first fails.
1787        let stop_calls = Arc::new(AtomicUsize::new(0));
1788        for name in ["failing-a", "failing-b"] {
1789            let stop_calls = stop_calls.clone();
1790            let (engine, _) = MockRegionEngine::with_custom_apply_fn(name, move |engine| {
1791                engine.handle_stop_mock_fn = Some(Box::new(move || {
1792                    stop_calls.fetch_add(1, Ordering::Relaxed);
1793                    error::UnexpectedSnafu {
1794                        violated: "stop failed",
1795                    }
1796                    .fail()
1797                }));
1798            });
1799            datanode.region_server().register_engine(engine);
1800        }
1801
1802        let err = datanode.shutdown().await.unwrap_err();
1803
1804        assert_matches!(err, Error::StopRegionEngine { .. });
1805        assert_eq!(2, stop_calls.load(Ordering::Relaxed));
1806        assert!(is_stopped(&log_store).await);
1807    }
1808
1809    #[tokio::test]
1810    async fn test_shutdown_closes_raft_engine_wal_of_written_region() {
1811        common_telemetry::init_default_ut_logging();
1812        let data_home = create_temp_dir("raft-engine-wal-written-region");
1813        let mut datanode = datanode_builder(
1814            data_home.path().to_str().unwrap(),
1815            DatanodeWalConfig::RaftEngine(RaftEngineConfig::default()),
1816            Arc::new(MemoryKvBackend::new()),
1817        )
1818        .build()
1819        .await
1820        .unwrap();
1821        let region_server = datanode.region_server();
1822        let request = CreateRequestBuilder::new().build();
1823        let schema = rows_schema(&request);
1824        region_server
1825            .handle_request(REGION, RegionRequest::Create(request))
1826            .await
1827            .unwrap();
1828        put_rows(&region_server, &schema, 0..3, false).await;
1829
1830        datanode.shutdown().await.unwrap();
1831
1832        // The open region and its writes must not keep the WAL referenced: if the datanode does
1833        // not hold the last reference, it only stops the WAL and the reopen below fails.
1834        assert_raft_engine_wal_reopens(data_home.path().to_str().unwrap()).await;
1835    }
1836
1837    #[tokio::test]
1838    async fn test_shutdown_closes_raft_engine_wal_after_engine_failures() {
1839        common_telemetry::init_default_ut_logging();
1840        let data_home = create_temp_dir("raft-engine-wal-shutdown-failure");
1841        let builder = datanode_builder(
1842            data_home.path().to_str().unwrap(),
1843            DatanodeWalConfig::RaftEngine(RaftEngineConfig::default()),
1844            Arc::new(MemoryKvBackend::new()),
1845        );
1846        let mut datanode = builder.build().await.unwrap();
1847        let (engine, _) = MockRegionEngine::with_custom_apply_fn("failing", |engine| {
1848            engine.handle_stop_mock_fn = Some(Box::new(|| {
1849                error::UnexpectedSnafu {
1850                    violated: "stop failed",
1851                }
1852                .fail()
1853            }));
1854        });
1855        datanode.region_server().register_engine(engine);
1856
1857        let err = datanode.shutdown().await.unwrap_err();
1858
1859        assert_matches!(err, Error::StopRegionEngine { .. });
1860        assert_raft_engine_wal_reopens(data_home.path().to_str().unwrap()).await;
1861    }
1862
1863    #[tokio::test]
1864    async fn test_failed_build_stops_object_store_log_store() {
1865        common_telemetry::init_default_ut_logging();
1866        let data_home = create_temp_dir("object-store-wal-failed-build");
1867        let kv_backend = Arc::new(MemoryKvBackend::new());
1868        // Regions registered in the metadata but missing on disk fail the region open
1869        // that runs after the log store and the engines are built.
1870        setup_table_datanode(&(kv_backend.clone() as _)).await;
1871        let mut builder = datanode_builder(
1872            data_home.path().to_str().unwrap(),
1873            DatanodeWalConfig::ObjectStore(ObjectStoreWalConfig::default()),
1874            kv_backend,
1875        );
1876
1877        assert!(builder.try_build().await.is_err());
1878
1879        let log_store = builder.object_store_log_store.as_ref().unwrap();
1880        assert!(is_stopped(log_store).await);
1881    }
1882
1883    #[tokio::test]
1884    async fn test_failed_build_closes_raft_engine_log_store() {
1885        common_telemetry::init_default_ut_logging();
1886        let data_home = create_temp_dir("raft-engine-wal-failed-build");
1887        let kv_backend = Arc::new(MemoryKvBackend::new());
1888        // Regions registered in the metadata but missing on disk fail the region open
1889        // that runs after the log store and the engines are built.
1890        setup_table_datanode(&(kv_backend.clone() as _)).await;
1891        let mut builder = datanode_builder(
1892            data_home.path().to_str().unwrap(),
1893            DatanodeWalConfig::RaftEngine(RaftEngineConfig::default()),
1894            kv_backend,
1895        );
1896
1897        assert!(builder.try_build().await.is_err());
1898
1899        assert!(builder.region_server.is_none());
1900        assert!(builder.raft_engine_log_store.is_none());
1901        assert_raft_engine_wal_reopens(data_home.path().to_str().unwrap()).await;
1902    }
1903
1904    const REGION: RegionId = RegionId::new(1024, 1);
1905
1906    fn enqueued_wal_config(flush_interval: Duration) -> ObjectStoreWalConfig {
1907        ObjectStoreWalConfig {
1908            ack_mode: AckMode::Enqueued,
1909            flush_interval,
1910            ..Default::default()
1911        }
1912    }
1913
1914    /// The options that put a region on the WAL of a datanode with `config`.
1915    fn region_wal_options(config: &ObjectStoreWalConfig) -> HashMap<String, String> {
1916        let prefix = config.node_prefix(0, STANDALONE_GENERATION);
1917        let wal_options =
1918            serde_json::to_string(&WalOptions::ObjectStore(ObjectStoreWalOptions::new(prefix)))
1919                .unwrap();
1920        HashMap::from([(WAL_OPTIONS_KEY.to_string(), wal_options)])
1921    }
1922
1923    async fn put_rows(
1924        region_server: &RegionServer,
1925        schema: &[ColumnSchema],
1926        rows: Range<usize>,
1927        skip_wal: bool,
1928    ) {
1929        if rows.is_empty() {
1930            return;
1931        }
1932        region_server
1933            .handle_request(
1934                REGION,
1935                RegionRequest::Put(RegionPutRequest {
1936                    rows: Rows {
1937                        schema: schema.to_vec(),
1938                        rows: build_rows(rows.start, rows.end),
1939                    },
1940                    skip_wal,
1941                    hint: None,
1942                    partition_expr_version: None,
1943                }),
1944            )
1945            .await
1946            .unwrap();
1947    }
1948
1949    /// Builds a datanode on an `enqueued` object store WAL with
1950    /// `flush_interval` and creates [`REGION`] on it with `durable_rows`
1951    /// written through the WAL and uploaded. Then makes every later create
1952    /// of the WAL fail with a transient error if `uploads_fail`, and writes
1953    /// `rows` through the WAL and `skipped_rows` past it. Returns the datanode
1954    /// and its log store.
1955    async fn datanode_with_object_store_wal(
1956        data_home: &str,
1957        flush_interval: Duration,
1958        durable_rows: usize,
1959        uploads_fail: bool,
1960        rows: usize,
1961        skipped_rows: usize,
1962    ) -> (Datanode, Arc<ObjectStoreLogStore>) {
1963        let config = enqueued_wal_config(flush_interval);
1964        let prefix = config.node_prefix(0, STANDALONE_GENERATION);
1965        let mut request = CreateRequestBuilder::new().build();
1966        request.options.extend(region_wal_options(&config));
1967        let datanode = object_store_wal_builder(data_home, config)
1968            .build()
1969            .await
1970            .unwrap();
1971        let log_store = datanode.object_store_log_store.clone().unwrap();
1972        let region_server = datanode.region_server();
1973        let schema = rows_schema(&request);
1974        region_server
1975            .handle_request(REGION, RegionRequest::Create(request))
1976            .await
1977            .unwrap();
1978
1979        put_rows(&region_server, &schema, 0..durable_rows, false).await;
1980        if durable_rows > 0 {
1981            let provider = Provider::object_store_provider(REGION, prefix);
1982            tokio::time::timeout(WAIT, async {
1983                while log_store.latest_entry_id(&provider).unwrap() == 0 {
1984                    tokio::time::sleep(Duration::from_millis(10)).await;
1985                }
1986            })
1987            .await
1988            .expect("the rows must become durable");
1989        }
1990        if uploads_fail {
1991            log_store.fail_creates();
1992        }
1993        let skipped = durable_rows + rows;
1994        put_rows(&region_server, &schema, durable_rows..skipped, false).await;
1995        put_rows(
1996            &region_server,
1997            &schema,
1998            skipped..skipped + skipped_rows,
1999            true,
2000        )
2001        .await;
2002        (datanode, log_store)
2003    }
2004
2005    /// Restarts a datanode on the files of `data_home`, opens [`REGION`] and
2006    /// returns the number of rows it holds after replaying the WAL.
2007    async fn rows_after_restart(data_home: &str) -> u64 {
2008        let config = enqueued_wal_config(ObjectStoreWalConfig::default().flush_interval);
2009        let options = region_wal_options(&config);
2010        let mut datanode = object_store_wal_builder(data_home, config)
2011            .build()
2012            .await
2013            .unwrap();
2014        let region_server = datanode.region_server();
2015        region_server
2016            .handle_request(
2017                REGION,
2018                RegionRequest::Open(RegionOpenRequest {
2019                    engine: MITO_ENGINE_NAME.to_string(),
2020                    table_dir: CreateRequestBuilder::new().build().table_dir,
2021                    path_type: PathType::Bare,
2022                    options,
2023                    skip_wal_replay: false,
2024                    checkpoint: None,
2025                    requirements: Default::default(),
2026                }),
2027            )
2028            .await
2029            .unwrap();
2030        let rows = region_server.region_statistic(REGION).unwrap().num_rows;
2031        datanode.shutdown().await.unwrap();
2032        rows
2033    }
2034
2035    /// Returns the Parquet files under `dir` and its subdirectories.
2036    fn sst_files(dir: &Path) -> Vec<String> {
2037        let mut files = Vec::new();
2038        for entry in std::fs::read_dir(dir).unwrap() {
2039            let path = entry.unwrap().path();
2040            if path.is_dir() {
2041                files.extend(sst_files(&path));
2042            } else if path.extension().is_some_and(|ext| ext == "parquet") {
2043                files.push(path.display().to_string());
2044            }
2045        }
2046        files
2047    }
2048
2049    /// A flush that waits for WAL uploads that keep failing does not hold up
2050    /// the shutdown: the WAL stops first and drops its backlog, which fails
2051    /// the flush with the stopped error before it writes its manifest edit,
2052    /// and the shutdown reports the dropped backlog. After a restart the rows
2053    /// that were durable before the outage replay once.
2054    #[tokio::test]
2055    async fn test_shutdown_while_a_flush_waits_for_failing_wal_uploads() {
2056        common_telemetry::init_default_ut_logging();
2057        let data_home = create_temp_dir("object-store-wal-shutdown-flush");
2058        let data_home = data_home.path().to_str().unwrap();
2059        let (mut datanode, log_store) = datanode_with_object_store_wal(
2060            data_home,
2061            ObjectStoreWalConfig::default().flush_interval,
2062            3,
2063            true,
2064            2,
2065            0,
2066        )
2067        .await;
2068        let flush = {
2069            let region_server = datanode.region_server();
2070            tokio::spawn(async move {
2071                region_server
2072                    .handle_request(REGION, RegionRequest::Flush(Default::default()))
2073                    .await
2074            })
2075        };
2076        tokio::time::timeout(WAIT, log_store.wait_for_durability_waits(1))
2077            .await
2078            .unwrap()
2079            .unwrap();
2080
2081        let err = tokio::time::timeout(WAIT, datanode.shutdown())
2082            .await
2083            .expect("the shutdown must finish")
2084            .unwrap_err();
2085
2086        assert_matches!(err, Error::ShutdownInstance { .. });
2087        assert!(
2088            format!("{err:?}").contains("injected create failure"),
2089            "{err:?}"
2090        );
2091        let flush_err = flush.await.unwrap().unwrap_err();
2092        assert!(
2093            format!("{flush_err:?}").contains("Object store WAL log store is stopped"),
2094            "{flush_err:?}"
2095        );
2096        assert_eq!(Vec::<String>::new(), sst_files(Path::new(data_home)));
2097        assert!(is_stopped(&log_store).await);
2098        drop(datanode);
2099
2100        assert_eq!(3, rows_after_restart(data_home).await);
2101    }
2102
2103    /// A region that holds writes that skipped the WAL is flushed when it
2104    /// closes. When the WAL's uploads keep failing, the WAL stops first and
2105    /// drops its backlog, so that flush fails at once instead of waiting: the
2106    /// close fails, the writes that skipped the WAL are lost, and the shutdown
2107    /// returns the dropped backlog.
2108    #[tokio::test]
2109    async fn test_shutdown_closes_a_region_after_the_wal_dropped_its_backlog() {
2110        common_telemetry::init_default_ut_logging();
2111        let data_home = create_temp_dir("object-store-wal-shutdown-closing-flush");
2112        let (mut datanode, log_store) = datanode_with_object_store_wal(
2113            data_home.path().to_str().unwrap(),
2114            ObjectStoreWalConfig::default().flush_interval,
2115            0,
2116            true,
2117            3,
2118            2,
2119        )
2120        .await;
2121
2122        let err = tokio::time::timeout(WAIT, datanode.shutdown())
2123            .await
2124            .expect("the shutdown must finish")
2125            .unwrap_err();
2126
2127        assert_matches!(err, Error::ShutdownInstance { .. });
2128        assert!(
2129            format!("{err:?}").contains("injected create failure"),
2130            "{err:?}"
2131        );
2132        assert_eq!(Vec::<String>::new(), sst_files(data_home.path()));
2133        assert!(is_stopped(&log_store).await);
2134    }
2135
2136    /// With a healthy WAL whose seal interval is longer than any test, the
2137    /// close-time flush of a region that holds writes that skipped the WAL
2138    /// commits: the WAL's stop seals and uploads its open batch before the
2139    /// region closes, so the flush finds its entries durable, and every row
2140    /// survives a restart.
2141    #[tokio::test]
2142    async fn test_shutdown_flushes_a_region_while_the_wal_is_healthy() {
2143        common_telemetry::init_default_ut_logging();
2144        let data_home = create_temp_dir("object-store-wal-shutdown-healthy-flush");
2145        let data_home = data_home.path().to_str().unwrap();
2146        let (mut datanode, log_store) =
2147            datanode_with_object_store_wal(data_home, Duration::from_secs(3600), 0, false, 3, 2)
2148                .await;
2149
2150        tokio::time::timeout(WAIT, datanode.shutdown())
2151            .await
2152            .expect("the shutdown must finish")
2153            .unwrap();
2154        assert_eq!(1, sst_files(Path::new(data_home)).len());
2155        assert!(is_stopped(&log_store).await);
2156        drop(datanode);
2157
2158        assert_eq!(5, rows_after_restart(data_home).await);
2159    }
2160
2161    #[test]
2162    fn test_node_prefix_passes_validation() {
2163        let config = ObjectStoreWalConfig::default();
2164        let derived = ObjectStoreWalConfig {
2165            prefix: config.node_prefix(0, STANDALONE_GENERATION),
2166            ..config
2167        };
2168        validate_object_store_wal_config(&derived).unwrap();
2169    }
2170
2171    #[tokio::test]
2172    async fn test_build_store_engines_registers_relational_engine() {
2173        common_telemetry::init_default_ut_logging();
2174        let kv_backend = Arc::new(MemoryKvBackend::new());
2175        let cache_registry = Arc::new(build_datanode_layered_cache_registry(kv_backend.clone()));
2176        let schema_cache: SchemaCacheRef = cache_registry.get().unwrap();
2177        let table_id_schema_cache: TableSchemaCacheRef = cache_registry.get().unwrap();
2178        let schema_metadata_manager = Arc::new(SchemaMetadataManager::new(
2179            table_id_schema_cache,
2180            schema_cache,
2181        ));
2182        let file_ref_manager = Arc::new(FileReferenceManager::with_gc_enabled(Some(0), false));
2183
2184        let data_home = create_temp_dir("build-store-engines");
2185        let mut opts = DatanodeOptions {
2186            node_id: Some(0),
2187            wal: DatanodeWalConfig::Noop,
2188            enable_telemetry: false,
2189            ..Default::default()
2190        };
2191        opts.storage.data_home = data_home.path().to_str().unwrap().to_string();
2192
2193        let object_store_manager =
2194            DatanodeBuilder::build_object_store_manager(&opts.storage, opts.node_id.unwrap())
2195                .await
2196                .unwrap();
2197
2198        let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend);
2199        let engines = builder
2200            .build_store_engines(
2201                object_store_manager,
2202                schema_metadata_manager,
2203                file_ref_manager,
2204                Plugins::default(),
2205            )
2206            .await
2207            .unwrap();
2208
2209        let names: Vec<&str> = engines.iter().map(|engine| engine.name()).collect();
2210        assert_eq!(engines.len(), 4, "expected 4 engines, got: {names:?}");
2211        for expected in [MITO_ENGINE_NAME, "metric", "file", "relational"] {
2212            assert!(
2213                names.contains(&expected),
2214                "engine {expected:?} not registered, got: {names:?}"
2215            );
2216        }
2217    }
2218}