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

store_api/
region_engine.rs
