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

common_meta/
instruction.rs
