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 to yield an empty batch to distinguish partition ranges.
310    pub distinguish_partition_range: bool,
311
312    /// The target partitions of the scanner. 0 indicates using the number of partitions as target partitions.
313    target_partitions: usize,
314
315    /// Whether the scanner is scanning a logical region.
316    logical_region: bool,
317
318    /// Region id that should receive query-load metrics for this scanner.
319    query_load_region_id: Option<RegionId>,
320    /// Counters that should receive query-load metrics for this scanner.
321    query_stat_counters: Option<RegionQueryStatCounters>,
322}
323
324impl ScannerProperties {
325    /// Sets append mode for scanner.
326    pub fn with_append_mode(mut self, append_mode: bool) -> Self {
327        self.append_mode = append_mode;
328        self
329    }
330
331    /// Sets total rows for scanner.
332    pub fn with_total_rows(mut self, total_rows: usize) -> Self {
333        self.total_rows = total_rows;
334        self
335    }
336
337    /// Creates a new [`ScannerProperties`] with the given partitioning.
338    pub fn new(partitions: Vec<Vec<PartitionRange>>, append_mode: bool, total_rows: usize) -> Self {
339        Self {
340            partitions,
341            append_mode,
342            total_rows,
343            distinguish_partition_range: false,
344            target_partitions: 0,
345            logical_region: false,
346            query_load_region_id: None,
347            query_stat_counters: None,
348        }
349    }
350
351    /// Updates the properties with the given [PrepareRequest].
352    pub fn prepare(&mut self, request: PrepareRequest) {
353        if let Some(ranges) = request.ranges {
354            self.partitions = ranges;
355        }
356        if let Some(distinguish_partition_range) = request.distinguish_partition_range {
357            self.distinguish_partition_range = distinguish_partition_range;
358        }
359        if let Some(target_partitions) = request.target_partitions {
360            self.target_partitions = target_partitions;
361        }
362    }
363
364    /// Returns the number of actual partitions.
365    pub fn num_partitions(&self) -> usize {
366        self.partitions.len()
367    }
368
369    pub fn append_mode(&self) -> bool {
370        self.append_mode
371    }
372
373    pub fn total_rows(&self) -> usize {
374        self.total_rows
375    }
376
377    /// Returns whether the scanner is scanning a logical region.
378    pub fn is_logical_region(&self) -> bool {
379        self.logical_region
380    }
381
382    /// Returns the target partitions of the scanner. If it is not set, returns the number of partitions.
383    pub fn target_partitions(&self) -> usize {
384        if self.target_partitions == 0 {
385            self.num_partitions()
386        } else {
387            self.target_partitions
388        }
389    }
390
391    /// Sets whether the scanner is reading a logical region.
392    pub fn set_logical_region(&mut self, logical_region: bool) {
393        self.logical_region = logical_region;
394    }
395
396    /// Returns the region id that should receive query-load metrics.
397    pub fn query_load_region_id(&self) -> Option<RegionId> {
398        self.query_load_region_id
399    }
400
401    /// Returns the counters that should receive query-load metrics.
402    pub fn query_stat_counters(&self) -> Option<RegionQueryStatCounters> {
403        self.query_stat_counters.clone()
404    }
405
406    /// Sets the region id that should receive query-load metrics.
407    pub fn set_query_load_region_id(&mut self, region_id: RegionId) {
408        self.query_load_region_id = Some(region_id);
409    }
410
411    /// Sets the counters that should receive query-load metrics.
412    pub fn set_query_stat_counters(&mut self, counters: RegionQueryStatCounters) {
413        self.query_stat_counters = Some(counters);
414    }
415}
416
417/// Request to override the scanner properties.
418#[derive(Default)]
419pub struct PrepareRequest {
420    /// Assigned partition ranges.
421    pub ranges: Option<Vec<Vec<PartitionRange>>>,
422    /// Distringuishes partition range by empty batches.
423    pub distinguish_partition_range: Option<bool>,
424    /// The expected number of target partitions.
425    pub target_partitions: Option<usize>,
426}
427
428impl PrepareRequest {
429    /// Sets the ranges.
430    pub fn with_ranges(mut self, ranges: Vec<Vec<PartitionRange>>) -> Self {
431        self.ranges = Some(ranges);
432        self
433    }
434
435    /// Sets the distinguish partition range flag.
436    pub fn with_distinguish_partition_range(mut self, distinguish_partition_range: bool) -> Self {
437        self.distinguish_partition_range = Some(distinguish_partition_range);
438        self
439    }
440
441    /// Sets the target partitions.
442    pub fn with_target_partitions(mut self, target_partitions: usize) -> Self {
443        self.target_partitions = Some(target_partitions);
444        self
445    }
446}
447
448/// Necessary context of the query for the scanner.
449#[derive(Clone, Default)]
450pub struct QueryScanContext {
451    /// Whether the query is EXPLAIN ANALYZE VERBOSE.
452    pub explain_verbose: bool,
453}
454
455/// A scanner that provides a way to scan the region concurrently.
456///
457/// The scanner splits the region into partitions so that each partition can be scanned concurrently.
458/// You can use this trait to implement an [`ExecutionPlan`](datafusion_physical_plan::ExecutionPlan).
459pub trait RegionScanner: Debug + DisplayAs + Send {
460    fn name(&self) -> &str;
461
462    /// Returns the properties of the scanner.
463    fn properties(&self) -> &ScannerProperties;
464
465    /// Returns the schema of the record batches.
466    fn schema(&self) -> SchemaRef;
467
468    /// Returns the metadata of the region.
469    fn metadata(&self) -> RegionMetadataRef;
470
471    /// Prepares the scanner with the given partition ranges.
472    ///
473    /// This method is for the planner to adjust the scanner's behavior based on the partition ranges.
474    fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError>;
475
476    /// Scans the partition and returns a stream of record batches.
477    ///
478    /// # Panics
479    /// Panics if the `partition` is out of bound.
480    fn scan_partition(
481        &self,
482        ctx: &QueryScanContext,
483        metrics_set: &ExecutionPlanMetricsSet,
484        partition: usize,
485    ) -> Result<SendableRecordBatchStream, BoxedError>;
486
487    /// Check if there is any predicate exclude region partition exprs that may be executed in this scanner.
488    fn has_predicate_without_region(&self) -> bool;
489
490    /// Add the given dynamic filter expressions to the predicate of the scanner.
491    /// Returns a vector of booleans indicating which filter expressions were applied.
492    /// true indicates the filter expression was applied(will be use by scanner to prune by stat for row group),
493    /// false otherwise.
494    fn add_dyn_filter_to_predicate(
495        &mut self,
496        filter_exprs: Vec<Arc<dyn PhysicalExpr>>,
497    ) -> Vec<bool>;
498
499    /// Sets whether the scanner is reading a logical region.
500    fn set_logical_region(&mut self, logical_region: bool);
501
502    /// Sets the region id that should receive query-load metrics.
503    fn set_query_load_region_id(&mut self, region_id: RegionId);
504
505    fn snapshot_sequence(&self) -> Option<SequenceNumber> {
506        None
507    }
508}
509
510pub type RegionScannerRef = Box<dyn RegionScanner>;
511
512pub type BatchResponses = Vec<(RegionId, Result<RegionResponse, BoxedError>)>;
513
514/// Represents the statistics of a region.
515#[derive(Debug, Deserialize, Serialize, Default)]
516pub struct RegionStatistic {
517    /// The number of rows stored in SST files owned by this region plus rows in memtables.
518    ///
519    /// Rows from SST files referenced from other regions, for example after repartition,
520    /// are not counted to avoid table-level double counting when summing region statistics.
521    #[serde(default)]
522    pub num_rows: u64,
523    /// The size of memtable in bytes.
524    pub memtable_size: u64,
525    /// The size of WAL in bytes.
526    pub wal_size: u64,
527    /// The size of manifest in bytes.
528    pub manifest_size: u64,
529    /// The size of SST data files owned by this region in bytes.
530    ///
531    /// SST files referenced from other regions, for example after repartition, are not counted.
532    pub sst_size: u64,
533    /// The number of SST files owned by this region.
534    ///
535    /// SST files referenced from other regions, for example after repartition, are not counted.
536    pub sst_num: u64,
537    /// The size of SST index files owned by this region in bytes.
538    ///
539    /// SST index files referenced from other regions, for example after repartition, are not counted.
540    #[serde(default)]
541    pub index_size: u64,
542    /// The details of the region.
543    #[serde(default)]
544    pub manifest: RegionManifestInfo,
545    #[serde(default)]
546    /// The total bytes written of the region since region opened.
547    pub written_bytes: u64,
548    /// The total query CPU time of the region since region opened.
549    ///
550    /// Unit: nanoseconds.
551    #[serde(default)]
552    pub query_cpu_time: u64,
553    /// The total scanned bytes of the region since region opened.
554    #[serde(default)]
555    pub query_scanned_bytes: u64,
556    /// The latest entry id of the region's remote WAL since last flush.
557    /// For metric engine, there're two latest entry ids, one for data and one for metadata.
558    /// TODO(weny): remove this two fields and use single instead.
559    #[serde(default)]
560    pub data_topic_latest_entry_id: u64,
561    #[serde(default)]
562    pub metadata_topic_latest_entry_id: u64,
563}
564
565/// The manifest info of a region.
566#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
567pub enum RegionManifestInfo {
568    Mito {
569        manifest_version: u64,
570        flushed_entry_id: u64,
571        /// Number of files removed in the manifest's `removed_files` field.
572        file_removed_cnt: u64,
573    },
574    Metric {
575        data_manifest_version: u64,
576        data_flushed_entry_id: u64,
577        metadata_manifest_version: u64,
578        metadata_flushed_entry_id: u64,
579    },
580}
581
582impl RegionManifestInfo {
583    /// Creates a new [RegionManifestInfo] for mito2 engine.
584    pub fn mito(manifest_version: u64, flushed_entry_id: u64, file_removal_rate: u64) -> Self {
585        Self::Mito {
586            manifest_version,
587            flushed_entry_id,
588            file_removed_cnt: file_removal_rate,
589        }
590    }
591
592    /// Creates a new [RegionManifestInfo] for metric engine.
593    pub fn metric(
594        data_manifest_version: u64,
595        data_flushed_entry_id: u64,
596        metadata_manifest_version: u64,
597        metadata_flushed_entry_id: u64,
598    ) -> Self {
599        Self::Metric {
600            data_manifest_version,
601            data_flushed_entry_id,
602            metadata_manifest_version,
603            metadata_flushed_entry_id,
604        }
605    }
606
607    /// Returns true if the region is a mito2 region.
608    pub fn is_mito(&self) -> bool {
609        matches!(self, RegionManifestInfo::Mito { .. })
610    }
611
612    /// Returns true if the region is a metric region.
613    pub fn is_metric(&self) -> bool {
614        matches!(self, RegionManifestInfo::Metric { .. })
615    }
616
617    /// Returns the flushed entry id of the data region.
618    pub fn data_flushed_entry_id(&self) -> u64 {
619        match self {
620            RegionManifestInfo::Mito {
621                flushed_entry_id, ..
622            } => *flushed_entry_id,
623            RegionManifestInfo::Metric {
624                data_flushed_entry_id,
625                ..
626            } => *data_flushed_entry_id,
627        }
628    }
629
630    /// Returns the manifest version of the data region.
631    pub fn data_manifest_version(&self) -> u64 {
632        match self {
633            RegionManifestInfo::Mito {
634                manifest_version, ..
635            } => *manifest_version,
636            RegionManifestInfo::Metric {
637                data_manifest_version,
638                ..
639            } => *data_manifest_version,
640        }
641    }
642
643    /// Returns the manifest version of the metadata region.
644    pub fn metadata_manifest_version(&self) -> Option<u64> {
645        match self {
646            RegionManifestInfo::Mito { .. } => None,
647            RegionManifestInfo::Metric {
648                metadata_manifest_version,
649                ..
650            } => Some(*metadata_manifest_version),
651        }
652    }
653
654    /// Returns the flushed entry id of the metadata region.
655    pub fn metadata_flushed_entry_id(&self) -> Option<u64> {
656        match self {
657            RegionManifestInfo::Mito { .. } => None,
658            RegionManifestInfo::Metric {
659                metadata_flushed_entry_id,
660                ..
661            } => Some(*metadata_flushed_entry_id),
662        }
663    }
664
665    /// Encodes a list of ([RegionId], [RegionManifestInfo]) to a byte array.
666    pub fn encode_list(manifest_infos: &[(RegionId, Self)]) -> serde_json::Result<Vec<u8>> {
667        serde_json::to_vec(manifest_infos)
668    }
669
670    /// Decodes a list of ([RegionId], [RegionManifestInfo]) from a byte array.
671    pub fn decode_list(value: &[u8]) -> serde_json::Result<Vec<(RegionId, Self)>> {
672        serde_json::from_slice(value)
673    }
674}
675
676impl Default for RegionManifestInfo {
677    fn default() -> Self {
678        Self::Mito {
679            manifest_version: 0,
680            flushed_entry_id: 0,
681            file_removed_cnt: 0,
682        }
683    }
684}
685
686impl RegionStatistic {
687    /// Deserializes the region statistic to a byte array.
688    ///
689    /// Returns None if the deserialization fails.
690    pub fn deserialize_from_slice(value: &[u8]) -> Option<RegionStatistic> {
691        serde_json::from_slice(value).ok()
692    }
693
694    /// Serializes the region statistic to a byte array.
695    ///
696    /// Returns None if the serialization fails.
697    pub fn serialize_to_vec(&self) -> Option<Vec<u8>> {
698        serde_json::to_vec(self).ok()
699    }
700}
701
702impl RegionStatistic {
703    /// Returns the estimated disk size of the region.
704    pub fn estimated_disk_size(&self) -> u64 {
705        self.wal_size + self.sst_size + self.manifest_size + self.index_size
706    }
707}
708
709#[cfg(test)]
710mod tests {
711    use serde_json::json;
712
713    use super::*;
714
715    #[test]
716    fn region_statistic_deserializes_without_query_stats() {
717        let statistic: RegionStatistic = serde_json::from_value(json!({
718            "num_rows": 1,
719            "memtable_size": 2,
720            "wal_size": 3,
721            "manifest_size": 4,
722            "sst_size": 5,
723            "sst_num": 6,
724            "index_size": 7,
725            "manifest": {
726                "Mito": {
727                    "manifest_version": 8,
728                    "flushed_entry_id": 9,
729                    "file_removed_cnt": 10
730                }
731            },
732            "written_bytes": 11
733        }))
734        .unwrap();
735
736        assert_eq!(statistic.query_cpu_time, 0);
737        assert_eq!(statistic.query_scanned_bytes, 0);
738    }
739}
740
741/// Request to sync the region from a manifest or a region.
742#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
743pub enum SyncRegionFromRequest {
744    /// Syncs the region using manifest information.
745    /// Used in leader-follower manifest sync scenarios.
746    FromManifest(RegionManifestInfo),
747    /// Syncs the region from another region.
748    ///
749    /// Used by the metric engine to sync logical regions from a source physical region
750    /// to a target physical region. This copies metadata region SST files and transforms
751    /// logical region entries to use the target's region number.
752    FromRegion {
753        /// The [`RegionId`] of the source region.
754        source_region_id: RegionId,
755        /// The parallelism of the sync operation.
756        parallelism: usize,
757    },
758}
759
760impl From<RegionManifestInfo> for SyncRegionFromRequest {
761    fn from(manifest_info: RegionManifestInfo) -> Self {
762        SyncRegionFromRequest::FromManifest(manifest_info)
763    }
764}
765
766impl SyncRegionFromRequest {
767    /// Creates a new request from a manifest info.
768    pub fn from_manifest(manifest_info: RegionManifestInfo) -> Self {
769        SyncRegionFromRequest::FromManifest(manifest_info)
770    }
771
772    /// Creates a new request from a region.
773    pub fn from_region(source_region_id: RegionId, parallelism: usize) -> Self {
774        SyncRegionFromRequest::FromRegion {
775            source_region_id,
776            parallelism,
777        }
778    }
779
780    /// Returns true if the request is from a manifest.
781    pub fn is_from_manifest(&self) -> bool {
782        matches!(self, SyncRegionFromRequest::FromManifest { .. })
783    }
784
785    /// Converts the request to a region manifest info.
786    ///
787    /// Returns None if the request is not from a manifest.
788    pub fn into_region_manifest_info(self) -> Option<RegionManifestInfo> {
789        match self {
790            SyncRegionFromRequest::FromManifest(manifest_info) => Some(manifest_info),
791            SyncRegionFromRequest::FromRegion { .. } => None,
792        }
793    }
794}
795
796/// The response of syncing the region.
797#[derive(Debug)]
798pub enum SyncRegionFromResponse {
799    NotSupported,
800    Mito {
801        /// Indicates if the data region was synced.
802        synced: bool,
803    },
804    Metric {
805        /// Indicates if the metadata region was synced.
806        metadata_synced: bool,
807        /// Indicates if the data region was synced.
808        data_synced: bool,
809        /// The logical regions that were newly opened during the sync operation.
810        /// This only occurs after the metadata region has been successfully synced.
811        new_opened_logical_region_ids: Vec<RegionId>,
812    },
813}
814
815impl SyncRegionFromResponse {
816    /// Returns true if data region is synced.
817    pub fn is_data_synced(&self) -> bool {
818        match self {
819            SyncRegionFromResponse::NotSupported => false,
820            SyncRegionFromResponse::Mito { synced } => *synced,
821            SyncRegionFromResponse::Metric { data_synced, .. } => *data_synced,
822        }
823    }
824
825    /// Returns true if the engine is a mito2 engine.
826    pub fn is_mito(&self) -> bool {
827        matches!(self, SyncRegionFromResponse::Mito { .. })
828    }
829
830    /// Returns true if the engine is a metric engine.
831    pub fn is_metric(&self) -> bool {
832        matches!(self, SyncRegionFromResponse::Metric { .. })
833    }
834
835    /// Returns the new opened logical region ids.
836    pub fn new_opened_logical_region_ids(self) -> Option<Vec<RegionId>> {
837        match self {
838            SyncRegionFromResponse::Metric {
839                new_opened_logical_region_ids,
840                ..
841            } => Some(new_opened_logical_region_ids),
842            _ => None,
843        }
844    }
845}
846
847/// Request to remap manifests from old regions to new regions.
848#[derive(Debug, Clone)]
849pub struct RemapManifestsRequest {
850    /// The [`RegionId`] of a staging region used to obtain table directory and storage configuration for the remap operation.
851    pub region_id: RegionId,
852    /// Regions to remap manifests from.
853    pub input_regions: Vec<RegionId>,
854    /// For each old region, which new regions should receive its files
855    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
856    /// New partition expressions for the new regions.
857    pub new_partition_exprs: HashMap<RegionId, String>,
858}
859
860/// Response to remap manifests from old regions to new regions.
861#[derive(Debug, Clone)]
862pub struct RemapManifestsResponse {
863    /// Maps region id to its staging manifest path.
864    ///
865    /// These paths are relative paths within the central region's staging blob storage,
866    /// and should be passed to [`ApplyStagingManifestRequest`](RegionRequest::ApplyStagingManifest) to finalize the repartition.
867    pub manifest_paths: HashMap<RegionId, String>,
868}
869
870/// Request to copy files from a source region to a target region.
871#[derive(Debug, Clone)]
872pub struct MitoCopyRegionFromRequest {
873    /// The [`RegionId`] of the source region.
874    pub source_region_id: RegionId,
875    /// The parallelism of the copy operation.
876    pub parallelism: usize,
877}
878
879#[derive(Debug, Clone)]
880pub struct MitoCopyRegionFromResponse {
881    /// The file ids that were copied from the source region to the target region.
882    pub copied_file_ids: Vec<FileId>,
883}
884
885#[async_trait]
886pub trait RegionEngine: Send + Sync {
887    /// Name of this engine
888    fn name(&self) -> &str;
889
890    /// Handles batch open region requests.
891    async fn handle_batch_open_requests(
892        &self,
893        parallelism: usize,
894        requests: Vec<(RegionId, RegionOpenRequest)>,
895    ) -> Result<BatchResponses, BoxedError> {
896        let semaphore = Arc::new(Semaphore::new(parallelism));
897        let mut tasks = Vec::with_capacity(requests.len());
898
899        for (region_id, request) in requests {
900            let semaphore_moved = semaphore.clone();
901
902            tasks.push(async move {
903                // Safety: semaphore must exist
904                let _permit = semaphore_moved.acquire().await.unwrap();
905                let result = self
906                    .handle_request(region_id, RegionRequest::Open(request))
907                    .await;
908                (region_id, result)
909            });
910        }
911
912        Ok(join_all(tasks).await)
913    }
914
915    async fn handle_batch_catchup_requests(
916        &self,
917        parallelism: usize,
918        requests: Vec<(RegionId, RegionCatchupRequest)>,
919    ) -> Result<BatchResponses, BoxedError> {
920        let semaphore = Arc::new(Semaphore::new(parallelism));
921        let mut tasks = Vec::with_capacity(requests.len());
922
923        for (region_id, request) in requests {
924            let semaphore_moved = semaphore.clone();
925
926            tasks.push(async move {
927                // Safety: semaphore must exist
928                let _permit = semaphore_moved.acquire().await.unwrap();
929                let result = self
930                    .handle_request(region_id, RegionRequest::Catchup(request))
931                    .await;
932                (region_id, result)
933            });
934        }
935
936        Ok(join_all(tasks).await)
937    }
938
939    async fn handle_batch_ddl_requests(
940        &self,
941        request: BatchRegionDdlRequest,
942    ) -> Result<RegionResponse, BoxedError> {
943        let requests = request.into_region_requests();
944
945        let mut affected_rows = 0;
946        let mut extensions = HashMap::new();
947
948        for (region_id, request) in requests {
949            let result = self.handle_request(region_id, request).await?;
950            affected_rows += result.affected_rows;
951            extensions.extend(result.extensions);
952        }
953
954        Ok(RegionResponse {
955            affected_rows,
956            extensions,
957            metadata: Vec::new(),
958        })
959    }
960
961    /// Handles non-query request to the region. Returns the count of affected rows.
962    async fn handle_request(
963        &self,
964        region_id: RegionId,
965        request: RegionRequest,
966    ) -> Result<RegionResponse, BoxedError>;
967
968    /// Returns the committed sequence (sequence of latest written data).
969    async fn get_committed_sequence(
970        &self,
971        region_id: RegionId,
972    ) -> Result<SequenceNumber, BoxedError>;
973
974    /// Handles query and return a scanner that can be used to scan the region concurrently.
975    async fn handle_query(
976        &self,
977        region_id: RegionId,
978        request: ScanRequest,
979    ) -> Result<RegionScannerRef, BoxedError>;
980
981    /// Returns the query memory tracker for scan execution.
982    fn query_memory_tracker(&self) -> Option<QueryMemoryTracker> {
983        None
984    }
985
986    /// Retrieves region's metadata.
987    async fn get_metadata(&self, region_id: RegionId) -> Result<RegionMetadataRef, BoxedError>;
988
989    /// Retrieves region's statistic.
990    fn region_statistic(&self, region_id: RegionId) -> Option<RegionStatistic>;
991
992    /// Stops the engine
993    async fn stop(&self) -> Result<(), BoxedError>;
994
995    /// Sets [RegionRole] for a region.
996    ///
997    /// The engine checks whether the region is writable before writing to the region. Setting
998    /// the region as readonly doesn't guarantee that write operations in progress will not
999    /// take effect.
1000    fn set_region_role(&self, region_id: RegionId, role: RegionRole) -> Result<(), BoxedError>;
1001
1002    /// Syncs the region manifest to the given manifest version.
1003    async fn sync_region(
1004        &self,
1005        region_id: RegionId,
1006        request: SyncRegionFromRequest,
1007    ) -> Result<SyncRegionFromResponse, BoxedError>;
1008
1009    /// Remaps manifests from old regions to new regions.
1010    async fn remap_manifests(
1011        &self,
1012        request: RemapManifestsRequest,
1013    ) -> Result<RemapManifestsResponse, BoxedError>;
1014
1015    /// Sets region role state gracefully.
1016    ///
1017    /// After the call returns, the engine ensures no more write operations will succeed in the region.
1018    async fn set_region_role_state_gracefully(
1019        &self,
1020        region_id: RegionId,
1021        region_role_state: SettableRegionRoleState,
1022    ) -> Result<SetRegionRoleStateResponse, BoxedError>;
1023
1024    /// Indicates region role.
1025    ///
1026    /// Returns the `None` if the region is not found.
1027    fn role(&self, region_id: RegionId) -> Option<RegionRole>;
1028
1029    fn as_any(&self) -> &dyn Any;
1030}
1031
1032pub type RegionEngineRef = Arc<dyn RegionEngine>;
1033
1034/// A [RegionScanner] that only scans a single partition.
1035pub struct SinglePartitionScanner {
1036    stream: Mutex<Option<SendableRecordBatchStream>>,
1037    schema: SchemaRef,
1038    properties: ScannerProperties,
1039    metadata: RegionMetadataRef,
1040    snapshot_sequence: Option<SequenceNumber>,
1041}
1042
1043impl SinglePartitionScanner {
1044    /// Creates a new [SinglePartitionScanner] with the given stream and metadata.
1045    pub fn new(
1046        stream: SendableRecordBatchStream,
1047        append_mode: bool,
1048        metadata: RegionMetadataRef,
1049        snapshot_sequence: Option<SequenceNumber>,
1050    ) -> Self {
1051        let schema = stream.schema();
1052        Self {
1053            stream: Mutex::new(Some(stream)),
1054            schema,
1055            properties: ScannerProperties::default().with_append_mode(append_mode),
1056            metadata,
1057            snapshot_sequence,
1058        }
1059    }
1060}
1061
1062impl Debug for SinglePartitionScanner {
1063    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1064        write!(f, "SinglePartitionScanner: <SendableRecordBatchStream>")
1065    }
1066}
1067
1068impl RegionScanner for SinglePartitionScanner {
1069    fn name(&self) -> &str {
1070        "SinglePartition"
1071    }
1072
1073    fn properties(&self) -> &ScannerProperties {
1074        &self.properties
1075    }
1076
1077    fn schema(&self) -> SchemaRef {
1078        self.schema.clone()
1079    }
1080
1081    fn prepare(&mut self, request: PrepareRequest) -> Result<(), BoxedError> {
1082        self.properties.prepare(request);
1083        Ok(())
1084    }
1085
1086    fn scan_partition(
1087        &self,
1088        _ctx: &QueryScanContext,
1089        _metrics_set: &ExecutionPlanMetricsSet,
1090        _partition: usize,
1091    ) -> Result<SendableRecordBatchStream, BoxedError> {
1092        let mut stream = self.stream.lock().unwrap();
1093        let result = stream
1094            .take()
1095            .or_else(|| Some(Box::pin(EmptyRecordBatchStream::new(self.schema.clone()))));
1096        Ok(result.unwrap())
1097    }
1098
1099    fn has_predicate_without_region(&self) -> bool {
1100        false
1101    }
1102
1103    fn add_dyn_filter_to_predicate(
1104        &mut self,
1105        filter_exprs: Vec<Arc<dyn datafusion_physical_plan::PhysicalExpr>>,
1106    ) -> Vec<bool> {
1107        vec![false; filter_exprs.len()]
1108    }
1109
1110    fn metadata(&self) -> RegionMetadataRef {
1111        self.metadata.clone()
1112    }
1113
1114    fn set_logical_region(&mut self, logical_region: bool) {
1115        self.properties.set_logical_region(logical_region);
1116    }
1117
1118    fn set_query_load_region_id(&mut self, region_id: RegionId) {
1119        self.properties.set_query_load_region_id(region_id);
1120    }
1121
1122    fn snapshot_sequence(&self) -> Option<SequenceNumber> {
1123        self.snapshot_sequence
1124    }
1125}
1126
1127impl DisplayAs for SinglePartitionScanner {
1128    fn fmt_as(&self, _t: DisplayFormatType, f: &mut std::fmt::Formatter) -> std::fmt::Result {
1129        write!(f, "{:?}", self)
1130    }
1131}
store_api/region_engine.rs

store_api/
region_engine.rs
