mito2/region/
opener.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//! Region opener.
16
17use std::any::TypeId;
18use std::collections::HashMap;
19use std::sync::atomic::{AtomicI64, AtomicU64};
20use std::sync::{Arc, Mutex};
21use std::time::Instant;
22
23use common_telemetry::{debug, error, info, warn};
24use common_wal::options::WalOptions;
25use futures::StreamExt;
26use futures::future::BoxFuture;
27use log_store::kafka::log_store::KafkaLogStore;
28use log_store::noop::log_store::NoopLogStore;
29use log_store::raft_engine::log_store::RaftEngineLogStore;
30use object_store::manager::ObjectStoreManagerRef;
31use object_store::util::normalize_dir;
32use snafu::{OptionExt, ResultExt, ensure};
33use store_api::logstore::LogStore;
34use store_api::logstore::provider::Provider;
35use store_api::metadata::{
36    ColumnMetadata, RegionMetadata, RegionMetadataBuilder, RegionMetadataRef,
37};
38use store_api::region_engine::RegionRole;
39use store_api::region_request::PathType;
40use store_api::storage::{ColumnId, RegionId};
41
42use crate::access_layer::AccessLayer;
43use crate::cache::CacheManagerRef;
44use crate::cache::file_cache::{FileCache, FileType, IndexKey};
45use crate::config::MitoConfig;
46use crate::error;
47use crate::error::{
48    EmptyRegionDirSnafu, InvalidMetadataSnafu, ObjectStoreNotFoundSnafu, RegionCorruptedSnafu,
49    Result, StaleLogEntrySnafu,
50};
51use crate::manifest::action::RegionManifest;
52use crate::manifest::manager::{RegionManifestManager, RegionManifestOptions};
53use crate::memtable::MemtableBuilderProvider;
54use crate::memtable::bulk::part::BulkPart;
55use crate::memtable::time_partition::{TimePartitions, TimePartitionsRef};
56use crate::metrics::{CACHE_FILL_DOWNLOADED_FILES, CACHE_FILL_PENDING_FILES};
57use crate::region::options::RegionOptions;
58use crate::region::version::{VersionBuilder, VersionControl, VersionControlRef};
59use crate::region::{
60    ManifestContext, ManifestStats, MitoRegion, MitoRegionRef, RegionLeaderState, RegionRoleState,
61};
62use crate::region_write_ctx::RegionWriteCtx;
63use crate::request::OptionOutputTx;
64use crate::schedule::scheduler::SchedulerRef;
65use crate::sst::FormatType;
66use crate::sst::file::{RegionFileId, RegionIndexId};
67use crate::sst::file_purger::{FilePurgerRef, create_file_purger};
68use crate::sst::file_ref::FileReferenceManagerRef;
69use crate::sst::index::intermediate::IntermediateManager;
70use crate::sst::index::puffin_manager::PuffinManagerFactory;
71use crate::sst::location::{self, region_dir_from_table_dir};
72use crate::time_provider::TimeProviderRef;
73use crate::wal::entry_reader::WalEntryReader;
74use crate::wal::{EntryId, Wal};
75
76/// A fetcher to retrieve partition expr for a region.
77///
78/// Compatibility: older regions didn't persist `partition_expr` in engine metadata,
79/// while newer ones do. On open, we backfill it via this fetcher and persist it
80/// to the manifest so future opens don't need refetching.
81#[async_trait::async_trait]
82pub trait PartitionExprFetcher {
83    async fn fetch_expr(&self, region_id: RegionId) -> Option<String>;
84}
85
86pub type PartitionExprFetcherRef = Arc<dyn PartitionExprFetcher + Send + Sync>;
87
88/// Builder to create a new [MitoRegion] or open an existing one.
89pub(crate) struct RegionOpener {
90    region_id: RegionId,
91    metadata_builder: Option<RegionMetadataBuilder>,
92    memtable_builder_provider: MemtableBuilderProvider,
93    object_store_manager: ObjectStoreManagerRef,
94    table_dir: String,
95    path_type: PathType,
96    purge_scheduler: SchedulerRef,
97    options: Option<RegionOptions>,
98    cache_manager: Option<CacheManagerRef>,
99    skip_wal_replay: bool,
100    puffin_manager_factory: PuffinManagerFactory,
101    intermediate_manager: IntermediateManager,
102    time_provider: TimeProviderRef,
103    stats: ManifestStats,
104    wal_entry_reader: Option<Box<dyn WalEntryReader>>,
105    replay_checkpoint: Option<u64>,
106    file_ref_manager: FileReferenceManagerRef,
107    partition_expr_fetcher: PartitionExprFetcherRef,
108}
109
110impl RegionOpener {
111    /// Returns a new opener.
112    // TODO(LFC): Reduce the number of arguments.
113    #[allow(clippy::too_many_arguments)]
114    pub(crate) fn new(
115        region_id: RegionId,
116        table_dir: &str,
117        path_type: PathType,
118        memtable_builder_provider: MemtableBuilderProvider,
119        object_store_manager: ObjectStoreManagerRef,
120        purge_scheduler: SchedulerRef,
121        puffin_manager_factory: PuffinManagerFactory,
122        intermediate_manager: IntermediateManager,
123        time_provider: TimeProviderRef,
124        file_ref_manager: FileReferenceManagerRef,
125        partition_expr_fetcher: PartitionExprFetcherRef,
126    ) -> RegionOpener {
127        RegionOpener {
128            region_id,
129            metadata_builder: None,
130            memtable_builder_provider,
131            object_store_manager,
132            table_dir: normalize_dir(table_dir),
133            path_type,
134            purge_scheduler,
135            options: None,
136            cache_manager: None,
137            skip_wal_replay: false,
138            puffin_manager_factory,
139            intermediate_manager,
140            time_provider,
141            stats: Default::default(),
142            wal_entry_reader: None,
143            replay_checkpoint: None,
144            file_ref_manager,
145            partition_expr_fetcher,
146        }
147    }
148
149    /// Sets metadata builder of the region to create.
150    pub(crate) fn metadata_builder(mut self, builder: RegionMetadataBuilder) -> Self {
151        self.metadata_builder = Some(builder);
152        self
153    }
154
155    /// Computes the region directory from table_dir and region_id.
156    fn region_dir(&self) -> String {
157        region_dir_from_table_dir(&self.table_dir, self.region_id, self.path_type)
158    }
159
160    /// Builds the region metadata.
161    ///
162    /// # Panics
163    /// - Panics if `options` is not set.
164    /// - Panics if `metadata_builder` is not set.
165    fn build_metadata(&mut self) -> Result<RegionMetadata> {
166        let options = self.options.as_ref().unwrap();
167        let mut metadata_builder = self.metadata_builder.take().unwrap();
168        metadata_builder.primary_key_encoding(options.primary_key_encoding());
169        metadata_builder.build().context(InvalidMetadataSnafu)
170    }
171
172    /// Parses and sets options for the region.
173    pub(crate) fn parse_options(self, options: HashMap<String, String>) -> Result<Self> {
174        self.options(RegionOptions::try_from(&options)?)
175    }
176
177    /// Sets the replay checkpoint for the region.
178    pub(crate) fn replay_checkpoint(mut self, replay_checkpoint: Option<u64>) -> Self {
179        self.replay_checkpoint = replay_checkpoint;
180        self
181    }
182
183    /// If a [WalEntryReader] is set, the [RegionOpener] will use [WalEntryReader] instead of
184    /// constructing a new one from scratch.
185    pub(crate) fn wal_entry_reader(
186        mut self,
187        wal_entry_reader: Option<Box<dyn WalEntryReader>>,
188    ) -> Self {
189        self.wal_entry_reader = wal_entry_reader;
190        self
191    }
192
193    /// Sets options for the region.
194    pub(crate) fn options(mut self, options: RegionOptions) -> Result<Self> {
195        options.validate()?;
196        self.options = Some(options);
197        Ok(self)
198    }
199
200    /// Sets the cache manager for the region.
201    pub(crate) fn cache(mut self, cache_manager: Option<CacheManagerRef>) -> Self {
202        self.cache_manager = cache_manager;
203        self
204    }
205
206    /// Sets the `skip_wal_replay`.
207    pub(crate) fn skip_wal_replay(mut self, skip: bool) -> Self {
208        self.skip_wal_replay = skip;
209        self
210    }
211
212    /// Writes region manifest and creates a new region if it does not exist.
213    /// Opens the region if it already exists.
214    ///
215    /// # Panics
216    /// - Panics if `metadata_builder` is not set.
217    /// - Panics if `options` is not set.
218    pub(crate) async fn create_or_open<S: LogStore>(
219        mut self,
220        config: &MitoConfig,
221        wal: &Wal<S>,
222    ) -> Result<MitoRegionRef> {
223        let region_id = self.region_id;
224        let region_dir = self.region_dir();
225        let metadata = self.build_metadata()?;
226        // Tries to open the region.
227        match self.maybe_open(config, wal).await {
228            Ok(Some(region)) => {
229                let recovered = region.metadata();
230                // Checks the schema of the region.
231                let expect = &metadata;
232                check_recovered_region(
233                    &recovered,
234                    expect.region_id,
235                    &expect.column_metadatas,
236                    &expect.primary_key,
237                )?;
238                // To keep consistency with Create behavior, set the opened Region to RegionRole::Leader.
239                region.set_role(RegionRole::Leader);
240
241                return Ok(region);
242            }
243            Ok(None) => {
244                debug!(
245                    "No data under directory {}, region_id: {}",
246                    region_dir, self.region_id
247                );
248            }
249            Err(e) => {
250                warn!(e;
251                    "Failed to open region {} before creating it, region_dir: {}",
252                    self.region_id, region_dir
253                );
254            }
255        }
256        // Safety: must be set before calling this method.
257        let mut options = self.options.take().unwrap();
258        let object_store = get_object_store(&options.storage, &self.object_store_manager)?;
259        let provider = self.provider::<S>(&options.wal_options)?;
260        let metadata = Arc::new(metadata);
261        // Sets the sst_format based on options or flat_format flag
262        let sst_format = if let Some(format) = options.sst_format {
263            format
264        } else if config.default_experimental_flat_format {
265            options.sst_format = Some(FormatType::Flat);
266            FormatType::Flat
267        } else {
268            // Default to PrimaryKeyParquet for newly created regions
269            options.sst_format = Some(FormatType::PrimaryKey);
270            FormatType::PrimaryKey
271        };
272        // Create a manifest manager for this region and writes regions to the manifest file.
273        let mut region_manifest_options =
274            RegionManifestOptions::new(config, &region_dir, &object_store);
275        // Set manifest cache if available
276        region_manifest_options.manifest_cache = self
277            .cache_manager
278            .as_ref()
279            .and_then(|cm| cm.write_cache())
280            .and_then(|wc| wc.manifest_cache());
281        // For remote WAL, we need to set flushed_entry_id to current topic's latest entry id.
282        let flushed_entry_id = provider.initial_flushed_entry_id::<S>(wal.store());
283        let manifest_manager = RegionManifestManager::new(
284            metadata.clone(),
285            flushed_entry_id,
286            region_manifest_options,
287            sst_format,
288            &self.stats,
289        )
290        .await?;
291
292        let memtable_builder = self.memtable_builder_provider.builder_for_options(&options);
293        let part_duration = options.compaction.time_window();
294        // Initial memtable id is 0.
295        let mutable = Arc::new(TimePartitions::new(
296            metadata.clone(),
297            memtable_builder.clone(),
298            0,
299            part_duration,
300        ));
301
302        debug!(
303            "Create region {} with options: {:?}, default_flat_format: {}",
304            region_id, options, config.default_experimental_flat_format
305        );
306
307        let version = VersionBuilder::new(metadata, mutable)
308            .options(options)
309            .build();
310        let version_control = Arc::new(VersionControl::new(version));
311        let access_layer = Arc::new(AccessLayer::new(
312            self.table_dir.clone(),
313            self.path_type,
314            object_store,
315            self.puffin_manager_factory,
316            self.intermediate_manager,
317        ));
318        let now = self.time_provider.current_time_millis();
319
320        Ok(Arc::new(MitoRegion {
321            region_id,
322            version_control,
323            access_layer: access_layer.clone(),
324            // Region is writable after it is created.
325            manifest_ctx: Arc::new(ManifestContext::new(
326                manifest_manager,
327                RegionRoleState::Leader(RegionLeaderState::Writable),
328            )),
329            file_purger: create_file_purger(
330                config.gc.enable,
331                self.path_type,
332                self.purge_scheduler,
333                access_layer,
334                self.cache_manager,
335                self.file_ref_manager.clone(),
336            ),
337            provider,
338            last_flush_millis: AtomicI64::new(now),
339            last_compaction_millis: AtomicI64::new(now),
340            time_provider: self.time_provider.clone(),
341            topic_latest_entry_id: AtomicU64::new(0),
342            written_bytes: Arc::new(AtomicU64::new(0)),
343            stats: self.stats,
344            staging_partition_info: Mutex::new(None),
345        }))
346    }
347
348    /// Opens an existing region in read only mode.
349    ///
350    /// Returns error if the region doesn't exist.
351    pub(crate) async fn open<S: LogStore>(
352        mut self,
353        config: &MitoConfig,
354        wal: &Wal<S>,
355    ) -> Result<MitoRegionRef> {
356        let region_id = self.region_id;
357        let region_dir = self.region_dir();
358        let region = self
359            .maybe_open(config, wal)
360            .await?
361            .with_context(|| EmptyRegionDirSnafu {
362                region_id,
363                region_dir: &region_dir,
364            })?;
365
366        ensure!(
367            region.region_id == self.region_id,
368            RegionCorruptedSnafu {
369                region_id: self.region_id,
370                reason: format!(
371                    "recovered region has different region id {}",
372                    region.region_id
373                ),
374            }
375        );
376
377        Ok(region)
378    }
379
380    fn provider<S: LogStore>(&self, wal_options: &WalOptions) -> Result<Provider> {
381        match wal_options {
382            WalOptions::RaftEngine => {
383                ensure!(
384                    TypeId::of::<RaftEngineLogStore>() == TypeId::of::<S>()
385                        || TypeId::of::<NoopLogStore>() == TypeId::of::<S>(),
386                    error::IncompatibleWalProviderChangeSnafu {
387                        global: "`kafka`",
388                        region: "`raft_engine`",
389                    }
390                );
391                Ok(Provider::raft_engine_provider(self.region_id.as_u64()))
392            }
393            WalOptions::Kafka(options) => {
394                ensure!(
395                    TypeId::of::<KafkaLogStore>() == TypeId::of::<S>()
396                        || TypeId::of::<NoopLogStore>() == TypeId::of::<S>(),
397                    error::IncompatibleWalProviderChangeSnafu {
398                        global: "`raft_engine`",
399                        region: "`kafka`",
400                    }
401                );
402                Ok(Provider::kafka_provider(options.topic.clone()))
403            }
404            WalOptions::Noop => Ok(Provider::noop_provider()),
405        }
406    }
407
408    /// Tries to open the region and returns `None` if the region directory is empty.
409    async fn maybe_open<S: LogStore>(
410        &mut self,
411        config: &MitoConfig,
412        wal: &Wal<S>,
413    ) -> Result<Option<MitoRegionRef>> {
414        let now = Instant::now();
415        let mut region_options = self.options.as_ref().unwrap().clone();
416        let object_storage = get_object_store(&region_options.storage, &self.object_store_manager)?;
417        let mut region_manifest_options =
418            RegionManifestOptions::new(config, &self.region_dir(), &object_storage);
419        // Set manifest cache if available
420        region_manifest_options.manifest_cache = self
421            .cache_manager
422            .as_ref()
423            .and_then(|cm| cm.write_cache())
424            .and_then(|wc| wc.manifest_cache());
425        let Some(manifest_manager) =
426            RegionManifestManager::open(region_manifest_options, &self.stats).await?
427        else {
428            return Ok(None);
429        };
430
431        // Backfill `partition_expr` if missing. Use the backfilled metadata in-memory during this open.
432        let manifest = manifest_manager.manifest();
433        let metadata = if manifest.metadata.partition_expr.is_none()
434            && let Some(expr_json) = self.partition_expr_fetcher.fetch_expr(self.region_id).await
435        {
436            let metadata = manifest.metadata.as_ref().clone();
437            let mut builder = RegionMetadataBuilder::from_existing(metadata);
438            builder.partition_expr_json(Some(expr_json));
439            Arc::new(builder.build().context(InvalidMetadataSnafu)?)
440        } else {
441            manifest.metadata.clone()
442        };
443        // Updates the region options with the manifest.
444        sanitize_region_options(&manifest, &mut region_options);
445
446        let region_id = self.region_id;
447        let provider = self.provider::<S>(&region_options.wal_options)?;
448        let wal_entry_reader = self
449            .wal_entry_reader
450            .take()
451            .unwrap_or_else(|| wal.wal_entry_reader(&provider, region_id, None));
452        let on_region_opened = wal.on_region_opened();
453        let object_store = get_object_store(&region_options.storage, &self.object_store_manager)?;
454
455        debug!(
456            "Open region {} at {} with options: {:?}",
457            region_id, self.table_dir, self.options
458        );
459
460        let access_layer = Arc::new(AccessLayer::new(
461            self.table_dir.clone(),
462            self.path_type,
463            object_store,
464            self.puffin_manager_factory.clone(),
465            self.intermediate_manager.clone(),
466        ));
467        let file_purger = create_file_purger(
468            config.gc.enable,
469            self.path_type,
470            self.purge_scheduler.clone(),
471            access_layer.clone(),
472            self.cache_manager.clone(),
473            self.file_ref_manager.clone(),
474        );
475        // We should sanitize the region options before creating a new memtable.
476        let memtable_builder = self
477            .memtable_builder_provider
478            .builder_for_options(&region_options);
479        // Use compaction time window in the manifest if region doesn't provide
480        // the time window option.
481        let part_duration = region_options
482            .compaction
483            .time_window()
484            .or(manifest.compaction_time_window);
485        // Initial memtable id is 0.
486        let mutable = Arc::new(TimePartitions::new(
487            metadata.clone(),
488            memtable_builder.clone(),
489            0,
490            part_duration,
491        ));
492
493        // Updates region options by manifest before creating version.
494        let version_builder = version_builder_from_manifest(
495            &manifest,
496            metadata,
497            file_purger.clone(),
498            mutable,
499            region_options,
500        );
501        let version = version_builder.build();
502        let flushed_entry_id = version.flushed_entry_id;
503        let version_control = Arc::new(VersionControl::new(version));
504
505        let topic_latest_entry_id = if !self.skip_wal_replay {
506            let replay_from_entry_id = self
507                .replay_checkpoint
508                .unwrap_or_default()
509                .max(flushed_entry_id);
510            info!(
511                "Start replaying memtable at replay_from_entry_id: {} for region {}, manifest version: {}, flushed entry id: {}, elapsed: {:?}",
512                replay_from_entry_id,
513                region_id,
514                manifest.manifest_version,
515                flushed_entry_id,
516                now.elapsed()
517            );
518            replay_memtable(
519                &provider,
520                wal_entry_reader,
521                region_id,
522                replay_from_entry_id,
523                &version_control,
524                config.allow_stale_entries,
525                on_region_opened,
526            )
527            .await?;
528            // For remote WAL, we need to set topic_latest_entry_id to current topic's latest entry id.
529            // Only set after the WAL replay is completed.
530
531            if provider.is_remote_wal() && version_control.current().version.memtables.is_empty() {
532                wal.store().latest_entry_id(&provider).unwrap_or(0)
533            } else {
534                0
535            }
536        } else {
537            info!(
538                "Skip the WAL replay for region: {}, manifest version: {}, flushed_entry_id: {}, elapsed: {:?}",
539                region_id,
540                manifest.manifest_version,
541                flushed_entry_id,
542                now.elapsed()
543            );
544
545            0
546        };
547
548        if let Some(committed_in_manifest) = manifest.committed_sequence {
549            let committed_after_replay = version_control.committed_sequence();
550            if committed_in_manifest > committed_after_replay {
551                info!(
552                    "Overriding committed sequence, region: {}, flushed_sequence: {}, committed_sequence: {} -> {}",
553                    self.region_id,
554                    version_control.current().version.flushed_sequence,
555                    version_control.committed_sequence(),
556                    committed_in_manifest
557                );
558                version_control.set_committed_sequence(committed_in_manifest);
559            }
560        }
561
562        let now = self.time_provider.current_time_millis();
563
564        let region = MitoRegion {
565            region_id: self.region_id,
566            version_control: version_control.clone(),
567            access_layer: access_layer.clone(),
568            // Region is always opened in read only mode.
569            manifest_ctx: Arc::new(ManifestContext::new(
570                manifest_manager,
571                RegionRoleState::Follower,
572            )),
573            file_purger,
574            provider: provider.clone(),
575            last_flush_millis: AtomicI64::new(now),
576            last_compaction_millis: AtomicI64::new(now),
577            time_provider: self.time_provider.clone(),
578            topic_latest_entry_id: AtomicU64::new(topic_latest_entry_id),
579            written_bytes: Arc::new(AtomicU64::new(0)),
580            stats: self.stats.clone(),
581            // TODO(weny): reload the staging partition info from the manifest.
582            staging_partition_info: Mutex::new(None),
583        };
584
585        let region = Arc::new(region);
586
587        maybe_load_cache(&region, config, &self.cache_manager);
588
589        Ok(Some(region))
590    }
591}
592
593/// Creates a version builder from a region manifest.
594pub(crate) fn version_builder_from_manifest(
595    manifest: &RegionManifest,
596    metadata: RegionMetadataRef,
597    file_purger: FilePurgerRef,
598    mutable: TimePartitionsRef,
599    region_options: RegionOptions,
600) -> VersionBuilder {
601    VersionBuilder::new(metadata, mutable)
602        .add_files(file_purger, manifest.files.values().cloned())
603        .flushed_entry_id(manifest.flushed_entry_id)
604        .flushed_sequence(manifest.flushed_sequence)
605        .truncated_entry_id(manifest.truncated_entry_id)
606        .compaction_time_window(manifest.compaction_time_window)
607        .options(region_options)
608}
609
610/// Updates region options by persistent options.
611pub(crate) fn sanitize_region_options(manifest: &RegionManifest, options: &mut RegionOptions) {
612    let option_format = options.sst_format.unwrap_or_default();
613    if option_format != manifest.sst_format {
614        common_telemetry::warn!(
615            "Overriding SST format from {:?} to {:?} for region {}",
616            option_format,
617            manifest.sst_format,
618            manifest.metadata.region_id,
619        );
620        options.sst_format = Some(manifest.sst_format);
621    }
622    if let Some(manifest_append_mode) = manifest.append_mode
623        && options.append_mode != manifest_append_mode
624    {
625        common_telemetry::warn!(
626            "Overriding append_mode from {} to {} for region {}",
627            options.append_mode,
628            manifest_append_mode,
629            manifest.metadata.region_id,
630        );
631        options.append_mode = manifest_append_mode;
632    }
633    if options.append_mode && options.merge_mode.take().is_some() {
634        common_telemetry::warn!(
635            "Ignoring merge_mode because append_mode is enabled for region {}",
636            manifest.metadata.region_id,
637        );
638    }
639}
640
641/// Sanitizes open request options before parsing.
642pub(crate) fn sanitize_open_request_options(options: &mut HashMap<String, String>) {
643    let append_mode_enabled = options
644        .get("append_mode")
645        .is_some_and(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "true" | "1"));
646
647    if append_mode_enabled && options.remove("merge_mode").is_some() {
648        common_telemetry::warn!(
649            "Ignoring merge_mode in open request options because append_mode is enabled"
650        );
651    }
652}
653
654/// Returns an object store corresponding to `name`. If `name` is `None`, this method returns the default object store.
655pub fn get_object_store(
656    name: &Option<String>,
657    object_store_manager: &ObjectStoreManagerRef,
658) -> Result<object_store::ObjectStore> {
659    if let Some(name) = name {
660        Ok(object_store_manager
661            .find(name)
662            .with_context(|| ObjectStoreNotFoundSnafu {
663                object_store: name.clone(),
664            })?
665            .clone())
666    } else {
667        Ok(object_store_manager.default_object_store().clone())
668    }
669}
670
671/// Checks whether the recovered region has the same schema as region to create.
672pub(crate) fn check_recovered_region(
673    recovered: &RegionMetadata,
674    region_id: RegionId,
675    column_metadatas: &[ColumnMetadata],
676    primary_key: &[ColumnId],
677) -> Result<()> {
678    if recovered.region_id != region_id {
679        error!(
680            "Recovered region {}, expect region {}",
681            recovered.region_id, region_id
682        );
683        return RegionCorruptedSnafu {
684            region_id,
685            reason: format!(
686                "recovered metadata has different region id {}",
687                recovered.region_id
688            ),
689        }
690        .fail();
691    }
692    if recovered.column_metadatas != column_metadatas {
693        error!(
694            "Unexpected schema in recovered region {}, recovered: {:?}, expect: {:?}",
695            recovered.region_id, recovered.column_metadatas, column_metadatas
696        );
697
698        return RegionCorruptedSnafu {
699            region_id,
700            reason: "recovered metadata has different schema",
701        }
702        .fail();
703    }
704    if recovered.primary_key != primary_key {
705        error!(
706            "Unexpected primary key in recovered region {}, recovered: {:?}, expect: {:?}",
707            recovered.region_id, recovered.primary_key, primary_key
708        );
709
710        return RegionCorruptedSnafu {
711            region_id,
712            reason: "recovered metadata has different primary key",
713        }
714        .fail();
715    }
716
717    Ok(())
718}
719
720/// Replays the mutations from WAL and inserts mutations to memtable of given region.
721pub(crate) async fn replay_memtable<F>(
722    provider: &Provider,
723    mut wal_entry_reader: Box<dyn WalEntryReader>,
724    region_id: RegionId,
725    flushed_entry_id: EntryId,
726    version_control: &VersionControlRef,
727    allow_stale_entries: bool,
728    on_region_opened: F,
729) -> Result<EntryId>
730where
731    F: FnOnce(RegionId, EntryId, &Provider) -> BoxFuture<Result<()>> + Send,
732{
733    let now = Instant::now();
734    let mut rows_replayed = 0;
735    // Last entry id should start from flushed entry id since there might be no
736    // data in the WAL.
737    let mut last_entry_id = flushed_entry_id;
738    let replay_from_entry_id = flushed_entry_id + 1;
739
740    let mut wal_stream = wal_entry_reader.read(provider, replay_from_entry_id)?;
741    while let Some(res) = wal_stream.next().await {
742        let (entry_id, entry) = res?;
743        if entry_id <= flushed_entry_id {
744            warn!(
745                "Stale WAL entries read during replay, region id: {}, flushed entry id: {}, entry id read: {}",
746                region_id, flushed_entry_id, entry_id
747            );
748            ensure!(
749                allow_stale_entries,
750                StaleLogEntrySnafu {
751                    region_id,
752                    flushed_entry_id,
753                    unexpected_entry_id: entry_id,
754                }
755            );
756        }
757        last_entry_id = last_entry_id.max(entry_id);
758
759        let mut region_write_ctx = RegionWriteCtx::new(
760            region_id,
761            version_control,
762            provider.clone(),
763            // For WAL replay, we don't need to track the write bytes rate.
764            None,
765        );
766        for mutation in entry.mutations {
767            rows_replayed += mutation
768                .rows
769                .as_ref()
770                .map(|rows| rows.rows.len())
771                .unwrap_or(0);
772            region_write_ctx.push_mutation(
773                mutation.op_type,
774                mutation.rows,
775                mutation.write_hint,
776                OptionOutputTx::none(),
777                // We should respect the sequence in WAL during replay.
778                Some(mutation.sequence),
779            );
780        }
781
782        for bulk_entry in entry.bulk_entries {
783            let part = BulkPart::try_from(bulk_entry)?;
784            rows_replayed += part.num_rows();
785            // During replay, we should adopt the sequence from WAL.
786            let bulk_sequence_from_wal = part.sequence;
787            ensure!(
788                region_write_ctx.push_bulk(
789                    OptionOutputTx::none(),
790                    part,
791                    Some(bulk_sequence_from_wal)
792                ),
793                RegionCorruptedSnafu {
794                    region_id,
795                    reason: "unable to replay memtable with bulk entries",
796                }
797            );
798        }
799
800        // set next_entry_id and write to memtable.
801        region_write_ctx.set_next_entry_id(last_entry_id + 1);
802        region_write_ctx.write_memtable().await;
803        region_write_ctx.write_bulk().await;
804    }
805
806    // TODO(weny): We need to update `flushed_entry_id` in the region manifest
807    // to avoid reading potentially incomplete entries in the future.
808    (on_region_opened)(region_id, flushed_entry_id, provider).await?;
809
810    let series_count = version_control.current().series_count();
811    info!(
812        "Replay WAL for region: {}, provider: {:?}, rows recovered: {}, replay from entry id: {}, last entry id: {}, total timeseries replayed: {}, elapsed: {:?}",
813        region_id,
814        provider,
815        rows_replayed,
816        replay_from_entry_id,
817        last_entry_id,
818        series_count,
819        now.elapsed()
820    );
821    Ok(last_entry_id)
822}
823
824/// A task to load and fill the region file cache.
825pub(crate) struct RegionLoadCacheTask {
826    region: MitoRegionRef,
827}
828
829impl RegionLoadCacheTask {
830    pub(crate) fn new(region: MitoRegionRef) -> Self {
831        Self { region }
832    }
833
834    /// Fills the file cache with index files from the region.
835    pub(crate) async fn fill_cache(&self, file_cache: &FileCache) {
836        let region_id = self.region.region_id;
837        let table_dir = self.region.access_layer.table_dir();
838        let path_type = self.region.access_layer.path_type();
839        let object_store = self.region.access_layer.object_store();
840        let version_control = &self.region.version_control;
841
842        // Collects IndexKeys, file sizes, and max timestamps for files that need to be downloaded
843        let mut files_to_download = Vec::new();
844        let mut files_already_cached = 0;
845
846        {
847            let version = version_control.current().version;
848            for level in version.ssts.levels() {
849                for file_handle in level.files.values() {
850                    let file_meta = file_handle.meta_ref();
851                    if file_meta.exists_index() {
852                        let puffin_key = IndexKey::new(
853                            file_meta.region_id,
854                            file_meta.file_id,
855                            FileType::Puffin(file_meta.index_version),
856                        );
857
858                        if !file_cache.contains_key(&puffin_key) {
859                            files_to_download.push((
860                                puffin_key,
861                                file_meta.index_file_size,
862                                file_meta.time_range.1, // max timestamp
863                            ));
864                        } else {
865                            files_already_cached += 1;
866                        }
867                    }
868                }
869            }
870            // Releases the Version after the scope to avoid holding the memtables and file handles
871            // for a long time.
872        }
873
874        // Sorts files by max timestamp in descending order to loads latest files first
875        files_to_download.sort_by(|a, b| b.2.cmp(&a.2));
876
877        let total_files = files_to_download.len() as i64;
878
879        info!(
880            "Starting background index cache preload for region {}, total_files_to_download: {}, files_already_cached: {}",
881            region_id, total_files, files_already_cached
882        );
883
884        CACHE_FILL_PENDING_FILES.add(total_files);
885
886        let mut files_downloaded = 0;
887        let mut files_skipped = 0;
888
889        for (puffin_key, file_size, max_timestamp) in files_to_download {
890            let current_size = file_cache.puffin_cache_size();
891            let capacity = file_cache.puffin_cache_capacity();
892            let region_state = self.region.state();
893            if !can_load_cache(region_state) {
894                info!(
895                    "Stopping index cache by state: {:?}, region: {}, current_size: {}, capacity: {}",
896                    region_state, region_id, current_size, capacity
897                );
898                break;
899            }
900
901            // Checks if adding this file would exceed capacity
902            if current_size + file_size > capacity {
903                info!(
904                    "Stopping index cache preload due to capacity limit, region: {}, file_id: {}, current_size: {}, file_size: {}, capacity: {}, file_timestamp: {:?}",
905                    region_id, puffin_key.file_id, current_size, file_size, capacity, max_timestamp
906                );
907                files_skipped = (total_files - files_downloaded) as usize;
908                CACHE_FILL_PENDING_FILES.sub(total_files - files_downloaded);
909                break;
910            }
911
912            let index_version = if let FileType::Puffin(version) = puffin_key.file_type {
913                version
914            } else {
915                unreachable!("`files_to_download` should only contains Puffin files");
916            };
917            let index_id = RegionIndexId::new(
918                RegionFileId::new(puffin_key.region_id, puffin_key.file_id),
919                index_version,
920            );
921
922            let index_remote_path = location::index_file_path(table_dir, index_id, path_type);
923
924            match file_cache
925                .download(puffin_key, &index_remote_path, object_store, file_size)
926                .await
927            {
928                Ok(_) => {
929                    debug!(
930                        "Downloaded index file to write cache, region: {}, file_id: {}",
931                        region_id, puffin_key.file_id
932                    );
933                    files_downloaded += 1;
934                    CACHE_FILL_DOWNLOADED_FILES.inc_by(1);
935                    CACHE_FILL_PENDING_FILES.dec();
936                }
937                Err(e) => {
938                    warn!(
939                        e; "Failed to download index file to write cache, region: {}, file_id: {}",
940                        region_id, puffin_key.file_id
941                    );
942                    CACHE_FILL_PENDING_FILES.dec();
943                }
944            }
945        }
946
947        info!(
948            "Completed background cache fill task for region {}, total_files: {}, files_downloaded: {}, files_already_cached: {}, files_skipped: {}",
949            region_id, total_files, files_downloaded, files_already_cached, files_skipped
950        );
951    }
952}
953
954/// Loads all index (Puffin) files from the version into the write cache.
955fn maybe_load_cache(
956    region: &MitoRegionRef,
957    config: &MitoConfig,
958    cache_manager: &Option<CacheManagerRef>,
959) {
960    let Some(cache_manager) = cache_manager else {
961        return;
962    };
963    let Some(write_cache) = cache_manager.write_cache() else {
964        return;
965    };
966
967    let preload_enabled = config.preload_index_cache;
968    if !preload_enabled {
969        return;
970    }
971
972    let task = RegionLoadCacheTask::new(region.clone());
973    write_cache.load_region_cache(task);
974}
975
976fn can_load_cache(state: RegionRoleState) -> bool {
977    match state {
978        RegionRoleState::Leader(RegionLeaderState::Writable)
979        | RegionRoleState::Leader(RegionLeaderState::Staging)
980        | RegionRoleState::Leader(RegionLeaderState::Altering)
981        | RegionRoleState::Leader(RegionLeaderState::EnteringStaging)
982        | RegionRoleState::Leader(RegionLeaderState::Editing)
983        | RegionRoleState::Follower => true,
984        // The region will be closed soon if it is downgrading.
985        RegionRoleState::Leader(RegionLeaderState::Downgrading)
986        | RegionRoleState::Leader(RegionLeaderState::Dropping)
987        | RegionRoleState::Leader(RegionLeaderState::Truncating) => false,
988    }
989}
mito2/region/opener.rs

mito2/region/
opener.rs
