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