common_meta/
instruction.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
15use std::collections::{HashMap, HashSet};
16use std::fmt::{Display, Formatter};
17use std::time::Duration;
18
19use serde::{Deserialize, Deserializer, Serialize, Serializer};
20use store_api::region_engine::SyncRegionFromRequest;
21use store_api::region_request::{RegionFlushReason, RegionRequirements};
22use store_api::storage::{FileRefsManifest, GcReport, RegionId, RegionNumber};
23use strum::Display;
24use table::metadata::TableId;
25use table::table_name::TableName;
26
27use crate::flow_name::FlowName;
28use crate::key::schema_name::SchemaName;
29use crate::key::{FlowId, FlowPartitionId};
30use crate::peer::Peer;
31use crate::{DatanodeId, FlownodeId};
32
33#[derive(Eq, Hash, PartialEq, Clone, Debug, Serialize, Deserialize)]
34pub struct RegionIdent {
35    pub datanode_id: DatanodeId,
36    pub table_id: TableId,
37    pub region_number: RegionNumber,
38    pub engine: String,
39}
40
41impl RegionIdent {
42    pub fn get_region_id(&self) -> RegionId {
43        RegionId::new(self.table_id, self.region_number)
44    }
45}
46
47impl Display for RegionIdent {
48    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
49        write!(
50            f,
51            "RegionIdent(datanode_id='{}', table_id={}, region_number={}, engine = {})",
52            self.datanode_id, self.table_id, self.region_number, self.engine
53        )
54    }
55}
56
57/// The result of downgrade leader region.
58#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
59pub struct DowngradeRegionReply {
60    /// The [RegionId].
61    /// For compatibility, it is defaulted to [RegionId::new(0, 0)].
62    #[serde(default)]
63    pub region_id: RegionId,
64    /// Returns the `last_entry_id` if available.
65    pub last_entry_id: Option<u64>,
66    /// Returns the `metadata_last_entry_id` if available (Only available for metric engine).
67    pub metadata_last_entry_id: Option<u64>,
68    /// Indicates whether the region exists.
69    pub exists: bool,
70    /// Return error if any during the operation.
71    pub error: Option<String>,
72}
73
74impl Display for DowngradeRegionReply {
75    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
76        write!(
77            f,
78            "(last_entry_id={:?}, exists={}, error={:?})",
79            self.last_entry_id, self.exists, self.error
80        )
81    }
82}
83
84#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
85pub struct SimpleReply {
86    pub result: bool,
87    pub error: Option<String>,
88}
89
90/// Reply for flush region operations with support for batch results.
91#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
92pub struct FlushRegionReply {
93    /// Results for each region that was attempted to be flushed.
94    /// For single region flushes, this will contain one result.
95    /// For batch flushes, this contains results for all attempted regions.
96    pub results: Vec<(RegionId, Result<(), String>)>,
97    /// Overall success: true if all regions were flushed successfully.
98    pub overall_success: bool,
99}
100
101impl FlushRegionReply {
102    /// Create a successful single region reply.
103    pub fn success_single(region_id: RegionId) -> Self {
104        Self {
105            results: vec![(region_id, Ok(()))],
106            overall_success: true,
107        }
108    }
109
110    /// Create a failed single region reply.
111    pub fn error_single(region_id: RegionId, error: String) -> Self {
112        Self {
113            results: vec![(region_id, Err(error))],
114            overall_success: false,
115        }
116    }
117
118    /// Create a batch reply from individual results.
119    pub fn from_results(results: Vec<(RegionId, Result<(), String>)>) -> Self {
120        let overall_success = results.iter().all(|(_, result)| result.is_ok());
121        Self {
122            results,
123            overall_success,
124        }
125    }
126
127    /// Convert to SimpleReply for backward compatibility.
128    pub fn to_simple_reply(&self) -> SimpleReply {
129        if self.overall_success {
130            SimpleReply {
131                result: true,
132                error: None,
133            }
134        } else {
135            let errors: Vec<String> = self
136                .results
137                .iter()
138                .filter_map(|(region_id, result)| {
139                    result
140                        .as_ref()
141                        .err()
142                        .map(|err| format!("{}: {}", region_id, err))
143                })
144                .collect();
145            SimpleReply {
146                result: false,
147                error: Some(errors.join("; ")),
148            }
149        }
150    }
151}
152
153impl Display for SimpleReply {
154    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
155        write!(f, "(result={}, error={:?})", self.result, self.error)
156    }
157}
158
159impl Display for FlushRegionReply {
160    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
161        let results_str = self
162            .results
163            .iter()
164            .map(|(region_id, result)| match result {
165                Ok(()) => format!("{}:OK", region_id),
166                Err(err) => format!("{}:ERR({})", region_id, err),
167            })
168            .collect::<Vec<_>>()
169            .join(", ");
170        write!(
171            f,
172            "(overall_success={}, results=[{}])",
173            self.overall_success, results_str
174        )
175    }
176}
177
178impl Display for OpenRegion {
179    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
180        write!(
181            f,
182            "OpenRegion(region_ident={}, region_storage_path={}, reason={:?})",
183            self.region_ident, self.region_storage_path, self.reason
184        )
185    }
186}
187
188/// The reason why an open region instruction is triggered.
189#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
190pub enum OpenRegionReason {
191    /// Open triggered before region migration.
192    RegionMigration,
193    /// Open triggered by region failover.
194    RegionFailover,
195    /// Open triggered when adding a follower region.
196    #[cfg(feature = "enterprise")]
197    RegionFollower,
198}
199
200#[serde_with::serde_as]
201#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
202pub struct OpenRegion {
203    pub region_ident: RegionIdent,
204    pub region_storage_path: String,
205    pub region_options: HashMap<String, String>,
206    #[serde(default)]
207    #[serde_as(as = "HashMap<serde_with::DisplayFromStr, _>")]
208    pub region_wal_options: HashMap<RegionNumber, String>,
209    #[serde(default)]
210    pub skip_wal_replay: bool,
211    #[serde(default, skip_serializing_if = "Option::is_none")]
212    pub reason: Option<OpenRegionReason>,
213    #[serde(default)]
214    pub requirements: RegionRequirements,
215}
216
217impl OpenRegion {
218    pub fn new(
219        region_ident: RegionIdent,
220        path: &str,
221        region_options: HashMap<String, String>,
222        region_wal_options: HashMap<RegionNumber, String>,
223        skip_wal_replay: bool,
224        reason: Option<OpenRegionReason>,
225        requirements: RegionRequirements,
226    ) -> Self {
227        Self {
228            region_ident,
229            region_storage_path: path.to_string(),
230            region_options,
231            region_wal_options,
232            skip_wal_replay,
233            reason,
234            requirements,
235        }
236    }
237}
238
239/// The instruction of downgrading leader region.
240#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
241pub struct DowngradeRegion {
242    /// The [RegionId].
243    pub region_id: RegionId,
244    /// The timeout of waiting for flush the region.
245    ///
246    /// `None` stands for don't flush before downgrading the region.
247    #[serde(default)]
248    pub flush_timeout: Option<Duration>,
249}
250
251impl Display for DowngradeRegion {
252    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
253        write!(
254            f,
255            "DowngradeRegion(region_id={}, flush_timeout={:?})",
256            self.region_id, self.flush_timeout,
257        )
258    }
259}
260
261/// Upgrades a follower region to leader region.
262#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
263pub struct UpgradeRegion {
264    /// The [RegionId].
265    pub region_id: RegionId,
266    /// The `last_entry_id` of old leader region.
267    pub last_entry_id: Option<u64>,
268    /// The `last_entry_id` of old leader metadata region (Only used for metric engine).
269    pub metadata_last_entry_id: Option<u64>,
270    /// The timeout of waiting for a wal replay.
271    ///
272    /// `None` stands for no wait,
273    /// it's helpful to verify whether the leader region is ready.
274    #[serde(with = "humantime_serde")]
275    pub replay_timeout: Duration,
276    /// The hint for replaying memtable.
277    #[serde(default)]
278    pub location_id: Option<u64>,
279    #[serde(default, skip_serializing_if = "Option::is_none")]
280    pub replay_entry_id: Option<u64>,
281    #[serde(default, skip_serializing_if = "Option::is_none")]
282    pub metadata_replay_entry_id: Option<u64>,
283}
284
285impl UpgradeRegion {
286    /// Sets the replay entry id.
287    pub fn with_replay_entry_id(mut self, replay_entry_id: Option<u64>) -> Self {
288        self.replay_entry_id = replay_entry_id;
289        self
290    }
291
292    /// Sets the metadata replay entry id.
293    pub fn with_metadata_replay_entry_id(mut self, metadata_replay_entry_id: Option<u64>) -> Self {
294        self.metadata_replay_entry_id = metadata_replay_entry_id;
295        self
296    }
297}
298
299#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
300/// The identifier of cache.
301pub enum CacheIdent {
302    FlowId(FlowId),
303    /// Indicate change of address of flownode.
304    FlowNodeAddressChange(u64),
305    FlowName(FlowName),
306    TableId(TableId),
307    TableName(TableName),
308    SchemaName(SchemaName),
309    CreateFlow(CreateFlow),
310    DropFlow(DropFlow),
311    /// Indicate change of user metadata.
312    User(UserCacheIdent),
313}
314
315#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
316pub struct UserCacheIdent {
317    pub catalog: String,
318    pub username: String,
319}
320
321#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
322pub struct CreateFlow {
323    /// The unique identifier for the flow.
324    pub flow_id: FlowId,
325    pub source_table_ids: Vec<TableId>,
326    /// Mapping of flow partition to peer information
327    pub partition_to_peer_mapping: Vec<(FlowPartitionId, Peer)>,
328}
329
330#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
331pub struct DropFlow {
332    pub flow_id: FlowId,
333    pub source_table_ids: Vec<TableId>,
334    /// Mapping of flow partition to flownode id
335    pub flow_part2node_id: Vec<(FlowPartitionId, FlownodeId)>,
336}
337
338/// Strategy for executing flush operations.
339#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
340pub enum FlushStrategy {
341    /// Synchronous operation that waits for completion and expects a reply
342    #[default]
343    Sync,
344    /// Asynchronous hint operation (fire-and-forget, no reply expected)
345    Async,
346}
347
348/// Error handling strategy for batch flush operations.
349#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
350pub enum FlushErrorStrategy {
351    /// Abort on first error (fail-fast)
352    #[default]
353    FailFast,
354    /// Attempt to flush all regions and collect all errors
355    TryAll,
356}
357
358/// Unified flush instruction supporting both single and batch operations
359/// with configurable execution strategies and error handling.
360#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
361pub struct FlushRegions {
362    /// List of region IDs to flush. Can contain a single region or multiple regions.
363    pub region_ids: Vec<RegionId>,
364    /// Execution strategy: Sync (expects reply) or Async (fire-and-forget hint).
365    #[serde(default)]
366    pub strategy: FlushStrategy,
367    /// Error handling strategy for batch operations (only applies when multiple regions and sync strategy).
368    #[serde(default)]
369    pub error_strategy: FlushErrorStrategy,
370    /// The source that triggered this flush.
371    #[serde(default, skip_serializing_if = "Option::is_none")]
372    pub reason: Option<RegionFlushReason>,
373}
374
375impl Display for FlushRegions {
376    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
377        write!(
378            f,
379            "FlushRegions(region_ids={:?}, strategy={:?}, error_strategy={:?}, reason={:?})",
380            self.region_ids, self.strategy, self.error_strategy, self.reason
381        )
382    }
383}
384
385impl FlushRegions {
386    /// Create synchronous single-region flush
387    pub fn sync_single(region_id: RegionId) -> Self {
388        Self {
389            region_ids: vec![region_id],
390            strategy: FlushStrategy::Sync,
391            error_strategy: FlushErrorStrategy::FailFast,
392            reason: None,
393        }
394    }
395
396    /// Create asynchronous batch flush (fire-and-forget)
397    pub fn async_batch(region_ids: Vec<RegionId>) -> Self {
398        Self {
399            region_ids,
400            strategy: FlushStrategy::Async,
401            error_strategy: FlushErrorStrategy::TryAll,
402            reason: None,
403        }
404    }
405
406    /// Create synchronous batch flush with error strategy
407    pub fn sync_batch(region_ids: Vec<RegionId>, error_strategy: FlushErrorStrategy) -> Self {
408        Self {
409            region_ids,
410            strategy: FlushStrategy::Sync,
411            error_strategy,
412            reason: None,
413        }
414    }
415
416    pub fn with_reason(mut self, reason: RegionFlushReason) -> Self {
417        self.reason = Some(reason);
418        self
419    }
420
421    /// Check if this is a single region flush.
422    pub fn is_single_region(&self) -> bool {
423        self.region_ids.len() == 1
424    }
425
426    /// Get the single region ID if this is a single region flush.
427    pub fn single_region_id(&self) -> Option<RegionId> {
428        if self.is_single_region() {
429            self.region_ids.first().copied()
430        } else {
431            None
432        }
433    }
434
435    /// Check if this is a hint (asynchronous) operation.
436    pub fn is_hint(&self) -> bool {
437        matches!(self.strategy, FlushStrategy::Async)
438    }
439
440    /// Check if this is a synchronous operation.
441    pub fn is_sync(&self) -> bool {
442        matches!(self.strategy, FlushStrategy::Sync)
443    }
444}
445
446impl From<RegionId> for FlushRegions {
447    fn from(region_id: RegionId) -> Self {
448        Self::sync_single(region_id)
449    }
450}
451
452#[derive(Debug, Deserialize)]
453#[serde(untagged)]
454enum SingleOrMultiple<T> {
455    Single(T),
456    Multiple(Vec<T>),
457}
458
459fn single_or_multiple_from<'de, D, T>(deserializer: D) -> Result<Vec<T>, D::Error>
460where
461    D: Deserializer<'de>,
462    T: Deserialize<'de>,
463{
464    let helper = SingleOrMultiple::<T>::deserialize(deserializer)?;
465    Ok(match helper {
466        SingleOrMultiple::Single(x) => vec![x],
467        SingleOrMultiple::Multiple(xs) => xs,
468    })
469}
470
471/// Instruction to get file references for specified regions.
472#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
473pub struct GetFileRefs {
474    /// List of region IDs to get file references from active FileHandles (in-memory).
475    pub query_regions: Vec<RegionId>,
476    /// Mapping from the src region IDs (whose file references to look for) to
477    /// the dst region IDs (where to read the manifests).
478    /// Key: The source region IDs (where files originally came from).
479    /// Value: The set of destination region IDs (whose manifests need to be read).
480    pub related_regions: HashMap<RegionId, HashSet<RegionId>>,
481}
482
483impl Display for GetFileRefs {
484    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
485        write!(f, "GetFileRefs(region_ids={:?})", self.query_regions)
486    }
487}
488
489/// Instruction to trigger garbage collection for a region.
490#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
491pub struct GcRegions {
492    /// The region ID to perform GC on, only regions that are currently on the given datanode can be garbage collected, regions not on the datanode will report errors.
493    pub regions: Vec<RegionId>,
494    /// The file references manifest containing temporary file references.
495    pub file_refs_manifest: FileRefsManifest,
496    /// Whether to perform a full file listing to find orphan files.
497    pub full_file_listing: bool,
498}
499
500impl Display for GcRegions {
501    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
502        write!(
503            f,
504            "GcRegion(regions={:?}, file_refs_count={}, full_file_listing={})",
505            self.regions,
506            self.file_refs_manifest.file_refs.len(),
507            self.full_file_listing
508        )
509    }
510}
511
512/// Reply for GetFileRefs instruction.
513#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
514pub struct GetFileRefsReply {
515    /// The file references manifest.
516    pub file_refs_manifest: FileRefsManifest,
517    /// Whether the operation was successful.
518    pub success: bool,
519    /// Error message if any.
520    pub error: Option<String>,
521}
522
523impl Display for GetFileRefsReply {
524    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
525        write!(
526            f,
527            "GetFileRefsReply(success={}, file_refs_count={}, error={:?})",
528            self.success,
529            self.file_refs_manifest.file_refs.len(),
530            self.error
531        )
532    }
533}
534
535/// Reply for GC instruction.
536#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
537pub struct GcRegionsReply {
538    pub result: Result<GcReport, String>,
539}
540
541impl Display for GcRegionsReply {
542    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
543        write!(
544            f,
545            "GcReply(result={})",
546            match &self.result {
547                Ok(report) => format!(
548                    "GcReport(deleted_files_count={}, need_retry_regions_count={})",
549                    report.deleted_files.len(),
550                    report.need_retry_regions.len()
551                ),
552                Err(err) => format!("Err({})", err),
553            }
554        )
555    }
556}
557
558#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
559pub struct EnterStagingRegion {
560    pub region_id: RegionId,
561    #[serde(
562        alias = "partition_expr",
563        deserialize_with = "deserialize_enter_staging_partition_directive",
564        serialize_with = "serialize_enter_staging_partition_directive"
565    )]
566    pub partition_directive: StagingPartitionDirective,
567}
568
569impl Display for EnterStagingRegion {
570    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
571        write!(
572            f,
573            "EnterStagingRegion(region_id={}, partition_directive={})",
574            self.region_id, self.partition_directive
575        )
576    }
577}
578
579#[derive(Debug, Clone, PartialEq, Eq)]
580pub enum StagingPartitionDirective {
581    UpdatePartitionExpr(String),
582    RejectAllWrites,
583}
584
585impl StagingPartitionDirective {
586    /// Returns the partition expression carried by this directive, if any.
587    pub fn as_partition_expr(&self) -> Option<&str> {
588        match self {
589            Self::UpdatePartitionExpr(expr) => Some(expr),
590            Self::RejectAllWrites => None,
591        }
592    }
593}
594
595impl Display for StagingPartitionDirective {
596    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
597        match self {
598            Self::UpdatePartitionExpr(expr) => write!(f, "UpdatePartitionExpr({})", expr),
599            Self::RejectAllWrites => write!(f, "RejectAllWrites"),
600        }
601    }
602}
603
604fn serialize_enter_staging_partition_directive<S>(
605    rule: &StagingPartitionDirective,
606    serializer: S,
607) -> std::result::Result<S::Ok, S::Error>
608where
609    S: Serializer,
610{
611    match rule {
612        StagingPartitionDirective::UpdatePartitionExpr(expr) => serializer.serialize_str(expr),
613        StagingPartitionDirective::RejectAllWrites => {
614            #[derive(Serialize)]
615            struct RejectAllWritesSer<'a> {
616                r#type: &'a str,
617            }
618
619            RejectAllWritesSer {
620                r#type: "reject_all_writes",
621            }
622            .serialize(serializer)
623        }
624    }
625}
626
627fn deserialize_enter_staging_partition_directive<'de, D>(
628    deserializer: D,
629) -> std::result::Result<StagingPartitionDirective, D::Error>
630where
631    D: Deserializer<'de>,
632{
633    #[derive(Deserialize)]
634    #[serde(untagged)]
635    enum Compat {
636        Legacy(String),
637        TypeTagged { r#type: String },
638    }
639
640    match Compat::deserialize(deserializer)? {
641        Compat::Legacy(expr) => Ok(StagingPartitionDirective::UpdatePartitionExpr(expr)),
642        Compat::TypeTagged { r#type } if r#type == "reject_all_writes" => {
643            Ok(StagingPartitionDirective::RejectAllWrites)
644        }
645        Compat::TypeTagged { r#type } => Err(serde::de::Error::custom(format!(
646            "Unknown enter staging partition directive type: {}",
647            r#type
648        ))),
649    }
650}
651
652/// Instruction payload for syncing a region from a manifest or another region.
653#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
654pub struct SyncRegion {
655    /// Region id to sync.
656    pub region_id: RegionId,
657    /// Request to sync the region.
658    pub request: SyncRegionFromRequest,
659}
660
661impl Display for SyncRegion {
662    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
663        write!(
664            f,
665            "SyncRegion(region_id={}, request={:?})",
666            self.region_id, self.request
667        )
668    }
669}
670
671#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
672pub struct RemapManifest {
673    pub region_id: RegionId,
674    /// Regions to remap manifests from.
675    pub input_regions: Vec<RegionId>,
676    /// For each old region, which new regions should receive its files
677    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
678    /// New partition expressions for the new regions.
679    pub new_partition_exprs: HashMap<RegionId, String>,
680}
681
682impl Display for RemapManifest {
683    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
684        write!(
685            f,
686            "RemapManifest(region_id={}, input_regions={:?}, region_mapping={:?}, new_partition_exprs={:?})",
687            self.region_id, self.input_regions, self.region_mapping, self.new_partition_exprs
688        )
689    }
690}
691
692#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
693pub struct ApplyStagingManifest {
694    /// The region ID to apply the staging manifest to.
695    pub region_id: RegionId,
696    /// The partition expression of the staging region.
697    pub partition_expr: String,
698    /// The region that stores the staging manifests in its staging blob storage.
699    pub central_region_id: RegionId,
700    /// The relative path to the staging manifest within the central region's staging blob storage.
701    pub manifest_path: String,
702}
703
704impl Display for ApplyStagingManifest {
705    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
706        write!(
707            f,
708            "ApplyStagingManifest(region_id={}, partition_expr={}, central_region_id={}, manifest_path={})",
709            self.region_id, self.partition_expr, self.central_region_id, self.manifest_path
710        )
711    }
712}
713
714#[derive(Debug, Clone, Serialize, Deserialize, Display, PartialEq)]
715pub enum Instruction {
716    /// Opens regions.
717    #[serde(deserialize_with = "single_or_multiple_from", alias = "OpenRegion")]
718    OpenRegions(Vec<OpenRegion>),
719    /// Closes regions.
720    #[serde(deserialize_with = "single_or_multiple_from", alias = "CloseRegion")]
721    CloseRegions(Vec<RegionIdent>),
722    /// Upgrades regions.
723    #[serde(deserialize_with = "single_or_multiple_from", alias = "UpgradeRegion")]
724    UpgradeRegions(Vec<UpgradeRegion>),
725    #[serde(
726        deserialize_with = "single_or_multiple_from",
727        alias = "DowngradeRegion"
728    )]
729    /// Downgrades regions.
730    DowngradeRegions(Vec<DowngradeRegion>),
731    /// Invalidates batch cache.
732    InvalidateCaches(Vec<CacheIdent>),
733    /// Flushes regions.
734    FlushRegions(FlushRegions),
735    /// Gets file references for regions.
736    GetFileRefs(GetFileRefs),
737    /// Triggers garbage collection for a region.
738    GcRegions(GcRegions),
739    /// Temporary suspend serving reads or writes
740    Suspend,
741    /// Makes regions enter staging state.
742    EnterStagingRegions(Vec<EnterStagingRegion>),
743    /// Syncs regions.
744    SyncRegions(Vec<SyncRegion>),
745    /// Remaps manifests for a region.
746    RemapManifest(RemapManifest),
747
748    /// Applies staging manifests for a region.
749    ApplyStagingManifests(Vec<ApplyStagingManifest>),
750}
751
752impl Instruction {
753    /// Converts the instruction into a vector of [OpenRegion].
754    pub fn into_open_regions(self) -> Option<Vec<OpenRegion>> {
755        match self {
756            Self::OpenRegions(open_regions) => Some(open_regions),
757            _ => None,
758        }
759    }
760
761    /// Converts the instruction into a vector of [RegionIdent].
762    pub fn into_close_regions(self) -> Option<Vec<RegionIdent>> {
763        match self {
764            Self::CloseRegions(close_regions) => Some(close_regions),
765            _ => None,
766        }
767    }
768
769    /// Converts the instruction into a [FlushRegions].
770    pub fn into_flush_regions(self) -> Option<FlushRegions> {
771        match self {
772            Self::FlushRegions(flush_regions) => Some(flush_regions),
773            _ => None,
774        }
775    }
776
777    /// Converts the instruction into a [DowngradeRegion].
778    pub fn into_downgrade_regions(self) -> Option<Vec<DowngradeRegion>> {
779        match self {
780            Self::DowngradeRegions(downgrade_region) => Some(downgrade_region),
781            _ => None,
782        }
783    }
784
785    /// Converts the instruction into a [UpgradeRegion].
786    pub fn into_upgrade_regions(self) -> Option<Vec<UpgradeRegion>> {
787        match self {
788            Self::UpgradeRegions(upgrade_region) => Some(upgrade_region),
789            _ => None,
790        }
791    }
792
793    pub fn into_get_file_refs(self) -> Option<GetFileRefs> {
794        match self {
795            Self::GetFileRefs(get_file_refs) => Some(get_file_refs),
796            _ => None,
797        }
798    }
799
800    pub fn into_gc_regions(self) -> Option<GcRegions> {
801        match self {
802            Self::GcRegions(gc_regions) => Some(gc_regions),
803            _ => None,
804        }
805    }
806
807    pub fn into_enter_staging_regions(self) -> Option<Vec<EnterStagingRegion>> {
808        match self {
809            Self::EnterStagingRegions(enter_staging) => Some(enter_staging),
810            _ => None,
811        }
812    }
813
814    pub fn into_sync_regions(self) -> Option<Vec<SyncRegion>> {
815        match self {
816            Self::SyncRegions(sync_regions) => Some(sync_regions),
817            _ => None,
818        }
819    }
820}
821
822/// The reply of [UpgradeRegion].
823#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
824pub struct UpgradeRegionReply {
825    /// The [RegionId].
826    /// For compatibility, it is defaulted to [RegionId::new(0, 0)].
827    #[serde(default)]
828    pub region_id: RegionId,
829    /// Returns true if `last_entry_id` has been replayed to the latest.
830    pub ready: bool,
831    /// Indicates whether the region exists.
832    pub exists: bool,
833    /// Returns error if any.
834    pub error: Option<String>,
835}
836
837impl Display for UpgradeRegionReply {
838    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
839        write!(
840            f,
841            "(ready={}, exists={}, error={:?})",
842            self.ready, self.exists, self.error
843        )
844    }
845}
846
847#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
848pub struct DowngradeRegionsReply {
849    pub replies: Vec<DowngradeRegionReply>,
850}
851
852impl DowngradeRegionsReply {
853    pub fn new(replies: Vec<DowngradeRegionReply>) -> Self {
854        Self { replies }
855    }
856
857    pub fn single(reply: DowngradeRegionReply) -> Self {
858        Self::new(vec![reply])
859    }
860}
861
862#[derive(Deserialize)]
863#[serde(untagged)]
864enum DowngradeRegionsCompat {
865    Single(DowngradeRegionReply),
866    Multiple(DowngradeRegionsReply),
867}
868
869fn downgrade_regions_compat_from<'de, D>(deserializer: D) -> Result<DowngradeRegionsReply, D::Error>
870where
871    D: Deserializer<'de>,
872{
873    let helper = DowngradeRegionsCompat::deserialize(deserializer)?;
874    Ok(match helper {
875        DowngradeRegionsCompat::Single(x) => DowngradeRegionsReply::new(vec![x]),
876        DowngradeRegionsCompat::Multiple(reply) => reply,
877    })
878}
879
880#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
881pub struct UpgradeRegionsReply {
882    pub replies: Vec<UpgradeRegionReply>,
883}
884
885impl UpgradeRegionsReply {
886    pub fn new(replies: Vec<UpgradeRegionReply>) -> Self {
887        Self { replies }
888    }
889
890    pub fn single(reply: UpgradeRegionReply) -> Self {
891        Self::new(vec![reply])
892    }
893}
894
895#[derive(Deserialize)]
896#[serde(untagged)]
897enum UpgradeRegionsCompat {
898    Single(UpgradeRegionReply),
899    Multiple(UpgradeRegionsReply),
900}
901
902fn upgrade_regions_compat_from<'de, D>(deserializer: D) -> Result<UpgradeRegionsReply, D::Error>
903where
904    D: Deserializer<'de>,
905{
906    let helper = UpgradeRegionsCompat::deserialize(deserializer)?;
907    Ok(match helper {
908        UpgradeRegionsCompat::Single(x) => UpgradeRegionsReply::new(vec![x]),
909        UpgradeRegionsCompat::Multiple(reply) => reply,
910    })
911}
912
913#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
914pub struct EnterStagingRegionReply {
915    pub region_id: RegionId,
916    /// Returns true if the region has entered staging with the target directive.
917    pub ready: bool,
918    /// Indicates whether the region exists.
919    pub exists: bool,
920    /// Return error if any during the operation.
921    pub error: Option<String>,
922}
923
924#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
925pub struct EnterStagingRegionsReply {
926    pub replies: Vec<EnterStagingRegionReply>,
927}
928
929impl EnterStagingRegionsReply {
930    pub fn new(replies: Vec<EnterStagingRegionReply>) -> Self {
931        Self { replies }
932    }
933}
934
935/// Reply for a single region sync request.
936#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
937pub struct SyncRegionReply {
938    /// Region id of the synced region.
939    pub region_id: RegionId,
940    /// Returns true if the region is successfully synced and ready.
941    pub ready: bool,
942    /// Indicates whether the region exists.
943    pub exists: bool,
944    /// Return error message if any during the operation.
945    pub error: Option<String>,
946}
947
948/// Reply for a batch of region sync requests.
949#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
950pub struct SyncRegionsReply {
951    pub replies: Vec<SyncRegionReply>,
952}
953
954impl SyncRegionsReply {
955    pub fn new(replies: Vec<SyncRegionReply>) -> Self {
956        Self { replies }
957    }
958}
959
960#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
961pub struct RemapManifestReply {
962    /// Returns false if the region does not exist.
963    pub exists: bool,
964    /// A map from region IDs to their corresponding remapped manifest paths.
965    pub manifest_paths: HashMap<RegionId, String>,
966    /// Return error if any during the operation.
967    pub error: Option<String>,
968}
969
970impl Display for RemapManifestReply {
971    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
972        write!(
973            f,
974            "RemapManifestReply(manifest_paths={:?}, error={:?})",
975            self.manifest_paths, self.error
976        )
977    }
978}
979
980#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
981pub struct ApplyStagingManifestsReply {
982    pub replies: Vec<ApplyStagingManifestReply>,
983}
984
985impl ApplyStagingManifestsReply {
986    pub fn new(replies: Vec<ApplyStagingManifestReply>) -> Self {
987        Self { replies }
988    }
989}
990
991#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
992pub struct ApplyStagingManifestReply {
993    pub region_id: RegionId,
994    /// Returns true if the region is ready to serve reads and writes.
995    pub ready: bool,
996    /// Indicates whether the region exists.
997    pub exists: bool,
998    /// Return error if any during the operation.
999    pub error: Option<String>,
1000}
1001
1002#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1003#[serde(tag = "type", rename_all = "snake_case")]
1004pub enum InstructionReply {
1005    #[serde(alias = "open_region")]
1006    OpenRegions(SimpleReply),
1007    #[serde(alias = "close_region")]
1008    CloseRegions(SimpleReply),
1009    #[serde(
1010        deserialize_with = "upgrade_regions_compat_from",
1011        alias = "upgrade_region"
1012    )]
1013    UpgradeRegions(UpgradeRegionsReply),
1014    #[serde(
1015        alias = "downgrade_region",
1016        deserialize_with = "downgrade_regions_compat_from"
1017    )]
1018    DowngradeRegions(DowngradeRegionsReply),
1019    FlushRegions(FlushRegionReply),
1020    GetFileRefs(GetFileRefsReply),
1021    GcRegions(GcRegionsReply),
1022    EnterStagingRegions(EnterStagingRegionsReply),
1023    SyncRegions(SyncRegionsReply),
1024    RemapManifest(RemapManifestReply),
1025    ApplyStagingManifests(ApplyStagingManifestsReply),
1026}
1027
1028impl Display for InstructionReply {
1029    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1030        match self {
1031            Self::OpenRegions(reply) => write!(f, "InstructionReply::OpenRegions({})", reply),
1032            Self::CloseRegions(reply) => write!(f, "InstructionReply::CloseRegions({})", reply),
1033            Self::UpgradeRegions(reply) => {
1034                write!(f, "InstructionReply::UpgradeRegions({:?})", reply.replies)
1035            }
1036            Self::DowngradeRegions(reply) => {
1037                write!(f, "InstructionReply::DowngradeRegions({:?})", reply.replies)
1038            }
1039            Self::FlushRegions(reply) => write!(f, "InstructionReply::FlushRegions({})", reply),
1040            Self::GetFileRefs(reply) => write!(f, "InstructionReply::GetFileRefs({})", reply),
1041            Self::GcRegions(reply) => write!(f, "InstructionReply::GcRegion({})", reply),
1042            Self::EnterStagingRegions(reply) => {
1043                write!(
1044                    f,
1045                    "InstructionReply::EnterStagingRegions({:?})",
1046                    reply.replies
1047                )
1048            }
1049            Self::SyncRegions(reply) => {
1050                write!(f, "InstructionReply::SyncRegions({:?})", reply.replies)
1051            }
1052            Self::RemapManifest(reply) => write!(f, "InstructionReply::RemapManifest({})", reply),
1053            Self::ApplyStagingManifests(reply) => write!(
1054                f,
1055                "InstructionReply::ApplyStagingManifests({:?})",
1056                reply.replies
1057            ),
1058        }
1059    }
1060}
1061
1062#[cfg(any(test, feature = "testing"))]
1063impl InstructionReply {
1064    pub fn expect_close_regions_reply(self) -> SimpleReply {
1065        match self {
1066            Self::CloseRegions(reply) => reply,
1067            _ => panic!("Expected CloseRegions reply"),
1068        }
1069    }
1070
1071    pub fn expect_open_regions_reply(self) -> SimpleReply {
1072        match self {
1073            Self::OpenRegions(reply) => reply,
1074            _ => panic!("Expected OpenRegions reply"),
1075        }
1076    }
1077
1078    pub fn expect_upgrade_regions_reply(self) -> Vec<UpgradeRegionReply> {
1079        match self {
1080            Self::UpgradeRegions(reply) => reply.replies,
1081            _ => panic!("Expected UpgradeRegion reply"),
1082        }
1083    }
1084
1085    pub fn expect_downgrade_regions_reply(self) -> Vec<DowngradeRegionReply> {
1086        match self {
1087            Self::DowngradeRegions(reply) => reply.replies,
1088            _ => panic!("Expected DowngradeRegion reply"),
1089        }
1090    }
1091
1092    pub fn expect_flush_regions_reply(self) -> FlushRegionReply {
1093        match self {
1094            Self::FlushRegions(reply) => reply,
1095            _ => panic!("Expected FlushRegions reply"),
1096        }
1097    }
1098
1099    pub fn expect_enter_staging_regions_reply(self) -> Vec<EnterStagingRegionReply> {
1100        match self {
1101            Self::EnterStagingRegions(reply) => reply.replies,
1102            _ => panic!("Expected EnterStagingRegion reply"),
1103        }
1104    }
1105
1106    pub fn expect_sync_regions_reply(self) -> Vec<SyncRegionReply> {
1107        match self {
1108            Self::SyncRegions(reply) => reply.replies,
1109            _ => panic!("Expected SyncRegion reply"),
1110        }
1111    }
1112
1113    pub fn expect_remap_manifest_reply(self) -> RemapManifestReply {
1114        match self {
1115            Self::RemapManifest(reply) => reply,
1116            _ => panic!("Expected RemapManifest reply"),
1117        }
1118    }
1119
1120    pub fn expect_apply_staging_manifests_reply(self) -> Vec<ApplyStagingManifestReply> {
1121        match self {
1122            Self::ApplyStagingManifests(reply) => reply.replies,
1123            _ => panic!("Expected ApplyStagingManifest reply"),
1124        }
1125    }
1126}
1127
1128#[cfg(test)]
1129mod tests {
1130    use std::collections::HashSet;
1131
1132    use store_api::storage::{FileId, FileRef};
1133
1134    use super::*;
1135
1136    #[test]
1137    fn test_serialize_instruction() {
1138        let open_region = Instruction::OpenRegions(vec![OpenRegion::new(
1139            RegionIdent {
1140                datanode_id: 2,
1141                table_id: 1024,
1142                region_number: 1,
1143                engine: "mito2".to_string(),
1144            },
1145            "test/foo",
1146            HashMap::new(),
1147            HashMap::new(),
1148            false,
1149            None,
1150            RegionRequirements::empty(),
1151        )]);
1152
1153        let serialized = serde_json::to_string(&open_region).unwrap();
1154        assert_eq!(
1155            r#"{"OpenRegions":[{"region_ident":{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"},"region_storage_path":"test/foo","region_options":{},"region_wal_options":{},"skip_wal_replay":false,"requirements":{"object_storage":false}}]}"#,
1156            serialized
1157        );
1158
1159        let close_region = Instruction::CloseRegions(vec![RegionIdent {
1160            datanode_id: 2,
1161            table_id: 1024,
1162            region_number: 1,
1163            engine: "mito2".to_string(),
1164        }]);
1165
1166        let serialized = serde_json::to_string(&close_region).unwrap();
1167        assert_eq!(
1168            r#"{"CloseRegions":[{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"}]}"#,
1169            serialized
1170        );
1171
1172        let upgrade_region = Instruction::UpgradeRegions(vec![UpgradeRegion {
1173            region_id: RegionId::new(1024, 1),
1174            last_entry_id: None,
1175            metadata_last_entry_id: None,
1176            replay_timeout: Duration::from_millis(1000),
1177            location_id: None,
1178            replay_entry_id: None,
1179            metadata_replay_entry_id: None,
1180        }]);
1181
1182        let serialized = serde_json::to_string(&upgrade_region).unwrap();
1183        assert_eq!(
1184            r#"{"UpgradeRegions":[{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"replay_timeout":"1s","location_id":null}]}"#,
1185            serialized
1186        );
1187    }
1188
1189    #[test]
1190    fn test_serialize_instruction_reply() {
1191        let downgrade_region_reply = InstructionReply::DowngradeRegions(
1192            DowngradeRegionsReply::single(DowngradeRegionReply {
1193                region_id: RegionId::new(1024, 1),
1194                last_entry_id: None,
1195                metadata_last_entry_id: None,
1196                exists: true,
1197                error: None,
1198            }),
1199        );
1200
1201        let serialized = serde_json::to_string(&downgrade_region_reply).unwrap();
1202        assert_eq!(
1203            r#"{"type":"downgrade_regions","replies":[{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"exists":true,"error":null}]}"#,
1204            serialized
1205        );
1206
1207        let upgrade_region_reply =
1208            InstructionReply::UpgradeRegions(UpgradeRegionsReply::single(UpgradeRegionReply {
1209                region_id: RegionId::new(1024, 1),
1210                ready: true,
1211                exists: true,
1212                error: None,
1213            }));
1214        let serialized = serde_json::to_string(&upgrade_region_reply).unwrap();
1215        assert_eq!(
1216            r#"{"type":"upgrade_regions","replies":[{"region_id":4398046511105,"ready":true,"exists":true,"error":null}]}"#,
1217            serialized
1218        );
1219    }
1220
1221    #[test]
1222    fn test_deserialize_instruction() {
1223        // legacy open region instruction
1224        let open_region_instruction = r#"{"OpenRegion":{"region_ident":{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"},"region_storage_path":"test/foo","region_options":{},"region_wal_options":{},"skip_wal_replay":false}}"#;
1225        let open_region_instruction: Instruction =
1226            serde_json::from_str(open_region_instruction).unwrap();
1227        let open_region = Instruction::OpenRegions(vec![OpenRegion::new(
1228            RegionIdent {
1229                datanode_id: 2,
1230                table_id: 1024,
1231                region_number: 1,
1232                engine: "mito2".to_string(),
1233            },
1234            "test/foo",
1235            HashMap::new(),
1236            HashMap::new(),
1237            false,
1238            None,
1239            RegionRequirements::empty(),
1240        )]);
1241        assert_eq!(open_region_instruction, open_region);
1242
1243        // legacy close region instruction
1244        let close_region_instruction = r#"{"CloseRegion":{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"}}"#;
1245        let close_region_instruction: Instruction =
1246            serde_json::from_str(close_region_instruction).unwrap();
1247        let close_region = Instruction::CloseRegions(vec![RegionIdent {
1248            datanode_id: 2,
1249            table_id: 1024,
1250            region_number: 1,
1251            engine: "mito2".to_string(),
1252        }]);
1253        assert_eq!(close_region_instruction, close_region);
1254
1255        // legacy downgrade region instruction
1256        let downgrade_region_instruction = r#"{"DowngradeRegions":{"region_id":4398046511105,"flush_timeout":{"secs":1,"nanos":0}}}"#;
1257        let downgrade_region_instruction: Instruction =
1258            serde_json::from_str(downgrade_region_instruction).unwrap();
1259        let downgrade_region = Instruction::DowngradeRegions(vec![DowngradeRegion {
1260            region_id: RegionId::new(1024, 1),
1261            flush_timeout: Some(Duration::from_millis(1000)),
1262        }]);
1263        assert_eq!(downgrade_region_instruction, downgrade_region);
1264
1265        // legacy upgrade region instruction
1266        let upgrade_region_instruction = r#"{"UpgradeRegion":{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"replay_timeout":"1s","location_id":null,"replay_entry_id":null,"metadata_replay_entry_id":null}}"#;
1267        let upgrade_region_instruction: Instruction =
1268            serde_json::from_str(upgrade_region_instruction).unwrap();
1269        let upgrade_region = Instruction::UpgradeRegions(vec![UpgradeRegion {
1270            region_id: RegionId::new(1024, 1),
1271            last_entry_id: None,
1272            metadata_last_entry_id: None,
1273            replay_timeout: Duration::from_millis(1000),
1274            location_id: None,
1275            replay_entry_id: None,
1276            metadata_replay_entry_id: None,
1277        }]);
1278        assert_eq!(upgrade_region_instruction, upgrade_region);
1279    }
1280
1281    #[test]
1282    fn test_deserialize_instruction_reply() {
1283        // legacy close region reply
1284        let close_region_instruction_reply =
1285            r#"{"result":true,"error":null,"type":"close_region"}"#;
1286        let close_region_instruction_reply: InstructionReply =
1287            serde_json::from_str(close_region_instruction_reply).unwrap();
1288        let close_region_reply = InstructionReply::CloseRegions(SimpleReply {
1289            result: true,
1290            error: None,
1291        });
1292        assert_eq!(close_region_instruction_reply, close_region_reply);
1293
1294        // legacy open region reply
1295        let open_region_instruction_reply = r#"{"result":true,"error":null,"type":"open_region"}"#;
1296        let open_region_instruction_reply: InstructionReply =
1297            serde_json::from_str(open_region_instruction_reply).unwrap();
1298        let open_region_reply = InstructionReply::OpenRegions(SimpleReply {
1299            result: true,
1300            error: None,
1301        });
1302        assert_eq!(open_region_instruction_reply, open_region_reply);
1303
1304        // legacy downgrade region reply
1305        let downgrade_region_instruction_reply = r#"{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"exists":true,"error":null,"type":"downgrade_region"}"#;
1306        let downgrade_region_instruction_reply: InstructionReply =
1307            serde_json::from_str(downgrade_region_instruction_reply).unwrap();
1308        let downgrade_region_reply = InstructionReply::DowngradeRegions(
1309            DowngradeRegionsReply::single(DowngradeRegionReply {
1310                region_id: RegionId::new(1024, 1),
1311                last_entry_id: None,
1312                metadata_last_entry_id: None,
1313                exists: true,
1314                error: None,
1315            }),
1316        );
1317        assert_eq!(downgrade_region_instruction_reply, downgrade_region_reply);
1318
1319        // legacy upgrade region reply
1320        let upgrade_region_instruction_reply = r#"{"region_id":4398046511105,"ready":true,"exists":true,"error":null,"type":"upgrade_region"}"#;
1321        let upgrade_region_instruction_reply: InstructionReply =
1322            serde_json::from_str(upgrade_region_instruction_reply).unwrap();
1323        let upgrade_region_reply =
1324            InstructionReply::UpgradeRegions(UpgradeRegionsReply::single(UpgradeRegionReply {
1325                region_id: RegionId::new(1024, 1),
1326                ready: true,
1327                exists: true,
1328                error: None,
1329            }));
1330        assert_eq!(upgrade_region_instruction_reply, upgrade_region_reply);
1331    }
1332
1333    #[test]
1334    fn test_enter_staging_partition_rule_compatibility() {
1335        let legacy = r#"{"region_id":4398046511105,"partition_expr":"{\"Expr\":{\"lhs\":{\"Column\":\"x\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"Int32\":0}}}}"}"#;
1336        let enter: EnterStagingRegion = serde_json::from_str(legacy).unwrap();
1337        assert_eq!(enter.region_id, RegionId::new(1024, 1));
1338        assert_eq!(
1339            enter.partition_directive,
1340            StagingPartitionDirective::UpdatePartitionExpr(
1341                "{\"Expr\":{\"lhs\":{\"Column\":\"x\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"Int32\":0}}}}"
1342                    .to_string()
1343            )
1344        );
1345
1346        let serialized = serde_json::to_string(&enter).unwrap();
1347        assert!(serialized.contains("\"partition_directive\":\""));
1348        assert!(!serialized.contains("partition_expr"));
1349
1350        let reject = r#"{"region_id":4398046511105,"partition_expr":{"type":"reject_all_writes"}}"#;
1351        let enter: EnterStagingRegion = serde_json::from_str(reject).unwrap();
1352        assert_eq!(
1353            enter.partition_directive,
1354            StagingPartitionDirective::RejectAllWrites
1355        );
1356    }
1357
1358    #[derive(Debug, Clone, Serialize, Deserialize)]
1359    struct LegacyOpenRegion {
1360        region_ident: RegionIdent,
1361        region_storage_path: String,
1362        region_options: HashMap<String, String>,
1363    }
1364
1365    #[test]
1366    fn test_compatible_serialize_open_region() {
1367        let region_ident = RegionIdent {
1368            datanode_id: 2,
1369            table_id: 1024,
1370            region_number: 1,
1371            engine: "mito2".to_string(),
1372        };
1373        let region_storage_path = "test/foo".to_string();
1374        let region_options = HashMap::from([
1375            ("a".to_string(), "aa".to_string()),
1376            ("b".to_string(), "bb".to_string()),
1377        ]);
1378
1379        // Serialize a legacy OpenRegion.
1380        let legacy_open_region = LegacyOpenRegion {
1381            region_ident: region_ident.clone(),
1382            region_storage_path: region_storage_path.clone(),
1383            region_options: region_options.clone(),
1384        };
1385        let serialized = serde_json::to_string(&legacy_open_region).unwrap();
1386
1387        // Deserialize to OpenRegion.
1388        let deserialized = serde_json::from_str(&serialized).unwrap();
1389        let expected = OpenRegion {
1390            region_ident,
1391            region_storage_path,
1392            region_options,
1393            region_wal_options: HashMap::new(),
1394            skip_wal_replay: false,
1395            reason: None,
1396            requirements: RegionRequirements::empty(),
1397        };
1398        assert_eq!(expected, deserialized);
1399    }
1400
1401    #[test]
1402    fn test_serialize_open_region_with_reason_and_requirements() {
1403        let open_region = OpenRegion::new(
1404            RegionIdent {
1405                datanode_id: 2,
1406                table_id: 1024,
1407                region_number: 1,
1408                engine: "mito2".to_string(),
1409            },
1410            "test/foo",
1411            HashMap::new(),
1412            HashMap::new(),
1413            false,
1414            Some(OpenRegionReason::RegionMigration),
1415            RegionRequirements::object_storage(),
1416        );
1417
1418        let serialized = serde_json::to_string(&open_region).unwrap();
1419        assert!(serialized.contains(r#""reason":"RegionMigration""#));
1420        assert!(serialized.contains(r#""object_storage":true"#));
1421
1422        let deserialized: OpenRegion = serde_json::from_str(&serialized).unwrap();
1423        assert_eq!(Some(OpenRegionReason::RegionMigration), deserialized.reason);
1424        assert_eq!(
1425            RegionRequirements::object_storage(),
1426            deserialized.requirements
1427        );
1428    }
1429
1430    #[test]
1431    fn test_flush_regions_creation() {
1432        let region_id = RegionId::new(1024, 1);
1433
1434        // Single region sync flush
1435        let single_sync = FlushRegions::sync_single(region_id);
1436        assert_eq!(single_sync.region_ids, vec![region_id]);
1437        assert_eq!(single_sync.strategy, FlushStrategy::Sync);
1438        assert!(!single_sync.is_hint());
1439        assert!(single_sync.is_sync());
1440        assert_eq!(single_sync.error_strategy, FlushErrorStrategy::FailFast);
1441        assert_eq!(single_sync.reason, None);
1442        assert!(single_sync.is_single_region());
1443        assert_eq!(single_sync.single_region_id(), Some(region_id));
1444
1445        // Batch async flush (hint)
1446        let region_ids = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
1447        let batch_async = FlushRegions::async_batch(region_ids.clone());
1448        assert_eq!(batch_async.region_ids, region_ids);
1449        assert_eq!(batch_async.strategy, FlushStrategy::Async);
1450        assert!(batch_async.is_hint());
1451        assert!(!batch_async.is_sync());
1452        assert_eq!(batch_async.error_strategy, FlushErrorStrategy::TryAll);
1453        assert_eq!(batch_async.reason, None);
1454        assert!(!batch_async.is_single_region());
1455        assert_eq!(batch_async.single_region_id(), None);
1456
1457        // Batch sync flush
1458        let batch_sync = FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::FailFast);
1459        assert_eq!(batch_sync.region_ids, region_ids);
1460        assert_eq!(batch_sync.strategy, FlushStrategy::Sync);
1461        assert!(!batch_sync.is_hint());
1462        assert!(batch_sync.is_sync());
1463        assert_eq!(batch_sync.error_strategy, FlushErrorStrategy::FailFast);
1464        assert_eq!(batch_sync.reason, None);
1465
1466        let with_reason = batch_sync.with_reason(RegionFlushReason::RemoteWalPrune);
1467        assert_eq!(with_reason.reason, Some(RegionFlushReason::RemoteWalPrune));
1468    }
1469
1470    #[test]
1471    fn test_flush_regions_conversion() {
1472        let region_id = RegionId::new(1024, 1);
1473
1474        let from_region_id: FlushRegions = region_id.into();
1475        assert_eq!(from_region_id.region_ids, vec![region_id]);
1476        assert_eq!(from_region_id.strategy, FlushStrategy::Sync);
1477        assert!(!from_region_id.is_hint());
1478        assert!(from_region_id.is_sync());
1479
1480        // Test default construction
1481        let flush_regions = FlushRegions {
1482            region_ids: vec![region_id],
1483            strategy: FlushStrategy::Async,
1484            error_strategy: FlushErrorStrategy::TryAll,
1485            reason: None,
1486        };
1487        assert_eq!(flush_regions.region_ids, vec![region_id]);
1488        assert_eq!(flush_regions.strategy, FlushStrategy::Async);
1489        assert!(flush_regions.is_hint());
1490        assert!(!flush_regions.is_sync());
1491    }
1492
1493    #[test]
1494    fn test_flush_region_reply() {
1495        let region_id = RegionId::new(1024, 1);
1496
1497        // Successful single region reply
1498        let success_reply = FlushRegionReply::success_single(region_id);
1499        assert!(success_reply.overall_success);
1500        assert_eq!(success_reply.results.len(), 1);
1501        assert_eq!(success_reply.results[0].0, region_id);
1502        assert!(success_reply.results[0].1.is_ok());
1503
1504        // Failed single region reply
1505        let error_reply = FlushRegionReply::error_single(region_id, "test error".to_string());
1506        assert!(!error_reply.overall_success);
1507        assert_eq!(error_reply.results.len(), 1);
1508        assert_eq!(error_reply.results[0].0, region_id);
1509        assert!(error_reply.results[0].1.is_err());
1510
1511        // Batch reply
1512        let region_id2 = RegionId::new(1024, 2);
1513        let results = vec![
1514            (region_id, Ok(())),
1515            (region_id2, Err("flush failed".to_string())),
1516        ];
1517        let batch_reply = FlushRegionReply::from_results(results);
1518        assert!(!batch_reply.overall_success);
1519        assert_eq!(batch_reply.results.len(), 2);
1520
1521        // Conversion to SimpleReply
1522        let simple_reply = batch_reply.to_simple_reply();
1523        assert!(!simple_reply.result);
1524        assert!(simple_reply.error.is_some());
1525        assert!(simple_reply.error.unwrap().contains("flush failed"));
1526    }
1527
1528    #[test]
1529    fn test_serialize_flush_regions_instruction() {
1530        let region_id = RegionId::new(1024, 1);
1531        let flush_regions = FlushRegions::sync_single(region_id);
1532        let instruction = Instruction::FlushRegions(flush_regions.clone());
1533
1534        let serialized = serde_json::to_string(&instruction).unwrap();
1535        assert!(!serialized.contains("reason"));
1536        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1537
1538        match deserialized {
1539            Instruction::FlushRegions(fr) => {
1540                assert_eq!(fr.region_ids, vec![region_id]);
1541                assert_eq!(fr.strategy, FlushStrategy::Sync);
1542                assert_eq!(fr.error_strategy, FlushErrorStrategy::FailFast);
1543                assert_eq!(fr.reason, None);
1544            }
1545            _ => panic!("Expected FlushRegions instruction"),
1546        }
1547
1548        let legacy = r#"{"FlushRegions":{"region_ids":[4398046511105],"strategy":"Sync","error_strategy":"FailFast"}}"#;
1549        let deserialized: Instruction = serde_json::from_str(legacy).unwrap();
1550        match deserialized {
1551            Instruction::FlushRegions(fr) => {
1552                assert_eq!(fr.region_ids, vec![region_id]);
1553                assert_eq!(fr.strategy, FlushStrategy::Sync);
1554                assert_eq!(fr.error_strategy, FlushErrorStrategy::FailFast);
1555                assert_eq!(fr.reason, None);
1556            }
1557            _ => panic!("Expected FlushRegions instruction"),
1558        }
1559
1560        let flush_regions = FlushRegions::async_batch(vec![region_id])
1561            .with_reason(RegionFlushReason::RemoteWalPrune);
1562        let instruction = Instruction::FlushRegions(flush_regions);
1563        let serialized = serde_json::to_string(&instruction).unwrap();
1564        assert!(serialized.contains(r#""reason":"RemoteWalPrune""#));
1565        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1566        match deserialized {
1567            Instruction::FlushRegions(fr) => {
1568                assert_eq!(fr.reason, Some(RegionFlushReason::RemoteWalPrune));
1569            }
1570            _ => panic!("Expected FlushRegions instruction"),
1571        }
1572    }
1573
1574    #[test]
1575    fn test_serialize_flush_regions_batch_instruction() {
1576        let region_ids = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
1577        let flush_regions =
1578            FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::TryAll);
1579        let instruction = Instruction::FlushRegions(flush_regions);
1580
1581        let serialized = serde_json::to_string(&instruction).unwrap();
1582        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1583
1584        match deserialized {
1585            Instruction::FlushRegions(fr) => {
1586                assert_eq!(fr.region_ids, region_ids);
1587                assert_eq!(fr.strategy, FlushStrategy::Sync);
1588                assert!(!fr.is_hint());
1589                assert!(fr.is_sync());
1590                assert_eq!(fr.error_strategy, FlushErrorStrategy::TryAll);
1591                assert_eq!(fr.reason, None);
1592            }
1593            _ => panic!("Expected FlushRegions instruction"),
1594        }
1595    }
1596
1597    #[test]
1598    fn test_serialize_get_file_refs_instruction_reply() {
1599        let mut manifest = FileRefsManifest::default();
1600        let r0 = RegionId::new(1024, 1);
1601        let r1 = RegionId::new(1024, 2);
1602        manifest.file_refs.insert(
1603            r0,
1604            HashSet::from([FileRef::new(r0, FileId::random(), None)]),
1605        );
1606        manifest.file_refs.insert(
1607            r1,
1608            HashSet::from([FileRef::new(r1, FileId::random(), None)]),
1609        );
1610        manifest.manifest_version.insert(r0, 10);
1611        manifest.manifest_version.insert(r1, 20);
1612
1613        let instruction_reply = InstructionReply::GetFileRefs(GetFileRefsReply {
1614            file_refs_manifest: manifest,
1615            success: true,
1616            error: None,
1617        });
1618
1619        let serialized = serde_json::to_string(&instruction_reply).unwrap();
1620        let deserialized = serde_json::from_str(&serialized).unwrap();
1621
1622        assert_eq!(instruction_reply, deserialized);
1623    }
1624}
common_meta/instruction.rs

common_meta/
instruction.rs
