Skip to main content

metric_engine/
engine.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod 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/// # Metric Engine
72///
73/// ## Regions
74///
75/// Regions in this metric engine has several roles. There is `PhysicalRegion`,
76/// which refer to the region that actually stores the data. And `LogicalRegion`
77/// that is "simulated" over physical regions. Each logical region is associated
78/// with one physical region group, which is a group of two physical regions.
79/// Their relationship is illustrated below:
80///
81/// ```mermaid
82/// erDiagram
83///     LogicalRegion ||--o{ PhysicalRegionGroup : corresponds
84///     PhysicalRegionGroup ||--|| DataRegion : contains
85///     PhysicalRegionGroup ||--|| MetadataRegion : contains
86/// ```
87///
88/// Metric engine uses two region groups. One is for data region
89/// ([METRIC_DATA_REGION_GROUP](crate::consts::METRIC_DATA_REGION_GROUP)), and the
90/// other is for metadata region ([METRIC_METADATA_REGION_GROUP](crate::consts::METRIC_METADATA_REGION_GROUP)).
91/// From the definition of [`RegionId`], we can convert between these two physical
92/// region ids easily. Thus in the code base we usually refer to one "physical
93/// region id", and convert it to the other one when necessary.
94///
95/// The logical region, in contrast, is a virtual region. It doesn't has dedicated
96/// storage or region group. Only a region id that is allocated by meta server.
97/// And all other things is shared with other logical region that are associated
98/// with the same physical region group.
99///
100/// For more document about physical regions, please refer to [`MetadataRegion`]
101/// and [`DataRegion`].
102///
103/// ## Operations
104///
105/// Both physical and logical region are accessible to user. But the operation
106/// they support are different. List below:
107///
108/// | Operations | Logical Region | Physical Region |
109/// | ---------- | -------------- | --------------- |
110/// |   Create   |       ✅        |        ✅        |
111/// |    Drop    |       ✅        |        ❓*       |
112/// |   Write    |       ✅        |        ❌        |
113/// |    Read    |       ✅        |        ✅        |
114/// |   Close    |       ✅        |        ✅        |
115/// |    Open    |       ✅        |        ✅        |
116/// |   Alter    |       ✅        |        ❓*       |
117///
118/// *: Physical region can be dropped only when all related logical regions are dropped.
119/// *: Alter: Physical regions only support altering region options.
120///
121/// ## Internal Columns
122///
123/// The physical data region contains two internal columns. Should
124/// mention that "internal" here is for metric engine itself. Mito
125/// engine will add it's internal columns to the region as well.
126///
127/// Their column id is registered in [`ReservedColumnId`]. And column name is
128/// defined in [`DATA_SCHEMA_TSID_COLUMN_NAME`] and [`DATA_SCHEMA_TABLE_ID_COLUMN_NAME`].
129///
130/// Tsid is generated by hashing all tags. And table id is retrieved from logical region
131/// id to distinguish data from different logical tables.
132#[derive(Clone)]
133pub struct MetricEngine {
134    inner: Arc<MetricEngineInner>,
135}
136
137#[async_trait]
138impl RegionEngine for MetricEngine {
139    /// Name of this engine
140    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    /// Handles non-query request to the region. Returns the count of affected rows.
211    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    /// Retrieves region's metadata.
348    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    /// Retrieves region's disk usage.
356    ///
357    /// Note: Returns `None` if it's a logical region.
358    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    /// Stops the engine
367    async fn stop(&self) -> Result<(), BoxedError> {
368        // don't need to stop the underlying mito engine
369        Ok(())
370    }
371
372    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> {
373        // ignore the region not found error
374        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        // ignore the region not found error
394        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    /// Returns the physical region role.
500    ///
501    /// Note: Returns `None` if it's a logical region.
502    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/// Outcome of a group of input requests.
516///
517/// Requests rejected before physical-region grouping produce singleton failures.
518#[derive(Debug)]
519pub struct PutGroupResult {
520    /// Positions in the input batch, including separate entries for the same region.
521    pub request_indices: Vec<usize>,
522    /// Aggregate rows on success; no other physical group's outcome is implied.
523    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    /// Writes all physical-region groups and retains each group's outcome.
559    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    /// Returns all logical regions associated with the physical region.
567    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    /// Handles substrait query and return a stream of record batches
575    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
606/// Shared definitions must agree before CREATE/ALTER mutates either region.
607fn 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        // Default Mito config disables series indexes: reaching that error proves forwarding.
683        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        // This maintenance operation applies to a physical region group only.
766        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        // close physical region
782        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        // reopen physical region
792        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, // Use Bare path type for engine regions
799            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        // close nonexistent region won't report error
810        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        // open nonexistent region won't report error
820        let invalid_open_request = RegionOpenRequest {
821            engine: METRIC_ENGINE_NAME.to_string(),
822            table_dir: TestEnv::default_table_dir(),
823            path_type: PathType::Bare, // Use Bare path type for engine regions
824            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        // Delete parquet files in metadata region
937        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        // Opening an already opened region should succeed.
960        // Since the region is already open, no metadata recovery operations will be performed.
961        metric_engine
962            .handle_request(physical_region_id, RegionRequest::Open(open_request))
963            .await
964            .unwrap();
965
966        // Close the region
967        metric_engine
968            .handle_request(
969                physical_region_id,
970                RegionRequest::Close(RegionCloseRequest::default()),
971            )
972            .await
973            .unwrap();
974
975        // Try to reopen region.
976        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        // Failed to open region because of missing parquet files.
993        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        // The metadata/data region should be closed.
999        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        // Creates physical regions
1035        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            // Creates logical regions for each physical region
1050            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        // Closes all regions
1060        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        // Opens all regions and skip the wal
1071        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}