store_api/
region_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//! Region Engine's definition
16
17use std::any::Any;
18use std::collections::HashMap;
19use std::fmt::{Debug, Display};
20use std::sync::{Arc, Mutex};
21
22use api::greptime_proto::v1::meta::{GrantedRegion as PbGrantedRegion, RegionRole as PbRegionRole};
23use api::region::RegionResponse;
24use async_trait::async_trait;
25use common_error::ext::BoxedError;
26use common_recordbatch::adapter::RegionQueryStatCounters;
27use common_recordbatch::{EmptyRecordBatchStream, QueryMemoryTracker, SendableRecordBatchStream};
28use common_time::Timestamp;
29use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
30use datafusion_physical_plan::{DisplayAs, DisplayFormatType, PhysicalExpr};
31use datatypes::schema::SchemaRef;
32use futures::future::join_all;
33use serde::{Deserialize, Serialize};
34use tokio::sync::Semaphore;
35
36use crate::logstore::entry;
37use crate::metadata::RegionMetadataRef;
38use crate::region_request::{
39    BatchRegionDdlRequest, RegionCatchupRequest, RegionOpenRequest, RegionRequest,
40};
41use crate::storage::{FileId, RegionId, ScanRequest, SequenceNumber};
42
43/// The settable region role state.
44#[derive(Debug, PartialEq, Eq, Clone, Copy)]
45pub enum SettableRegionRoleState {
46    Follower,
47    DowngradingLeader,
48    /// Exit staging mode and return to normal leader state. Only allowed from staging state.
49    Leader,
50    /// Enter staging mode. Region remains writable but disables checkpoint and compaction. Only allowed from normal leader state.
51    StagingLeader,
52}
53
54impl Display for SettableRegionRoleState {
55    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        match self {
57            SettableRegionRoleState::Follower => write!(f, "Follower"),
58            SettableRegionRoleState::DowngradingLeader => write!(f, "Leader(Downgrading)"),
59            SettableRegionRoleState::Leader => write!(f, "Leader"),
60            SettableRegionRoleState::StagingLeader => write!(f, "Leader(Staging)"),
61        }
62    }
63}
64
65impl From<SettableRegionRoleState> for RegionRole {
66    fn from(value: SettableRegionRoleState) -> Self {
67        match value {
68            SettableRegionRoleState::Follower => RegionRole::Follower,
69            SettableRegionRoleState::DowngradingLeader => RegionRole::DowngradingLeader,
70            SettableRegionRoleState::Leader => RegionRole::Leader,
71            SettableRegionRoleState::StagingLeader => RegionRole::StagingLeader,
72        }
73    }
74}
75
76/// The request to set region role state.
77#[derive(Debug, PartialEq, Eq)]
78pub struct SetRegionRoleStateRequest {
79    region_id: RegionId,
80    region_role_state: SettableRegionRoleState,
81}
82
83/// The success response of setting region role state.
84#[derive(Debug, PartialEq, Eq)]
85pub enum SetRegionRoleStateSuccess {
86    File,
87    Mito {
88        last_entry_id: entry::Id,
89    },
90    Metric {
91        last_entry_id: entry::Id,
92        metadata_last_entry_id: entry::Id,
93    },
94}
95
96impl SetRegionRoleStateSuccess {
97    /// Returns a [SetRegionRoleStateSuccess::File].
98    pub fn file() -> Self {
99        Self::File
100    }
101
102    /// Returns a [SetRegionRoleStateSuccess::Mito] with the `last_entry_id`.
103    pub fn mito(last_entry_id: entry::Id) -> Self {
104        SetRegionRoleStateSuccess::Mito { last_entry_id }
105    }
106
107    /// Returns a [SetRegionRoleStateSuccess::Metric] with the `last_entry_id` and `metadata_last_entry_id`.
108    pub fn metric(last_entry_id: entry::Id, metadata_last_entry_id: entry::Id) -> Self {
109        SetRegionRoleStateSuccess::Metric {
110            last_entry_id,
111            metadata_last_entry_id,
112        }
113    }
114}
115
116impl SetRegionRoleStateSuccess {
117    /// Returns the last entry id of the region.
118    pub fn last_entry_id(&self) -> Option<entry::Id> {
119        match self {
120            SetRegionRoleStateSuccess::File => None,
121            SetRegionRoleStateSuccess::Mito { last_entry_id } => Some(*last_entry_id),
122            SetRegionRoleStateSuccess::Metric { last_entry_id, .. } => Some(*last_entry_id),
123        }
124    }
125
126    /// Returns the last entry id of the metadata of the region.
127    pub fn metadata_last_entry_id(&self) -> Option<entry::Id> {
128        match self {
129            SetRegionRoleStateSuccess::File => None,
130            SetRegionRoleStateSuccess::Mito { .. } => None,
131            SetRegionRoleStateSuccess::Metric {
132                metadata_last_entry_id,
133                ..
134            } => Some(*metadata_last_entry_id),
135        }
136    }
137}
138
139/// The response of setting region role state.
140#[derive(Debug)]
141pub enum SetRegionRoleStateResponse {
142    Success(SetRegionRoleStateSuccess),
143    NotFound,
144    InvalidTransition(BoxedError),
145}
146
147impl SetRegionRoleStateResponse {
148    /// Returns a [SetRegionRoleStateResponse::Success] with the `File` success.
149    pub fn success(success: SetRegionRoleStateSuccess) -> Self {
150        Self::Success(success)
151    }
152
153    /// Returns a [SetRegionRoleStateResponse::InvalidTransition] with the error.
154    pub fn invalid_transition(error: BoxedError) -> Self {
155        Self::InvalidTransition(error)
156    }
157
158    /// Returns true if the response is a [SetRegionRoleStateResponse::NotFound].
159    pub fn is_not_found(&self) -> bool {
160        matches!(self, SetRegionRoleStateResponse::NotFound)
161    }
162
163    /// Returns true if the response is a [SetRegionRoleStateResponse::InvalidTransition].
164    pub fn is_invalid_transition(&self) -> bool {
165        matches!(self, SetRegionRoleStateResponse::InvalidTransition(_))
166    }
167}
168
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub struct GrantedRegion {
171    pub region_id: RegionId,
172    pub region_role: RegionRole,
173    pub extensions: HashMap<String, Vec<u8>>,
174}
175
176impl GrantedRegion {
177    pub fn new(region_id: RegionId, region_role: RegionRole) -> Self {
178        Self {
179            region_id,
180            region_role,
181            extensions: HashMap::new(),
182        }
183    }
184}
185
186impl From<GrantedRegion> for PbGrantedRegion {
187    fn from(value: GrantedRegion) -> Self {
188        PbGrantedRegion {
189            region_id: value.region_id.as_u64(),
190            role: PbRegionRole::from(value.region_role).into(),
191            extensions: value.extensions,
192        }
193    }
194}
195
196impl From<PbGrantedRegion> for GrantedRegion {
197    fn from(value: PbGrantedRegion) -> Self {
198        GrantedRegion {
199            region_id: RegionId::from_u64(value.region_id),
200            region_role: value.role().into(),
201            extensions: value.extensions,
202        }
203    }
204}
205
206/// The role of the region.
207/// TODO(weny): rename it to `RegionRoleState`
208#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
209pub enum RegionRole {
210    // Readonly region(mito2)
211    Follower,
212    // Writable region(mito2), Readonly region(file).
213    Leader,
214    // Leader is in staging mode.
215    //
216    // This is leader-like and writable, but it follows the staging workflow
217    // semantics instead of a normal leader's steady state.
218    StagingLeader,
219    // Leader is downgrading to follower.
220    //
221    // This state is used to prevent new write requests.
222    DowngradingLeader,
223}
224
225impl Display for RegionRole {
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        match self {
228            RegionRole::Follower => write!(f, "Follower"),
229            RegionRole::Leader => write!(f, "Leader"),
230            RegionRole::StagingLeader => write!(f, "Leader(Staging)"),
231            RegionRole::DowngradingLeader => write!(f, "Leader(Downgrading)"),
232        }
233    }
234}
235
236impl RegionRole {
237    pub fn writable(&self) -> bool {
238        matches!(self, RegionRole::Leader | RegionRole::StagingLeader)
239    }
240}
241
242impl From<RegionRole> for PbRegionRole {
243    fn from(value: RegionRole) -> Self {
244        match value {
245            RegionRole::Follower => PbRegionRole::Follower,
246            RegionRole::Leader => PbRegionRole::Leader,
247            RegionRole::StagingLeader => PbRegionRole::StagingLeader,
248            RegionRole::DowngradingLeader => PbRegionRole::DowngradingLeader,
249        }
250    }
251}
252
253impl From<PbRegionRole> for RegionRole {
254    fn from(value: PbRegionRole) -> Self {
255        match value {
256            PbRegionRole::Leader => RegionRole::Leader,
257            PbRegionRole::StagingLeader => RegionRole::StagingLeader,
258            PbRegionRole::Follower => RegionRole::Follower,
259            PbRegionRole::DowngradingLeader => RegionRole::DowngradingLeader,
260        }
261    }
262}
263
264/// Output partition properties of the [RegionScanner].
265#[derive(Debug)]
266pub enum ScannerPartitioning {
267    /// Unknown partitioning scheme with a known number of partitions
268    Unknown(usize),
269}
270
271impl ScannerPartitioning {
272    /// Returns the number of partitions.
273    pub fn num_partitions(&self) -> usize {
274        match self {
275            ScannerPartitioning::Unknown(num_partitions) => *num_partitions,
276        }
277    }
278}
279
280/// Represents one data range within a partition
281#[derive(Debug, Clone, Copy, PartialEq, Eq)]
282pub struct PartitionRange {
283    /// Start time of time index column. Inclusive.
284    pub start: Timestamp,
285    /// End time of time index column. Exclusive.
286    pub end: Timestamp,
287    /// Number of rows in this range. Is used to balance ranges between partitions.
288    pub num_rows: usize,
289    /// Identifier to this range. Assigned by storage engine.
290    pub identifier: usize,
291}
292
293/// Properties of the [RegionScanner].
294#[derive(Debug, Default)]
295pub struct ScannerProperties {
296    /// A 2-dim partition ranges.
297    ///
298    /// The first dim vector's length represents the output partition number. The second
299    /// dim is ranges within one partition.
300    pub partitions: Vec<Vec<PartitionRange>>,
301
302    /// Whether scanner is in append-only mode.
303    append_mode: bool,
304
305    /// Total rows that **may** return by scanner. This field is only read iff
306    /// [ScannerProperties::append_mode] is true.
307    total_rows: usize,
308
309    /// Whether total_rows is exact before applying query predicates.
310    total_rows_is_exact: bool,
311
312    /// Whether to yield an empty batch to distinguish partition ranges.
313    pub distinguish_partition_range: bool,
314
315    /// The target partitions of the scanner. 0 indicates using the number of partitions as target partitions.
316    target_partitions: usize,
317
318    /// Whether the scanner is scanning a logical region.
319    logical_region: bool,
320
321    /// Region id that should receive query-load metrics for this scanner.
322    query_load_region_id: Option<RegionId>,
323    /// Counters that should receive query-load metrics for this scanner.
324    query_stat_counters: Option<RegionQueryStatCounters>,
325}
326
327impl ScannerProperties {
328    /// Sets append mode for scanner.
329    pub fn with_append_mode(mut self, append_mode: bool) -> Self {
330        self.append_mode = append_mode;
331        self
332    }
333
334    /// Sets total rows for scanner.
335    pub fn with_total_rows(mut self, total_rows: usize) -> Self {
336        self.total_rows = total_rows;
337        self
338    }
339
340    /// Sets whether the source row count is exact before query predicates.
341    pub fn with_total_rows_is_exact(mut self, exact: bool) -> Self {
342        self.total_rows_is_exact = exact;
343        self
344    }
345
346    /// Creates a new [`ScannerProperties`] with the given partitioning.
347    pub fn new(partitions: Vec<Vec<PartitionRange>>, append_mode: bool, total_rows: usize) -> Self {
348        Self {
349            partitions,
350            append_mode,
351            total_rows,
352            total_rows_is_exact: false,
353            distinguish_partition_range: false,
354            target_partitions: 0,
355            logical_region: false,
356            query_load_region_id: None,
357            query_stat_counters: None,
358        }
359    }
360
361    /// Updates the properties with the given [PrepareRequest].
362    pub fn prepare(&mut self, request: PrepareRequest) {
363        if let Some(ranges) = request.ranges {
364            self.partitions = ranges;
365        }
366        if let Some(distinguish_partition_range) = request.distinguish_partition_range {
367            self.distinguish_partition_range = distinguish_partition_range;
368        }
369        if let Some(target_partitions) = request.target_partitions {
370            self.target_partitions = target_partitions;
371        }
372    }
373
374    /// Returns the number of actual partitions.
375    pub fn num_partitions(&self) -> usize {
376        self.partitions.len()
377    }
378
379    pub fn append_mode(&self) -> bool {
380        self.append_mode
381    }
382
383    pub fn total_rows(&self) -> usize {
384        self.total_rows
385    }
386
387    /// Returns whether the source row count is exact before query predicates.
388    pub fn total_rows_is_exact(&self) -> bool {
389        self.total_rows_is_exact
390    }
391
392    /// Returns whether the scanner is scanning a logical region.
393    pub fn is_logical_region(&self) -> bool {
394        self.logical_region
395    }
396
397    /// Returns the target partitions of the scanner. If it is not set, returns the number of partitions.
398    pub fn target_partitions(&self) -> usize {
399        if self.target_partitions == 0 {
400            self.num_partitions()
401        } else {
402            self.target_partitions
403        }
404    }
405
406    /// Sets whether the scanner is reading a logical region.
407    pub fn set_logical_region(&mut self, logical_region: bool) {
408        self.logical_region = logical_region;
409    }
410
411    /// Returns the region id that should receive query-load metrics.
412    pub fn query_load_region_id(&self) -> Option<RegionId> {
413        self.query_load_region_id
414    }
415
416    /// Returns the counters that should receive query-load metrics.
417    pub fn query_stat_counters(&self) -> Option<RegionQueryStatCounters> {
418        self.query_stat_counters.clone()
419    }
420
421    /// Sets the region id that should receive query-load metrics.
422    pub fn set_query_load_region_id(&mut self, region_id: RegionId) {
423        self.query_load_region_id = Some(region_id);
424    }
425
426    /// Sets the counters that should receive query-load metrics.
427    pub fn set_query_stat_counters(&mut self, counters: RegionQueryStatCounters) {
428        self.query_stat_counters = Some(counters);
429    }
430}
431
432/// Request to override the scanner properties.
433#[derive(Default)]
434pub struct PrepareRequest {
435    /// Assigned partition ranges.
436    pub ranges: Option<Vec<Vec<PartitionRange>>>,
437    /// Distringuishes partition range by empty batches.
438    pub distinguish_partition_range: Option<bool>,
439    /// The expected number of target partitions.
440    pub target_partitions: Option<usize>,
441}
442
443impl PrepareRequest {
444    /// Sets the ranges.
445    pub fn with_ranges(mut self, ranges: Vec<Vec<PartitionRange>>) -> Self {
446        self.ranges = Some(ranges);
447        self
448    }
449
450    /// Sets the distinguish partition range flag.
451    pub fn with_distinguish_partition_range(mut self, distinguish_partition_range: bool) -> Self {
452        self.distinguish_partition_range = Some(distinguish_partition_range);
453        self
454    }
455
456    /// Sets the target partitions.
457    pub fn with_target_partitions(mut self, target_partitions: usize) -> Self {
458        self.target_partitions = Some(target_partitions);
459        self
460    }
461}
462
463/// Necessary context of the query for the scanner.
464#[derive(Clone, Default)]
465pub struct QueryScanContext {
466    /// Whether the query is EXPLAIN ANALYZE VERBOSE.
467    pub explain_verbose: bool,
468}
469
470/// A scanner that provides a way to scan the region concurrently.
471///
472/// The scanner splits the region into partitions so that each partition can be scanned concurrently.
473/// You can use this trait to implement an [`ExecutionPlan`](datafusion_physical_plan::ExecutionPlan).
474pub trait RegionScanner: Debug + DisplayAs + Send {
475    fn name(&self) -> &str;
476
477    /// Returns the properties of the scanner.
478    fn properties(&self) -> &ScannerProperties;
479
480    /// Returns the schema of the record batches.
481    fn schema(&self) -> SchemaRef;
482
483    /// Returns the metadata of the region.
484    fn metadata(&self) -> RegionMetadataRef;
485
486    /// Prepares the scanner with the given partition ranges.
487    ///
488    /// This method is for the planner to adjust the scanner's behavior based on the partition ranges.
489    fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError>;
490
491    /// Scans the partition and returns a stream of record batches.
492    ///
493    /// # Panics
494    /// Panics if the `partition` is out of bound.
495    fn scan_partition(
496        &self,
497        ctx: &QueryScanContext,
498        metrics_set: &ExecutionPlanMetricsSet,
499        partition: usize,
500    ) -> Result<SendableRecordBatchStream, BoxedError>;
501
502    /// Check if there is any predicate exclude region partition exprs that may be executed in this scanner.
503    fn has_predicate_without_region(&self) -> bool;
504
505    /// Add the given dynamic filter expressions to the predicate of the scanner.
506    /// Returns a vector of booleans indicating which filter expressions were applied.
507    /// true indicates the filter expression was applied(will be use by scanner to prune by stat for row group),
508    /// false otherwise.
509    fn add_dyn_filter_to_predicate(
510        &mut self,
511        filter_exprs: Vec<Arc<dyn PhysicalExpr>>,
512    ) -> Vec<bool>;
513
514    /// Resets execution state after streams from the previous execution have been dropped.
515    ///
516    /// Callers must not concurrently execute, reset, or install filters through scanner aliases.
517    /// Implementations retaining dynamic predicates must discard them and derived pruning state,
518    /// while preserving static predicates, the source snapshot, and prepared properties.
519    /// The default is a no-op and does not make inherently one-shot sources replayable.
520    fn reset_state(&mut self) {}
521
522    /// Sets whether the scanner is reading a logical region.
523    fn set_logical_region(&mut self, logical_region: bool);
524
525    /// Sets the region id that should receive query-load metrics.
526    fn set_query_load_region_id(&mut self, region_id: RegionId);
527
528    fn snapshot_sequence(&self) -> Option<SequenceNumber> {
529        None
530    }
531}
532
533pub type RegionScannerRef = Box<dyn RegionScanner>;
534
535pub type BatchResponses = Vec<(RegionId, Result<RegionResponse, BoxedError>)>;
536
537/// Represents the statistics of a region.
538#[derive(Debug, Deserialize, Serialize, Default)]
539pub struct RegionStatistic {
540    /// The number of rows stored in SST files owned by this region plus rows in memtables.
541    ///
542    /// Rows from SST files referenced from other regions, for example after repartition,
543    /// are not counted to avoid table-level double counting when summing region statistics.
544    #[serde(default)]
545    pub num_rows: u64,
546    /// The size of memtable in bytes.
547    pub memtable_size: u64,
548    /// The size of WAL in bytes.
549    pub wal_size: u64,
550    /// The size of manifest in bytes.
551    pub manifest_size: u64,
552    /// The size of SST data files owned by this region in bytes.
553    ///
554    /// SST files referenced from other regions, for example after repartition, are not counted.
555    pub sst_size: u64,
556    /// The number of SST files owned by this region.
557    ///
558    /// SST files referenced from other regions, for example after repartition, are not counted.
559    pub sst_num: u64,
560    /// The size of SST index files owned by this region in bytes.
561    ///
562    /// SST index files referenced from other regions, for example after repartition, are not counted.
563    #[serde(default)]
564    pub index_size: u64,
565    /// The details of the region.
566    #[serde(default)]
567    pub manifest: RegionManifestInfo,
568    #[serde(default)]
569    /// The total bytes written of the region since region opened.
570    pub written_bytes: u64,
571    /// The total query CPU time of the region since region opened.
572    ///
573    /// Unit: nanoseconds.
574    #[serde(default)]
575    pub query_cpu_time: u64,
576    /// The total scanned bytes of the region since region opened.
577    #[serde(default)]
578    pub query_scanned_bytes: u64,
579    /// The latest entry id of the region's remote WAL since last flush.
580    /// For metric engine, there're two latest entry ids, one for data and one for metadata.
581    /// TODO(weny): remove this two fields and use single instead.
582    #[serde(default)]
583    pub data_topic_latest_entry_id: u64,
584    #[serde(default)]
585    pub metadata_topic_latest_entry_id: u64,
586}
587
588/// The manifest info of a region.
589#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
590pub enum RegionManifestInfo {
591    Mito {
592        manifest_version: u64,
593        flushed_entry_id: u64,
594        /// Number of files removed in the manifest's `removed_files` field.
595        file_removed_cnt: u64,
596    },
597    Metric {
598        data_manifest_version: u64,
599        data_flushed_entry_id: u64,
600        metadata_manifest_version: u64,
601        metadata_flushed_entry_id: u64,
602    },
603}
604
605impl RegionManifestInfo {
606    /// Creates a new [RegionManifestInfo] for mito2 engine.
607    pub fn mito(manifest_version: u64, flushed_entry_id: u64, file_removal_rate: u64) -> Self {
608        Self::Mito {
609            manifest_version,
610            flushed_entry_id,
611            file_removed_cnt: file_removal_rate,
612        }
613    }
614
615    /// Creates a new [RegionManifestInfo] for metric engine.
616    pub fn metric(
617        data_manifest_version: u64,
618        data_flushed_entry_id: u64,
619        metadata_manifest_version: u64,
620        metadata_flushed_entry_id: u64,
621    ) -> Self {
622        Self::Metric {
623            data_manifest_version,
624            data_flushed_entry_id,
625            metadata_manifest_version,
626            metadata_flushed_entry_id,
627        }
628    }
629
630    /// Returns true if the region is a mito2 region.
631    pub fn is_mito(&self) -> bool {
632        matches!(self, RegionManifestInfo::Mito { .. })
633    }
634
635    /// Returns true if the region is a metric region.
636    pub fn is_metric(&self) -> bool {
637        matches!(self, RegionManifestInfo::Metric { .. })
638    }
639
640    /// Returns the flushed entry id of the data region.
641    pub fn data_flushed_entry_id(&self) -> u64 {
642        match self {
643            RegionManifestInfo::Mito {
644                flushed_entry_id, ..
645            } => *flushed_entry_id,
646            RegionManifestInfo::Metric {
647                data_flushed_entry_id,
648                ..
649            } => *data_flushed_entry_id,
650        }
651    }
652
653    /// Returns the manifest version of the data region.
654    pub fn data_manifest_version(&self) -> u64 {
655        match self {
656            RegionManifestInfo::Mito {
657                manifest_version, ..
658            } => *manifest_version,
659            RegionManifestInfo::Metric {
660                data_manifest_version,
661                ..
662            } => *data_manifest_version,
663        }
664    }
665
666    /// Returns the manifest version of the metadata region.
667    pub fn metadata_manifest_version(&self) -> Option<u64> {
668        match self {
669            RegionManifestInfo::Mito { .. } => None,
670            RegionManifestInfo::Metric {
671                metadata_manifest_version,
672                ..
673            } => Some(*metadata_manifest_version),
674        }
675    }
676
677    /// Returns the flushed entry id of the metadata region.
678    pub fn metadata_flushed_entry_id(&self) -> Option<u64> {
679        match self {
680            RegionManifestInfo::Mito { .. } => None,
681            RegionManifestInfo::Metric {
682                metadata_flushed_entry_id,
683                ..
684            } => Some(*metadata_flushed_entry_id),
685        }
686    }
687
688    /// Encodes a list of ([RegionId], [RegionManifestInfo]) to a byte array.
689    pub fn encode_list(manifest_infos: &[(RegionId, Self)]) -> serde_json::Result<Vec<u8>> {
690        serde_json::to_vec(manifest_infos)
691    }
692
693    /// Decodes a list of ([RegionId], [RegionManifestInfo]) from a byte array.
694    pub fn decode_list(value: &[u8]) -> serde_json::Result<Vec<(RegionId, Self)>> {
695        serde_json::from_slice(value)
696    }
697}
698
699impl Default for RegionManifestInfo {
700    fn default() -> Self {
701        Self::Mito {
702            manifest_version: 0,
703            flushed_entry_id: 0,
704            file_removed_cnt: 0,
705        }
706    }
707}
708
709impl RegionStatistic {
710    /// Deserializes the region statistic to a byte array.
711    ///
712    /// Returns None if the deserialization fails.
713    pub fn deserialize_from_slice(value: &[u8]) -> Option<RegionStatistic> {
714        serde_json::from_slice(value).ok()
715    }
716
717    /// Serializes the region statistic to a byte array.
718    ///
719    /// Returns None if the serialization fails.
720    pub fn serialize_to_vec(&self) -> Option<Vec<u8>> {
721        serde_json::to_vec(self).ok()
722    }
723}
724
725impl RegionStatistic {
726    /// Returns the estimated disk size of the region.
727    pub fn estimated_disk_size(&self) -> u64 {
728        self.wal_size + self.sst_size + self.manifest_size + self.index_size
729    }
730}
731
732#[cfg(test)]
733mod tests {
734    use serde_json::json;
735
736    use super::*;
737
738    #[test]
739    fn region_statistic_deserializes_without_query_stats() {
740        let statistic: RegionStatistic = serde_json::from_value(json!({
741            "num_rows": 1,
742            "memtable_size": 2,
743            "wal_size": 3,
744            "manifest_size": 4,
745            "sst_size": 5,
746            "sst_num": 6,
747            "index_size": 7,
748            "manifest": {
749                "Mito": {
750                    "manifest_version": 8,
751                    "flushed_entry_id": 9,
752                    "file_removed_cnt": 10
753                }
754            },
755            "written_bytes": 11
756        }))
757        .unwrap();
758
759        assert_eq!(statistic.query_cpu_time, 0);
760        assert_eq!(statistic.query_scanned_bytes, 0);
761    }
762}
763
764/// Request to sync the region from a manifest or a region.
765#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
766pub enum SyncRegionFromRequest {
767    /// Syncs the region using manifest information.
768    /// Used in leader-follower manifest sync scenarios.
769    FromManifest(RegionManifestInfo),
770    /// Syncs the region from another region.
771    ///
772    /// Used by the metric engine to sync logical regions from a source physical region
773    /// to a target physical region. This copies metadata region SST files and transforms
774    /// logical region entries to use the target's region number.
775    FromRegion {
776        /// The [`RegionId`] of the source region.
777        source_region_id: RegionId,
778        /// The parallelism of the sync operation.
779        parallelism: usize,
780    },
781}
782
783impl From<RegionManifestInfo> for SyncRegionFromRequest {
784    fn from(manifest_info: RegionManifestInfo) -> Self {
785        SyncRegionFromRequest::FromManifest(manifest_info)
786    }
787}
788
789impl SyncRegionFromRequest {
790    /// Creates a new request from a manifest info.
791    pub fn from_manifest(manifest_info: RegionManifestInfo) -> Self {
792        SyncRegionFromRequest::FromManifest(manifest_info)
793    }
794
795    /// Creates a new request from a region.
796    pub fn from_region(source_region_id: RegionId, parallelism: usize) -> Self {
797        SyncRegionFromRequest::FromRegion {
798            source_region_id,
799            parallelism,
800        }
801    }
802
803    /// Returns true if the request is from a manifest.
804    pub fn is_from_manifest(&self) -> bool {
805        matches!(self, SyncRegionFromRequest::FromManifest { .. })
806    }
807
808    /// Converts the request to a region manifest info.
809    ///
810    /// Returns None if the request is not from a manifest.
811    pub fn into_region_manifest_info(self) -> Option<RegionManifestInfo> {
812        match self {
813            SyncRegionFromRequest::FromManifest(manifest_info) => Some(manifest_info),
814            SyncRegionFromRequest::FromRegion { .. } => None,
815        }
816    }
817}
818
819/// The response of syncing the region.
820#[derive(Debug)]
821pub enum SyncRegionFromResponse {
822    NotSupported,
823    Mito {
824        /// Indicates if the data region was synced.
825        synced: bool,
826    },
827    Metric {
828        /// Indicates if the metadata region was synced.
829        metadata_synced: bool,
830        /// Indicates if the data region was synced.
831        data_synced: bool,
832        /// The logical regions that were newly opened during the sync operation.
833        /// This only occurs after the metadata region has been successfully synced.
834        new_opened_logical_region_ids: Vec<RegionId>,
835    },
836}
837
838impl SyncRegionFromResponse {
839    /// Returns true if data region is synced.
840    pub fn is_data_synced(&self) -> bool {
841        match self {
842            SyncRegionFromResponse::NotSupported => false,
843            SyncRegionFromResponse::Mito { synced } => *synced,
844            SyncRegionFromResponse::Metric { data_synced, .. } => *data_synced,
845        }
846    }
847
848    /// Returns true if the engine is a mito2 engine.
849    pub fn is_mito(&self) -> bool {
850        matches!(self, SyncRegionFromResponse::Mito { .. })
851    }
852
853    /// Returns true if the engine is a metric engine.
854    pub fn is_metric(&self) -> bool {
855        matches!(self, SyncRegionFromResponse::Metric { .. })
856    }
857
858    /// Returns the new opened logical region ids.
859    pub fn new_opened_logical_region_ids(self) -> Option<Vec<RegionId>> {
860        match self {
861            SyncRegionFromResponse::Metric {
862                new_opened_logical_region_ids,
863                ..
864            } => Some(new_opened_logical_region_ids),
865            _ => None,
866        }
867    }
868}
869
870/// Request to remap manifests from old regions to new regions.
871#[derive(Debug, Clone)]
872pub struct RemapManifestsRequest {
873    /// The [`RegionId`] of a staging region used to obtain table directory and storage configuration for the remap operation.
874    pub region_id: RegionId,
875    /// Regions to remap manifests from.
876    pub input_regions: Vec<RegionId>,
877    /// For each old region, which new regions should receive its files
878    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
879    /// New partition expressions for the new regions.
880    pub new_partition_exprs: HashMap<RegionId, String>,
881}
882
883/// Response to remap manifests from old regions to new regions.
884#[derive(Debug, Clone)]
885pub struct RemapManifestsResponse {
886    /// Maps region id to its staging manifest path.
887    ///
888    /// These paths are relative paths within the central region's staging blob storage,
889    /// and should be passed to [`ApplyStagingManifestRequest`](RegionRequest::ApplyStagingManifest) to finalize the repartition.
890    pub manifest_paths: HashMap<RegionId, String>,
891}
892
893/// Request to copy files from a source region to a target region.
894#[derive(Debug, Clone)]
895pub struct MitoCopyRegionFromRequest {
896    /// The [`RegionId`] of the source region.
897    pub source_region_id: RegionId,
898    /// The parallelism of the copy operation.
899    pub parallelism: usize,
900}
901
902#[derive(Debug, Clone)]
903pub struct MitoCopyRegionFromResponse {
904    /// The file ids that were copied from the source region to the target region.
905    pub copied_file_ids: Vec<FileId>,
906}
907
908#[async_trait]
909pub trait RegionEngine: Send + Sync {
910    /// Name of this engine
911    fn name(&self) -> &str;
912
913    /// Handles batch open region requests.
914    async fn handle_batch_open_requests(
915        &self,
916        parallelism: usize,
917        requests: Vec<(RegionId, RegionOpenRequest)>,
918    ) -> Result<BatchResponses, BoxedError> {
919        let semaphore = Arc::new(Semaphore::new(parallelism));
920        let mut tasks = Vec::with_capacity(requests.len());
921
922        for (region_id, request) in requests {
923            let semaphore_moved = semaphore.clone();
924
925            tasks.push(async move {
926                // Safety: semaphore must exist
927                let _permit = semaphore_moved.acquire().await.unwrap();
928                let result = self
929                    .handle_request(region_id, RegionRequest::Open(request))
930                    .await;
931                (region_id, result)
932            });
933        }
934
935        Ok(join_all(tasks).await)
936    }
937
938    async fn handle_batch_catchup_requests(
939        &self,
940        parallelism: usize,
941        requests: Vec<(RegionId, RegionCatchupRequest)>,
942    ) -> Result<BatchResponses, BoxedError> {
943        let semaphore = Arc::new(Semaphore::new(parallelism));
944        let mut tasks = Vec::with_capacity(requests.len());
945
946        for (region_id, request) in requests {
947            let semaphore_moved = semaphore.clone();
948
949            tasks.push(async move {
950                // Safety: semaphore must exist
951                let _permit = semaphore_moved.acquire().await.unwrap();
952                let result = self
953                    .handle_request(region_id, RegionRequest::Catchup(request))
954                    .await;
955                (region_id, result)
956            });
957        }
958
959        Ok(join_all(tasks).await)
960    }
961
962    async fn handle_batch_ddl_requests(
963        &self,
964        request: BatchRegionDdlRequest,
965    ) -> Result<RegionResponse, BoxedError> {
966        let requests = request.into_region_requests();
967
968        let mut affected_rows = 0;
969        let mut extensions = HashMap::new();
970
971        for (region_id, request) in requests {
972            let result = self.handle_request(region_id, request).await?;
973            affected_rows += result.affected_rows;
974            extensions.extend(result.extensions);
975        }
976
977        Ok(RegionResponse {
978            affected_rows,
979            extensions,
980            metadata: Vec::new(),
981        })
982    }
983
984    /// Handles non-query request to the region. Returns the count of affected rows.
985    async fn handle_request(
986        &self,
987        region_id: RegionId,
988        request: RegionRequest,
989    ) -> Result<RegionResponse, BoxedError>;
990
991    /// Returns the committed sequence (sequence of latest written data).
992    async fn get_committed_sequence(
993        &self,
994        region_id: RegionId,
995    ) -> Result<SequenceNumber, BoxedError>;
996
997    /// Handles query and return a scanner that can be used to scan the region concurrently.
998    async fn handle_query(
999        &self,
1000        region_id: RegionId,
1001        request: ScanRequest,
1002    ) -> Result<RegionScannerRef, BoxedError>;
1003
1004    /// Returns the query memory tracker for scan execution.
1005    fn query_memory_tracker(&self) -> Option<QueryMemoryTracker> {
1006        None
1007    }
1008
1009    /// Retrieves region's metadata.
1010    async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError>;
1011
1012    /// Retrieves region's statistic.
1013    fn region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic>;
1014
1015    /// Stops the engine
1016    async fn stop(&self) -> Result<(), BoxedError>;
1017
1018    /// Sets [RegionRole] for a region.
1019    ///
1020    /// The engine checks whether the region is writable before writing to the region. Setting
1021    /// the region as readonly doesn't guarantee that write operations in progress will not
1022    /// take effect.
1023    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>;
1024
1025    /// Syncs the region manifest to the given manifest version.
1026    async fn sync_region(
1027        &self,
1028        region_id: RegionId,
1029        request: SyncRegionFromRequest,
1030    ) -> Result<SyncRegionFromResponse, BoxedError>;
1031
1032    /// Remaps manifests from old regions to new regions.
1033    async fn remap_manifests(
1034        &self,
1035        request: RemapManifestsRequest,
1036    ) -> Result<RemapManifestsResponse, BoxedError>;
1037
1038    /// Sets region role state gracefully.
1039    ///
1040    /// After the call returns, the engine ensures no more write operations will succeed in the region.
1041    async fn set_region_role_state_gracefully(
1042        &self,
1043        region_id: RegionId,
1044        region_role_state: SettableRegionRoleState,
1045    ) -> Result<SetRegionRoleStateResponse, BoxedError>;
1046
1047    /// Indicates region role.
1048    ///
1049    /// Returns the `None` if the region is not found.
1050    fn role(&self, region_id: RegionId) -> Option<RegionRole>;
1051
1052    fn as_any(&self) -> &dyn Any;
1053}
1054
1055pub type RegionEngineRef = Arc<dyn RegionEngine>;
1056
1057/// A [RegionScanner] that only scans a single partition.
1058pub struct SinglePartitionScanner {
1059    stream: Mutex<Option<SendableRecordBatchStream>>,
1060    schema: SchemaRef,
1061    properties: ScannerProperties,
1062    metadata: RegionMetadataRef,
1063    snapshot_sequence: Option<SequenceNumber>,
1064}
1065
1066impl SinglePartitionScanner {
1067    /// Creates a new [SinglePartitionScanner] with the given stream and metadata.
1068    pub fn new(
1069        stream: SendableRecordBatchStream,
1070        append_mode: bool,
1071        metadata: RegionMetadataRef,
1072        snapshot_sequence: Option<SequenceNumber>,
1073    ) -> Self {
1074        let schema = stream.schema();
1075        Self {
1076            stream: Mutex::new(Some(stream)),
1077            schema,
1078            properties: ScannerProperties::default().with_append_mode(append_mode),
1079            metadata,
1080            snapshot_sequence,
1081        }
1082    }
1083}
1084
1085impl Debug for SinglePartitionScanner {
1086    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1087        write!(f, "SinglePartitionScanner: <SendableRecordBatchStream>")
1088    }
1089}
1090
1091impl RegionScanner for SinglePartitionScanner {
1092    fn name(&self) -> &str {
1093        "SinglePartition"
1094    }
1095
1096    fn properties(&self) -> &ScannerProperties {
1097        &self.properties
1098    }
1099
1100    fn schema(&self) -> SchemaRef {
1101        self.schema.clone()
1102    }
1103
1104    fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
1105        self.properties.prepare(request);
1106        Ok(())
1107    }
1108
1109    fn scan_partition(
1110        &self,
1111        _ctx: &QueryScanContext,
1112        _metrics_set: &ExecutionPlanMetricsSet,
1113        _partition: usize,
1114    ) -> Result<SendableRecordBatchStream, BoxedError> {
1115        let mut stream = self.stream.lock().unwrap();
1116        let result = stream
1117            .take()
1118            .or_else(|| Some(Box::pin(EmptyRecordBatchStream::new(self.schema.clone()))));
1119        Ok(result.unwrap())
1120    }
1121
1122    fn has_predicate_without_region(&self) -> bool {
1123        false
1124    }
1125
1126    fn add_dyn_filter_to_predicate(
1127        &mut self,
1128        filter_exprs: Vec<Arc<dyn datafusion_physical_plan::PhysicalExpr>>,
1129    ) -> Vec<bool> {
1130        vec![false; filter_exprs.len()]
1131    }
1132
1133    fn metadata(&self) -> RegionMetadataRef {
1134        self.metadata.clone()
1135    }
1136
1137    fn set_logical_region(&mut self, logical_region: bool) {
1138        self.properties.set_logical_region(logical_region);
1139    }
1140
1141    fn set_query_load_region_id(&mut self, region_id: RegionId) {
1142        self.properties.set_query_load_region_id(region_id);
1143    }
1144
1145    fn snapshot_sequence(&self) -> Option<SequenceNumber> {
1146        self.snapshot_sequence
1147    }
1148}
1149
1150impl DisplayAs for SinglePartitionScanner {
1151    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1152        write!(f, "{:?}", self)
1153    }
1154}
store_api/region_engine.rs

store_api/
region_engine.rs
