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

store_api/
region_engine.rs
