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 Ok(())
784}
785
786async fn open_all_regions(
788 region_server: RegionServer,
789 region_open_requests: RegionOpenRequests,
790 open_with_writable: bool,
791 init_regions_parallelism: usize,
792 ignore_nonexistent_region: bool,
793) -> Result<()> {
794 let RegionOpenRequests {
795 leader_regions,
796 #[cfg(feature = "enterprise")]
797 follower_regions,
798 } = region_open_requests;
799
800 let leader_region_num = leader_regions.len();
801 info!("going to open {} region(s)", leader_region_num);
802 let now = Instant::now();
803 let open_regions = region_server
804 .handle_batch_open_requests(
805 init_regions_parallelism,
806 leader_regions,
807 ignore_nonexistent_region,
808 )
809 .await?;
810 info!(
811 "Opened {} regions in {:?}",
812 open_regions.len(),
813 now.elapsed()
814 );
815 if !ignore_nonexistent_region {
816 ensure!(
817 open_regions.len() == leader_region_num,
818 error::UnexpectedSnafu {
819 violated: format!(
820 "Expected to open {} of regions, only {} of regions has opened",
821 leader_region_num,
822 open_regions.len()
823 )
824 }
825 );
826 } else if open_regions.len() != leader_region_num {
827 warn!(
828 "ignore nonexistent region, expected to open {} of regions, only {} of regions has opened",
829 leader_region_num,
830 open_regions.len()
831 );
832 }
833
834 for region_id in open_regions {
835 if open_with_writable {
836 let res = region_server.set_region_role(region_id, RegionRole::Leader);
837 match res {
838 Ok(_) => {
839 if let SetRegionRoleStateResponse::InvalidTransition(err) = region_server
841 .set_region_role_state_gracefully(
842 region_id,
843 SettableRegionRoleState::Leader,
844 )
845 .await?
846 {
847 error!(err; "failed to convert region {region_id} to leader");
848 }
849 }
850 Err(e) => {
851 error!(e; "failed to convert region {region_id} to leader");
852 }
853 }
854 }
855 }
856
857 #[cfg(feature = "enterprise")]
858 if !follower_regions.is_empty() {
859 use tokio::time::Instant;
860
861 let follower_region_num = follower_regions.len();
862 info!("going to open {} follower region(s)", follower_region_num);
863
864 let now = Instant::now();
865 let open_regions = region_server
866 .handle_batch_open_requests(
867 init_regions_parallelism,
868 follower_regions,
869 ignore_nonexistent_region,
870 )
871 .await?;
872 info!(
873 "Opened {} follower regions in {:?}",
874 open_regions.len(),
875 now.elapsed()
876 );
877
878 if !ignore_nonexistent_region {
879 ensure!(
880 open_regions.len() == follower_region_num,
881 error::UnexpectedSnafu {
882 violated: format!(
883 "Expected to open {} of follower regions, only {} of regions has opened",
884 follower_region_num,
885 open_regions.len()
886 )
887 }
888 );
889 } else if open_regions.len() != follower_region_num {
890 warn!(
891 "ignore nonexistent region, expected to open {} of follower regions, only {} of regions has opened",
892 follower_region_num,
893 open_regions.len()
894 );
895 }
896 }
897
898 info!("all regions are opened");
899
900 Ok(())
901}
902
903#[cfg(test)]
904mod tests {
905 use std::assert_matches;
906 use std::collections::{BTreeMap, HashMap};
907 use std::io::Write;
908 use std::path::Path;
909 use std::sync::Arc;
910 use std::time::Duration;
911
912 use cache::build_datanode_cache_registry;
913 use common_base::Plugins;
914 use common_base::readable_size::ReadableSize;
915 use common_config::Configurable;
916 use common_error::ext::ErrorExt;
917 use common_error::status_code::StatusCode;
918 use common_meta::cache::LayeredCacheRegistryBuilder;
919 use common_meta::key::RegionRoleSet;
920 use common_meta::key::datanode_table::DatanodeTableManager;
921 use common_meta::kv_backend::KvBackendRef;
922 use common_meta::kv_backend::memory::MemoryKvBackend;
923 use common_test_util::temp_dir::{create_named_temp_file, create_temp_dir};
924 use common_wal::config::DatanodeWalConfig;
925 use common_wal::config::object_store::{ObjectStoreWalConfig, STANDALONE_GENERATION};
926 use mito2::engine::MITO_ENGINE_NAME;
927 use store_api::region_request::RegionRequest;
928 use store_api::storage::RegionId;
929
930 use crate::config::{DatanodeOptions, RegionEngineConfig};
931 use crate::datanode::{
932 DatanodeBuilder, validate_object_store_wal_config, validate_region_engine_config,
933 };
934 use crate::error::Error;
935 use crate::tests::{MockRegionEngine, mock_region_server};
936
937 async fn setup_table_datanode(kv: &KvBackendRef) {
938 let mgr = DatanodeTableManager::new(kv.clone());
939 let txn = mgr
940 .build_create_txn(
941 1028,
942 MITO_ENGINE_NAME,
943 "foo/bar/weny",
944 HashMap::from([("foo".to_string(), "bar".to_string())]),
945 HashMap::default(),
946 BTreeMap::from([(0, RegionRoleSet::new(vec![0, 1, 2], vec![]))]),
947 )
948 .unwrap();
949
950 let r = kv.txn(txn).await.unwrap();
951 assert!(r.succeeded);
952 }
953
954 #[tokio::test]
955 async fn test_initialize_region_server() {
956 common_telemetry::init_default_ut_logging();
957 let mut mock_region_server = mock_region_server();
958 let (mock_region, mut mock_region_handler) = MockRegionEngine::new(MITO_ENGINE_NAME);
959
960 mock_region_server.register_engine(mock_region.clone());
961
962 let kv_backend = Arc::new(MemoryKvBackend::new());
963 let layered_cache_registry = Arc::new(
964 LayeredCacheRegistryBuilder::default()
965 .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
966 .build(),
967 );
968
969 let mut builder = DatanodeBuilder::new(
970 DatanodeOptions {
971 node_id: Some(0),
972 ..Default::default()
973 },
974 Plugins::default(),
975 kv_backend.clone(),
976 );
977 builder.with_cache_registry(layered_cache_registry);
978 setup_table_datanode(&(kv_backend as _)).await;
979
980 builder
981 .initialize_region_server(&mock_region_server, false)
982 .await
983 .unwrap();
984
985 for i in 0..3 {
986 let (region_id, req) = mock_region_handler.recv().await.unwrap();
987 assert_eq!(region_id, RegionId::new(1028, i));
988 if let RegionRequest::Open(req) = req {
989 assert_eq!(
990 req.options,
991 HashMap::from([("foo".to_string(), "bar".to_string())])
992 )
993 } else {
994 unreachable!()
995 }
996 }
997
998 assert_matches!(
999 mock_region_handler.try_recv(),
1000 Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1001 );
1002 }
1003
1004 fn datanode_builder(
1005 data_home: &str,
1006 wal: DatanodeWalConfig,
1007 kv_backend: KvBackendRef,
1008 ) -> DatanodeBuilder {
1009 let layered_cache_registry = Arc::new(
1010 LayeredCacheRegistryBuilder::default()
1011 .add_cache_registry(build_datanode_cache_registry(kv_backend.clone()))
1012 .build(),
1013 );
1014 let mut opts = DatanodeOptions {
1015 node_id: Some(0),
1016 wal,
1017 enable_telemetry: false,
1018 ..Default::default()
1019 };
1020 opts.storage.data_home = data_home.to_string();
1021
1022 let mut builder = DatanodeBuilder::new(opts, Plugins::default(), kv_backend);
1023 builder.with_cache_registry(layered_cache_registry);
1024 builder
1025 }
1026
1027 fn object_store_wal_builder(data_home: &str, config: ObjectStoreWalConfig) -> DatanodeBuilder {
1028 datanode_builder(
1029 data_home,
1030 DatanodeWalConfig::ObjectStore(config),
1031 Arc::new(MemoryKvBackend::new()),
1032 )
1033 }
1034
1035 async fn build_err(builder: DatanodeBuilder) -> Error {
1036 match builder.build().await {
1037 Ok(_) => panic!("build must fail"),
1038 Err(err) => err,
1039 }
1040 }
1041
1042 fn is_empty_dir(dir: &Path) -> bool {
1043 std::fs::read_dir(dir).unwrap().next().is_none()
1044 }
1045
1046 #[test]
1047 fn test_validate_unique_region_engine_config() {
1048 let mut opts = DatanodeOptions::default();
1049 validate_region_engine_config(&opts.region_engine).unwrap();
1050
1051 opts.region_engine
1052 .push(RegionEngineConfig::Metric(Default::default()));
1053 validate_region_engine_config(&opts.region_engine).unwrap();
1054
1055 validate_region_engine_config(&[]).unwrap();
1057 }
1058
1059 #[tokio::test]
1060 async fn test_build_rejects_duplicate_region_engine_config() {
1061 for (engine, first_index) in [("mito", 0), ("file", 1), ("metric", 2)] {
1062 let mut file = create_named_temp_file();
1063 write!(
1064 file,
1065 r#"
1066 [[region_engine]]
1067 [region_engine.mito]
1068 global_write_buffer_size = "4GiB"
1069 min_compaction_interval = "600s"
1070
1071 [[region_engine]]
1072 [region_engine.file]
1073
1074 [[region_engine]]
1075 [region_engine.metric]
1076
1077 [[region_engine]]
1078 [region_engine.{engine}]
1079 "#
1080 )
1081 .unwrap();
1082 let opts = DatanodeOptions::load_layered_options(
1083 Some(file.path().to_str().unwrap()),
1084 "DATANODE_DUPLICATE_REGION_ENGINE_UT",
1085 )
1086 .unwrap();
1087 let data_home = create_temp_dir("duplicate-region-engine-config");
1088 let mut builder = datanode_builder(
1089 data_home.path().to_str().unwrap(),
1090 DatanodeWalConfig::default(),
1091 Arc::new(MemoryKvBackend::new()),
1092 );
1093 builder.opts.region_engine = opts.region_engine;
1094 let err = build_err(builder).await;
1095
1096 let Error::DuplicateRegionEngineConfig {
1097 engine: actual_engine,
1098 first_index: actual_first_index,
1099 duplicate_index,
1100 ..
1101 } = &err
1102 else {
1103 panic!("unexpected error for {engine}: {err:?}");
1104 };
1105 assert_eq!(*actual_engine, engine);
1106 assert_eq!(*actual_first_index, first_index);
1107 assert_eq!(*duplicate_index, 3);
1108 assert_eq!(err.status_code(), StatusCode::InvalidArguments);
1109 assert!(is_empty_dir(data_home.path()));
1111 }
1112 }
1113
1114 #[tokio::test]
1115 async fn test_build_rejects_invalid_object_store_wal_config() {
1116 let cases = [
1117 (
1118 ObjectStoreWalConfig {
1119 prefix: String::new(),
1120 ..Default::default()
1121 },
1122 "prefix",
1123 ),
1124 (
1125 ObjectStoreWalConfig {
1126 prefix: "/wal".to_string(),
1127 ..Default::default()
1128 },
1129 "prefix",
1130 ),
1131 (
1132 ObjectStoreWalConfig {
1133 prefix: "wal/../other".to_string(),
1134 ..Default::default()
1135 },
1136 "prefix",
1137 ),
1138 (
1139 ObjectStoreWalConfig {
1140 prefix: "wal//objects".to_string(),
1141 ..Default::default()
1142 },
1143 "prefix",
1144 ),
1145 (
1146 ObjectStoreWalConfig {
1147 prefix: "wal/".to_string(),
1148 ..Default::default()
1149 },
1150 "prefix",
1151 ),
1152 (
1153 ObjectStoreWalConfig {
1154 flush_interval: Duration::from_millis(9),
1155 ..Default::default()
1156 },
1157 "flush_interval",
1158 ),
1159 (
1160 ObjectStoreWalConfig {
1161 max_batch_bytes: ReadableSize(0),
1162 ..Default::default()
1163 },
1164 "max_batch_bytes",
1165 ),
1166 ];
1167
1168 for (config, expected_field) in cases {
1169 let expected_value = match expected_field {
1170 "prefix" => config.prefix.clone(),
1171 "flush_interval" => format!("{:?}", config.flush_interval),
1172 _ => config.max_batch_bytes.to_string(),
1173 };
1174 let data_home = create_temp_dir("object-store-wal-invalid-config");
1175 let builder = object_store_wal_builder(data_home.path().to_str().unwrap(), config);
1176 let err = build_err(builder).await;
1177
1178 let Error::InvalidObjectStoreWalConfig { field, value, .. } = &err else {
1179 panic!("unexpected error for {expected_field}: {err:?}");
1180 };
1181 assert_eq!(*field, expected_field);
1182 assert_eq!(*value, expected_value);
1183 assert_eq!(StatusCode::InvalidArguments, err.status_code());
1184 assert!(is_empty_dir(data_home.path()));
1186 }
1187 }
1188
1189 #[tokio::test]
1190 async fn test_build_rejects_object_store_wal_as_not_supported() {
1191 let data_home = create_temp_dir("object-store-wal-not-supported");
1192 let builder = object_store_wal_builder(
1193 data_home.path().to_str().unwrap(),
1194 ObjectStoreWalConfig::default(),
1195 );
1196 let err = build_err(builder).await;
1197
1198 assert_matches!(err, Error::ObjectStoreWalNotSupported { .. });
1199 assert_eq!(StatusCode::Unsupported, err.status_code());
1200 assert!(!data_home.path().join("wal").exists());
1202 }
1203
1204 #[test]
1205 fn test_node_prefix_passes_validation() {
1206 let config = ObjectStoreWalConfig::default();
1207 let derived = ObjectStoreWalConfig {
1208 prefix: config.node_prefix(0, STANDALONE_GENERATION),
1209 ..config
1210 };
1211 validate_object_store_wal_config(&derived).unwrap();
1212 }
1213}