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 earliest timestamp of the region's current data, or `None` if it holds
580    /// none. Unlike the size and row counters this covers files referenced from
581    /// other regions too, since a time range cannot double count.
582    #[serde(default)]
583    pub min_timestamp: Option<Timestamp>,
584    /// The latest timestamp of the region's current data. See [`Self::min_timestamp`].
585    #[serde(default)]
586    pub max_timestamp: Option<Timestamp>,
587    /// The latest entry id of the region's remote WAL since last flush.
588    /// For metric engine, there're two latest entry ids, one for data and one for metadata.
589    /// TODO(weny): remove this two fields and use single instead.
590    #[serde(default)]
591    pub data_topic_latest_entry_id: u64,
592    #[serde(default)]
593    pub metadata_topic_latest_entry_id: u64,
594}
595
596/// The manifest info of a region.
597#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
598pub enum RegionManifestInfo {
599    Mito {
600        manifest_version: u64,
601        flushed_entry_id: u64,
602        /// Number of files removed in the manifest's `removed_files` field.
603        file_removed_cnt: u64,
604    },
605    Metric {
606        data_manifest_version: u64,
607        data_flushed_entry_id: u64,
608        metadata_manifest_version: u64,
609        metadata_flushed_entry_id: u64,
610    },
611}
612
613impl RegionManifestInfo {
614    /// Creates a new [RegionManifestInfo] for mito2 engine.
615    pub fn mito(manifest_version: u64, flushed_entry_id: u64, file_removal_rate: u64) -> Self {
616        Self::Mito {
617            manifest_version,
618            flushed_entry_id,
619            file_removed_cnt: file_removal_rate,
620        }
621    }
622
623    /// Creates a new [RegionManifestInfo] for metric engine.
624    pub fn metric(
625        data_manifest_version: u64,
626        data_flushed_entry_id: u64,
627        metadata_manifest_version: u64,
628        metadata_flushed_entry_id: u64,
629    ) -> Self {
630        Self::Metric {
631            data_manifest_version,
632            data_flushed_entry_id,
633            metadata_manifest_version,
634            metadata_flushed_entry_id,
635        }
636    }
637
638    /// Returns true if the region is a mito2 region.
639    pub fn is_mito(&self) -> bool {
640        matches!(self, RegionManifestInfo::Mito { .. })
641    }
642
643    /// Returns true if the region is a metric region.
644    pub fn is_metric(&self) -> bool {
645        matches!(self, RegionManifestInfo::Metric { .. })
646    }
647
648    /// Returns the flushed entry id of the data region.
649    pub fn data_flushed_entry_id(&self) -> u64 {
650        match self {
651            RegionManifestInfo::Mito {
652                flushed_entry_id, ..
653            } => *flushed_entry_id,
654            RegionManifestInfo::Metric {
655                data_flushed_entry_id,
656                ..
657            } => *data_flushed_entry_id,
658        }
659    }
660
661    /// Returns the manifest version of the data region.
662    pub fn data_manifest_version(&self) -> u64 {
663        match self {
664            RegionManifestInfo::Mito {
665                manifest_version, ..
666            } => *manifest_version,
667            RegionManifestInfo::Metric {
668                data_manifest_version,
669                ..
670            } => *data_manifest_version,
671        }
672    }
673
674    /// Returns the manifest version of the metadata region.
675    pub fn metadata_manifest_version(&self) -> Option<u64> {
676        match self {
677            RegionManifestInfo::Mito { .. } => None,
678            RegionManifestInfo::Metric {
679                metadata_manifest_version,
680                ..
681            } => Some(*metadata_manifest_version),
682        }
683    }
684
685    /// Returns the flushed entry id of the metadata region.
686    pub fn metadata_flushed_entry_id(&self) -> Option<u64> {
687        match self {
688            RegionManifestInfo::Mito { .. } => None,
689            RegionManifestInfo::Metric {
690                metadata_flushed_entry_id,
691                ..
692            } => Some(*metadata_flushed_entry_id),
693        }
694    }
695
696    /// Encodes a list of ([RegionId], [RegionManifestInfo]) to a byte array.
697    pub fn encode_list(manifest_infos: &[(RegionId, Self)]) -> serde_json::Result<Vec<u8>> {
698        serde_json::to_vec(manifest_infos)
699    }
700
701    /// Decodes a list of ([RegionId], [RegionManifestInfo]) from a byte array.
702    pub fn decode_list(value: &[u8]) -> serde_json::Result<Vec<(RegionId, Self)>> {
703        serde_json::from_slice(value)
704    }
705}
706
707impl Default for RegionManifestInfo {
708    fn default() -> Self {
709        Self::Mito {
710            manifest_version: 0,
711            flushed_entry_id: 0,
712            file_removed_cnt: 0,
713        }
714    }
715}
716
717impl RegionStatistic {
718    /// Deserializes the region statistic to a byte array.
719    ///
720    /// Returns None if the deserialization fails.
721    pub fn deserialize_from_slice(value: &[u8]) -> Option<RegionStatistic> {
722        serde_json::from_slice(value).ok()
723    }
724
725    /// Serializes the region statistic to a byte array.
726    ///
727    /// Returns None if the serialization fails.
728    pub fn serialize_to_vec(&self) -> Option<Vec<u8>> {
729        serde_json::to_vec(self).ok()
730    }
731}
732
733impl RegionStatistic {
734    /// Returns the estimated disk size of the region.
735    pub fn estimated_disk_size(&self) -> u64 {
736        self.wal_size + self.sst_size + self.manifest_size + self.index_size
737    }
738}
739
740#[cfg(test)]
741mod tests {
742    use serde_json::json;
743
744    use super::*;
745
746    #[test]
747    fn region_statistic_deserializes_without_query_stats() {
748        let statistic: RegionStatistic = serde_json::from_value(json!({
749            "num_rows": 1,
750            "memtable_size": 2,
751            "wal_size": 3,
752            "manifest_size": 4,
753            "sst_size": 5,
754            "sst_num": 6,
755            "index_size": 7,
756            "manifest": {
757                "Mito": {
758                    "manifest_version": 8,
759                    "flushed_entry_id": 9,
760                    "file_removed_cnt": 10
761                }
762            },
763            "written_bytes": 11
764        }))
765        .unwrap();
766
767        assert_eq!(statistic.query_cpu_time, 0);
768        assert_eq!(statistic.query_scanned_bytes, 0);
769    }
770}
771
772/// Request to sync the region from a manifest or a region.
773#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
774pub enum SyncRegionFromRequest {
775    /// Syncs the region using manifest information.
776    /// Used in leader-follower manifest sync scenarios.
777    FromManifest(RegionManifestInfo),
778    /// Syncs the region from another region.
779    ///
780    /// Used by the metric engine to sync logical regions from a source physical region
781    /// to a target physical region. This copies metadata region SST files and transforms
782    /// logical region entries to use the target's region number.
783    FromRegion {
784        /// The [`RegionId`] of the source region.
785        source_region_id: RegionId,
786        /// The parallelism of the sync operation.
787        parallelism: usize,
788    },
789}
790
791impl From<RegionManifestInfo> for SyncRegionFromRequest {
792    fn from(manifest_info: RegionManifestInfo) -> Self {
793        SyncRegionFromRequest::FromManifest(manifest_info)
794    }
795}
796
797impl SyncRegionFromRequest {
798    /// Creates a new request from a manifest info.
799    pub fn from_manifest(manifest_info: RegionManifestInfo) -> Self {
800        SyncRegionFromRequest::FromManifest(manifest_info)
801    }
802
803    /// Creates a new request from a region.
804    pub fn from_region(source_region_id: RegionId, parallelism: usize) -> Self {
805        SyncRegionFromRequest::FromRegion {
806            source_region_id,
807            parallelism,
808        }
809    }
810
811    /// Returns true if the request is from a manifest.
812    pub fn is_from_manifest(&self) -> bool {
813        matches!(self, SyncRegionFromRequest::FromManifest { .. })
814    }
815
816    /// Converts the request to a region manifest info.
817    ///
818    /// Returns None if the request is not from a manifest.
819    pub fn into_region_manifest_info(self) -> Option<RegionManifestInfo> {
820        match self {
821            SyncRegionFromRequest::FromManifest(manifest_info) => Some(manifest_info),
822            SyncRegionFromRequest::FromRegion { .. } => None,
823        }
824    }
825}
826
827/// The response of syncing the region.
828#[derive(Debug)]
829pub enum SyncRegionFromResponse {
830    NotSupported,
831    Mito {
832        /// Indicates if the data region was synced.
833        synced: bool,
834    },
835    Metric {
836        /// Indicates if the metadata region was synced.
837        metadata_synced: bool,
838        /// Indicates if the data region was synced.
839        data_synced: bool,
840        /// The logical regions that were newly opened during the sync operation.
841        /// This only occurs after the metadata region has been successfully synced.
842        new_opened_logical_region_ids: Vec<RegionId>,
843    },
844}
845
846impl SyncRegionFromResponse {
847    /// Returns true if data region is synced.
848    pub fn is_data_synced(&self) -> bool {
849        match self {
850            SyncRegionFromResponse::NotSupported => false,
851            SyncRegionFromResponse::Mito { synced } => *synced,
852            SyncRegionFromResponse::Metric { data_synced, .. } => *data_synced,
853        }
854    }
855
856    /// Returns true if the engine is a mito2 engine.
857    pub fn is_mito(&self) -> bool {
858        matches!(self, SyncRegionFromResponse::Mito { .. })
859    }
860
861    /// Returns true if the engine is a metric engine.
862    pub fn is_metric(&self) -> bool {
863        matches!(self, SyncRegionFromResponse::Metric { .. })
864    }
865
866    /// Returns the new opened logical region ids.
867    pub fn new_opened_logical_region_ids(self) -> Option<Vec<RegionId>> {
868        match self {
869            SyncRegionFromResponse::Metric {
870                new_opened_logical_region_ids,
871                ..
872            } => Some(new_opened_logical_region_ids),
873            _ => None,
874        }
875    }
876}
877
878/// Request to remap manifests from old regions to new regions.
879#[derive(Debug, Clone)]
880pub struct RemapManifestsRequest {
881    /// The [`RegionId`] of a staging region used to obtain table directory and storage configuration for the remap operation.
882    pub region_id: RegionId,
883    /// Regions to remap manifests from.
884    pub input_regions: Vec<RegionId>,
885    /// For each old region, which new regions should receive its files
886    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
887    /// New partition expressions for the new regions.
888    pub new_partition_exprs: HashMap<RegionId, String>,
889}
890
891/// Response to remap manifests from old regions to new regions.
892#[derive(Debug, Clone)]
893pub struct RemapManifestsResponse {
894    /// Maps region id to its staging manifest path.
895    ///
896    /// These paths are relative paths within the central region's staging blob storage,
897    /// and should be passed to [`ApplyStagingManifestRequest`](RegionRequest::ApplyStagingManifest) to finalize the repartition.
898    pub manifest_paths: HashMap<RegionId, String>,
899}
900
901/// Request to copy files from a source region to a target region.
902#[derive(Debug, Clone)]
903pub struct MitoCopyRegionFromRequest {
904    /// The [`RegionId`] of the source region.
905    pub source_region_id: RegionId,
906    /// The parallelism of the copy operation.
907    pub parallelism: usize,
908}
909
910#[derive(Debug, Clone)]
911pub struct MitoCopyRegionFromResponse {
912    /// The file ids that were copied from the source region to the target region.
913    pub copied_file_ids: Vec<FileId>,
914}
915
916#[async_trait]
917pub trait RegionEngine: Send + Sync {
918    /// Name of this engine
919    fn name(&self) -> &str;
920
921    /// Handles batch open region requests.
922    async fn handle_batch_open_requests(
923        &self,
924        parallelism: usize,
925        requests: Vec<(RegionId, RegionOpenRequest)>,
926    ) -> Result<BatchResponses, BoxedError> {
927        let semaphore = Arc::new(Semaphore::new(parallelism));
928        let mut tasks = Vec::with_capacity(requests.len());
929
930        for (region_id, request) in requests {
931            let semaphore_moved = semaphore.clone();
932
933            tasks.push(async move {
934                // Safety: semaphore must exist
935                let _permit = semaphore_moved.acquire().await.unwrap();
936                let result = self
937                    .handle_request(region_id, RegionRequest::Open(request))
938                    .await;
939                (region_id, result)
940            });
941        }
942
943        Ok(join_all(tasks).await)
944    }
945
946    async fn handle_batch_catchup_requests(
947        &self,
948        parallelism: usize,
949        requests: Vec<(RegionId, RegionCatchupRequest)>,
950    ) -> Result<BatchResponses, BoxedError> {
951        let semaphore = Arc::new(Semaphore::new(parallelism));
952        let mut tasks = Vec::with_capacity(requests.len());
953
954        for (region_id, request) in requests {
955            let semaphore_moved = semaphore.clone();
956
957            tasks.push(async move {
958                // Safety: semaphore must exist
959                let _permit = semaphore_moved.acquire().await.unwrap();
960                let result = self
961                    .handle_request(region_id, RegionRequest::Catchup(request))
962                    .await;
963                (region_id, result)
964            });
965        }
966
967        Ok(join_all(tasks).await)
968    }
969
970    async fn handle_batch_ddl_requests(
971        &self,
972        request: BatchRegionDdlRequest,
973    ) -> Result<RegionResponse, BoxedError> {
974        let requests = request.into_region_requests();
975
976        let mut affected_rows = 0;
977        let mut extensions = HashMap::new();
978
979        for (region_id, request) in requests {
980            let result = self.handle_request(region_id, request).await?;
981            affected_rows += result.affected_rows;
982            extensions.extend(result.extensions);
983        }
984
985        Ok(RegionResponse {
986            affected_rows,
987            extensions,
988            metadata: Vec::new(),
989        })
990    }
991
992    /// Handles non-query request to the region. Returns the count of affected rows.
993    async fn handle_request(
994        &self,
995        region_id: RegionId,
996        request: RegionRequest,
997    ) -> Result<RegionResponse, BoxedError>;
998
999    /// Returns the committed sequence (sequence of latest written data).
1000    async fn get_committed_sequence(
1001        &self,
1002        region_id: RegionId,
1003    ) -> Result<SequenceNumber, BoxedError>;
1004
1005    /// Handles query and return a scanner that can be used to scan the region concurrently.
1006    async fn handle_query(
1007        &self,
1008        region_id: RegionId,
1009        request: ScanRequest,
1010    ) -> Result<RegionScannerRef, BoxedError>;
1011
1012    /// Returns the query memory tracker for scan execution.
1013    fn query_memory_tracker(&self) -> Option<QueryMemoryTracker> {
1014        None
1015    }
1016
1017    /// Retrieves region's metadata.
1018    async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError>;
1019
1020    /// Retrieves region's statistic.
1021    fn region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic>;
1022
1023    /// Stops the engine
1024    async fn stop(&self) -> Result<(), BoxedError>;
1025
1026    /// Sets [RegionRole] for a region.
1027    ///
1028    /// The engine checks whether the region is writable before writing to the region. Setting
1029    /// the region as readonly doesn't guarantee that write operations in progress will not
1030    /// take effect.
1031    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>;
1032
1033    /// Syncs the region manifest to the given manifest version.
1034    async fn sync_region(
1035        &self,
1036        region_id: RegionId,
1037        request: SyncRegionFromRequest,
1038    ) -> Result<SyncRegionFromResponse, BoxedError>;
1039
1040    /// Remaps manifests from old regions to new regions.
1041    async fn remap_manifests(
1042        &self,
1043        request: RemapManifestsRequest,
1044    ) -> Result<RemapManifestsResponse, BoxedError>;
1045
1046    /// Sets region role state gracefully.
1047    ///
1048    /// After the call returns, the engine ensures no more write operations will succeed in the region.
1049    async fn set_region_role_state_gracefully(
1050        &self,
1051        region_id: RegionId,
1052        region_role_state: SettableRegionRoleState,
1053    ) -> Result<SetRegionRoleStateResponse, BoxedError>;
1054
1055    /// Indicates region role.
1056    ///
1057    /// Returns the `None` if the region is not found.
1058    fn role(&self, region_id: RegionId) -> Option<RegionRole>;
1059
1060    fn as_any(&self) -> &dyn Any;
1061}
1062
1063pub type RegionEngineRef = Arc<dyn RegionEngine>;
1064
1065/// A [RegionScanner] that only scans a single partition.
1066pub struct SinglePartitionScanner {
1067    stream: Mutex<Option<SendableRecordBatchStream>>,
1068    schema: SchemaRef,
1069    properties: ScannerProperties,
1070    metadata: RegionMetadataRef,
1071    snapshot_sequence: Option<SequenceNumber>,
1072}
1073
1074impl SinglePartitionScanner {
1075    /// Creates a new [SinglePartitionScanner] with the given stream and metadata.
1076    pub fn new(
1077        stream: SendableRecordBatchStream,
1078        append_mode: bool,
1079        metadata: RegionMetadataRef,
1080        snapshot_sequence: Option<SequenceNumber>,
1081    ) -> Self {
1082        let schema = stream.schema();
1083        Self {
1084            stream: Mutex::new(Some(stream)),
1085            schema,
1086            properties: ScannerProperties::default().with_append_mode(append_mode),
1087            metadata,
1088            snapshot_sequence,
1089        }
1090    }
1091}
1092
1093impl Debug for SinglePartitionScanner {
1094    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1095        write!(f, "SinglePartitionScanner: <SendableRecordBatchStream>")
1096    }
1097}
1098
1099impl RegionScanner for SinglePartitionScanner {
1100    fn name(&self) -> &str {
1101        "SinglePartition"
1102    }
1103
1104    fn properties(&self) -> &ScannerProperties {
1105        &self.properties
1106    }
1107
1108    fn schema(&self) -> SchemaRef {
1109        self.schema.clone()
1110    }
1111
1112    fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
1113        self.properties.prepare(request);
1114        Ok(())
1115    }
1116
1117    fn scan_partition(
1118        &self,
1119        _ctx: &QueryScanContext,
1120        _metrics_set: &ExecutionPlanMetricsSet,
1121        _partition: usize,
1122    ) -> Result<SendableRecordBatchStream, BoxedError> {
1123        let mut stream = self.stream.lock().unwrap();
1124        let result = stream
1125            .take()
1126            .or_else(|| Some(Box::pin(EmptyRecordBatchStream::new(self.schema.clone()))));
1127        Ok(result.unwrap())
1128    }
1129
1130    fn has_predicate_without_region(&self) -> bool {
1131        false
1132    }
1133
1134    fn add_dyn_filter_to_predicate(
1135        &mut self,
1136        filter_exprs: Vec<Arc<dyn datafusion_physical_plan::PhysicalExpr>>,
1137    ) -> Vec<bool> {
1138        vec![false; filter_exprs.len()]
1139    }
1140
1141    fn metadata(&self) -> RegionMetadataRef {
1142        self.metadata.clone()
1143    }
1144
1145    fn set_logical_region(&mut self, logical_region: bool) {
1146        self.properties.set_logical_region(logical_region);
1147    }
1148
1149    fn set_query_load_region_id(&mut self, region_id: RegionId) {
1150        self.properties.set_query_load_region_id(region_id);
1151    }
1152
1153    fn snapshot_sequence(&self) -> Option<SequenceNumber> {
1154        self.snapshot_sequence
1155    }
1156}
1157
1158impl DisplayAs for SinglePartitionScanner {
1159    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1160        write!(f, "{:?}", self)
1161    }
1162}
store_api/region_engine.rs

store_api/
region_engine.rs
