Skip to main content

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 base64::Engine;
20use common_error::ext::{ErrorExt, RetryHint};
21use common_error::status_code::StatusCode;
22use serde::{Deserialize, Deserializer, Serialize, Serializer};
23use snafu::ResultExt as _;
24use store_api::region_engine::SyncRegionFromRequest;
25use store_api::region_request::{RegionFlushReason, RegionRequirements};
26use store_api::storage::{FileId, FileRef, FileRefsManifest, GcReport, RegionId, RegionNumber};
27use strum::Display;
28use table::metadata::TableId;
29use table::table_name::TableName;
30
31use crate::error::{DecodePackedFileRefsSnafu, InvalidPackedFileRefsSnafu};
32use crate::flow_name::FlowName;
33use crate::key::schema_name::SchemaName;
34use crate::key::{FlowId, FlowPartitionId};
35use crate::peer::Peer;
36use crate::wal_provider::{RegionWalOptions, region_wal_options_serde};
37use crate::{DatanodeId, FlownodeId};
38
39/// A structured error returned by instruction replies.
40#[derive(Debug, Serialize, PartialEq, Eq, Clone)]
41pub struct InstructionError {
42    /// Numeric status code aligned with [`StatusCode`].
43    #[serde(
44        serialize_with = "StatusCode::serialize_as_u32",
45        deserialize_with = "StatusCode::deserialize_from_u32"
46    )]
47    pub code: StatusCode,
48    /// User-facing error message aligned with [`ErrorExt::output_msg`].
49    pub message: String,
50    /// Retry hint serialized as a string.
51    #[serde(
52        serialize_with = "RetryHint::serialize_as_str",
53        deserialize_with = "RetryHint::deserialize_from_str"
54    )]
55    pub retry_hint: RetryHint,
56}
57
58impl<'de> Deserialize<'de> for InstructionError {
59    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
60    where
61        D: Deserializer<'de>,
62    {
63        #[derive(Deserialize)]
64        #[serde(untagged)]
65        enum Compat {
66            Structured {
67                #[serde(deserialize_with = "StatusCode::deserialize_from_u32")]
68                code: StatusCode,
69                message: String,
70                #[serde(deserialize_with = "RetryHint::deserialize_from_str")]
71                retry_hint: RetryHint,
72            },
73            Legacy(String),
74        }
75
76        match Compat::deserialize(deserializer)? {
77            Compat::Structured {
78                code,
79                message,
80                retry_hint,
81            } => Ok(Self {
82                code,
83                message,
84                retry_hint,
85            }),
86            Compat::Legacy(message) => Ok(Self::legacy_internal_retryable(message)),
87        }
88    }
89}
90
91impl InstructionError {
92    pub fn new(code: StatusCode, message: impl Into<String>, retry_hint: RetryHint) -> Self {
93        Self {
94            code,
95            message: message.into(),
96            retry_hint,
97        }
98    }
99
100    pub fn legacy_internal_retryable(message: impl Into<String>) -> Self {
101        Self::new(StatusCode::Internal, message, RetryHint::Retryable)
102    }
103
104    pub fn from_error<E: ErrorExt>(error: &E) -> Self {
105        Self {
106            code: error.status_code(),
107            message: error.output_msg(),
108            retry_hint: error.retry_hint(),
109        }
110    }
111}
112
113impl Display for InstructionError {
114    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
115        write!(
116            f,
117            "InstructionError(code={}, retry_hint={}, message={})",
118            self.code as u32,
119            self.retry_hint.as_str(),
120            self.message
121        )
122    }
123}
124
125pub type InstructionResult<T> = std::result::Result<T, InstructionError>;
126
127#[derive(Eq, Hash, PartialEq, Clone, Debug, Serialize, Deserialize)]
128pub struct RegionIdent {
129    pub datanode_id: DatanodeId,
130    pub table_id: TableId,
131    pub region_number: RegionNumber,
132    pub engine: String,
133}
134
135impl RegionIdent {
136    pub fn get_region_id(&self) -> RegionId {
137        RegionId::new(self.table_id, self.region_number)
138    }
139}
140
141impl Display for RegionIdent {
142    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
143        write!(
144            f,
145            "RegionIdent(datanode_id='{}', table_id={}, region_number={}, engine = {})",
146            self.datanode_id, self.table_id, self.region_number, self.engine
147        )
148    }
149}
150
151/// The result of downgrade leader region.
152#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
153pub struct DowngradeRegionReply {
154    /// The [RegionId].
155    /// For compatibility, it is defaulted to [RegionId::new(0, 0)].
156    #[serde(default)]
157    pub region_id: RegionId,
158    /// Returns the `last_entry_id` if available.
159    pub last_entry_id: Option<u64>,
160    /// Returns the `metadata_last_entry_id` if available (Only available for metric engine).
161    pub metadata_last_entry_id: Option<u64>,
162    /// Indicates whether the region exists.
163    pub exists: bool,
164    /// Return error if any during the operation.
165    pub error: Option<InstructionError>,
166}
167
168impl Display for DowngradeRegionReply {
169    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
170        write!(
171            f,
172            "(last_entry_id={:?}, exists={}, error={:?})",
173            self.last_entry_id, self.exists, self.error
174        )
175    }
176}
177
178#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
179pub struct SimpleReply {
180    pub result: bool,
181    pub error: Option<InstructionError>,
182}
183
184/// Reply for flush region operations with support for batch results.
185#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
186pub struct FlushRegionReply {
187    /// Results for each region that was attempted to be flushed.
188    /// For single region flushes, this will contain one result.
189    /// For batch flushes, this contains results for all attempted regions.
190    pub results: Vec<(RegionId, InstructionResult<()>)>,
191    /// Overall success: true if all regions were flushed successfully.
192    pub overall_success: bool,
193}
194
195impl FlushRegionReply {
196    /// Create a successful single region reply.
197    pub fn success_single(region_id: RegionId) -> Self {
198        Self {
199            results: vec![(region_id, Ok(()))],
200            overall_success: true,
201        }
202    }
203
204    /// Create a failed single region reply.
205    pub fn error_single(region_id: RegionId, error: InstructionError) -> Self {
206        Self {
207            results: vec![(region_id, Err(error))],
208            overall_success: false,
209        }
210    }
211
212    /// Create a batch reply from individual results.
213    pub fn from_results(results: Vec<(RegionId, InstructionResult<()>)>) -> Self {
214        let overall_success = results.iter().all(|(_, result)| result.is_ok());
215        Self {
216            results,
217            overall_success,
218        }
219    }
220
221    /// Convert to SimpleReply for backward compatibility.
222    pub fn to_simple_reply(&self) -> SimpleReply {
223        if self.overall_success {
224            SimpleReply {
225                result: true,
226                error: None,
227            }
228        } else {
229            let errors: Vec<String> = self
230                .results
231                .iter()
232                .filter_map(|(region_id, result)| {
233                    result
234                        .as_ref()
235                        .err()
236                        .map(|err| format!("{}: {}", region_id, err))
237                })
238                .collect();
239            SimpleReply {
240                result: false,
241                error: Some(InstructionError::legacy_internal_retryable(
242                    errors.join("; "),
243                )),
244            }
245        }
246    }
247}
248
249impl Display for SimpleReply {
250    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
251        write!(f, "(result={}, error={:?})", self.result, self.error)
252    }
253}
254
255impl Display for FlushRegionReply {
256    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
257        let results_str = self
258            .results
259            .iter()
260            .map(|(region_id, result)| match result {
261                Ok(()) => format!("{}:OK", region_id),
262                Err(err) => format!("{}:ERR({})", region_id, err),
263            })
264            .collect::<Vec<_>>()
265            .join(", ");
266        write!(
267            f,
268            "(overall_success={}, results=[{}])",
269            self.overall_success, results_str
270        )
271    }
272}
273
274impl Display for OpenRegion {
275    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
276        write!(
277            f,
278            "OpenRegion(region_ident={}, region_storage_path={}, reason={:?})",
279            self.region_ident, self.region_storage_path, self.reason
280        )
281    }
282}
283
284/// The reason why an open region instruction is triggered.
285#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
286pub enum OpenRegionReason {
287    /// Open triggered before region migration.
288    RegionMigration,
289    /// Open triggered by region failover.
290    RegionFailover,
291    /// Open triggered when adding a follower region.
292    #[cfg(feature = "enterprise")]
293    RegionFollower,
294}
295
296#[serde_with::serde_as]
297#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
298pub struct OpenRegion {
299    pub region_ident: RegionIdent,
300    pub region_storage_path: String,
301    pub region_options: HashMap<String, String>,
302    #[serde(default)]
303    #[serde(with = "region_wal_options_serde")]
304    pub region_wal_options: RegionWalOptions,
305    #[serde(default)]
306    pub skip_wal_replay: bool,
307    #[serde(default, skip_serializing_if = "Option::is_none")]
308    pub reason: Option<OpenRegionReason>,
309    #[serde(default)]
310    pub requirements: RegionRequirements,
311}
312
313impl OpenRegion {
314    pub fn new(
315        region_ident: RegionIdent,
316        path: &str,
317        region_options: HashMap<String, String>,
318        region_wal_options: RegionWalOptions,
319        skip_wal_replay: bool,
320        reason: Option<OpenRegionReason>,
321        requirements: RegionRequirements,
322    ) -> Self {
323        Self {
324            region_ident,
325            region_storage_path: path.to_string(),
326            region_options,
327            region_wal_options,
328            skip_wal_replay,
329            reason,
330            requirements,
331        }
332    }
333}
334
335/// The instruction of downgrading leader region.
336#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
337pub struct DowngradeRegion {
338    /// The [RegionId].
339    pub region_id: RegionId,
340    /// The timeout of waiting for flush the region.
341    ///
342    /// `None` stands for don't flush before downgrading the region.
343    #[serde(default)]
344    pub flush_timeout: Option<Duration>,
345}
346
347impl Display for DowngradeRegion {
348    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
349        write!(
350            f,
351            "DowngradeRegion(region_id={}, flush_timeout={:?})",
352            self.region_id, self.flush_timeout,
353        )
354    }
355}
356
357/// Upgrades a follower region to leader region.
358#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
359pub struct UpgradeRegion {
360    /// The [RegionId].
361    pub region_id: RegionId,
362    /// The `last_entry_id` of old leader region.
363    pub last_entry_id: Option<u64>,
364    /// The `last_entry_id` of old leader metadata region (Only used for metric engine).
365    pub metadata_last_entry_id: Option<u64>,
366    /// The timeout of waiting for a wal replay.
367    ///
368    /// `None` stands for no wait,
369    /// it's helpful to verify whether the leader region is ready.
370    #[serde(with = "humantime_serde")]
371    pub replay_timeout: Duration,
372    /// The hint for replaying memtable.
373    #[serde(default)]
374    pub location_id: Option<u64>,
375    #[serde(default, skip_serializing_if = "Option::is_none")]
376    pub replay_entry_id: Option<u64>,
377    #[serde(default, skip_serializing_if = "Option::is_none")]
378    pub metadata_replay_entry_id: Option<u64>,
379}
380
381impl UpgradeRegion {
382    /// Sets the replay entry id.
383    pub fn with_replay_entry_id(mut self, replay_entry_id: Option<u64>) -> Self {
384        self.replay_entry_id = replay_entry_id;
385        self
386    }
387
388    /// Sets the metadata replay entry id.
389    pub fn with_metadata_replay_entry_id(mut self, metadata_replay_entry_id: Option<u64>) -> Self {
390        self.metadata_replay_entry_id = metadata_replay_entry_id;
391        self
392    }
393}
394
395#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
396/// The identifier of cache.
397pub enum CacheIdent {
398    FlowId(FlowId),
399    /// Indicate change of address of flownode.
400    FlowNodeAddressChange(u64),
401    FlowName(FlowName),
402    TableId(TableId),
403    TableName(TableName),
404    SchemaName(SchemaName),
405    CreateFlow(CreateFlow),
406    DropFlow(DropFlow),
407    /// Indicate change of user metadata.
408    User(UserCacheIdent),
409}
410
411#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
412pub struct UserCacheIdent {
413    pub catalog: String,
414    pub username: String,
415}
416
417#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
418pub struct CreateFlow {
419    /// The unique identifier for the flow.
420    pub flow_id: FlowId,
421    pub source_table_ids: Vec<TableId>,
422    /// Mapping of flow partition to peer information
423    pub partition_to_peer_mapping: Vec<(FlowPartitionId, Peer)>,
424}
425
426#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
427pub struct DropFlow {
428    pub flow_id: FlowId,
429    pub source_table_ids: Vec<TableId>,
430    /// Mapping of flow partition to flownode id
431    pub flow_part2node_id: Vec<(FlowPartitionId, FlownodeId)>,
432}
433
434/// Strategy for executing flush operations.
435#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
436pub enum FlushStrategy {
437    /// Synchronous operation that waits for completion and expects a reply
438    #[default]
439    Sync,
440    /// Asynchronous hint operation (fire-and-forget, no reply expected)
441    Async,
442}
443
444/// Error handling strategy for batch flush operations.
445#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
446pub enum FlushErrorStrategy {
447    /// Abort on first error (fail-fast)
448    #[default]
449    FailFast,
450    /// Attempt to flush all regions and collect all errors
451    TryAll,
452}
453
454/// Unified flush instruction supporting both single and batch operations
455/// with configurable execution strategies and error handling.
456#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
457pub struct FlushRegions {
458    /// List of region IDs to flush. Can contain a single region or multiple regions.
459    pub region_ids: Vec<RegionId>,
460    /// Execution strategy: Sync (expects reply) or Async (fire-and-forget hint).
461    #[serde(default)]
462    pub strategy: FlushStrategy,
463    /// Error handling strategy for batch operations (only applies when multiple regions and sync strategy).
464    #[serde(default)]
465    pub error_strategy: FlushErrorStrategy,
466    /// The source that triggered this flush.
467    #[serde(default, skip_serializing_if = "Option::is_none")]
468    pub reason: Option<RegionFlushReason>,
469}
470
471impl Display for FlushRegions {
472    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
473        write!(
474            f,
475            "FlushRegions(region_ids={:?}, strategy={:?}, error_strategy={:?}, reason={:?})",
476            self.region_ids, self.strategy, self.error_strategy, self.reason
477        )
478    }
479}
480
481impl FlushRegions {
482    /// Create synchronous single-region flush
483    pub fn sync_single(region_id: RegionId) -> Self {
484        Self {
485            region_ids: vec![region_id],
486            strategy: FlushStrategy::Sync,
487            error_strategy: FlushErrorStrategy::FailFast,
488            reason: None,
489        }
490    }
491
492    /// Create asynchronous batch flush (fire-and-forget)
493    pub fn async_batch(region_ids: Vec<RegionId>) -> Self {
494        Self {
495            region_ids,
496            strategy: FlushStrategy::Async,
497            error_strategy: FlushErrorStrategy::TryAll,
498            reason: None,
499        }
500    }
501
502    /// Create synchronous batch flush with error strategy
503    pub fn sync_batch(region_ids: Vec<RegionId>, error_strategy: FlushErrorStrategy) -> Self {
504        Self {
505            region_ids,
506            strategy: FlushStrategy::Sync,
507            error_strategy,
508            reason: None,
509        }
510    }
511
512    pub fn with_reason(mut self, reason: RegionFlushReason) -> Self {
513        self.reason = Some(reason);
514        self
515    }
516
517    /// Check if this is a single region flush.
518    pub fn is_single_region(&self) -> bool {
519        self.region_ids.len() == 1
520    }
521
522    /// Get the single region ID if this is a single region flush.
523    pub fn single_region_id(&self) -> Option<RegionId> {
524        if self.is_single_region() {
525            self.region_ids.first().copied()
526        } else {
527            None
528        }
529    }
530
531    /// Check if this is a hint (asynchronous) operation.
532    pub fn is_hint(&self) -> bool {
533        matches!(self.strategy, FlushStrategy::Async)
534    }
535
536    /// Check if this is a synchronous operation.
537    pub fn is_sync(&self) -> bool {
538        matches!(self.strategy, FlushStrategy::Sync)
539    }
540}
541
542impl From<RegionId> for FlushRegions {
543    fn from(region_id: RegionId) -> Self {
544        Self::sync_single(region_id)
545    }
546}
547
548#[derive(Debug, Deserialize)]
549#[serde(untagged)]
550enum SingleOrMultiple<T> {
551    Single(T),
552    Multiple(Vec<T>),
553}
554
555fn single_or_multiple_from<'de, D, T>(deserializer: D) -> Result<Vec<T>, D::Error>
556where
557    D: Deserializer<'de>,
558    T: Deserialize<'de>,
559{
560    let helper = SingleOrMultiple::<T>::deserialize(deserializer)?;
561    Ok(match helper {
562        SingleOrMultiple::Single(x) => vec![x],
563        SingleOrMultiple::Multiple(xs) => xs,
564    })
565}
566
567/// Instruction to get file references for specified regions.
568#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
569pub struct GetFileRefs {
570    /// List of region IDs to get file references from active FileHandles (in-memory).
571    pub query_regions: Vec<RegionId>,
572    /// Mapping from the src region IDs (whose file references to look for) to
573    /// the dst region IDs (where to read the manifests).
574    /// Key: The source region IDs (where files originally came from).
575    /// Value: The set of destination region IDs (whose manifests need to be read).
576    pub related_regions: HashMap<RegionId, HashSet<RegionId>>,
577}
578
579impl Display for GetFileRefs {
580    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
581        write!(f, "GetFileRefs(region_ids={:?})", self.query_regions)
582    }
583}
584
585/// Instruction to get file references using the packed manifest format.
586#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
587pub struct GetPackedFileRefs {
588    /// List of region IDs to get file references from active FileHandles (in-memory).
589    pub query_regions: Vec<RegionId>,
590    /// Mapping from the src region IDs (whose file references to look for) to
591    /// the dst region IDs (where to read the manifests).
592    /// Key: The source region IDs (where files originally came from).
593    /// Value: The set of destination region IDs (whose manifests need to be read).
594    pub related_regions: HashMap<RegionId, HashSet<RegionId>>,
595}
596
597impl Display for GetPackedFileRefs {
598    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
599        write!(f, "GetPackedFileRefs(region_ids={:?})", self.query_regions)
600    }
601}
602
603/// Instruction to trigger garbage collection for a region.
604#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
605pub struct GcRegions {
606    /// 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.
607    pub regions: Vec<RegionId>,
608    /// The file references manifest containing temporary file references.
609    pub file_refs_manifest: FileRefsManifest,
610    /// Whether to perform a full file listing to find orphan files.
611    pub full_file_listing: bool,
612}
613
614/// Instruction to trigger garbage collection with a packed file-reference manifest.
615#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
616pub struct PackedGcRegions {
617    /// The region ID to perform GC on.
618    pub regions: Vec<RegionId>,
619    /// Packed manifest used by the packed-only GC protocol.
620    pub packed_file_refs_manifest: PackedFileRefsManifest,
621    /// Whether to perform a full file listing to find orphan files.
622    pub full_file_listing: bool,
623}
624
625impl Display for PackedGcRegions {
626    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
627        let file_refs_count = self.packed_file_refs_manifest.file_refs.len();
628        write!(
629            f,
630            "PackedGcRegions(regions={:?}, file_refs_count={}, full_file_listing={})",
631            self.regions, file_refs_count, self.full_file_listing
632        )
633    }
634}
635
636/// A packed, JSON-compatible encoding of a file reference manifest.
637#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
638pub struct PackedFileRefsManifest {
639    /// Required packed file references for each region.
640    pub file_refs: HashMap<RegionId, PackedRegionFileRefs>,
641    #[serde(default)]
642    pub manifest_version: HashMap<RegionId, u64>,
643    #[serde(default)]
644    pub cross_region_refs: HashMap<RegionId, HashSet<RegionId>>,
645}
646
647/// Packed file references for one region, encoded with STANDARD Base64.
648///
649/// `files` is required and contains 16-byte UUID records for references whose
650/// `index_version` is `None`. `indexed` is required and contains 24-byte records
651/// for references whose `index_version` is `Some`: a 16-byte UUID followed by
652/// an 8-byte big-endian `u64` version. Decoding fails when either field is not
653/// valid STANDARD Base64 or its decoded bytes do not align to the corresponding
654/// record width.
655#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
656pub struct PackedRegionFileRefs {
657    /// Required STANDARD Base64-encoded 16-byte UUID records.
658    pub files: String,
659    /// Required STANDARD Base64-encoded 16-byte UUID and 8-byte version records.
660    pub indexed: String,
661}
662
663impl PackedFileRefsManifest {
664    /// Packs a file reference manifest using required [`PackedRegionFileRefs`]
665    /// STANDARD Base64 fields.
666    ///
667    /// References whose `index_version` is `None` are encoded as 16-byte UUID
668    /// `files` records; references whose `index_version` is `Some` are encoded
669    /// as 24-byte `indexed` records: a 16-byte UUID followed by an 8-byte
670    /// big-endian `u64` version.
671    pub fn from_manifest(manifest: &FileRefsManifest) -> Self {
672        Self {
673            file_refs: manifest
674                .file_refs
675                .iter()
676                .map(|(region, refs)| (*region, PackedRegionFileRefs::from_refs(refs)))
677                .collect(),
678            manifest_version: manifest.manifest_version.clone(),
679            cross_region_refs: manifest.cross_region_refs.clone(),
680        }
681    }
682
683    /// Decodes this packed representation into a file reference manifest.
684    ///
685    /// Returns a typed error for invalid STANDARD Base64 or malformed record
686    /// framing without retaining encoded payloads.
687    pub fn into_manifest(self) -> crate::error::Result<FileRefsManifest> {
688        let mut file_refs = HashMap::new();
689        for (region, encoded) in self.file_refs {
690            file_refs.insert(region, encoded.into_refs(region)?);
691        }
692        Ok(FileRefsManifest {
693            file_refs,
694            manifest_version: self.manifest_version,
695            cross_region_refs: self.cross_region_refs,
696        })
697    }
698}
699
700impl PackedRegionFileRefs {
701    /// Packs file references as the documented STANDARD Base64 record streams.
702    pub fn from_refs(refs: &HashSet<FileRef>) -> Self {
703        let engine = base64::engine::general_purpose::STANDARD;
704        let mut files = Vec::new();
705        let mut indexed = Vec::new();
706        for file_ref in refs {
707            match file_ref.index_version {
708                None => files.extend_from_slice(file_ref.file_id.as_bytes()),
709                Some(version) => {
710                    indexed.extend_from_slice(file_ref.file_id.as_bytes());
711                    indexed.extend_from_slice(&version.to_be_bytes());
712                }
713            }
714        }
715        Self {
716            files: engine.encode(files),
717            indexed: engine.encode(indexed),
718        }
719    }
720
721    /// Decodes the documented STANDARD Base64 record streams for `region`.
722    fn into_refs(self, region: RegionId) -> crate::error::Result<HashSet<FileRef>> {
723        let engine = base64::engine::general_purpose::STANDARD;
724        let files = engine
725            .decode(self.files)
726            .context(DecodePackedFileRefsSnafu)?;
727        let indexed = engine
728            .decode(self.indexed)
729            .context(DecodePackedFileRefsSnafu)?;
730        if files.len() % 16 != 0 || indexed.len() % 24 != 0 {
731            return InvalidPackedFileRefsSnafu.fail();
732        }
733
734        let mut refs = HashSet::new();
735        for bytes in files.chunks_exact(16) {
736            let mut id = [0; 16];
737            id.copy_from_slice(bytes);
738            refs.insert(FileRef::new(region, FileId::from_bytes(id), None));
739        }
740        for bytes in indexed.chunks_exact(24) {
741            let mut id = [0; 16];
742            id.copy_from_slice(&bytes[..16]);
743            let mut version = [0; 8];
744            version.copy_from_slice(&bytes[16..]);
745            refs.insert(FileRef::new(
746                region,
747                FileId::from_bytes(id),
748                Some(u64::from_be_bytes(version)),
749            ));
750        }
751        Ok(refs)
752    }
753}
754
755impl Display for GcRegions {
756    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
757        write!(
758            f,
759            "GcRegion(regions={:?}, file_refs_count={}, full_file_listing={})",
760            self.regions,
761            self.file_refs_manifest.file_refs.len(),
762            self.full_file_listing
763        )
764    }
765}
766
767/// Reply for GetFileRefs instruction.
768#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
769pub struct GetFileRefsReply {
770    /// The file references manifest.
771    pub file_refs_manifest: FileRefsManifest,
772    /// Whether the operation was successful.
773    pub success: bool,
774    /// Error message if any.
775    pub error: Option<InstructionError>,
776}
777
778impl Display for GetFileRefsReply {
779    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
780        write!(
781            f,
782            "GetFileRefsReply(success={}, file_refs_count={}, error={:?})",
783            self.success,
784            self.file_refs_manifest.file_refs.len(),
785            self.error
786        )
787    }
788}
789
790/// Reply for GetPackedFileRefs instruction.
791#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
792pub struct GetPackedFileRefsReply {
793    pub packed_file_refs_manifest: PackedFileRefsManifest,
794    pub success: bool,
795    pub error: Option<InstructionError>,
796}
797
798impl Display for GetPackedFileRefsReply {
799    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
800        write!(
801            f,
802            "GetPackedFileRefsReply(success={}, file_refs_count={}, error={:?})",
803            self.success,
804            self.packed_file_refs_manifest.file_refs.len(),
805            self.error
806        )
807    }
808}
809
810/// Reply for GC instruction.
811#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
812pub struct GcRegionsReply {
813    pub result: InstructionResult<GcReport>,
814}
815
816impl Display for GcRegionsReply {
817    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
818        write!(
819            f,
820            "GcReply(result={})",
821            match &self.result {
822                Ok(report) => format!(
823                    "GcReport(deleted_files_count={}, need_retry_regions_count={})",
824                    report.deleted_files.len(),
825                    report.need_retry_regions.len()
826                ),
827                Err(err) => format!("Err({})", err),
828            }
829        )
830    }
831}
832
833#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
834pub struct EnterStagingRegion {
835    pub region_id: RegionId,
836    #[serde(
837        alias = "partition_expr",
838        deserialize_with = "deserialize_enter_staging_partition_directive",
839        serialize_with = "serialize_enter_staging_partition_directive"
840    )]
841    pub partition_directive: StagingPartitionDirective,
842}
843
844impl Display for EnterStagingRegion {
845    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
846        write!(
847            f,
848            "EnterStagingRegion(region_id={}, partition_directive={})",
849            self.region_id, self.partition_directive
850        )
851    }
852}
853
854#[derive(Debug, Clone, PartialEq, Eq)]
855pub enum StagingPartitionDirective {
856    UpdatePartitionExpr(String),
857    RejectAllWrites,
858}
859
860impl StagingPartitionDirective {
861    /// Returns the partition expression carried by this directive, if any.
862    pub fn as_partition_expr(&self) -> Option<&str> {
863        match self {
864            Self::UpdatePartitionExpr(expr) => Some(expr),
865            Self::RejectAllWrites => None,
866        }
867    }
868}
869
870impl Display for StagingPartitionDirective {
871    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
872        match self {
873            Self::UpdatePartitionExpr(expr) => write!(f, "UpdatePartitionExpr({})", expr),
874            Self::RejectAllWrites => write!(f, "RejectAllWrites"),
875        }
876    }
877}
878
879fn serialize_enter_staging_partition_directive<S>(
880    rule: &StagingPartitionDirective,
881    serializer: S,
882) -> std::result::Result<S::Ok, S::Error>
883where
884    S: Serializer,
885{
886    match rule {
887        StagingPartitionDirective::UpdatePartitionExpr(expr) => serializer.serialize_str(expr),
888        StagingPartitionDirective::RejectAllWrites => {
889            #[derive(Serialize)]
890            struct RejectAllWritesSer<'a> {
891                r#type: &'a str,
892            }
893
894            RejectAllWritesSer {
895                r#type: "reject_all_writes",
896            }
897            .serialize(serializer)
898        }
899    }
900}
901
902fn deserialize_enter_staging_partition_directive<'de, D>(
903    deserializer: D,
904) -> std::result::Result<StagingPartitionDirective, D::Error>
905where
906    D: Deserializer<'de>,
907{
908    #[derive(Deserialize)]
909    #[serde(untagged)]
910    enum Compat {
911        Legacy(String),
912        TypeTagged { r#type: String },
913    }
914
915    match Compat::deserialize(deserializer)? {
916        Compat::Legacy(expr) => Ok(StagingPartitionDirective::UpdatePartitionExpr(expr)),
917        Compat::TypeTagged { r#type } if r#type == "reject_all_writes" => {
918            Ok(StagingPartitionDirective::RejectAllWrites)
919        }
920        Compat::TypeTagged { r#type } => Err(serde::de::Error::custom(format!(
921            "Unknown enter staging partition directive type: {}",
922            r#type
923        ))),
924    }
925}
926
927/// Instruction payload for syncing a region from a manifest or another region.
928#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
929pub struct SyncRegion {
930    /// Region id to sync.
931    pub region_id: RegionId,
932    /// Request to sync the region.
933    pub request: SyncRegionFromRequest,
934}
935
936impl Display for SyncRegion {
937    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
938        write!(
939            f,
940            "SyncRegion(region_id={}, request={:?})",
941            self.region_id, self.request
942        )
943    }
944}
945
946#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
947pub struct RemapManifest {
948    pub region_id: RegionId,
949    /// Regions to remap manifests from.
950    pub input_regions: Vec<RegionId>,
951    /// For each old region, which new regions should receive its files
952    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
953    /// New partition expressions for the new regions.
954    pub new_partition_exprs: HashMap<RegionId, String>,
955}
956
957impl Display for RemapManifest {
958    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
959        write!(
960            f,
961            "RemapManifest(region_id={}, input_regions={:?}, region_mapping={:?}, new_partition_exprs={:?})",
962            self.region_id, self.input_regions, self.region_mapping, self.new_partition_exprs
963        )
964    }
965}
966
967#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
968pub struct ApplyStagingManifest {
969    /// The region ID to apply the staging manifest to.
970    pub region_id: RegionId,
971    /// The partition expression of the staging region.
972    pub partition_expr: String,
973    /// The region that stores the staging manifests in its staging blob storage.
974    pub central_region_id: RegionId,
975    /// The relative path to the staging manifest within the central region's staging blob storage.
976    pub manifest_path: String,
977}
978
979impl Display for ApplyStagingManifest {
980    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
981        write!(
982            f,
983            "ApplyStagingManifest(region_id={}, partition_expr={}, central_region_id={}, manifest_path={})",
984            self.region_id, self.partition_expr, self.central_region_id, self.manifest_path
985        )
986    }
987}
988
989#[derive(Debug, Clone, Serialize, Deserialize, Display, PartialEq)]
990pub enum Instruction {
991    /// Opens regions.
992    #[serde(deserialize_with = "single_or_multiple_from", alias = "OpenRegion")]
993    OpenRegions(Vec<OpenRegion>),
994    /// Closes regions.
995    #[serde(deserialize_with = "single_or_multiple_from", alias = "CloseRegion")]
996    CloseRegions(Vec<RegionIdent>),
997    /// Upgrades regions.
998    #[serde(deserialize_with = "single_or_multiple_from", alias = "UpgradeRegion")]
999    UpgradeRegions(Vec<UpgradeRegion>),
1000    #[serde(
1001        deserialize_with = "single_or_multiple_from",
1002        alias = "DowngradeRegion"
1003    )]
1004    /// Downgrades regions.
1005    DowngradeRegions(Vec<DowngradeRegion>),
1006    /// Invalidates batch cache.
1007    InvalidateCaches(Vec<CacheIdent>),
1008    /// Flushes regions.
1009    FlushRegions(FlushRegions),
1010    /// Gets file references for regions.
1011    GetFileRefs(GetFileRefs),
1012    /// Gets file references using the packed manifest format.
1013    GetPackedFileRefs(GetPackedFileRefs),
1014    /// Triggers garbage collection for a region.
1015    GcRegions(GcRegions),
1016    /// Triggers garbage collection with a packed file-reference manifest.
1017    PackedGcRegions(PackedGcRegions),
1018    /// Temporary suspend serving reads or writes
1019    Suspend,
1020    /// Makes regions enter staging state.
1021    EnterStagingRegions(Vec<EnterStagingRegion>),
1022    /// Syncs regions.
1023    SyncRegions(Vec<SyncRegion>),
1024    /// Remaps manifests for a region.
1025    RemapManifest(RemapManifest),
1026
1027    /// Applies staging manifests for a region.
1028    ApplyStagingManifests(Vec<ApplyStagingManifest>),
1029}
1030
1031impl Instruction {
1032    /// Converts the instruction into a vector of [OpenRegion].
1033    pub fn into_open_regions(self) -> Option<Vec<OpenRegion>> {
1034        match self {
1035            Self::OpenRegions(open_regions) => Some(open_regions),
1036            _ => None,
1037        }
1038    }
1039
1040    /// Converts the instruction into a vector of [RegionIdent].
1041    pub fn into_close_regions(self) -> Option<Vec<RegionIdent>> {
1042        match self {
1043            Self::CloseRegions(close_regions) => Some(close_regions),
1044            _ => None,
1045        }
1046    }
1047
1048    /// Converts the instruction into a [FlushRegions].
1049    pub fn into_flush_regions(self) -> Option<FlushRegions> {
1050        match self {
1051            Self::FlushRegions(flush_regions) => Some(flush_regions),
1052            _ => None,
1053        }
1054    }
1055
1056    /// Converts the instruction into a [DowngradeRegion].
1057    pub fn into_downgrade_regions(self) -> Option<Vec<DowngradeRegion>> {
1058        match self {
1059            Self::DowngradeRegions(downgrade_region) => Some(downgrade_region),
1060            _ => None,
1061        }
1062    }
1063
1064    /// Converts the instruction into a [UpgradeRegion].
1065    pub fn into_upgrade_regions(self) -> Option<Vec<UpgradeRegion>> {
1066        match self {
1067            Self::UpgradeRegions(upgrade_region) => Some(upgrade_region),
1068            _ => None,
1069        }
1070    }
1071
1072    pub fn into_get_file_refs(self) -> Option<GetFileRefs> {
1073        match self {
1074            Self::GetFileRefs(get_file_refs) => Some(get_file_refs),
1075            _ => None,
1076        }
1077    }
1078
1079    pub fn into_gc_regions(self) -> Option<GcRegions> {
1080        match self {
1081            Self::GcRegions(gc_regions) => Some(gc_regions),
1082            _ => None,
1083        }
1084    }
1085
1086    pub fn into_enter_staging_regions(self) -> Option<Vec<EnterStagingRegion>> {
1087        match self {
1088            Self::EnterStagingRegions(enter_staging) => Some(enter_staging),
1089            _ => None,
1090        }
1091    }
1092
1093    pub fn into_sync_regions(self) -> Option<Vec<SyncRegion>> {
1094        match self {
1095            Self::SyncRegions(sync_regions) => Some(sync_regions),
1096            _ => None,
1097        }
1098    }
1099}
1100
1101/// The reply of [UpgradeRegion].
1102#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1103pub struct UpgradeRegionReply {
1104    /// The [RegionId].
1105    /// For compatibility, it is defaulted to [RegionId::new(0, 0)].
1106    #[serde(default)]
1107    pub region_id: RegionId,
1108    /// Returns true if `last_entry_id` has been replayed to the latest.
1109    pub ready: bool,
1110    /// Indicates whether the region exists.
1111    pub exists: bool,
1112    /// Returns error if any.
1113    pub error: Option<InstructionError>,
1114}
1115
1116impl Display for UpgradeRegionReply {
1117    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1118        write!(
1119            f,
1120            "(ready={}, exists={}, error={:?})",
1121            self.ready, self.exists, self.error
1122        )
1123    }
1124}
1125
1126#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1127pub struct DowngradeRegionsReply {
1128    pub replies: Vec<DowngradeRegionReply>,
1129}
1130
1131impl DowngradeRegionsReply {
1132    pub fn new(replies: Vec<DowngradeRegionReply>) -> Self {
1133        Self { replies }
1134    }
1135
1136    pub fn single(reply: DowngradeRegionReply) -> Self {
1137        Self::new(vec![reply])
1138    }
1139}
1140
1141#[derive(Deserialize)]
1142#[serde(untagged)]
1143enum DowngradeRegionsCompat {
1144    Single(DowngradeRegionReply),
1145    Multiple(DowngradeRegionsReply),
1146}
1147
1148fn downgrade_regions_compat_from<'de, D>(deserializer: D) -> Result<DowngradeRegionsReply, D::Error>
1149where
1150    D: Deserializer<'de>,
1151{
1152    let helper = DowngradeRegionsCompat::deserialize(deserializer)?;
1153    Ok(match helper {
1154        DowngradeRegionsCompat::Single(x) => DowngradeRegionsReply::new(vec![x]),
1155        DowngradeRegionsCompat::Multiple(reply) => reply,
1156    })
1157}
1158
1159#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1160pub struct UpgradeRegionsReply {
1161    pub replies: Vec<UpgradeRegionReply>,
1162}
1163
1164impl UpgradeRegionsReply {
1165    pub fn new(replies: Vec<UpgradeRegionReply>) -> Self {
1166        Self { replies }
1167    }
1168
1169    pub fn single(reply: UpgradeRegionReply) -> Self {
1170        Self::new(vec![reply])
1171    }
1172}
1173
1174#[derive(Deserialize)]
1175#[serde(untagged)]
1176enum UpgradeRegionsCompat {
1177    Single(UpgradeRegionReply),
1178    Multiple(UpgradeRegionsReply),
1179}
1180
1181fn upgrade_regions_compat_from<'de, D>(deserializer: D) -> Result<UpgradeRegionsReply, D::Error>
1182where
1183    D: Deserializer<'de>,
1184{
1185    let helper = UpgradeRegionsCompat::deserialize(deserializer)?;
1186    Ok(match helper {
1187        UpgradeRegionsCompat::Single(x) => UpgradeRegionsReply::new(vec![x]),
1188        UpgradeRegionsCompat::Multiple(reply) => reply,
1189    })
1190}
1191
1192#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1193pub struct EnterStagingRegionReply {
1194    pub region_id: RegionId,
1195    /// Returns true if the region has entered staging with the target directive.
1196    pub ready: bool,
1197    /// Indicates whether the region exists.
1198    pub exists: bool,
1199    /// Return error if any during the operation.
1200    pub error: Option<InstructionError>,
1201}
1202
1203#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1204pub struct EnterStagingRegionsReply {
1205    pub replies: Vec<EnterStagingRegionReply>,
1206}
1207
1208impl EnterStagingRegionsReply {
1209    pub fn new(replies: Vec<EnterStagingRegionReply>) -> Self {
1210        Self { replies }
1211    }
1212}
1213
1214/// Reply for a single region sync request.
1215#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1216pub struct SyncRegionReply {
1217    /// Region id of the synced region.
1218    pub region_id: RegionId,
1219    /// Returns true if the region is successfully synced and ready.
1220    pub ready: bool,
1221    /// Indicates whether the region exists.
1222    pub exists: bool,
1223    /// Return error message if any during the operation.
1224    pub error: Option<InstructionError>,
1225}
1226
1227/// Reply for a batch of region sync requests.
1228#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1229pub struct SyncRegionsReply {
1230    pub replies: Vec<SyncRegionReply>,
1231}
1232
1233impl SyncRegionsReply {
1234    pub fn new(replies: Vec<SyncRegionReply>) -> Self {
1235        Self { replies }
1236    }
1237}
1238
1239#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1240pub struct RemapManifestReply {
1241    /// Returns false if the region does not exist.
1242    pub exists: bool,
1243    /// A map from region IDs to their corresponding remapped manifest paths.
1244    pub manifest_paths: HashMap<RegionId, String>,
1245    /// Return error if any during the operation.
1246    pub error: Option<InstructionError>,
1247}
1248
1249impl Display for RemapManifestReply {
1250    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1251        write!(
1252            f,
1253            "RemapManifestReply(manifest_paths={:?}, error={:?})",
1254            self.manifest_paths, self.error
1255        )
1256    }
1257}
1258
1259#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1260pub struct ApplyStagingManifestsReply {
1261    pub replies: Vec<ApplyStagingManifestReply>,
1262}
1263
1264impl ApplyStagingManifestsReply {
1265    pub fn new(replies: Vec<ApplyStagingManifestReply>) -> Self {
1266        Self { replies }
1267    }
1268}
1269
1270#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1271pub struct ApplyStagingManifestReply {
1272    pub region_id: RegionId,
1273    /// Returns true if the region is ready to serve reads and writes.
1274    pub ready: bool,
1275    /// Indicates whether the region exists.
1276    pub exists: bool,
1277    /// Return error if any during the operation.
1278    pub error: Option<InstructionError>,
1279}
1280
1281#[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
1282#[serde(tag = "type", rename_all = "snake_case")]
1283pub enum InstructionReply {
1284    #[serde(alias = "open_region")]
1285    OpenRegions(SimpleReply),
1286    #[serde(alias = "close_region")]
1287    CloseRegions(SimpleReply),
1288    #[serde(
1289        deserialize_with = "upgrade_regions_compat_from",
1290        alias = "upgrade_region"
1291    )]
1292    UpgradeRegions(UpgradeRegionsReply),
1293    #[serde(
1294        alias = "downgrade_region",
1295        deserialize_with = "downgrade_regions_compat_from"
1296    )]
1297    DowngradeRegions(DowngradeRegionsReply),
1298    FlushRegions(FlushRegionReply),
1299    GetFileRefs(GetFileRefsReply),
1300    GetPackedFileRefs(GetPackedFileRefsReply),
1301    GcRegions(GcRegionsReply),
1302    EnterStagingRegions(EnterStagingRegionsReply),
1303    SyncRegions(SyncRegionsReply),
1304    RemapManifest(RemapManifestReply),
1305    ApplyStagingManifests(ApplyStagingManifestsReply),
1306}
1307
1308impl Display for InstructionReply {
1309    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
1310        match self {
1311            Self::OpenRegions(reply) => write!(f, "InstructionReply::OpenRegions({})", reply),
1312            Self::CloseRegions(reply) => write!(f, "InstructionReply::CloseRegions({})", reply),
1313            Self::UpgradeRegions(reply) => {
1314                write!(f, "InstructionReply::UpgradeRegions({:?})", reply.replies)
1315            }
1316            Self::DowngradeRegions(reply) => {
1317                write!(f, "InstructionReply::DowngradeRegions({:?})", reply.replies)
1318            }
1319            Self::FlushRegions(reply) => write!(f, "InstructionReply::FlushRegions({})", reply),
1320            Self::GetFileRefs(reply) => write!(f, "InstructionReply::GetFileRefs({})", reply),
1321            Self::GetPackedFileRefs(reply) => {
1322                write!(f, "InstructionReply::GetPackedFileRefs({})", reply)
1323            }
1324            Self::GcRegions(reply) => write!(f, "InstructionReply::GcRegion({})", reply),
1325            Self::EnterStagingRegions(reply) => {
1326                write!(
1327                    f,
1328                    "InstructionReply::EnterStagingRegions({:?})",
1329                    reply.replies
1330                )
1331            }
1332            Self::SyncRegions(reply) => {
1333                write!(f, "InstructionReply::SyncRegions({:?})", reply.replies)
1334            }
1335            Self::RemapManifest(reply) => write!(f, "InstructionReply::RemapManifest({})", reply),
1336            Self::ApplyStagingManifests(reply) => write!(
1337                f,
1338                "InstructionReply::ApplyStagingManifests({:?})",
1339                reply.replies
1340            ),
1341        }
1342    }
1343}
1344
1345#[cfg(any(test, feature = "testing"))]
1346impl InstructionReply {
1347    pub fn expect_close_regions_reply(self) -> SimpleReply {
1348        match self {
1349            Self::CloseRegions(reply) => reply,
1350            _ => panic!("Expected CloseRegions reply"),
1351        }
1352    }
1353
1354    pub fn expect_open_regions_reply(self) -> SimpleReply {
1355        match self {
1356            Self::OpenRegions(reply) => reply,
1357            _ => panic!("Expected OpenRegions reply"),
1358        }
1359    }
1360
1361    pub fn expect_upgrade_regions_reply(self) -> Vec<UpgradeRegionReply> {
1362        match self {
1363            Self::UpgradeRegions(reply) => reply.replies,
1364            _ => panic!("Expected UpgradeRegion reply"),
1365        }
1366    }
1367
1368    pub fn expect_downgrade_regions_reply(self) -> Vec<DowngradeRegionReply> {
1369        match self {
1370            Self::DowngradeRegions(reply) => reply.replies,
1371            _ => panic!("Expected DowngradeRegion reply"),
1372        }
1373    }
1374
1375    pub fn expect_flush_regions_reply(self) -> FlushRegionReply {
1376        match self {
1377            Self::FlushRegions(reply) => reply,
1378            _ => panic!("Expected FlushRegions reply"),
1379        }
1380    }
1381
1382    pub fn expect_enter_staging_regions_reply(self) -> Vec<EnterStagingRegionReply> {
1383        match self {
1384            Self::EnterStagingRegions(reply) => reply.replies,
1385            _ => panic!("Expected EnterStagingRegion reply"),
1386        }
1387    }
1388
1389    pub fn expect_sync_regions_reply(self) -> Vec<SyncRegionReply> {
1390        match self {
1391            Self::SyncRegions(reply) => reply.replies,
1392            _ => panic!("Expected SyncRegion reply"),
1393        }
1394    }
1395
1396    pub fn expect_remap_manifest_reply(self) -> RemapManifestReply {
1397        match self {
1398            Self::RemapManifest(reply) => reply,
1399            _ => panic!("Expected RemapManifest reply"),
1400        }
1401    }
1402
1403    pub fn expect_apply_staging_manifests_reply(self) -> Vec<ApplyStagingManifestReply> {
1404        match self {
1405            Self::ApplyStagingManifests(reply) => reply.replies,
1406            _ => panic!("Expected ApplyStagingManifest reply"),
1407        }
1408    }
1409}
1410
1411#[cfg(test)]
1412mod tests {
1413    use std::collections::HashSet;
1414    use std::error::Error as _;
1415
1416    use common_error::mock::MockError;
1417    use common_wal::options::WalOptions;
1418    use store_api::storage::{FileId, FileRef};
1419
1420    use super::*;
1421
1422    #[test]
1423    fn test_instruction_error_serde() {
1424        let error = InstructionError::new(
1425            StatusCode::RegionNotFound,
1426            "region not found",
1427            RetryHint::Retryable,
1428        );
1429
1430        let serialized = serde_json::to_string(&error).unwrap();
1431        assert_eq!(
1432            r#"{"code":4005,"message":"region not found","retry_hint":"retryable"}"#,
1433            serialized
1434        );
1435
1436        let deserialized: InstructionError = serde_json::from_str(&serialized).unwrap();
1437        assert_eq!(error, deserialized);
1438
1439        assert!(
1440            serde_json::from_str::<InstructionError>(
1441                r#"{"code":999999,"message":"unknown","retry_hint":"retryable"}"#
1442            )
1443            .is_err()
1444        );
1445        assert!(
1446            serde_json::from_str::<InstructionError>(
1447                r#"{"code":4005,"message":"unknown","retry_hint":"unknown"}"#
1448            )
1449            .is_err()
1450        );
1451    }
1452
1453    #[test]
1454    fn test_instruction_error_from_error() {
1455        let error = MockError::new(StatusCode::RegionNotFound);
1456
1457        let instruction_error = InstructionError::from_error(&error);
1458
1459        assert_eq!(StatusCode::RegionNotFound, instruction_error.code);
1460        assert_eq!("RegionNotFound", instruction_error.message);
1461        assert_eq!(RetryHint::NonRetryable, instruction_error.retry_hint);
1462    }
1463
1464    #[test]
1465    fn test_instruction_result_serde() {
1466        let success: InstructionResult<bool> = Ok(true);
1467        let serialized = serde_json::to_string(&success).unwrap();
1468        assert_eq!(r#"{"Ok":true}"#, serialized);
1469        let deserialized: InstructionResult<bool> = serde_json::from_str(&serialized).unwrap();
1470        assert_eq!(success, deserialized);
1471
1472        let failure: InstructionResult<bool> = Err(InstructionError::new(
1473            StatusCode::RegionBusy,
1474            "region busy",
1475            RetryHint::Retryable,
1476        ));
1477        let serialized = serde_json::to_string(&failure).unwrap();
1478        assert_eq!(
1479            r#"{"Err":{"code":4009,"message":"region busy","retry_hint":"retryable"}}"#,
1480            serialized
1481        );
1482        let deserialized: InstructionResult<bool> = serde_json::from_str(&serialized).unwrap();
1483        assert_eq!(failure, deserialized);
1484    }
1485
1486    #[test]
1487    fn test_serialize_instruction() {
1488        let open_region = Instruction::OpenRegions(vec![OpenRegion::new(
1489            RegionIdent {
1490                datanode_id: 2,
1491                table_id: 1024,
1492                region_number: 1,
1493                engine: "mito2".to_string(),
1494            },
1495            "test/foo",
1496            HashMap::new(),
1497            HashMap::new(),
1498            false,
1499            None,
1500            RegionRequirements::empty(),
1501        )]);
1502
1503        let serialized = serde_json::to_string(&open_region).unwrap();
1504        assert_eq!(
1505            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}}]}"#,
1506            serialized
1507        );
1508
1509        let close_region = Instruction::CloseRegions(vec![RegionIdent {
1510            datanode_id: 2,
1511            table_id: 1024,
1512            region_number: 1,
1513            engine: "mito2".to_string(),
1514        }]);
1515
1516        let serialized = serde_json::to_string(&close_region).unwrap();
1517        assert_eq!(
1518            r#"{"CloseRegions":[{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"}]}"#,
1519            serialized
1520        );
1521
1522        let upgrade_region = Instruction::UpgradeRegions(vec![UpgradeRegion {
1523            region_id: RegionId::new(1024, 1),
1524            last_entry_id: None,
1525            metadata_last_entry_id: None,
1526            replay_timeout: Duration::from_millis(1000),
1527            location_id: None,
1528            replay_entry_id: None,
1529            metadata_replay_entry_id: None,
1530        }]);
1531
1532        let serialized = serde_json::to_string(&upgrade_region).unwrap();
1533        assert_eq!(
1534            r#"{"UpgradeRegions":[{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"replay_timeout":"1s","location_id":null}]}"#,
1535            serialized
1536        );
1537    }
1538
1539    #[test]
1540    fn test_serialize_instruction_reply() {
1541        let downgrade_region_reply = InstructionReply::DowngradeRegions(
1542            DowngradeRegionsReply::single(DowngradeRegionReply {
1543                region_id: RegionId::new(1024, 1),
1544                last_entry_id: None,
1545                metadata_last_entry_id: None,
1546                exists: true,
1547                error: None,
1548            }),
1549        );
1550
1551        let serialized = serde_json::to_string(&downgrade_region_reply).unwrap();
1552        assert_eq!(
1553            r#"{"type":"downgrade_regions","replies":[{"region_id":4398046511105,"last_entry_id":null,"metadata_last_entry_id":null,"exists":true,"error":null}]}"#,
1554            serialized
1555        );
1556
1557        let upgrade_region_reply =
1558            InstructionReply::UpgradeRegions(UpgradeRegionsReply::single(UpgradeRegionReply {
1559                region_id: RegionId::new(1024, 1),
1560                ready: true,
1561                exists: true,
1562                error: None,
1563            }));
1564        let serialized = serde_json::to_string(&upgrade_region_reply).unwrap();
1565        assert_eq!(
1566            r#"{"type":"upgrade_regions","replies":[{"region_id":4398046511105,"ready":true,"exists":true,"error":null}]}"#,
1567            serialized
1568        );
1569    }
1570
1571    #[test]
1572    fn test_deserialize_instruction() {
1573        // legacy open region instruction
1574        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}}"#;
1575        let open_region_instruction: Instruction =
1576            serde_json::from_str(open_region_instruction).unwrap();
1577        let open_region = Instruction::OpenRegions(vec![OpenRegion::new(
1578            RegionIdent {
1579                datanode_id: 2,
1580                table_id: 1024,
1581                region_number: 1,
1582                engine: "mito2".to_string(),
1583            },
1584            "test/foo",
1585            HashMap::new(),
1586            HashMap::new(),
1587            false,
1588            None,
1589            RegionRequirements::empty(),
1590        )]);
1591        assert_eq!(open_region_instruction, open_region);
1592
1593        // legacy close region instruction
1594        let close_region_instruction = r#"{"CloseRegion":{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"}}"#;
1595        let close_region_instruction: Instruction =
1596            serde_json::from_str(close_region_instruction).unwrap();
1597        let close_region = Instruction::CloseRegions(vec![RegionIdent {
1598            datanode_id: 2,
1599            table_id: 1024,
1600            region_number: 1,
1601            engine: "mito2".to_string(),
1602        }]);
1603        assert_eq!(close_region_instruction, close_region);
1604
1605        // legacy downgrade region instruction
1606        let downgrade_region_instruction = r#"{"DowngradeRegions":{"region_id":4398046511105,"flush_timeout":{"secs":1,"nanos":0}}}"#;
1607        let downgrade_region_instruction: Instruction =
1608            serde_json::from_str(downgrade_region_instruction).unwrap();
1609        let downgrade_region = Instruction::DowngradeRegions(vec![DowngradeRegion {
1610            region_id: RegionId::new(1024, 1),
1611            flush_timeout: Some(Duration::from_millis(1000)),
1612        }]);
1613        assert_eq!(downgrade_region_instruction, downgrade_region);
1614
1615        // legacy upgrade region instruction
1616        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}}"#;
1617        let upgrade_region_instruction: Instruction =
1618            serde_json::from_str(upgrade_region_instruction).unwrap();
1619        let upgrade_region = Instruction::UpgradeRegions(vec![UpgradeRegion {
1620            region_id: RegionId::new(1024, 1),
1621            last_entry_id: None,
1622            metadata_last_entry_id: None,
1623            replay_timeout: Duration::from_millis(1000),
1624            location_id: None,
1625            replay_entry_id: None,
1626            metadata_replay_entry_id: None,
1627        }]);
1628        assert_eq!(upgrade_region_instruction, upgrade_region);
1629    }
1630
1631    #[test]
1632    fn test_deserialize_instruction_reply() {
1633        // legacy close region reply
1634        let close_region_instruction_reply =
1635            r#"{"result":true,"error":null,"type":"close_region"}"#;
1636        let close_region_instruction_reply: InstructionReply =
1637            serde_json::from_str(close_region_instruction_reply).unwrap();
1638        let close_region_reply = InstructionReply::CloseRegions(SimpleReply {
1639            result: true,
1640            error: None,
1641        });
1642        assert_eq!(close_region_instruction_reply, close_region_reply);
1643
1644        // legacy open region reply
1645        let open_region_instruction_reply = r#"{"result":true,"error":null,"type":"open_region"}"#;
1646        let open_region_instruction_reply: InstructionReply =
1647            serde_json::from_str(open_region_instruction_reply).unwrap();
1648        let open_region_reply = InstructionReply::OpenRegions(SimpleReply {
1649            result: true,
1650            error: None,
1651        });
1652        assert_eq!(open_region_instruction_reply, open_region_reply);
1653
1654        // legacy downgrade region reply
1655        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"}"#;
1656        let downgrade_region_instruction_reply: InstructionReply =
1657            serde_json::from_str(downgrade_region_instruction_reply).unwrap();
1658        let downgrade_region_reply = InstructionReply::DowngradeRegions(
1659            DowngradeRegionsReply::single(DowngradeRegionReply {
1660                region_id: RegionId::new(1024, 1),
1661                last_entry_id: None,
1662                metadata_last_entry_id: None,
1663                exists: true,
1664                error: None,
1665            }),
1666        );
1667        assert_eq!(downgrade_region_instruction_reply, downgrade_region_reply);
1668
1669        // legacy upgrade region reply
1670        let upgrade_region_instruction_reply = r#"{"region_id":4398046511105,"ready":true,"exists":true,"error":null,"type":"upgrade_region"}"#;
1671        let upgrade_region_instruction_reply: InstructionReply =
1672            serde_json::from_str(upgrade_region_instruction_reply).unwrap();
1673        let upgrade_region_reply =
1674            InstructionReply::UpgradeRegions(UpgradeRegionsReply::single(UpgradeRegionReply {
1675                region_id: RegionId::new(1024, 1),
1676                ready: true,
1677                exists: true,
1678                error: None,
1679            }));
1680        assert_eq!(upgrade_region_instruction_reply, upgrade_region_reply);
1681    }
1682
1683    #[test]
1684    fn test_enter_staging_partition_rule_compatibility() {
1685        let legacy = r#"{"region_id":4398046511105,"partition_expr":"{\"Expr\":{\"lhs\":{\"Column\":\"x\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"Int32\":0}}}}"}"#;
1686        let enter: EnterStagingRegion = serde_json::from_str(legacy).unwrap();
1687        assert_eq!(enter.region_id, RegionId::new(1024, 1));
1688        assert_eq!(
1689            enter.partition_directive,
1690            StagingPartitionDirective::UpdatePartitionExpr(
1691                "{\"Expr\":{\"lhs\":{\"Column\":\"x\"},\"op\":\"GtEq\",\"rhs\":{\"Value\":{\"Int32\":0}}}}"
1692                    .to_string()
1693            )
1694        );
1695
1696        let serialized = serde_json::to_string(&enter).unwrap();
1697        assert!(serialized.contains("\"partition_directive\":\""));
1698        assert!(!serialized.contains("partition_expr"));
1699
1700        let reject = r#"{"region_id":4398046511105,"partition_expr":{"type":"reject_all_writes"}}"#;
1701        let enter: EnterStagingRegion = serde_json::from_str(reject).unwrap();
1702        assert_eq!(
1703            enter.partition_directive,
1704            StagingPartitionDirective::RejectAllWrites
1705        );
1706    }
1707
1708    #[derive(Debug, Clone, Serialize, Deserialize)]
1709    struct LegacyOpenRegion {
1710        region_ident: RegionIdent,
1711        region_storage_path: String,
1712        region_options: HashMap<String, String>,
1713    }
1714
1715    #[test]
1716    fn test_compatible_serialize_open_region() {
1717        let region_ident = RegionIdent {
1718            datanode_id: 2,
1719            table_id: 1024,
1720            region_number: 1,
1721            engine: "mito2".to_string(),
1722        };
1723        let region_storage_path = "test/foo".to_string();
1724        let region_options = HashMap::from([
1725            ("a".to_string(), "aa".to_string()),
1726            ("b".to_string(), "bb".to_string()),
1727        ]);
1728
1729        // Serialize a legacy OpenRegion.
1730        let legacy_open_region = LegacyOpenRegion {
1731            region_ident: region_ident.clone(),
1732            region_storage_path: region_storage_path.clone(),
1733            region_options: region_options.clone(),
1734        };
1735        let serialized = serde_json::to_string(&legacy_open_region).unwrap();
1736
1737        // Deserialize to OpenRegion.
1738        let deserialized = serde_json::from_str(&serialized).unwrap();
1739        let expected = OpenRegion {
1740            region_ident,
1741            region_storage_path,
1742            region_options,
1743            region_wal_options: HashMap::new(),
1744            skip_wal_replay: false,
1745            reason: None,
1746            requirements: RegionRequirements::empty(),
1747        };
1748        assert_eq!(expected, deserialized);
1749    }
1750
1751    #[test]
1752    fn test_deserialize_open_region_with_legacy_region_wal_options() {
1753        let open_region = r#"{"region_ident":{"datanode_id":2,"table_id":1024,"region_number":1,"engine":"mito2"},"region_storage_path":"test/foo","region_options":{},"region_wal_options":{"1":"{\"wal.provider\":\"raft_engine\"}"},"skip_wal_replay":false}"#;
1754
1755        let open_region: OpenRegion = serde_json::from_str(open_region).unwrap();
1756
1757        assert_eq!(
1758            open_region.region_wal_options,
1759            HashMap::from([(1, WalOptions::RaftEngine)])
1760        );
1761    }
1762
1763    #[test]
1764    fn test_serialize_open_region_with_reason_and_requirements() {
1765        let open_region = OpenRegion::new(
1766            RegionIdent {
1767                datanode_id: 2,
1768                table_id: 1024,
1769                region_number: 1,
1770                engine: "mito2".to_string(),
1771            },
1772            "test/foo",
1773            HashMap::new(),
1774            HashMap::new(),
1775            false,
1776            Some(OpenRegionReason::RegionMigration),
1777            RegionRequirements::object_storage(),
1778        );
1779
1780        let serialized = serde_json::to_string(&open_region).unwrap();
1781        assert!(serialized.contains(r#""reason":"RegionMigration""#));
1782        assert!(serialized.contains(r#""object_storage":true"#));
1783
1784        let deserialized: OpenRegion = serde_json::from_str(&serialized).unwrap();
1785        assert_eq!(Some(OpenRegionReason::RegionMigration), deserialized.reason);
1786        assert_eq!(
1787            RegionRequirements::object_storage(),
1788            deserialized.requirements
1789        );
1790    }
1791
1792    #[test]
1793    fn test_flush_regions_creation() {
1794        let region_id = RegionId::new(1024, 1);
1795
1796        // Single region sync flush
1797        let single_sync = FlushRegions::sync_single(region_id);
1798        assert_eq!(single_sync.region_ids, vec![region_id]);
1799        assert_eq!(single_sync.strategy, FlushStrategy::Sync);
1800        assert!(!single_sync.is_hint());
1801        assert!(single_sync.is_sync());
1802        assert_eq!(single_sync.error_strategy, FlushErrorStrategy::FailFast);
1803        assert_eq!(single_sync.reason, None);
1804        assert!(single_sync.is_single_region());
1805        assert_eq!(single_sync.single_region_id(), Some(region_id));
1806
1807        // Batch async flush (hint)
1808        let region_ids = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
1809        let batch_async = FlushRegions::async_batch(region_ids.clone());
1810        assert_eq!(batch_async.region_ids, region_ids);
1811        assert_eq!(batch_async.strategy, FlushStrategy::Async);
1812        assert!(batch_async.is_hint());
1813        assert!(!batch_async.is_sync());
1814        assert_eq!(batch_async.error_strategy, FlushErrorStrategy::TryAll);
1815        assert_eq!(batch_async.reason, None);
1816        assert!(!batch_async.is_single_region());
1817        assert_eq!(batch_async.single_region_id(), None);
1818
1819        // Batch sync flush
1820        let batch_sync = FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::FailFast);
1821        assert_eq!(batch_sync.region_ids, region_ids);
1822        assert_eq!(batch_sync.strategy, FlushStrategy::Sync);
1823        assert!(!batch_sync.is_hint());
1824        assert!(batch_sync.is_sync());
1825        assert_eq!(batch_sync.error_strategy, FlushErrorStrategy::FailFast);
1826        assert_eq!(batch_sync.reason, None);
1827
1828        let with_reason = batch_sync.with_reason(RegionFlushReason::RemoteWalPrune);
1829        assert_eq!(with_reason.reason, Some(RegionFlushReason::RemoteWalPrune));
1830    }
1831
1832    #[test]
1833    fn test_flush_regions_conversion() {
1834        let region_id = RegionId::new(1024, 1);
1835
1836        let from_region_id: FlushRegions = region_id.into();
1837        assert_eq!(from_region_id.region_ids, vec![region_id]);
1838        assert_eq!(from_region_id.strategy, FlushStrategy::Sync);
1839        assert!(!from_region_id.is_hint());
1840        assert!(from_region_id.is_sync());
1841
1842        // Test default construction
1843        let flush_regions = FlushRegions {
1844            region_ids: vec![region_id],
1845            strategy: FlushStrategy::Async,
1846            error_strategy: FlushErrorStrategy::TryAll,
1847            reason: None,
1848        };
1849        assert_eq!(flush_regions.region_ids, vec![region_id]);
1850        assert_eq!(flush_regions.strategy, FlushStrategy::Async);
1851        assert!(flush_regions.is_hint());
1852        assert!(!flush_regions.is_sync());
1853    }
1854
1855    #[test]
1856    fn test_flush_region_reply() {
1857        let region_id = RegionId::new(1024, 1);
1858
1859        // Successful single region reply
1860        let success_reply = FlushRegionReply::success_single(region_id);
1861        assert!(success_reply.overall_success);
1862        assert_eq!(success_reply.results.len(), 1);
1863        assert_eq!(success_reply.results[0].0, region_id);
1864        assert!(success_reply.results[0].1.is_ok());
1865
1866        // Failed single region reply
1867        let error_reply = FlushRegionReply::error_single(
1868            region_id,
1869            InstructionError::legacy_internal_retryable("test error"),
1870        );
1871        assert!(!error_reply.overall_success);
1872        assert_eq!(error_reply.results.len(), 1);
1873        assert_eq!(error_reply.results[0].0, region_id);
1874        assert!(error_reply.results[0].1.is_err());
1875
1876        // Batch reply
1877        let region_id2 = RegionId::new(1024, 2);
1878        let results = vec![
1879            (region_id, Ok(())),
1880            (
1881                region_id2,
1882                Err(InstructionError::legacy_internal_retryable("flush failed")),
1883            ),
1884        ];
1885        let batch_reply = FlushRegionReply::from_results(results);
1886        assert!(!batch_reply.overall_success);
1887        assert_eq!(batch_reply.results.len(), 2);
1888
1889        // Conversion to SimpleReply
1890        let simple_reply = batch_reply.to_simple_reply();
1891        assert!(!simple_reply.result);
1892        assert!(simple_reply.error.is_some());
1893        assert!(simple_reply.error.unwrap().message.contains("flush failed"));
1894    }
1895
1896    #[test]
1897    fn test_serialize_flush_regions_instruction() {
1898        let region_id = RegionId::new(1024, 1);
1899        let flush_regions = FlushRegions::sync_single(region_id);
1900        let instruction = Instruction::FlushRegions(flush_regions.clone());
1901
1902        let serialized = serde_json::to_string(&instruction).unwrap();
1903        assert!(!serialized.contains("reason"));
1904        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1905
1906        match deserialized {
1907            Instruction::FlushRegions(fr) => {
1908                assert_eq!(fr.region_ids, vec![region_id]);
1909                assert_eq!(fr.strategy, FlushStrategy::Sync);
1910                assert_eq!(fr.error_strategy, FlushErrorStrategy::FailFast);
1911                assert_eq!(fr.reason, None);
1912            }
1913            _ => panic!("Expected FlushRegions instruction"),
1914        }
1915
1916        let legacy = r#"{"FlushRegions":{"region_ids":[4398046511105],"strategy":"Sync","error_strategy":"FailFast"}}"#;
1917        let deserialized: Instruction = serde_json::from_str(legacy).unwrap();
1918        match deserialized {
1919            Instruction::FlushRegions(fr) => {
1920                assert_eq!(fr.region_ids, vec![region_id]);
1921                assert_eq!(fr.strategy, FlushStrategy::Sync);
1922                assert_eq!(fr.error_strategy, FlushErrorStrategy::FailFast);
1923                assert_eq!(fr.reason, None);
1924            }
1925            _ => panic!("Expected FlushRegions instruction"),
1926        }
1927
1928        let flush_regions = FlushRegions::async_batch(vec![region_id])
1929            .with_reason(RegionFlushReason::RemoteWalPrune);
1930        let instruction = Instruction::FlushRegions(flush_regions);
1931        let serialized = serde_json::to_string(&instruction).unwrap();
1932        assert!(serialized.contains(r#""reason":"RemoteWalPrune""#));
1933        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1934        match deserialized {
1935            Instruction::FlushRegions(fr) => {
1936                assert_eq!(fr.reason, Some(RegionFlushReason::RemoteWalPrune));
1937            }
1938            _ => panic!("Expected FlushRegions instruction"),
1939        }
1940    }
1941
1942    #[test]
1943    fn test_serialize_flush_regions_batch_instruction() {
1944        let region_ids = vec![RegionId::new(1024, 1), RegionId::new(1024, 2)];
1945        let flush_regions =
1946            FlushRegions::sync_batch(region_ids.clone(), FlushErrorStrategy::TryAll);
1947        let instruction = Instruction::FlushRegions(flush_regions);
1948
1949        let serialized = serde_json::to_string(&instruction).unwrap();
1950        let deserialized: Instruction = serde_json::from_str(&serialized).unwrap();
1951
1952        match deserialized {
1953            Instruction::FlushRegions(fr) => {
1954                assert_eq!(fr.region_ids, region_ids);
1955                assert_eq!(fr.strategy, FlushStrategy::Sync);
1956                assert!(!fr.is_hint());
1957                assert!(fr.is_sync());
1958                assert_eq!(fr.error_strategy, FlushErrorStrategy::TryAll);
1959                assert_eq!(fr.reason, None);
1960            }
1961            _ => panic!("Expected FlushRegions instruction"),
1962        }
1963    }
1964
1965    #[test]
1966    fn test_legacy_gc_file_refs_compatibility() {
1967        #[derive(Debug, Deserialize)]
1968        struct LegacyGetFileRefs {
1969            query_regions: Vec<RegionId>,
1970            related_regions: HashMap<RegionId, HashSet<RegionId>>,
1971        }
1972
1973        #[derive(Debug, Deserialize)]
1974        struct LegacyGetFileRefsReply {
1975            file_refs_manifest: FileRefsManifest,
1976            success: bool,
1977            error: Option<InstructionError>,
1978        }
1979
1980        #[derive(Debug, Deserialize)]
1981        enum LegacyInstruction {
1982            GetFileRefs(LegacyGetFileRefs),
1983        }
1984
1985        let get_file_refs = Instruction::GetFileRefs(GetFileRefs {
1986            query_regions: vec![RegionId::new(7, 3)],
1987            related_regions: HashMap::new(),
1988        });
1989        let get_file_refs_json = serde_json::to_string(&get_file_refs).unwrap();
1990        let legacy_get_file_refs: LegacyInstruction =
1991            serde_json::from_str(&get_file_refs_json).unwrap();
1992        let LegacyInstruction::GetFileRefs(legacy_get_file_refs) = legacy_get_file_refs;
1993        assert_eq!(legacy_get_file_refs.query_regions.len(), 1);
1994        assert!(legacy_get_file_refs.related_regions.is_empty());
1995
1996        let reply = GetFileRefsReply {
1997            file_refs_manifest: FileRefsManifest::default(),
1998            success: true,
1999            error: None,
2000        };
2001        let reply_json = serde_json::to_string(&reply).unwrap();
2002        let legacy_reply: LegacyGetFileRefsReply = serde_json::from_str(&reply_json).unwrap();
2003        let current_reply: GetFileRefsReply = serde_json::from_str(&reply_json).unwrap();
2004        assert!(legacy_reply.success);
2005        assert!(legacy_reply.error.is_none());
2006        assert!(legacy_reply.file_refs_manifest.file_refs.is_empty());
2007        assert_eq!(current_reply, reply);
2008
2009        let packed = Instruction::PackedGcRegions(PackedGcRegions {
2010            regions: vec![],
2011            packed_file_refs_manifest: PackedFileRefsManifest::default(),
2012            full_file_listing: false,
2013        });
2014        let packed_json = serde_json::to_string(&packed).unwrap();
2015        assert!(serde_json::from_str::<LegacyInstruction>(&packed_json).is_err());
2016    }
2017
2018    #[test]
2019    fn test_packed_gc_regions_round_trip_is_distinct() {
2020        let instruction = Instruction::PackedGcRegions(PackedGcRegions {
2021            regions: vec![RegionId::new(7, 3)],
2022            packed_file_refs_manifest: PackedFileRefsManifest::default(),
2023            full_file_listing: true,
2024        });
2025        let serialized = serde_json::to_string(&instruction).unwrap();
2026        assert!(serialized.contains("PackedGcRegions"));
2027        assert_eq!(
2028            serde_json::from_str::<Instruction>(&serialized).unwrap(),
2029            instruction
2030        );
2031    }
2032
2033    #[test]
2034    fn test_packed_file_refs_manifest_round_trip_and_validation() {
2035        let region = RegionId::new(7, 3);
2036        let file_id = FileId::parse_str("00000000-0000-0000-0000-000000000001").unwrap();
2037        let mut manifest = FileRefsManifest::default();
2038        manifest.file_refs.insert(
2039            region,
2040            HashSet::from([
2041                FileRef::new(region, file_id, None),
2042                FileRef::new(region, file_id, Some(0)),
2043                FileRef::new(region, file_id, Some(u64::MAX)),
2044            ]),
2045        );
2046        manifest.manifest_version.insert(region, 42);
2047        manifest
2048            .cross_region_refs
2049            .insert(region, HashSet::from([RegionId::new(7, 4)]));
2050        let packed = PackedFileRefsManifest::from_manifest(&manifest);
2051        assert_eq!(packed.clone().into_manifest().unwrap(), manifest);
2052        let mut missing_files = serde_json::to_value(&packed).unwrap();
2053        missing_files
2054            .get_mut("file_refs")
2055            .and_then(serde_json::Value::as_object_mut)
2056            .and_then(|file_refs| file_refs.values_mut().next())
2057            .and_then(serde_json::Value::as_object_mut)
2058            .unwrap()
2059            .remove("files");
2060        assert!(serde_json::from_value::<PackedFileRefsManifest>(missing_files).is_err());
2061
2062        let mut missing_indexed = serde_json::to_value(&packed).unwrap();
2063        missing_indexed
2064            .get_mut("file_refs")
2065            .and_then(serde_json::Value::as_object_mut)
2066            .and_then(|file_refs| file_refs.values_mut().next())
2067            .and_then(serde_json::Value::as_object_mut)
2068            .unwrap()
2069            .remove("indexed");
2070        assert!(serde_json::from_value::<PackedFileRefsManifest>(missing_indexed).is_err());
2071
2072        let mut malformed = packed.clone();
2073        malformed.file_refs.get_mut(&region).unwrap().indexed = "!".to_string();
2074        let err = malformed.into_manifest().unwrap_err();
2075        assert!(matches!(
2076            &err,
2077            crate::error::Error::DecodePackedFileRefs { .. }
2078        ));
2079        assert!(
2080            err.source()
2081                .is_some_and(|source| source.is::<base64::DecodeError>())
2082        );
2083        assert_eq!(err.status_code(), StatusCode::Unexpected);
2084        assert_eq!(err.retry_hint(), RetryHint::NonRetryable);
2085
2086        let mut malformed = packed;
2087        malformed.file_refs.get_mut(&region).unwrap().indexed =
2088            base64::engine::general_purpose::STANDARD.encode([0; 1]);
2089        let err = malformed.into_manifest().unwrap_err();
2090        assert!(matches!(
2091            &err,
2092            crate::error::Error::InvalidPackedFileRefs { .. }
2093        ));
2094        assert!(err.source().is_none());
2095        assert_eq!(err.status_code(), StatusCode::Unexpected);
2096        assert_eq!(err.retry_hint(), RetryHint::NonRetryable);
2097    }
2098
2099    #[test]
2100    fn test_packed_file_refs_is_smaller_than_legacy() {
2101        let region = RegionId::new(7, 3);
2102        let mut manifest = FileRefsManifest::default();
2103        manifest.file_refs.insert(
2104            region,
2105            (0..500)
2106                .map(|i| {
2107                    let id =
2108                        FileId::parse_str(&format!("00000000-0000-0000-0000-{i:012}")).unwrap();
2109                    FileRef::new(region, id, None)
2110                })
2111                .collect(),
2112        );
2113        let legacy = serde_json::to_vec(&manifest).unwrap().len();
2114        let packed = serde_json::to_vec(&PackedFileRefsManifest::from_manifest(&manifest))
2115            .unwrap()
2116            .len();
2117        assert!(packed * 2 < legacy, "packed={packed}, legacy={legacy}");
2118    }
2119
2120    #[test]
2121    fn test_serialize_get_file_refs_instruction_reply() {
2122        let mut manifest = FileRefsManifest::default();
2123        let r0 = RegionId::new(1024, 1);
2124        let r1 = RegionId::new(1024, 2);
2125        manifest.file_refs.insert(
2126            r0,
2127            HashSet::from([FileRef::new(r0, FileId::random(), None)]),
2128        );
2129        manifest.file_refs.insert(
2130            r1,
2131            HashSet::from([FileRef::new(r1, FileId::random(), None)]),
2132        );
2133        manifest.manifest_version.insert(r0, 10);
2134        manifest.manifest_version.insert(r1, 20);
2135
2136        let instruction_reply = InstructionReply::GetFileRefs(GetFileRefsReply {
2137            file_refs_manifest: manifest,
2138            success: true,
2139            error: None,
2140        });
2141
2142        let serialized = serde_json::to_string(&instruction_reply).unwrap();
2143        let deserialized = serde_json::from_str(&serialized).unwrap();
2144
2145        assert_eq!(instruction_reply, deserialized);
2146    }
2147}