1mod alter;
16mod bulk_insert;
17mod catchup;
18mod close;
19mod create;
20mod drop;
21mod flush;
22mod open;
23mod options;
24mod put;
25mod read;
26mod region_metadata;
27mod staging;
28mod state;
29mod sync;
30
31use std::any::Any;
32use std::collections::HashMap;
33use std::sync::{Arc, RwLock};
34use std::time::Instant;
35
36use api::region::RegionResponse;
37use async_trait::async_trait;
38use common_error::ext::{BoxedError, ErrorExt};
39use common_error::status_code::StatusCode;
40use common_runtime::RepeatedTask;
41use mito2::engine::MitoEngine;
42pub(crate) use options::IndexOptions;
43use snafu::{OptionExt, ResultExt};
44pub(crate) use state::MetricEngineState;
45use store_api::metadata::RegionMetadataRef;
46use store_api::metric_engine_consts::METRIC_ENGINE_NAME;
47use store_api::region_engine::{
48 BatchResponses, RegionEngine, RegionRole, RegionScannerRef, RegionStatistic,
49 RemapManifestsRequest, RemapManifestsResponse, SetRegionRoleStateResponse,
50 SetRegionRoleStateSuccess, SettableRegionRoleState, SyncRegionFromRequest,
51 SyncRegionFromResponse,
52};
53use store_api::region_request::{
54 AffectedRows, BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionPutRequest,
55 RegionRequest, RegionTruncateRequest,
56};
57use store_api::storage::{RegionId, ScanRequest, SequenceNumber};
58
59use crate::config::EngineConfig;
60use crate::data_region::DataRegion;
61use crate::error::{
62 self, Error, Result, StartRepeatedTaskSnafu, UnsupportedRegionRequestSnafu,
63 UnsupportedRemapManifestsRequestSnafu,
64};
65use crate::metadata_region::MetadataRegion;
66use crate::repeated_task::FlushMetadataRegionTask;
67use crate::row_modifier::RowModifier;
68use crate::utils::{self, get_region_statistic};
69
70#[cfg_attr(doc, aquamarine::aquamarine)]
71#[derive(Clone)]
133pub struct MetricEngine {
134 inner: Arc<MetricEngineInner>,
135}
136
137#[async_trait]
138impl RegionEngine for MetricEngine {
139 fn name(&self) -> &str {
141 METRIC_ENGINE_NAME
142 }
143
144 async fn handle_batch_open_requests(
145 &self,
146 parallelism: usize,
147 requests: Vec<(RegionId, RegionOpenRequest)>,
148 ) -> Result<BatchResponses, BoxedError> {
149 self.inner
150 .handle_batch_open_requests(parallelism, requests)
151 .await
152 .map_err(BoxedError::new)
153 }
154
155 async fn handle_batch_catchup_requests(
156 &self,
157 parallelism: usize,
158 requests: Vec<(RegionId, RegionCatchupRequest)>,
159 ) -> Result<BatchResponses, BoxedError> {
160 self.inner
161 .handle_batch_catchup_requests(parallelism, requests)
162 .await
163 .map_err(BoxedError::new)
164 }
165
166 async fn handle_batch_ddl_requests(
167 &self,
168 batch_request: BatchRegionDdlRequest,
169 ) -> Result<RegionResponse, BoxedError> {
170 match batch_request {
171 BatchRegionDdlRequest::Create(requests) => {
172 let mut extension_return_value = HashMap::new();
173 let rows = self
174 .inner
175 .create_regions(requests, &mut extension_return_value)
176 .await
177 .map_err(BoxedError::new)?;
178
179 Ok(RegionResponse {
180 affected_rows: rows,
181 extensions: extension_return_value,
182 metadata: Vec::new(),
183 })
184 }
185 BatchRegionDdlRequest::Alter(requests) => {
186 let mut extension_return_value = HashMap::new();
187 let rows = self
188 .inner
189 .alter_regions(requests, &mut extension_return_value)
190 .await
191 .map_err(BoxedError::new)?;
192
193 Ok(RegionResponse {
194 affected_rows: rows,
195 extensions: extension_return_value,
196 metadata: Vec::new(),
197 })
198 }
199 BatchRegionDdlRequest::Drop(requests) => {
200 self.handle_requests(
201 requests
202 .into_iter()
203 .map(|(region_id, req)| (region_id, RegionRequest::Drop(req))),
204 )
205 .await
206 }
207 }
208 }
209
210 async fn handle_request(
212 &self,
213 region_id: RegionId,
214 request: RegionRequest,
215 ) -> Result<RegionResponse, BoxedError> {
216 let mut extension_return_value = HashMap::new();
217
218 let result = match request {
219 RegionRequest::EnterStaging(_) => {
220 if self.inner.is_physical_region(region_id) {
221 self.handle_enter_staging_request(region_id, request).await
222 } else {
223 UnsupportedRegionRequestSnafu { request }.fail()
224 }
225 }
226 RegionRequest::ApplyStagingManifest(_) => {
227 if self.inner.is_physical_region(region_id) {
228 return self.inner.mito.handle_request(region_id, request).await;
229 } else {
230 UnsupportedRegionRequestSnafu { request }.fail()
231 }
232 }
233 RegionRequest::Put(put) => self.inner.put_region(region_id, put).await,
234 RegionRequest::Create(create) => {
235 self.inner
236 .create_regions(vec![(region_id, create)], &mut extension_return_value)
237 .await
238 }
239 RegionRequest::Drop(drop) => self.inner.drop_region(region_id, drop).await,
240 RegionRequest::Open(open) => self.inner.open_region(region_id, open).await,
241 RegionRequest::CleanUp(clean_up) => {
242 self.inner.clean_up_region(region_id, clean_up).await
243 }
244 RegionRequest::Close(close) => self.inner.close_region(region_id, close).await,
245 RegionRequest::Alter(alter) => {
246 self.inner
247 .alter_regions(vec![(region_id, alter)], &mut extension_return_value)
248 .await
249 }
250 RegionRequest::Compact(_) => {
251 if self.inner.is_physical_region(region_id) {
252 self.inner
253 .mito
254 .handle_request(region_id, request)
255 .await
256 .context(error::MitoFlushOperationSnafu)
257 .map(|response| response.affected_rows)
258 } else {
259 UnsupportedRegionRequestSnafu { request }.fail()
260 }
261 }
262 RegionRequest::Flush(req) => self.inner.flush_region(region_id, req).await,
263 RegionRequest::BuildIndex(_) => {
264 if self.inner.is_physical_region(region_id) {
265 self.inner
266 .mito
267 .handle_request(region_id, request)
268 .await
269 .context(error::MitoFlushOperationSnafu)
270 .map(|response| response.affected_rows)
271 } else {
272 UnsupportedRegionRequestSnafu { request }.fail()
273 }
274 }
275 RegionRequest::Truncate(RegionTruncateRequest::Unflushed) => {
276 if self.inner.is_physical_region(region_id) {
277 self.inner
278 .mito
279 .handle_request(
280 utils::to_data_region_id(region_id),
281 RegionRequest::Truncate(RegionTruncateRequest::Unflushed),
282 )
283 .await
284 .context(error::MitoTruncateOperationSnafu)
285 .map(|response| response.affected_rows)
286 } else {
287 UnsupportedRegionRequestSnafu {
288 request: RegionRequest::Truncate(RegionTruncateRequest::Unflushed),
289 }
290 .fail()
291 }
292 }
293 RegionRequest::Truncate(request) => UnsupportedRegionRequestSnafu {
294 request: RegionRequest::Truncate(request),
295 }
296 .fail(),
297 RegionRequest::Delete(delete) => self.inner.delete_region(region_id, delete).await,
298 RegionRequest::Catchup(_) => {
299 let mut response = self
300 .inner
301 .handle_batch_catchup_requests(
302 1,
303 vec![(region_id, RegionCatchupRequest::default())],
304 )
305 .await
306 .map_err(BoxedError::new)?;
307 debug_assert_eq!(response.len(), 1);
308 let (resp_region_id, response) = response
309 .pop()
310 .context(error::UnexpectedRequestSnafu {
311 reason: "expected 1 response, but got zero responses",
312 })
313 .map_err(BoxedError::new)?;
314 debug_assert_eq!(region_id, resp_region_id);
315 return response;
316 }
317 RegionRequest::BulkInserts(bulk) => {
318 self.inner.bulk_insert_region(region_id, bulk).await
319 }
320 };
321
322 result.map_err(BoxedError::new).map(|rows| RegionResponse {
323 affected_rows: rows,
324 extensions: extension_return_value,
325 metadata: Vec::new(),
326 })
327 }
328
329 async fn handle_query(
330 &self,
331 region_id: RegionId,
332 request: ScanRequest,
333 ) -> Result<RegionScannerRef, BoxedError> {
334 self.handle_query(region_id, request).await
335 }
336
337 async fn get_committed_sequence(
338 &self,
339 region_id: RegionId,
340 ) -> Result<SequenceNumber, BoxedError> {
341 self.inner
342 .get_last_seq_num(region_id)
343 .await
344 .map_err(BoxedError::new)
345 }
346
347 async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError> {
349 self.inner
350 .load_region_metadata(region_id)
351 .await
352 .map_err(BoxedError::new)
353 }
354
355 fn region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic> {
359 if self.inner.is_physical_region(region_id) {
360 get_region_statistic(&self.inner.mito, region_id)
361 } else {
362 None
363 }
364 }
365
366 async fn stop(&self) -> Result<(), BoxedError> {
368 Ok(())
370 }
371
372 fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> {
373 for x in [
375 utils::to_metadata_region_id(region_id),
376 utils::to_data_region_id(region_id),
377 ] {
378 if let Err(e) = self.inner.mito.set_region_role(x, role)
379 && e.status_code() != StatusCode::RegionNotFound
380 {
381 return Err(e);
382 }
383 }
384 Ok(())
385 }
386
387 fn set_region_role_with_lease(
388 &self,
389 region_id: RegionId,
390 role: RegionRole,
391 deadline: Instant,
392 ) -> Result<(), BoxedError> {
393 for x in [
395 utils::to_metadata_region_id(region_id),
396 utils::to_data_region_id(region_id),
397 ] {
398 if let Err(e) = self
399 .inner
400 .mito
401 .set_region_role_with_lease(x, role, deadline)
402 && e.status_code() != StatusCode::RegionNotFound
403 {
404 return Err(e);
405 }
406 }
407 Ok(())
408 }
409
410 async fn sync_region(
411 &self,
412 region_id: RegionId,
413 request: SyncRegionFromRequest,
414 ) -> Result<SyncRegionFromResponse, BoxedError> {
415 match request {
416 SyncRegionFromRequest::FromManifest(manifest_info) => self
417 .inner
418 .sync_region_from_manifest(region_id, manifest_info)
419 .await
420 .map_err(BoxedError::new),
421 SyncRegionFromRequest::FromRegion {
422 source_region_id,
423 parallelism,
424 } => {
425 if self.inner.is_physical_region(region_id) {
426 self.inner
427 .sync_region_from_region(region_id, source_region_id, parallelism)
428 .await
429 .map_err(BoxedError::new)
430 } else {
431 Err(BoxedError::new(
432 error::UnsupportedSyncRegionFromRequestSnafu { region_id }.build(),
433 ))
434 }
435 }
436 }
437 }
438
439 async fn remap_manifests(
440 &self,
441 request: RemapManifestsRequest,
442 ) -> Result<RemapManifestsResponse, BoxedError> {
443 let region_id = request.region_id;
444 if self.inner.is_physical_region(region_id) {
445 self.inner.mito.remap_manifests(request).await
446 } else {
447 Err(BoxedError::new(
448 UnsupportedRemapManifestsRequestSnafu { region_id }.build(),
449 ))
450 }
451 }
452
453 async fn set_region_role_state_gracefully(
454 &self,
455 region_id: RegionId,
456 region_role_state: SettableRegionRoleState,
457 ) -> std::result::Result<SetRegionRoleStateResponse, BoxedError> {
458 let metadata_result = match self
459 .inner
460 .mito
461 .set_region_role_state_gracefully(
462 utils::to_metadata_region_id(region_id),
463 region_role_state,
464 )
465 .await?
466 {
467 SetRegionRoleStateResponse::Success(success) => success,
468 SetRegionRoleStateResponse::NotFound => {
469 return Ok(SetRegionRoleStateResponse::NotFound);
470 }
471 SetRegionRoleStateResponse::InvalidTransition(error) => {
472 return Ok(SetRegionRoleStateResponse::InvalidTransition(error));
473 }
474 };
475
476 let data_result = match self
477 .inner
478 .mito
479 .set_region_role_state_gracefully(region_id, region_role_state)
480 .await?
481 {
482 SetRegionRoleStateResponse::Success(success) => success,
483 SetRegionRoleStateResponse::NotFound => {
484 return Ok(SetRegionRoleStateResponse::NotFound);
485 }
486 SetRegionRoleStateResponse::InvalidTransition(error) => {
487 return Ok(SetRegionRoleStateResponse::InvalidTransition(error));
488 }
489 };
490
491 Ok(SetRegionRoleStateResponse::success(
492 SetRegionRoleStateSuccess::metric(
493 data_result.last_entry_id().unwrap_or_default(),
494 metadata_result.last_entry_id().unwrap_or_default(),
495 ),
496 ))
497 }
498
499 fn role(&self, region_id: RegionId) -> Option<RegionRole> {
503 if self.inner.is_physical_region(region_id) {
504 self.inner.mito.role(region_id)
505 } else {
506 None
507 }
508 }
509
510 fn as_any(&self) -> &dyn Any {
511 self
512 }
513}
514
515#[derive(Debug)]
519pub struct PutGroupResult {
520 pub request_indices: Vec<usize>,
522 pub result: Result<AffectedRows>,
524}
525
526impl MetricEngine {
527 pub fn try_new(mito: MitoEngine, mut config: EngineConfig) -> Result<Self> {
528 let metadata_region = MetadataRegion::new(mito.clone());
529 let data_region = DataRegion::new(mito.clone());
530 let state = Arc::new(RwLock::default());
531 config.sanitize();
532 let flush_interval = config.flush_metadata_region_interval;
533 let inner = Arc::new(MetricEngineInner {
534 mito: mito.clone(),
535 metadata_region,
536 data_region,
537 state: state.clone(),
538 row_modifier: RowModifier::default(),
539 flush_task: RepeatedTask::new(
540 flush_interval,
541 Box::new(FlushMetadataRegionTask {
542 state: state.clone(),
543 mito: mito.clone(),
544 }),
545 ),
546 });
547 inner
548 .flush_task
549 .start(common_runtime::global_runtime())
550 .context(StartRepeatedTaskSnafu { name: "flush_task" })?;
551 Ok(Self { inner })
552 }
553
554 pub fn mito(&self) -> MitoEngine {
555 self.inner.mito.clone()
556 }
557
558 pub async fn put_regions_batch(
560 &self,
561 requests: impl ExactSizeIterator<Item = (RegionId, RegionPutRequest)>,
562 ) -> Vec<PutGroupResult> {
563 self.inner.put_regions_batch(requests).await
564 }
565
566 pub async fn logical_regions(&self, physical_region_id: RegionId) -> Result<Vec<RegionId>> {
568 self.inner
569 .metadata_region
570 .logical_regions(physical_region_id)
571 .await
572 }
573
574 async fn handle_query(
576 &self,
577 region_id: RegionId,
578 request: ScanRequest,
579 ) -> Result<RegionScannerRef, BoxedError> {
580 self.inner
581 .read_region(region_id, request)
582 .await
583 .map_err(BoxedError::new)
584 }
585
586 async fn handle_requests(
587 &self,
588 requests: impl IntoIterator<Item = (RegionId, RegionRequest)>,
589 ) -> Result<RegionResponse, BoxedError> {
590 let mut affected_rows = 0;
591 let mut extensions = HashMap::new();
592 for (region_id, request) in requests {
593 let response = self.handle_request(region_id, request).await?;
594 affected_rows += response.affected_rows;
595 extensions.extend(response.extensions);
596 }
597
598 Ok(RegionResponse {
599 affected_rows,
600 extensions,
601 metadata: Vec::new(),
602 })
603 }
604}
605
606fn validate_shared_column(
608 region_id: store_api::storage::RegionId,
609 existing: &store_api::metadata::ColumnMetadata,
610 proposed: &store_api::metadata::ColumnMetadata,
611) -> crate::error::Result<()> {
612 snafu::ensure!(
613 existing.column_schema.data_type == proposed.column_schema.data_type
614 && existing.semantic_type == proposed.semantic_type,
615 crate::error::InvalidRequestSnafu {
616 region_id,
617 reason: format!(
618 "incompatible shared physical column `{}`",
619 proposed.column_schema.name
620 ),
621 }
622 );
623 Ok(())
624}
625
626#[cfg(test)]
627impl MetricEngine {
628 pub async fn scan_to_stream(
629 &self,
630 region_id: RegionId,
631 request: ScanRequest,
632 ) -> Result<common_recordbatch::SendableRecordBatchStream, BoxedError> {
633 self.inner.scan_to_stream(region_id, request).await
634 }
635}
636
637struct MetricEngineInner {
638 mito: MitoEngine,
639 metadata_region: MetadataRegion,
640 data_region: DataRegion,
641 state: Arc<RwLock<MetricEngineState>>,
642 row_modifier: RowModifier,
643 flush_task: RepeatedTask<Error>,
644}
645
646#[cfg(test)]
647mod test {
648 use std::assert_matches;
649 use std::collections::HashMap;
650
651 use api::v1::Rows;
652 use common_recordbatch::RecordBatches;
653 use common_telemetry::info;
654 use common_wal::options::{KafkaWalOptions, WalOptions};
655 use mito2::sst::location::region_dir_from_table_dir;
656 use mito2::test_util::{kafka_log_store_factory, prepare_test_for_kafka_log_store};
657 use store_api::metric_engine_consts::PHYSICAL_TABLE_METADATA_KEY;
658 use store_api::mito_engine_options::WAL_OPTIONS_KEY;
659 use store_api::region_request::{
660 PathType, RegionCleanUpRequest, RegionCloseRequest, RegionDropRequest, RegionFlushRequest,
661 RegionOpenRequest, RegionPutRequest, RegionRequest, RegionTruncateRequest,
662 };
663
664 use super::*;
665 use crate::maybe_skip_kafka_log_store_integration_test;
666 use crate::test_util::{
667 TestEnv, build_rows, create_logical_region_request, row_schema_with_tags,
668 };
669
670 #[tokio::test]
671 async fn test_build_series_index_forwarding() {
672 use api::v1::region::build_index_request;
673 use store_api::region_request::RegionBuildIndexRequest;
674
675 let env = TestEnv::new().await;
676 env.init_metric_region().await;
677 let request = || {
678 RegionRequest::BuildIndex(RegionBuildIndexRequest {
679 options: Some(build_index_request::Options::SeriesIndex(Default::default())),
680 })
681 };
682 let error = env
684 .metric()
685 .handle_request(env.default_physical_region_id(), request())
686 .await
687 .unwrap_err();
688 let error = format!("{error:?}");
689 assert!(error.contains("series index is disabled"), "{error}");
690 let error = env
691 .metric()
692 .handle_request(env.default_logical_region_id(), request())
693 .await
694 .unwrap_err();
695 assert_eq!(
696 common_error::status_code::StatusCode::Unsupported,
697 error.status_code()
698 );
699 }
700
701 #[tokio::test]
702 async fn test_discard_unflushed_data_only() {
703 let env = TestEnv::new().await;
704 env.init_metric_region().await;
705
706 let engine = env.metric();
707 let mito = env.mito();
708 let physical_region_id = env.default_physical_region_id();
709 let logical_region_id = env.default_logical_region_id();
710 let data_region_id = utils::to_data_region_id(physical_region_id);
711 let metadata_region_id = utils::to_metadata_region_id(physical_region_id);
712
713 let metadata_memtable_size = mito
714 .region_statistic(metadata_region_id)
715 .unwrap()
716 .memtable_size;
717 assert!(metadata_memtable_size > 0);
718
719 engine
720 .handle_request(
721 logical_region_id,
722 RegionRequest::Put(RegionPutRequest {
723 skip_wal: false,
724 rows: Rows {
725 schema: row_schema_with_tags(&["job"]),
726 rows: build_rows(1, 5),
727 },
728 hint: None,
729 partition_expr_version: None,
730 }),
731 )
732 .await
733 .unwrap();
734 assert!(mito.region_statistic(data_region_id).unwrap().memtable_size > 0);
735
736 engine
737 .handle_request(
738 physical_region_id,
739 RegionRequest::Truncate(RegionTruncateRequest::Unflushed),
740 )
741 .await
742 .unwrap();
743
744 assert_eq!(
745 0,
746 mito.region_statistic(data_region_id).unwrap().memtable_size
747 );
748 assert_eq!(
749 metadata_memtable_size,
750 mito.region_statistic(metadata_region_id)
751 .unwrap()
752 .memtable_size
753 );
754
755 let stream = engine
756 .scan_to_stream(logical_region_id, ScanRequest::default())
757 .await
758 .unwrap();
759 let batches = RecordBatches::try_collect(stream).await.unwrap();
760 assert_eq!(
761 0,
762 batches.iter().map(|batch| batch.num_rows()).sum::<usize>()
763 );
764
765 engine
767 .handle_request(
768 logical_region_id,
769 RegionRequest::Truncate(RegionTruncateRequest::Unflushed),
770 )
771 .await
772 .unwrap_err();
773 }
774
775 #[tokio::test]
776 async fn close_open_regions() {
777 let env = TestEnv::new().await;
778 env.init_metric_region().await;
779 let engine = env.metric();
780
781 let physical_region_id = env.default_physical_region_id();
783 engine
784 .handle_request(
785 physical_region_id,
786 RegionRequest::Close(RegionCloseRequest::default()),
787 )
788 .await
789 .unwrap();
790
791 let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
793 .into_iter()
794 .collect();
795 let open_request = RegionOpenRequest {
796 engine: METRIC_ENGINE_NAME.to_string(),
797 table_dir: TestEnv::default_table_dir(),
798 path_type: PathType::Bare, options: physical_region_option,
800 skip_wal_replay: false,
801 checkpoint: None,
802 requirements: Default::default(),
803 };
804 engine
805 .handle_request(physical_region_id, RegionRequest::Open(open_request))
806 .await
807 .unwrap();
808
809 let nonexistent_region_id = RegionId::new(12313, 12);
811 engine
812 .handle_request(
813 nonexistent_region_id,
814 RegionRequest::Close(RegionCloseRequest::default()),
815 )
816 .await
817 .unwrap();
818
819 let invalid_open_request = RegionOpenRequest {
821 engine: METRIC_ENGINE_NAME.to_string(),
822 table_dir: TestEnv::default_table_dir(),
823 path_type: PathType::Bare, options: HashMap::new(),
825 skip_wal_replay: false,
826 checkpoint: None,
827 requirements: Default::default(),
828 };
829 engine
830 .handle_request(
831 nonexistent_region_id,
832 RegionRequest::Open(invalid_open_request),
833 )
834 .await
835 .unwrap();
836 }
837
838 #[tokio::test]
839 async fn test_offline_cleanup_physical_region() {
840 let env = TestEnv::new().await;
841 env.init_metric_region().await;
842 let engine = env.metric();
843 let mito = env.mito();
844 let physical_region_id = env.default_physical_region_id();
845 let metadata_region_id = crate::utils::to_metadata_region_id(physical_region_id);
846 let data_region_id = crate::utils::to_data_region_id(physical_region_id);
847
848 engine
849 .handle_request(
850 physical_region_id,
851 RegionRequest::Close(RegionCloseRequest::default()),
852 )
853 .await
854 .unwrap();
855
856 let object_store = env.get_object_store().unwrap();
857 let metadata_region_dir = region_dir_from_table_dir(
858 &TestEnv::default_table_dir(),
859 metadata_region_id,
860 PathType::Metadata,
861 );
862 let data_region_dir = region_dir_from_table_dir(
863 &TestEnv::default_table_dir(),
864 data_region_id,
865 PathType::Data,
866 );
867 assert!(object_store.exists(&metadata_region_dir).await.unwrap());
868 assert!(object_store.exists(&data_region_dir).await.unwrap());
869
870 let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
871 .into_iter()
872 .collect();
873 let clean_up_request = RegionCleanUpRequest {
874 engine: METRIC_ENGINE_NAME.to_string(),
875 table_dir: TestEnv::default_table_dir(),
876 path_type: PathType::Bare,
877 options: physical_region_option,
878 };
879 engine
880 .handle_request(physical_region_id, RegionRequest::CleanUp(clean_up_request))
881 .await
882 .unwrap();
883
884 assert!(!mito.is_region_exists(metadata_region_id));
885 assert!(!mito.is_region_exists(data_region_id));
886 assert!(!object_store.exists(&metadata_region_dir).await.unwrap());
887 assert!(!object_store.exists(&data_region_dir).await.unwrap());
888 }
889
890 #[tokio::test]
891 async fn test_role() {
892 let env = TestEnv::new().await;
893 env.init_metric_region().await;
894
895 let logical_region_id = env.default_logical_region_id();
896 let physical_region_id = env.default_physical_region_id();
897
898 assert!(env.metric().role(logical_region_id).is_none());
899 assert!(env.metric().role(physical_region_id).is_some());
900 }
901
902 #[tokio::test]
903 async fn test_region_disk_usage() {
904 let env = TestEnv::new().await;
905 env.init_metric_region().await;
906
907 let logical_region_id = env.default_logical_region_id();
908 let physical_region_id = env.default_physical_region_id();
909
910 assert!(env.metric().region_statistic(logical_region_id).is_none());
911 assert!(env.metric().region_statistic(physical_region_id).is_some());
912 }
913
914 #[tokio::test]
915 async fn test_open_region_failure() {
916 let env = TestEnv::new().await;
917 env.init_metric_region().await;
918 let physical_region_id = env.default_physical_region_id();
919
920 let metric_engine = env.metric();
921 metric_engine
922 .handle_request(
923 physical_region_id,
924 RegionRequest::Flush(RegionFlushRequest::default()),
925 )
926 .await
927 .unwrap();
928
929 let path = region_dir_from_table_dir(
930 &TestEnv::default_table_dir(),
931 physical_region_id,
932 PathType::Metadata,
933 );
934 let object_store = env.get_object_store().unwrap();
935 let list = object_store.list(&path).await.unwrap();
936 for entry in list {
938 if entry.metadata().is_dir() {
939 continue;
940 }
941 if entry.name().ends_with("parquet") {
942 info!("deleting {}", entry.path());
943 object_store.delete(entry.path()).await.unwrap();
944 }
945 }
946
947 let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
948 .into_iter()
949 .collect();
950 let open_request = RegionOpenRequest {
951 engine: METRIC_ENGINE_NAME.to_string(),
952 table_dir: TestEnv::default_table_dir(),
953 path_type: PathType::Bare,
954 options: physical_region_option,
955 skip_wal_replay: false,
956 checkpoint: None,
957 requirements: Default::default(),
958 };
959 metric_engine
962 .handle_request(physical_region_id, RegionRequest::Open(open_request))
963 .await
964 .unwrap();
965
966 metric_engine
968 .handle_request(
969 physical_region_id,
970 RegionRequest::Close(RegionCloseRequest::default()),
971 )
972 .await
973 .unwrap();
974
975 let physical_region_option = [(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new())]
977 .into_iter()
978 .collect();
979 let open_request = RegionOpenRequest {
980 engine: METRIC_ENGINE_NAME.to_string(),
981 table_dir: TestEnv::default_table_dir(),
982 path_type: PathType::Bare,
983 options: physical_region_option,
984 skip_wal_replay: false,
985 checkpoint: None,
986 requirements: Default::default(),
987 };
988 let err = metric_engine
989 .handle_request(physical_region_id, RegionRequest::Open(open_request))
990 .await
991 .unwrap_err();
992 assert_eq!(err.status_code(), StatusCode::StorageUnavailable);
994
995 let mito_engine = metric_engine.mito();
996 let data_region_id = utils::to_data_region_id(physical_region_id);
997 let metadata_region_id = utils::to_metadata_region_id(physical_region_id);
998 let err = mito_engine.get_metadata(data_region_id).await.unwrap_err();
1000 assert_eq!(err.status_code(), StatusCode::RegionNotFound);
1001 let err = mito_engine
1002 .get_metadata(metadata_region_id)
1003 .await
1004 .unwrap_err();
1005 assert_eq!(err.status_code(), StatusCode::RegionNotFound);
1006 }
1007
1008 #[tokio::test]
1009 async fn test_catchup_regions() {
1010 common_telemetry::init_default_ut_logging();
1011 maybe_skip_kafka_log_store_integration_test!();
1012 let kafka_log_store_factory = kafka_log_store_factory().unwrap();
1013 let mito_env = mito2::test_util::TestEnv::new()
1014 .await
1015 .with_log_store_factory(kafka_log_store_factory.clone());
1016 let env = TestEnv::with_mito_env(mito_env).await;
1017 let table_dir = |region_id| format!("table/{region_id}");
1018 let mut physical_region_ids = vec![];
1019 let mut logical_region_ids = vec![];
1020
1021 let num_topics = 3;
1022 let num_physical_regions = 8;
1023 let num_logical_regions = 16;
1024 let parallelism = 2;
1025 let mut topics = Vec::with_capacity(num_topics);
1026 for _ in 0..num_topics {
1027 let topic = prepare_test_for_kafka_log_store(&kafka_log_store_factory)
1028 .await
1029 .unwrap();
1030 topics.push(topic);
1031 }
1032
1033 let topic_idx = |id| (id as usize) % num_topics;
1034 for i in 0..num_physical_regions {
1036 let physical_region_id = RegionId::new(1, i);
1037 physical_region_ids.push(physical_region_id);
1038
1039 let wal_options = WalOptions::Kafka(KafkaWalOptions::new(topics[topic_idx(i)].clone()));
1040 env.create_physical_region(
1041 physical_region_id,
1042 &table_dir(physical_region_id),
1043 vec![(
1044 WAL_OPTIONS_KEY.to_string(),
1045 serde_json::to_string(&wal_options).unwrap(),
1046 )],
1047 )
1048 .await;
1049 for j in 0..num_logical_regions {
1051 let logical_region_id = RegionId::new(1024 + i, j);
1052 logical_region_ids.push(logical_region_id);
1053 env.create_logical_region(physical_region_id, logical_region_id)
1054 .await;
1055 }
1056 }
1057
1058 let metric_engine = env.metric();
1059 for region_id in logical_region_ids.iter().chain(physical_region_ids.iter()) {
1061 metric_engine
1062 .handle_request(
1063 *region_id,
1064 RegionRequest::Close(RegionCloseRequest::default()),
1065 )
1066 .await
1067 .unwrap();
1068 }
1069
1070 let requests = physical_region_ids
1072 .iter()
1073 .enumerate()
1074 .map(|(idx, region_id)| {
1075 let mut options = HashMap::new();
1076 let wal_options =
1077 WalOptions::Kafka(KafkaWalOptions::new(topics[topic_idx(idx as u32)].clone()));
1078 options.insert(PHYSICAL_TABLE_METADATA_KEY.to_string(), String::new());
1079 options.insert(
1080 WAL_OPTIONS_KEY.to_string(),
1081 serde_json::to_string(&wal_options).unwrap(),
1082 );
1083 (
1084 *region_id,
1085 RegionOpenRequest {
1086 engine: METRIC_ENGINE_NAME.to_string(),
1087 table_dir: table_dir(*region_id),
1088 path_type: PathType::Bare,
1089 options: options.clone(),
1090 skip_wal_replay: true,
1091 checkpoint: None,
1092 requirements: Default::default(),
1093 },
1094 )
1095 })
1096 .collect::<Vec<_>>();
1097 info!("Open batch regions with parallelism: {parallelism}");
1098 metric_engine
1099 .handle_batch_open_requests(parallelism, requests)
1100 .await
1101 .unwrap();
1102 {
1103 let state = metric_engine.inner.state.read().unwrap();
1104 for logical_region in &logical_region_ids {
1105 assert!(!state.logical_regions().contains_key(logical_region));
1106 }
1107 }
1108
1109 let catch_requests = physical_region_ids
1110 .iter()
1111 .map(|region_id| {
1112 (
1113 *region_id,
1114 RegionCatchupRequest {
1115 set_writable: true,
1116 ..Default::default()
1117 },
1118 )
1119 })
1120 .collect::<Vec<_>>();
1121 metric_engine
1122 .handle_batch_catchup_requests(parallelism, catch_requests)
1123 .await
1124 .unwrap();
1125 {
1126 let state = metric_engine.inner.state.read().unwrap();
1127 for logical_region in &logical_region_ids {
1128 assert!(state.logical_regions().contains_key(logical_region));
1129 }
1130 }
1131 }
1132
1133 #[tokio::test]
1134 async fn test_drop_region() {
1135 let env = TestEnv::new().await;
1136 let engine = env.metric();
1137 let physical_region_id1 = RegionId::new(1024, 0);
1138 let logical_region_id1 = RegionId::new(1025, 0);
1139 env.create_physical_region(physical_region_id1, "/test_dir1", vec![])
1140 .await;
1141 let region_create_request1 =
1142 create_logical_region_request(&["job"], physical_region_id1, "logical1");
1143 engine
1144 .handle_batch_ddl_requests(BatchRegionDdlRequest::Create(vec![(
1145 logical_region_id1,
1146 region_create_request1,
1147 )]))
1148 .await
1149 .unwrap();
1150 let err = engine
1151 .handle_request(
1152 physical_region_id1,
1153 RegionRequest::Drop(RegionDropRequest {
1154 fast_path: false,
1155 force: false,
1156 partial_drop: false,
1157 }),
1158 )
1159 .await
1160 .unwrap_err();
1161 assert_matches!(
1162 err.as_any().downcast_ref::<Error>().unwrap(),
1163 &Error::PhysicalRegionBusy { .. }
1164 );
1165
1166 engine
1167 .handle_request(
1168 physical_region_id1,
1169 RegionRequest::Drop(RegionDropRequest {
1170 fast_path: false,
1171 force: true,
1172 partial_drop: false,
1173 }),
1174 )
1175 .await
1176 .unwrap();
1177 assert!(
1178 engine
1179 .inner
1180 .state
1181 .read()
1182 .unwrap()
1183 .physical_region_states()
1184 .get(&physical_region_id1)
1185 .is_none()
1186 );
1187 assert!(
1188 engine
1189 .inner
1190 .state
1191 .read()
1192 .unwrap()
1193 .logical_regions()
1194 .get(&logical_region_id1)
1195 .is_none()
1196 );
1197 }
1198}