Skip to main content

mito2/
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
15//! Mito region engine.
16
17#[cfg(test)]
18mod alter_test;
19#[cfg(test)]
20mod append_mode_test;
21#[cfg(test)]
22mod basic_test;
23#[cfg(test)]
24mod batch_catchup_test;
25#[cfg(test)]
26mod batch_open_test;
27#[cfg(test)]
28mod bump_committed_sequence_test;
29#[cfg(test)]
30mod catchup_test;
31#[cfg(test)]
32mod close_test;
33#[cfg(test)]
34pub(crate) mod compaction_test;
35#[cfg(test)]
36mod create_test;
37#[cfg(test)]
38mod drop_test;
39#[cfg(test)]
40mod edit_region_test;
41#[cfg(test)]
42mod file_ref_test;
43#[cfg(test)]
44mod filter_deleted_test;
45#[cfg(test)]
46mod flush_test;
47#[cfg(test)]
48mod index_build_test;
49#[cfg(any(test, feature = "test"))]
50pub mod listener;
51#[cfg(test)]
52mod merge_mode_test;
53#[cfg(test)]
54mod open_test;
55#[cfg(test)]
56mod parallel_test;
57#[cfg(test)]
58mod projection_test;
59#[cfg(test)]
60mod prune_test;
61pub mod region_hook;
62#[cfg(test)]
63mod row_selector_test;
64#[cfg(test)]
65mod scan_corrupt;
66#[cfg(test)]
67mod scan_test;
68#[cfg(test)]
69mod set_role_state_test;
70#[cfg(test)]
71mod skip_wal_test;
72#[cfg(test)]
73mod staging_test;
74#[cfg(test)]
75mod sync_test;
76#[cfg(test)]
77mod truncate_test;
78
79#[cfg(test)]
80mod copy_region_from_test;
81#[cfg(test)]
82mod remap_manifests_test;
83
84#[cfg(test)]
85mod apply_staging_manifest_test;
86#[cfg(test)]
87mod partition_filter_test;
88mod puffin_index;
89
90use std::any::Any;
91use std::collections::{HashMap, HashSet};
92use std::sync::Arc;
93use std::time::Instant;
94
95use api::region::RegionResponse;
96use async_trait::async_trait;
97use common_base::Plugins;
98use common_error::ext::BoxedError;
99use common_meta::error::UnexpectedSnafu;
100use common_meta::key::SchemaMetadataManagerRef;
101use common_recordbatch::{QueryMemoryTracker, SendableRecordBatchStream};
102use common_stat::get_total_memory_bytes;
103use common_telemetry::{info, tracing, warn};
104use common_wal::options::WalOptions;
105use futures::future::{join_all, try_join_all};
106use futures::stream::{self, Stream, StreamExt};
107use object_store::manager::ObjectStoreManagerRef;
108use snafu::{OptionExt, ResultExt, ensure};
109use store_api::ManifestVersion;
110use store_api::codec::PrimaryKeyEncoding;
111use store_api::logstore::LogStore;
112use store_api::logstore::provider::{KafkaProvider, Provider};
113use store_api::metadata::{ColumnMetadata, RegionMetadataRef};
114use store_api::metric_engine_consts::{
115    MANIFEST_INFO_EXTENSION_KEY, TABLE_COLUMN_METADATA_EXTENSION_KEY,
116};
117use store_api::region_engine::{
118    BatchResponses, MitoCopyRegionFromRequest, MitoCopyRegionFromResponse, RegionEngine,
119    RegionManifestInfo, RegionRole, RegionScannerRef, RegionStatistic, RemapManifestsRequest,
120    RemapManifestsResponse, SetRegionRoleStateResponse, SettableRegionRoleState,
121    SyncRegionFromRequest, SyncRegionFromResponse,
122};
123use store_api::region_info::RegionInfoEntry;
124use store_api::region_request::{
125    AffectedRows, RegionCatchupRequest, RegionOpenRequest, RegionRequest,
126};
127use store_api::sst_entry::{ManifestSstEntry, PuffinIndexMetaEntry, StorageSstEntry};
128use store_api::storage::{FileId, FileRefsManifest, RegionId, ScanRequest, SequenceNumber};
129use tokio::sync::{Semaphore, oneshot};
130
131use crate::access_layer::RegionFilePathFactory;
132use crate::cache::{CacheManagerRef, CacheStrategy};
133use crate::config::MitoConfig;
134use crate::engine::puffin_index::{IndexEntryContext, collect_index_entries_from_puffin};
135use crate::error::{
136    IncrementalQueryStaleSnafu, InvalidRequestSnafu, JoinSnafu, MitoManifestInfoSnafu, RecvSnafu,
137    RegionNotFoundSnafu, Result, SerdeJsonSnafu, SerializeColumnMetadataSnafu,
138    SnapshotFenceStaleSnafu,
139};
140#[cfg(feature = "enterprise")]
141use crate::extension::BoxedExtensionRangeProviderFactory;
142use crate::gc::GcLimiterRef;
143use crate::manifest::action::RegionEdit;
144use crate::memtable::MemtableStats;
145use crate::metrics::{
146    HANDLE_REQUEST_ELAPSED, SCAN_MEMORY_EXHAUSTED_TOTAL, SCAN_MEMORY_USAGE_BYTES,
147    SCAN_REQUESTS_REJECTED_TOTAL,
148};
149use crate::read::scan_region::{ScanRegion, Scanner};
150use crate::read::stream::ScanBatchStream;
151use crate::region::MitoRegionRef;
152use crate::region::opener::PartitionExprFetcherRef;
153use crate::region::options::parse_wal_options;
154use crate::request::{RegionEditRequest, WorkerRequest};
155use crate::sst::file::{FileMeta, RegionFileId, RegionIndexId};
156use crate::sst::file_ref::FileReferenceManagerRef;
157use crate::sst::index::intermediate::IntermediateManager;
158use crate::sst::index::puffin_manager::PuffinManagerFactory;
159use crate::wal::entry_distributor::{
160    DEFAULT_ENTRY_RECEIVER_BUFFER_SIZE, build_wal_entry_distributor_and_receivers,
161};
162use crate::wal::raw_entry_reader::{LogStoreRawEntryReader, RawEntryReader};
163use crate::worker::WorkerGroup;
164
165pub const MITO_ENGINE_NAME: &str = "mito";
166
167pub struct MitoEngineBuilder<'a, S: LogStore> {
168    data_home: &'a str,
169    config: MitoConfig,
170    log_store: Arc<S>,
171    object_store_manager: ObjectStoreManagerRef,
172    schema_metadata_manager: SchemaMetadataManagerRef,
173    file_ref_manager: FileReferenceManagerRef,
174    partition_expr_fetcher: PartitionExprFetcherRef,
175    plugins: Plugins,
176    #[cfg(feature = "enterprise")]
177    extension_range_provider_factory: Option<BoxedExtensionRangeProviderFactory>,
178}
179
180impl<'a, S: LogStore> MitoEngineBuilder<'a, S> {
181    #[allow(clippy::too_many_arguments)]
182    pub fn new(
183        data_home: &'a str,
184        config: MitoConfig,
185        log_store: Arc<S>,
186        object_store_manager: ObjectStoreManagerRef,
187        schema_metadata_manager: SchemaMetadataManagerRef,
188        file_ref_manager: FileReferenceManagerRef,
189        partition_expr_fetcher: PartitionExprFetcherRef,
190        plugins: Plugins,
191    ) -> Self {
192        Self {
193            data_home,
194            config,
195            log_store,
196            object_store_manager,
197            schema_metadata_manager,
198            file_ref_manager,
199            plugins,
200            partition_expr_fetcher,
201            #[cfg(feature = "enterprise")]
202            extension_range_provider_factory: None,
203        }
204    }
205
206    #[cfg(feature = "enterprise")]
207    #[must_use]
208    pub fn with_extension_range_provider_factory(
209        self,
210        extension_range_provider_factory: Option<BoxedExtensionRangeProviderFactory>,
211    ) -> Self {
212        Self {
213            extension_range_provider_factory,
214            ..self
215        }
216    }
217
218    pub async fn try_build(mut self) -> Result<MitoEngine> {
219        self.config.sanitize(self.data_home)?;
220
221        let config = Arc::new(self.config);
222        let workers = WorkerGroup::start(
223            config.clone(),
224            self.log_store.clone(),
225            self.object_store_manager,
226            self.schema_metadata_manager,
227            self.file_ref_manager,
228            self.partition_expr_fetcher.clone(),
229            self.plugins,
230        )
231        .await?;
232        let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(self.log_store));
233        let total_memory = get_total_memory_bytes().max(0) as u64;
234        let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize;
235        let scan_memory_tracker =
236            QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted)
237                .on_update(|usage| {
238                    SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
239                })
240                .on_exhausted(|| {
241                    SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
242                })
243                .on_reject(|| {
244                    SCAN_REQUESTS_REJECTED_TOTAL.inc();
245                })
246                .build();
247
248        let inner = EngineInner {
249            workers,
250            config,
251            wal_raw_entry_reader,
252            scan_memory_tracker,
253            #[cfg(feature = "enterprise")]
254            extension_range_provider_factory: None,
255        };
256
257        #[cfg(feature = "enterprise")]
258        let inner =
259            inner.with_extension_range_provider_factory(self.extension_range_provider_factory);
260
261        Ok(MitoEngine {
262            inner: Arc::new(inner),
263        })
264    }
265}
266
267/// Region engine implementation for timeseries data.
268#[derive(Clone)]
269pub struct MitoEngine {
270    inner: Arc<EngineInner>,
271}
272
273impl MitoEngine {
274    /// Returns a new [MitoEngine] with specific `config`, `log_store` and `object_store`.
275    #[allow(clippy::too_many_arguments)]
276    pub async fn new<S: LogStore>(
277        data_home: &str,
278        config: MitoConfig,
279        log_store: Arc<S>,
280        object_store_manager: ObjectStoreManagerRef,
281        schema_metadata_manager: SchemaMetadataManagerRef,
282        file_ref_manager: FileReferenceManagerRef,
283        partition_expr_fetcher: PartitionExprFetcherRef,
284        plugins: Plugins,
285    ) -> Result<MitoEngine> {
286        let builder = MitoEngineBuilder::new(
287            data_home,
288            config,
289            log_store,
290            object_store_manager,
291            schema_metadata_manager,
292            file_ref_manager,
293            partition_expr_fetcher,
294            plugins,
295        );
296        builder.try_build().await
297    }
298
299    pub fn mito_config(&self) -> &MitoConfig {
300        &self.inner.config
301    }
302
303    pub fn cache_manager(&self) -> CacheManagerRef {
304        self.inner.workers.cache_manager()
305    }
306
307    pub fn file_ref_manager(&self) -> FileReferenceManagerRef {
308        self.inner.workers.file_ref_manager()
309    }
310
311    pub fn gc_limiter(&self) -> GcLimiterRef {
312        self.inner.workers.gc_limiter()
313    }
314
315    pub fn object_store_manager(&self) -> &ObjectStoreManagerRef {
316        self.inner.workers.object_store_manager()
317    }
318
319    pub fn puffin_manager_factory(&self) -> &PuffinManagerFactory {
320        self.inner.workers.puffin_manager_factory()
321    }
322
323    pub fn intermediate_manager(&self) -> &IntermediateManager {
324        self.inner.workers.intermediate_manager()
325    }
326
327    pub fn schema_metadata_manager(&self) -> &SchemaMetadataManagerRef {
328        self.inner.workers.schema_metadata_manager()
329    }
330
331    /// Get all tmp ref files for given region ids, excluding files that's already in manifest.
332    pub async fn get_snapshot_of_file_refs(
333        &self,
334        file_handle_regions: impl IntoIterator<Item = RegionId>,
335        related_regions: HashMap<RegionId, HashSet<RegionId>>,
336    ) -> Result<FileRefsManifest> {
337        let file_ref_mgr = self.file_ref_manager();
338
339        let file_handle_regions = file_handle_regions.into_iter().collect::<Vec<_>>();
340        // Convert region IDs to MitoRegionRef objects, ignore regions that do not exist on current datanode
341        // as regions on other datanodes are not managed by this engine.
342        let query_regions: Vec<MitoRegionRef> = file_handle_regions
343            .into_iter()
344            .filter_map(|region_id| self.find_region(region_id))
345            .collect();
346
347        let dst_region_to_src_regions: Vec<(MitoRegionRef, HashSet<RegionId>)> = {
348            let dst2src = related_regions
349                .into_iter()
350                .flat_map(|(src, dsts)| dsts.into_iter().map(move |dst| (dst, src)))
351                .fold(
352                    HashMap::<RegionId, HashSet<RegionId>>::new(),
353                    |mut acc, (k, v)| {
354                        let entry = acc.entry(k).or_default();
355                        entry.insert(v);
356                        acc
357                    },
358                );
359            let mut dst_region_to_src_regions = Vec::with_capacity(dst2src.len());
360            for (dst_region, srcs) in dst2src {
361                let Some(region) = self.find_region(dst_region) else {
362                    return RegionNotFoundSnafu {
363                        region_id: dst_region,
364                    }
365                    .fail();
366                };
367                dst_region_to_src_regions.push((region, srcs));
368            }
369            dst_region_to_src_regions
370        };
371
372        file_ref_mgr
373            .get_snapshot_of_file_refs(query_regions, dst_region_to_src_regions)
374            .await
375    }
376
377    /// Returns true if the specific region exists.
378    pub fn is_region_exists(&self, region_id: RegionId) -> bool {
379        self.inner.workers.is_region_exists(region_id)
380    }
381
382    /// Returns true if the specific region exists.
383    pub fn is_region_opening(&self, region_id: RegionId) -> bool {
384        self.inner.workers.is_region_opening(region_id)
385    }
386
387    /// Returns true if the specific region is catching up.
388    pub fn is_region_catching_up(&self, region_id: RegionId) -> bool {
389        self.inner.workers.is_region_catching_up(region_id)
390    }
391
392    /// Returns the region disk/memory statistic.
393    pub fn get_region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic> {
394        self.find_region(region_id)
395            .map(|region| region.region_statistic())
396    }
397
398    /// Returns primary key encoding of the region.
399    pub fn get_primary_key_encoding(&self, region_id: RegionId) -> Option<PrimaryKeyEncoding> {
400        self.find_region(region_id)
401            .map(|r| r.primary_key_encoding())
402    }
403
404    /// Handle substrait query and return a stream of record batches
405    ///
406    /// Notice that the output stream's ordering is not guranateed. If order
407    /// matter, please use [`scanner`] to build a [`Scanner`] to consume.
408    #[tracing::instrument(skip_all)]
409    pub async fn scan_to_stream(
410        &self,
411        region_id: RegionId,
412        request: ScanRequest,
413    ) -> Result<SendableRecordBatchStream, BoxedError> {
414        self.scanner(region_id, request)
415            .await
416            .map_err(BoxedError::new)?
417            .scan()
418            .await
419    }
420
421    /// Scan [`Batch`]es by [`ScanRequest`].
422    pub async fn scan_batch(
423        &self,
424        region_id: RegionId,
425        request: ScanRequest,
426        filter_deleted: bool,
427    ) -> Result<ScanBatchStream> {
428        let mut scan_region = self.scan_region(region_id, request)?;
429        scan_region.set_filter_deleted(filter_deleted);
430        scan_region.scanner().await?.scan_batch()
431    }
432
433    /// Returns a scanner to scan for `request`.
434    pub(crate) async fn scanner(
435        &self,
436        region_id: RegionId,
437        request: ScanRequest,
438    ) -> Result<Scanner> {
439        self.scan_region(region_id, request)?.scanner().await
440    }
441
442    /// Scans a region.
443    #[tracing::instrument(skip_all, fields(region_id = %region_id))]
444    fn scan_region(&self, region_id: RegionId, request: ScanRequest) -> Result<ScanRegion> {
445        self.inner.scan_region(region_id, request)
446    }
447
448    /// Edit region's metadata by [RegionEdit] directly. Use with care.
449    /// Now we only allow adding files or removing files from region (the [RegionEdit] struct can only contain a non-empty "files_to_add" or "files_to_remove" field).
450    /// Other region editing intention will result in an "invalid request" error.
451    /// Also note that if a region is to be edited directly, we MUST not write data to it thereafter.
452    pub async fn edit_region(&self, region_id: RegionId, edit: RegionEdit) -> Result<()> {
453        let _timer = HANDLE_REQUEST_ELAPSED
454            .with_label_values(&["edit_region"])
455            .start_timer();
456
457        ensure!(
458            is_valid_region_edit(&edit),
459            InvalidRequestSnafu {
460                region_id,
461                reason: "invalid region edit"
462            }
463        );
464
465        let (tx, rx) = oneshot::channel();
466        let request = WorkerRequest::EditRegion(RegionEditRequest::new(region_id, edit, true, tx));
467        self.inner
468            .workers
469            .submit_to_worker(region_id, request)
470            .await?;
471        rx.await.context(RecvSnafu)?
472    }
473
474    /// Handles copy region from request.
475    ///
476    /// This method is only supported for internal use and is not exposed in the trait implementation.
477    pub async fn copy_region_from(
478        &self,
479        region_id: RegionId,
480        request: MitoCopyRegionFromRequest,
481    ) -> Result<MitoCopyRegionFromResponse> {
482        self.inner.copy_region_from(region_id, request).await
483    }
484
485    #[cfg(test)]
486    pub(crate) fn get_region(&self, id: RegionId) -> Option<crate::region::MitoRegionRef> {
487        self.find_region(id)
488    }
489
490    pub fn find_region(&self, region_id: RegionId) -> Option<MitoRegionRef> {
491        self.inner.workers.get_region(region_id)
492    }
493
494    /// Returns all regions.
495    pub fn regions(&self) -> Vec<MitoRegionRef> {
496        self.inner.workers.all_regions().collect()
497    }
498
499    fn encode_manifest_info_to_extensions(
500        region_id: &RegionId,
501        manifest_info: RegionManifestInfo,
502        extensions: &mut HashMap<String, Vec<u8>>,
503    ) -> Result<()> {
504        let region_manifest_info = vec![(*region_id, manifest_info)];
505
506        extensions.insert(
507            MANIFEST_INFO_EXTENSION_KEY.to_string(),
508            RegionManifestInfo::encode_list(&region_manifest_info).context(SerdeJsonSnafu)?,
509        );
510        info!(
511            "Added manifest info: {:?} to extensions, region_id: {:?}",
512            region_manifest_info, region_id
513        );
514        Ok(())
515    }
516
517    fn encode_column_metadatas_to_extensions(
518        region_id: &RegionId,
519        column_metadatas: Vec<ColumnMetadata>,
520        extensions: &mut HashMap<String, Vec<u8>>,
521    ) -> Result<()> {
522        extensions.insert(
523            TABLE_COLUMN_METADATA_EXTENSION_KEY.to_string(),
524            ColumnMetadata::encode_list(&column_metadatas).context(SerializeColumnMetadataSnafu)?,
525        );
526        info!(
527            "Added column metadatas: {:?} to extensions, region_id: {:?}",
528            column_metadatas, region_id
529        );
530        Ok(())
531    }
532
533    /// Find the current version's memtables and SSTs stats by region_id.
534    /// The stats must be collected in one place one time to ensure data consistency.
535    pub fn find_memtable_and_sst_stats(
536        &self,
537        region_id: RegionId,
538    ) -> Result<(Vec<MemtableStats>, Vec<FileMeta>)> {
539        let region = self
540            .find_region(region_id)
541            .context(RegionNotFoundSnafu { region_id })?;
542
543        let version = region.version();
544        let memtable_stats = version
545            .memtables
546            .list_memtables()
547            .iter()
548            .map(|x| x.stats())
549            .collect::<Vec<_>>();
550
551        let sst_stats = version
552            .ssts
553            .levels()
554            .iter()
555            .flat_map(|level| level.files().map(|x| x.meta_ref()))
556            .cloned()
557            .collect::<Vec<_>>();
558        Ok((memtable_stats, sst_stats))
559    }
560
561    /// Lists all SSTs from the manifest of all regions in the engine.
562    pub async fn all_ssts_from_manifest(&self) -> Vec<ManifestSstEntry> {
563        let node_id = self.inner.workers.file_ref_manager().node_id();
564        let regions = self.inner.workers.all_regions();
565
566        let mut results = Vec::new();
567        for region in regions {
568            let mut entries = region.manifest_sst_entries().await;
569            for e in &mut entries {
570                e.node_id = node_id;
571            }
572            results.extend(entries);
573        }
574
575        results
576    }
577
578    /// Lists metadata about all puffin index targets stored in the engine.
579    pub async fn all_index_metas(&self) -> Vec<PuffinIndexMetaEntry> {
580        let node_id = self.inner.workers.file_ref_manager().node_id();
581        let cache_manager = self.inner.workers.cache_manager();
582        let puffin_metadata_cache = cache_manager.puffin_metadata_cache().cloned();
583        let bloom_filter_cache = cache_manager.bloom_filter_index_cache().cloned();
584        let inverted_index_cache = cache_manager.inverted_index_cache().cloned();
585
586        let mut results = Vec::new();
587
588        for region in self.inner.workers.all_regions() {
589            let manifest_entries = region.manifest_sst_entries().await;
590            let access_layer = region.access_layer.clone();
591            let table_dir = access_layer.table_dir().to_string();
592            let path_type = access_layer.path_type();
593            let object_store = access_layer.object_store().clone();
594            let puffin_factory = access_layer.puffin_manager_factory().clone();
595            let path_factory = RegionFilePathFactory::new(table_dir, path_type);
596
597            let entry_futures = manifest_entries.into_iter().map(|entry| {
598                let object_store = object_store.clone();
599                let path_factory = path_factory.clone();
600                let puffin_factory = puffin_factory.clone();
601                let puffin_metadata_cache = puffin_metadata_cache.clone();
602                let bloom_filter_cache = bloom_filter_cache.clone();
603                let inverted_index_cache = inverted_index_cache.clone();
604
605                async move {
606                    let Some(index_file_path) = entry.index_file_path.as_ref() else {
607                        return Vec::new();
608                    };
609
610                    let index_version = entry.index_version;
611                    let file_id = match FileId::parse_str(&entry.file_id) {
612                        Ok(file_id) => file_id,
613                        Err(err) => {
614                            warn!(
615                                err;
616                                "Failed to parse puffin index file id, table_dir: {}, file_id: {}",
617                                entry.table_dir,
618                                entry.file_id
619                            );
620                            return Vec::new();
621                        }
622                    };
623                    // The index file path is derived from the physical file owner. After
624                    // repartition, `entry.region_id` is only the referring region.
625                    let region_index_id = RegionIndexId::new(
626                        RegionFileId::new(entry.origin_region_id, file_id),
627                        index_version,
628                    );
629                    let context = IndexEntryContext {
630                        table_dir: &entry.table_dir,
631                        index_file_path: index_file_path.as_str(),
632                        region_id: entry.region_id,
633                        table_id: entry.table_id,
634                        region_number: entry.region_number,
635                        region_group: entry.region_group,
636                        region_sequence: entry.region_sequence,
637                        file_id: &entry.file_id,
638                        index_file_size: entry.index_file_size,
639                        node_id,
640                    };
641
642                    let manager = puffin_factory
643                        .build(object_store, path_factory)
644                        .with_puffin_metadata_cache(puffin_metadata_cache);
645
646                    collect_index_entries_from_puffin(
647                        manager,
648                        region_index_id,
649                        context,
650                        bloom_filter_cache,
651                        inverted_index_cache,
652                    )
653                    .await
654                }
655            });
656
657            let mut meta_stream = stream::iter(entry_futures).buffer_unordered(8); // Parallelism is 8.
658            while let Some(mut metas) = meta_stream.next().await {
659                results.append(&mut metas);
660            }
661        }
662
663        results
664    }
665
666    /// Lists region info entries of all regions in the engine.
667    pub async fn all_region_infos(&self) -> Vec<RegionInfoEntry> {
668        let node_id = self.inner.workers.file_ref_manager().node_id();
669        self.inner
670            .workers
671            .all_regions()
672            .map(|region| region.region_info_entry(node_id))
673            .collect()
674    }
675
676    /// Lists all SSTs from the storage layer of all regions in the engine.
677    pub fn all_ssts_from_storage(&self) -> impl Stream<Item = Result<StorageSstEntry>> {
678        let node_id = self.inner.workers.file_ref_manager().node_id();
679        let regions = self.inner.workers.all_regions();
680
681        let mut layers_distinct_table_dirs = HashMap::new();
682        for region in regions {
683            let table_dir = region.access_layer.table_dir();
684            if !layers_distinct_table_dirs.contains_key(table_dir) {
685                layers_distinct_table_dirs
686                    .insert(table_dir.to_string(), region.access_layer.clone());
687            }
688        }
689
690        stream::iter(layers_distinct_table_dirs)
691            .map(|(_, access_layer)| access_layer.storage_sst_entries())
692            .flatten()
693            .map(move |entry| {
694                entry.map(move |mut entry| {
695                    entry.node_id = node_id;
696                    entry
697                })
698            })
699    }
700}
701
702/// Check whether the region edit is valid.
703///
704/// Only adding or removing files to region is considered valid now.
705fn is_valid_region_edit(edit: &RegionEdit) -> bool {
706    (!edit.files_to_add.is_empty() || !edit.files_to_remove.is_empty())
707        && matches!(
708            edit,
709            RegionEdit {
710                files_to_add: _,
711                files_to_remove: _,
712                timestamp_ms: _,
713                compaction_time_window: None,
714                flushed_entry_id: None,
715                flushed_sequence: None,
716                ..
717            }
718        )
719}
720
721/// Inner struct of [MitoEngine].
722struct EngineInner {
723    /// Region workers group.
724    workers: WorkerGroup,
725    /// Config of the engine.
726    config: Arc<MitoConfig>,
727    /// The Wal raw entry reader.
728    wal_raw_entry_reader: Arc<dyn RawEntryReader>,
729    /// Memory tracker for table scans.
730    scan_memory_tracker: QueryMemoryTracker,
731    #[cfg(feature = "enterprise")]
732    extension_range_provider_factory: Option<BoxedExtensionRangeProviderFactory>,
733}
734
735type TopicGroupedRegionOpenRequests = HashMap<String, Vec<(RegionId, RegionOpenRequest)>>;
736
737/// Returns requests([TopicGroupedRegionOpenRequests]) grouped by topic and remaining requests.
738fn prepare_batch_open_requests(
739    requests: Vec<(RegionId, RegionOpenRequest)>,
740) -> Result<(
741    TopicGroupedRegionOpenRequests,
742    Vec<(RegionId, RegionOpenRequest)>,
743)> {
744    let mut topic_to_regions: HashMap<String, Vec<(RegionId, RegionOpenRequest)>> = HashMap::new();
745    let mut remaining_regions: Vec<(RegionId, RegionOpenRequest)> = Vec::new();
746    for (region_id, request) in requests {
747        match parse_wal_options(&request.options).context(SerdeJsonSnafu)? {
748            WalOptions::Kafka(options) => {
749                topic_to_regions
750                    .entry(options.topic)
751                    .or_default()
752                    .push((region_id, request));
753            }
754            WalOptions::RaftEngine | WalOptions::Noop => {
755                remaining_regions.push((region_id, request));
756            }
757        }
758    }
759
760    Ok((topic_to_regions, remaining_regions))
761}
762
763impl EngineInner {
764    #[cfg(feature = "enterprise")]
765    #[must_use]
766    fn with_extension_range_provider_factory(
767        self,
768        extension_range_provider_factory: Option<BoxedExtensionRangeProviderFactory>,
769    ) -> Self {
770        Self {
771            extension_range_provider_factory,
772            ..self
773        }
774    }
775
776    /// Stop the inner engine.
777    async fn stop(&self) -> Result<()> {
778        self.workers.stop().await
779    }
780
781    fn find_region(&self, region_id: RegionId) -> Result<MitoRegionRef> {
782        self.workers
783            .get_region(region_id)
784            .context(RegionNotFoundSnafu { region_id })
785    }
786
787    /// Get metadata of a region.
788    ///
789    /// Returns error if the region doesn't exist.
790    fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef> {
791        // Reading a region doesn't need to go through the region worker thread.
792        let region = self.find_region(region_id)?;
793        Ok(region.metadata())
794    }
795
796    async fn open_topic_regions(
797        &self,
798        topic: String,
799        region_requests: Vec<(RegionId, RegionOpenRequest)>,
800    ) -> Result<Vec<(RegionId, Result<AffectedRows>)>> {
801        let now = Instant::now();
802        let region_ids = region_requests
803            .iter()
804            .map(|(region_id, _)| *region_id)
805            .collect::<Vec<_>>();
806        let provider = Provider::kafka_provider(topic);
807        let (distributor, entry_receivers) = build_wal_entry_distributor_and_receivers(
808            provider.clone(),
809            self.wal_raw_entry_reader.clone(),
810            &region_ids,
811            DEFAULT_ENTRY_RECEIVER_BUFFER_SIZE,
812        );
813
814        let mut responses = Vec::with_capacity(region_requests.len());
815        for ((region_id, request), entry_receiver) in
816            region_requests.into_iter().zip(entry_receivers)
817        {
818            let (request, receiver) =
819                WorkerRequest::new_open_region_request(region_id, request, Some(entry_receiver));
820            self.workers.submit_to_worker(region_id, request).await?;
821            responses.push(async move { receiver.await.context(RecvSnafu)? });
822        }
823
824        // Waits for entries distribution.
825        let distribution =
826            common_runtime::spawn_global(async move { distributor.distribute().await });
827        // Waits for worker returns.
828        let responses = join_all(responses).await;
829        distribution.await.context(JoinSnafu)??;
830
831        let num_failure = responses.iter().filter(|r| r.is_err()).count();
832        info!(
833            "Opened {} regions for topic '{}', failures: {}, elapsed: {:?}",
834            region_ids.len() - num_failure,
835            // Safety: provider is kafka provider.
836            provider.as_kafka_provider().unwrap(),
837            num_failure,
838            now.elapsed(),
839        );
840        Ok(region_ids.into_iter().zip(responses).collect())
841    }
842
843    async fn handle_batch_open_requests(
844        &self,
845        parallelism: usize,
846        requests: Vec<(RegionId, RegionOpenRequest)>,
847    ) -> Result<Vec<(RegionId, Result<AffectedRows>)>> {
848        let semaphore = Arc::new(Semaphore::new(parallelism));
849        let (topic_to_region_requests, remaining_region_requests) =
850            prepare_batch_open_requests(requests)?;
851        let mut responses =
852            Vec::with_capacity(topic_to_region_requests.len() + remaining_region_requests.len());
853
854        if !topic_to_region_requests.is_empty() {
855            let mut tasks = Vec::with_capacity(topic_to_region_requests.len());
856            for (topic, region_requests) in topic_to_region_requests {
857                let semaphore_moved = semaphore.clone();
858                tasks.push(async move {
859                    // Safety: semaphore must exist
860                    let _permit = semaphore_moved.acquire().await.unwrap();
861                    self.open_topic_regions(topic, region_requests).await
862                })
863            }
864            let r = try_join_all(tasks).await?;
865            responses.extend(r.into_iter().flatten());
866        }
867
868        if !remaining_region_requests.is_empty() {
869            let mut tasks = Vec::with_capacity(remaining_region_requests.len());
870            let mut region_ids = Vec::with_capacity(remaining_region_requests.len());
871            for (region_id, request) in remaining_region_requests {
872                let semaphore_moved = semaphore.clone();
873                region_ids.push(region_id);
874                tasks.push(async move {
875                    // Safety: semaphore must exist
876                    let _permit = semaphore_moved.acquire().await.unwrap();
877                    let (request, receiver) =
878                        WorkerRequest::new_open_region_request(region_id, request, None);
879
880                    self.workers.submit_to_worker(region_id, request).await?;
881
882                    receiver.await.context(RecvSnafu)?
883                })
884            }
885
886            let results = join_all(tasks).await;
887            responses.extend(region_ids.into_iter().zip(results));
888        }
889
890        Ok(responses)
891    }
892
893    async fn catchup_topic_regions(
894        &self,
895        provider: Provider,
896        region_requests: Vec<(RegionId, RegionCatchupRequest)>,
897    ) -> Result<Vec<(RegionId, Result<AffectedRows>)>> {
898        let now = Instant::now();
899        let region_ids = region_requests
900            .iter()
901            .map(|(region_id, _)| *region_id)
902            .collect::<Vec<_>>();
903        let (distributor, entry_receivers) = build_wal_entry_distributor_and_receivers(
904            provider.clone(),
905            self.wal_raw_entry_reader.clone(),
906            &region_ids,
907            DEFAULT_ENTRY_RECEIVER_BUFFER_SIZE,
908        );
909
910        let mut responses = Vec::with_capacity(region_requests.len());
911        for ((region_id, request), entry_receiver) in
912            region_requests.into_iter().zip(entry_receivers)
913        {
914            let (request, receiver) =
915                WorkerRequest::new_catchup_region_request(region_id, request, Some(entry_receiver));
916            self.workers.submit_to_worker(region_id, request).await?;
917            responses.push(async move { receiver.await.context(RecvSnafu)? });
918        }
919
920        // Wait for entries distribution.
921        let distribution =
922            common_runtime::spawn_global(async move { distributor.distribute().await });
923        // Wait for worker returns.
924        let responses = join_all(responses).await;
925        distribution.await.context(JoinSnafu)??;
926
927        let num_failure = responses.iter().filter(|r| r.is_err()).count();
928        info!(
929            "Caught up {} regions for topic '{}', failures: {}, elapsed: {:?}",
930            region_ids.len() - num_failure,
931            // Safety: provider is kafka provider.
932            provider.as_kafka_provider().unwrap(),
933            num_failure,
934            now.elapsed(),
935        );
936
937        Ok(region_ids.into_iter().zip(responses).collect())
938    }
939
940    async fn handle_batch_catchup_requests(
941        &self,
942        parallelism: usize,
943        requests: Vec<(RegionId, RegionCatchupRequest)>,
944    ) -> Result<Vec<(RegionId, Result<AffectedRows>)>> {
945        let mut responses = Vec::with_capacity(requests.len());
946        let mut topic_regions: HashMap<Arc<KafkaProvider>, Vec<_>> = HashMap::new();
947        let mut remaining_region_requests = vec![];
948
949        for (region_id, request) in requests {
950            match self.workers.get_region(region_id) {
951                Some(region) => match region.provider.as_kafka_provider() {
952                    Some(provider) => {
953                        topic_regions
954                            .entry(provider.clone())
955                            .or_default()
956                            .push((region_id, request));
957                    }
958                    None => {
959                        remaining_region_requests.push((region_id, request));
960                    }
961                },
962                None => responses.push((region_id, RegionNotFoundSnafu { region_id }.fail())),
963            }
964        }
965
966        let semaphore = Arc::new(Semaphore::new(parallelism));
967
968        if !topic_regions.is_empty() {
969            let mut tasks = Vec::with_capacity(topic_regions.len());
970            for (provider, region_requests) in topic_regions {
971                let semaphore_moved = semaphore.clone();
972                tasks.push(async move {
973                    // Safety: semaphore must exist
974                    let _permit = semaphore_moved.acquire().await.unwrap();
975                    self.catchup_topic_regions(Provider::Kafka(provider), region_requests)
976                        .await
977                })
978            }
979
980            let r = try_join_all(tasks).await?;
981            responses.extend(r.into_iter().flatten());
982        }
983
984        if !remaining_region_requests.is_empty() {
985            let mut tasks = Vec::with_capacity(remaining_region_requests.len());
986            let mut region_ids = Vec::with_capacity(remaining_region_requests.len());
987            for (region_id, request) in remaining_region_requests {
988                let semaphore_moved = semaphore.clone();
989                region_ids.push(region_id);
990                tasks.push(async move {
991                    // Safety: semaphore must exist
992                    let _permit = semaphore_moved.acquire().await.unwrap();
993                    let (request, receiver) =
994                        WorkerRequest::new_catchup_region_request(region_id, request, None);
995
996                    self.workers.submit_to_worker(region_id, request).await?;
997
998                    receiver.await.context(RecvSnafu)?
999                })
1000            }
1001
1002            let results = join_all(tasks).await;
1003            responses.extend(region_ids.into_iter().zip(results));
1004        }
1005
1006        Ok(responses)
1007    }
1008
1009    /// Handles [RegionRequest] and return its executed result.
1010    async fn handle_request(
1011        &self,
1012        region_id: RegionId,
1013        request: RegionRequest,
1014    ) -> Result<AffectedRows> {
1015        let region_metadata = self.get_metadata(region_id).ok();
1016        let (request, receiver) =
1017            WorkerRequest::try_from_region_request(region_id, request, region_metadata)?;
1018        self.workers.submit_to_worker(region_id, request).await?;
1019
1020        receiver.await.context(RecvSnafu)?
1021    }
1022
1023    /// Returns the sequence of latest committed data.
1024    fn get_committed_sequence(&self, region_id: RegionId) -> Result<SequenceNumber> {
1025        // Reading a region doesn't need to go through the region worker thread.
1026        self.find_region(region_id)
1027            .map(|r| r.find_committed_sequence())
1028    }
1029
1030    /// Handles the scan `request` and returns a [ScanRegion].
1031    #[tracing::instrument(skip_all, fields(region_id = %region_id))]
1032    fn scan_region(&self, region_id: RegionId, mut request: ScanRequest) -> Result<ScanRegion> {
1033        let query_start = Instant::now();
1034        // Reading a region doesn't need to go through the region worker thread.
1035        let region = self.find_region(region_id)?;
1036        let version_data = region.version_control.current();
1037        let version = version_data.version;
1038
1039        if request.snapshot_on_scan && request.memtable_max_sequence.is_none() {
1040            request.memtable_max_sequence = Some(version_data.committed_sequence);
1041        }
1042
1043        if let Some(given_seq) = request.memtable_min_sequence {
1044            let min_readable_seq = version.flushed_sequence;
1045            ensure!(
1046                given_seq >= min_readable_seq,
1047                IncrementalQueryStaleSnafu {
1048                    region_id,
1049                    given_seq,
1050                    min_readable_seq,
1051                }
1052            );
1053        }
1054
1055        if let Some(given_seq) = request.memtable_max_sequence
1056            && !request.skip_sst_files
1057        {
1058            // Explicit snapshot fences that include SST reads are enforceable
1059            // only while the requested upper bound is not older than the
1060            // region's flushed frontier. If H has already been flushed into SST,
1061            // mito cannot apply a memtable-only sequence upper bound to that
1062            // SST scan, so fail and let Flow rebind the fenced repair instead
1063            // of reading rows beyond H.
1064            let min_enforceable_seq = version.flushed_sequence;
1065            ensure!(
1066                given_seq >= min_enforceable_seq,
1067                SnapshotFenceStaleSnafu {
1068                    region_id,
1069                    given_seq,
1070                    min_enforceable_seq,
1071                }
1072            );
1073        }
1074
1075        // Get cache.
1076        let cache_manager = self.workers.cache_manager();
1077
1078        let scan_region = ScanRegion::new(
1079            version,
1080            region.access_layer.clone(),
1081            request,
1082            CacheStrategy::EnableAll(cache_manager),
1083        )
1084        .with_query_stat_counters(region.region_stats.query_stat_counters())
1085        .with_max_concurrent_scan_files(self.config.max_concurrent_scan_files)
1086        .with_ignore_inverted_index(self.config.inverted_index.apply_on_query.disabled())
1087        .with_ignore_fulltext_index(self.config.fulltext_index.apply_on_query.disabled())
1088        .with_ignore_bloom_filter(self.config.bloom_filter_index.apply_on_query.disabled())
1089        .with_start_time(query_start);
1090
1091        #[cfg(feature = "enterprise")]
1092        let scan_region = self.maybe_fill_extension_range_provider(scan_region, region);
1093
1094        Ok(scan_region)
1095    }
1096
1097    #[cfg(feature = "enterprise")]
1098    fn maybe_fill_extension_range_provider(
1099        &self,
1100        mut scan_region: ScanRegion,
1101        region: MitoRegionRef,
1102    ) -> ScanRegion {
1103        if region.is_follower()
1104            && let Some(factory) = self.extension_range_provider_factory.as_ref()
1105        {
1106            scan_region
1107                .set_extension_range_provider(factory.create_extension_range_provider(region));
1108        }
1109        scan_region
1110    }
1111
1112    /// Converts the [`RegionRole`].
1113    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<()> {
1114        let region = self.find_region(region_id)?;
1115        region.set_role(role);
1116        Ok(())
1117    }
1118
1119    /// Sets read-only for a region and ensures no more writes in the region after it returns.
1120    async fn set_region_role_state_gracefully(
1121        &self,
1122        region_id: RegionId,
1123        region_role_state: SettableRegionRoleState,
1124    ) -> Result<SetRegionRoleStateResponse> {
1125        // Notes: It acquires the mutable ownership to ensure no other threads,
1126        // Therefore, we submit it to the worker.
1127        let (request, receiver) =
1128            WorkerRequest::new_set_readonly_gracefully(region_id, region_role_state);
1129        self.workers.submit_to_worker(region_id, request).await?;
1130
1131        receiver.await.context(RecvSnafu)
1132    }
1133
1134    async fn sync_region(
1135        &self,
1136        region_id: RegionId,
1137        manifest_info: RegionManifestInfo,
1138    ) -> Result<(ManifestVersion, bool)> {
1139        ensure!(manifest_info.is_mito(), MitoManifestInfoSnafu);
1140        let manifest_version = manifest_info.data_manifest_version();
1141        let (request, receiver) =
1142            WorkerRequest::new_sync_region_request(region_id, manifest_version);
1143        self.workers.submit_to_worker(region_id, request).await?;
1144
1145        receiver.await.context(RecvSnafu)?
1146    }
1147
1148    async fn remap_manifests(
1149        &self,
1150        request: RemapManifestsRequest,
1151    ) -> Result<RemapManifestsResponse> {
1152        let region_id = request.region_id;
1153        let (request, receiver) = WorkerRequest::try_from_remap_manifests_request(request)?;
1154        self.workers.submit_to_worker(region_id, request).await?;
1155        let manifest_paths = receiver.await.context(RecvSnafu)??;
1156        Ok(RemapManifestsResponse { manifest_paths })
1157    }
1158
1159    async fn copy_region_from(
1160        &self,
1161        region_id: RegionId,
1162        request: MitoCopyRegionFromRequest,
1163    ) -> Result<MitoCopyRegionFromResponse> {
1164        let (request, receiver) =
1165            WorkerRequest::try_from_copy_region_from_request(region_id, request)?;
1166        self.workers.submit_to_worker(region_id, request).await?;
1167        let response = receiver.await.context(RecvSnafu)??;
1168        Ok(response)
1169    }
1170
1171    fn role(&self, region_id: RegionId) -> Option<RegionRole> {
1172        self.workers
1173            .get_region(region_id)
1174            .map(|region| region.region_role())
1175    }
1176}
1177
1178fn map_batch_responses(responses: Vec<(RegionId, Result<AffectedRows>)>) -> BatchResponses {
1179    responses
1180        .into_iter()
1181        .map(|(region_id, response)| {
1182            (
1183                region_id,
1184                response.map(RegionResponse::new).map_err(BoxedError::new),
1185            )
1186        })
1187        .collect()
1188}
1189
1190#[async_trait]
1191impl RegionEngine for MitoEngine {
1192    fn name(&self) -> &str {
1193        MITO_ENGINE_NAME
1194    }
1195
1196    #[tracing::instrument(skip_all)]
1197    async fn handle_batch_open_requests(
1198        &self,
1199        parallelism: usize,
1200        requests: Vec<(RegionId, RegionOpenRequest)>,
1201    ) -> Result<BatchResponses, BoxedError> {
1202        // TODO(weny): add metrics.
1203        self.inner
1204            .handle_batch_open_requests(parallelism, requests)
1205            .await
1206            .map(map_batch_responses)
1207            .map_err(BoxedError::new)
1208    }
1209
1210    #[tracing::instrument(skip_all)]
1211    async fn handle_batch_catchup_requests(
1212        &self,
1213        parallelism: usize,
1214        requests: Vec<(RegionId, RegionCatchupRequest)>,
1215    ) -> Result<BatchResponses, BoxedError> {
1216        self.inner
1217            .handle_batch_catchup_requests(parallelism, requests)
1218            .await
1219            .map(map_batch_responses)
1220            .map_err(BoxedError::new)
1221    }
1222
1223    #[tracing::instrument(skip_all)]
1224    async fn handle_request(
1225        &self,
1226        region_id: RegionId,
1227        request: RegionRequest,
1228    ) -> Result<RegionResponse, BoxedError> {
1229        let _timer = HANDLE_REQUEST_ELAPSED
1230            .with_label_values(&[request.request_type()])
1231            .start_timer();
1232
1233        let is_alter = matches!(request, RegionRequest::Alter(_));
1234        let is_create = matches!(request, RegionRequest::Create(_));
1235        let mut response = self
1236            .inner
1237            .handle_request(region_id, request)
1238            .await
1239            .map(RegionResponse::new)
1240            .map_err(BoxedError::new)?;
1241
1242        if is_alter {
1243            self.handle_alter_response(region_id, &mut response)
1244                .map_err(BoxedError::new)?;
1245        } else if is_create {
1246            self.handle_create_response(region_id, &mut response)
1247                .map_err(BoxedError::new)?;
1248        }
1249
1250        Ok(response)
1251    }
1252
1253    #[tracing::instrument(skip_all)]
1254    async fn handle_query(
1255        &self,
1256        region_id: RegionId,
1257        request: ScanRequest,
1258    ) -> Result<RegionScannerRef, BoxedError> {
1259        self.scan_region(region_id, request)
1260            .map_err(BoxedError::new)?
1261            .region_scanner()
1262            .await
1263            .map_err(BoxedError::new)
1264    }
1265
1266    fn query_memory_tracker(&self) -> Option<QueryMemoryTracker> {
1267        Some(self.inner.scan_memory_tracker.clone())
1268    }
1269
1270    async fn get_committed_sequence(
1271        &self,
1272        region_id: RegionId,
1273    ) -> Result<SequenceNumber, BoxedError> {
1274        self.inner
1275            .get_committed_sequence(region_id)
1276            .map_err(BoxedError::new)
1277    }
1278
1279    /// Retrieve region's metadata.
1280    async fn get_metadata(
1281        &self,
1282        region_id: RegionId,
1283    ) -> std::result::Result<RegionMetadataRef, BoxedError> {
1284        self.inner.get_metadata(region_id).map_err(BoxedError::new)
1285    }
1286
1287    /// Stop the engine.
1288    ///
1289    /// Stopping the engine doesn't stop the underlying log store as other components might
1290    /// still use it. (When no other components are referencing the log store, it will
1291    /// automatically shutdown.)
1292    async fn stop(&self) -> std::result::Result<(), BoxedError> {
1293        self.inner.stop().await.map_err(BoxedError::new)
1294    }
1295
1296    fn region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic> {
1297        self.get_region_statistic(region_id)
1298    }
1299
1300    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError> {
1301        self.inner
1302            .set_region_role(region_id, role)
1303            .map_err(BoxedError::new)
1304    }
1305
1306    async fn set_region_role_state_gracefully(
1307        &self,
1308        region_id: RegionId,
1309        region_role_state: SettableRegionRoleState,
1310    ) -> Result<SetRegionRoleStateResponse, BoxedError> {
1311        let _timer = HANDLE_REQUEST_ELAPSED
1312            .with_label_values(&["set_region_role_state_gracefully"])
1313            .start_timer();
1314
1315        self.inner
1316            .set_region_role_state_gracefully(region_id, region_role_state)
1317            .await
1318            .map_err(BoxedError::new)
1319    }
1320
1321    async fn sync_region(
1322        &self,
1323        region_id: RegionId,
1324        request: SyncRegionFromRequest,
1325    ) -> Result<SyncRegionFromResponse, BoxedError> {
1326        let manifest_info = request
1327            .into_region_manifest_info()
1328            .context(UnexpectedSnafu {
1329                err_msg: "Expected a manifest info request",
1330            })
1331            .map_err(BoxedError::new)?;
1332        let (_, synced) = self
1333            .inner
1334            .sync_region(region_id, manifest_info)
1335            .await
1336            .map_err(BoxedError::new)?;
1337
1338        Ok(SyncRegionFromResponse::Mito { synced })
1339    }
1340
1341    async fn remap_manifests(
1342        &self,
1343        request: RemapManifestsRequest,
1344    ) -> Result<RemapManifestsResponse, BoxedError> {
1345        self.inner
1346            .remap_manifests(request)
1347            .await
1348            .map_err(BoxedError::new)
1349    }
1350
1351    fn role(&self, region_id: RegionId) -> Option<RegionRole> {
1352        self.inner.role(region_id)
1353    }
1354
1355    fn as_any(&self) -> &dyn Any {
1356        self
1357    }
1358}
1359
1360impl MitoEngine {
1361    fn handle_alter_response(
1362        &self,
1363        region_id: RegionId,
1364        response: &mut RegionResponse,
1365    ) -> Result<()> {
1366        if let Some(statistic) = self.region_statistic(region_id) {
1367            Self::encode_manifest_info_to_extensions(
1368                &region_id,
1369                statistic.manifest,
1370                &mut response.extensions,
1371            )?;
1372        }
1373        let column_metadatas = self
1374            .inner
1375            .find_region(region_id)
1376            .ok()
1377            .map(|r| r.metadata().column_metadatas.clone());
1378        if let Some(column_metadatas) = column_metadatas {
1379            Self::encode_column_metadatas_to_extensions(
1380                &region_id,
1381                column_metadatas,
1382                &mut response.extensions,
1383            )?;
1384        }
1385        Ok(())
1386    }
1387
1388    fn handle_create_response(
1389        &self,
1390        region_id: RegionId,
1391        response: &mut RegionResponse,
1392    ) -> Result<()> {
1393        let column_metadatas = self
1394            .inner
1395            .find_region(region_id)
1396            .ok()
1397            .map(|r| r.metadata().column_metadatas.clone());
1398        if let Some(column_metadatas) = column_metadatas {
1399            Self::encode_column_metadatas_to_extensions(
1400                &region_id,
1401                column_metadatas,
1402                &mut response.extensions,
1403            )?;
1404        }
1405        Ok(())
1406    }
1407}
1408
1409// Tests methods.
1410#[cfg(any(test, feature = "test"))]
1411#[allow(clippy::too_many_arguments)]
1412impl MitoEngine {
1413    /// Returns a new [MitoEngine] for tests.
1414    pub async fn new_for_test<S: LogStore>(
1415        data_home: &str,
1416        mut config: MitoConfig,
1417        log_store: Arc<S>,
1418        object_store_manager: ObjectStoreManagerRef,
1419        write_buffer_manager: Option<crate::flush::WriteBufferManagerRef>,
1420        listener: Option<crate::engine::listener::EventListenerRef>,
1421        time_provider: crate::time_provider::TimeProviderRef,
1422        schema_metadata_manager: SchemaMetadataManagerRef,
1423        file_ref_manager: FileReferenceManagerRef,
1424        partition_expr_fetcher: PartitionExprFetcherRef,
1425    ) -> Result<MitoEngine> {
1426        config.sanitize(data_home)?;
1427
1428        let config = Arc::new(config);
1429        let wal_raw_entry_reader = Arc::new(LogStoreRawEntryReader::new(log_store.clone()));
1430        let total_memory = get_total_memory_bytes().max(0) as u64;
1431        let scan_memory_limit = config.scan_memory_limit.resolve(total_memory) as usize;
1432        let scan_memory_tracker =
1433            QueryMemoryTracker::builder(scan_memory_limit, config.scan_memory_on_exhausted)
1434                .on_update(|usage| {
1435                    SCAN_MEMORY_USAGE_BYTES.set(usage as i64);
1436                })
1437                .on_exhausted(|| {
1438                    SCAN_MEMORY_EXHAUSTED_TOTAL.inc();
1439                })
1440                .on_reject(|| {
1441                    SCAN_REQUESTS_REJECTED_TOTAL.inc();
1442                })
1443                .build();
1444        Ok(MitoEngine {
1445            inner: Arc::new(EngineInner {
1446                workers: WorkerGroup::start_for_test(
1447                    config.clone(),
1448                    log_store,
1449                    object_store_manager,
1450                    write_buffer_manager,
1451                    listener,
1452                    schema_metadata_manager,
1453                    file_ref_manager,
1454                    time_provider,
1455                    partition_expr_fetcher,
1456                )
1457                .await?,
1458                config,
1459                wal_raw_entry_reader,
1460                scan_memory_tracker,
1461                #[cfg(feature = "enterprise")]
1462                extension_range_provider_factory: None,
1463            }),
1464        })
1465    }
1466
1467    /// Returns the purge scheduler.
1468    pub fn purge_scheduler(&self) -> &crate::schedule::scheduler::SchedulerRef {
1469        self.inner.workers.purge_scheduler()
1470    }
1471}
1472
1473#[cfg(test)]
1474mod tests {
1475    use std::time::Duration;
1476
1477    use super::*;
1478    use crate::sst::file::FileMeta;
1479
1480    #[test]
1481    fn test_is_valid_region_edit() {
1482        // Valid: has only "files_to_add"
1483        let edit = RegionEdit {
1484            files_to_add: vec![FileMeta::default()],
1485            files_to_remove: vec![],
1486            timestamp_ms: None,
1487            compaction_time_window: None,
1488            flushed_entry_id: None,
1489            flushed_sequence: None,
1490            committed_sequence: None,
1491        };
1492        assert!(is_valid_region_edit(&edit));
1493
1494        // Invalid: "files_to_add" and "files_to_remove" are both empty
1495        let edit = RegionEdit {
1496            files_to_add: vec![],
1497            files_to_remove: vec![],
1498            timestamp_ms: None,
1499            compaction_time_window: None,
1500            flushed_entry_id: None,
1501            flushed_sequence: None,
1502            committed_sequence: None,
1503        };
1504        assert!(!is_valid_region_edit(&edit));
1505
1506        // Valid: has only "files_to_remove"
1507        let edit = RegionEdit {
1508            files_to_add: vec![],
1509            files_to_remove: vec![FileMeta::default()],
1510            timestamp_ms: None,
1511            compaction_time_window: None,
1512            flushed_entry_id: None,
1513            flushed_sequence: None,
1514            committed_sequence: None,
1515        };
1516        assert!(is_valid_region_edit(&edit));
1517
1518        // Valid: both "files_to_add" and "files_to_remove" are not empty
1519        let edit = RegionEdit {
1520            files_to_add: vec![FileMeta::default()],
1521            files_to_remove: vec![FileMeta::default()],
1522            timestamp_ms: None,
1523            compaction_time_window: None,
1524            flushed_entry_id: None,
1525            flushed_sequence: None,
1526            committed_sequence: None,
1527        };
1528        assert!(is_valid_region_edit(&edit));
1529
1530        // Invalid: other fields are not all "None"s
1531        let edit = RegionEdit {
1532            files_to_add: vec![FileMeta::default()],
1533            files_to_remove: vec![],
1534            timestamp_ms: None,
1535            compaction_time_window: Some(Duration::from_secs(1)),
1536            flushed_entry_id: None,
1537            flushed_sequence: None,
1538            committed_sequence: None,
1539        };
1540        assert!(!is_valid_region_edit(&edit));
1541        let edit = RegionEdit {
1542            files_to_add: vec![FileMeta::default()],
1543            files_to_remove: vec![],
1544            timestamp_ms: None,
1545            compaction_time_window: None,
1546            flushed_entry_id: Some(1),
1547            flushed_sequence: None,
1548            committed_sequence: None,
1549        };
1550        assert!(!is_valid_region_edit(&edit));
1551        let edit = RegionEdit {
1552            files_to_add: vec![FileMeta::default()],
1553            files_to_remove: vec![],
1554            timestamp_ms: None,
1555            compaction_time_window: None,
1556            flushed_entry_id: None,
1557            flushed_sequence: Some(1),
1558            committed_sequence: None,
1559        };
1560        assert!(!is_valid_region_edit(&edit));
1561    }
1562}