meta_srv/procedure/repartition/
group.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
15pub(crate) mod apply_staging_manifest;
16pub(crate) mod enter_staging_region;
17pub(crate) mod remap_manifest;
18pub(crate) mod repartition_end;
19pub(crate) mod repartition_start;
20pub(crate) mod sync_region;
21pub(crate) mod update_metadata;
22pub(crate) mod utils;
23
24use std::any::Any;
25use std::collections::HashMap;
26use std::fmt::{Debug, Display};
27use std::time::{Duration, Instant};
28
29use common_error::ext::BoxedError;
30use common_meta::cache_invalidator::CacheInvalidatorRef;
31use common_meta::ddl::DdlContext;
32use common_meta::instruction::CacheIdent;
33use common_meta::key::datanode_table::{DatanodeTableValue, RegionInfo};
34use common_meta::key::table_route::TableRouteValue;
35use common_meta::key::{DeserializedValueWithBytes, TableMetadataManagerRef};
36use common_meta::lock_key::{CatalogLock, RegionLock, SchemaLock};
37use common_meta::peer::Peer;
38use common_meta::rpc::router::RegionRoute;
39use common_procedure::error::{FromJsonSnafu, ToJsonSnafu};
40use common_procedure::{
41    Context as ProcedureContext, Error as ProcedureError, EventContext, EventTrigger, LockKey,
42    Procedure, ProcedureId, Result as ProcedureResult, Status, StringKey,
43};
44use common_telemetry::{error, info};
45use serde::{Deserialize, Serialize};
46use snafu::{OptionExt, ResultExt};
47use store_api::storage::{RegionId, TableId};
48use uuid::Uuid;
49
50use crate::error::{self, Result};
51use crate::event::repartition::{REPARTITION_GROUP_EVENT_TYPE, RepartitionGroupEvent};
52use crate::procedure::repartition::group::repartition_start::RepartitionStart;
53use crate::procedure::repartition::plan::{SourceRegionDescriptor, TargetRegionDescriptor};
54use crate::procedure::repartition::utils::get_datanode_table_value;
55use crate::procedure::repartition::{self};
56use crate::service::mailbox::MailboxRef;
57
58#[derive(Debug, Clone, Default)]
59pub struct Metrics {
60    /// Elapsed time of flushing pending deallocate regions.
61    flush_pending_deallocate_regions_elapsed: Duration,
62    /// Elapsed time of entering staging region.
63    enter_staging_region_elapsed: Duration,
64    /// Elapsed time of applying staging manifest.
65    apply_staging_manifest_elapsed: Duration,
66    /// Elapsed time of remapping manifest.
67    remap_manifest_elapsed: Duration,
68    /// Elapsed time of updating metadata.
69    update_metadata_elapsed: Duration,
70}
71
72impl Display for Metrics {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        let total = self.flush_pending_deallocate_regions_elapsed
75            + self.enter_staging_region_elapsed
76            + self.apply_staging_manifest_elapsed
77            + self.remap_manifest_elapsed
78            + self.update_metadata_elapsed;
79        write!(f, "total: {:?}", total)?;
80        let mut parts = Vec::with_capacity(5);
81        if self.flush_pending_deallocate_regions_elapsed > Duration::ZERO {
82            parts.push(format!(
83                "flush_pending_deallocate_regions_elapsed: {:?}",
84                self.flush_pending_deallocate_regions_elapsed
85            ));
86        }
87        if self.enter_staging_region_elapsed > Duration::ZERO {
88            parts.push(format!(
89                "enter_staging_region_elapsed: {:?}",
90                self.enter_staging_region_elapsed
91            ));
92        }
93        if self.apply_staging_manifest_elapsed > Duration::ZERO {
94            parts.push(format!(
95                "apply_staging_manifest_elapsed: {:?}",
96                self.apply_staging_manifest_elapsed
97            ));
98        }
99        if self.remap_manifest_elapsed > Duration::ZERO {
100            parts.push(format!(
101                "remap_manifest_elapsed: {:?}",
102                self.remap_manifest_elapsed
103            ));
104        }
105        if self.update_metadata_elapsed > Duration::ZERO {
106            parts.push(format!(
107                "update_metadata_elapsed: {:?}",
108                self.update_metadata_elapsed
109            ));
110        }
111
112        if !parts.is_empty() {
113            write!(f, ", {}", parts.join(", "))?;
114        }
115        Ok(())
116    }
117}
118
119impl Metrics {
120    /// Updates the elapsed time of entering staging region.
121    pub fn update_enter_staging_region_elapsed(&mut self, elapsed: Duration) {
122        self.enter_staging_region_elapsed += elapsed;
123    }
124
125    pub fn update_flush_pending_deallocate_regions_elapsed(&mut self, elapsed: Duration) {
126        self.flush_pending_deallocate_regions_elapsed += elapsed;
127    }
128
129    /// Updates the elapsed time of applying staging manifest.
130    pub fn update_apply_staging_manifest_elapsed(&mut self, elapsed: Duration) {
131        self.apply_staging_manifest_elapsed += elapsed;
132    }
133
134    /// Updates the elapsed time of remapping manifest.
135    pub fn update_remap_manifest_elapsed(&mut self, elapsed: Duration) {
136        self.remap_manifest_elapsed += elapsed;
137    }
138
139    /// Updates the elapsed time of updating metadata.
140    pub fn update_update_metadata_elapsed(&mut self, elapsed: Duration) {
141        self.update_metadata_elapsed += elapsed;
142    }
143}
144
145pub type GroupId = Uuid;
146
147pub struct RepartitionGroupProcedure {
148    state: Box<dyn State>,
149    context: Context,
150}
151
152#[derive(Debug, Serialize)]
153struct RepartitionGroupData<'a> {
154    persistent_ctx: &'a PersistentContext,
155    state: &'a dyn State,
156}
157
158#[derive(Debug, Deserialize)]
159struct RepartitionGroupDataOwned {
160    persistent_ctx: PersistentContext,
161    state: Box<dyn State>,
162}
163
164impl RepartitionGroupProcedure {
165    pub(crate) const TYPE_NAME: &'static str = "metasrv-procedure::RepartitionGroup";
166
167    pub fn new(persistent_context: PersistentContext, context: &repartition::Context) -> Self {
168        let state = Box::new(RepartitionStart);
169
170        Self {
171            state,
172            context: Context {
173                persistent_ctx: persistent_context,
174                cache_invalidator: context.cache_invalidator.clone(),
175                table_metadata_manager: context.table_metadata_manager.clone(),
176                mailbox: context.mailbox.clone(),
177                server_addr: context.server_addr.clone(),
178                start_time: Instant::now(),
179                volatile_ctx: VolatileContext::default(),
180            },
181        }
182    }
183
184    pub fn from_json<F>(json: &str, ctx_factory: F) -> ProcedureResult<Self>
185    where
186        F: FnOnce(PersistentContext) -> Context,
187    {
188        let RepartitionGroupDataOwned {
189            state,
190            persistent_ctx,
191        } = serde_json::from_str(json).context(FromJsonSnafu)?;
192        let context = ctx_factory(persistent_ctx);
193
194        Ok(Self { state, context })
195    }
196}
197
198#[async_trait::async_trait]
199impl Procedure for RepartitionGroupProcedure {
200    fn type_name(&self) -> &str {
201        Self::TYPE_NAME
202    }
203
204    async fn rollback(&mut self, _ctx: &ProcedureContext) -> ProcedureResult<()> {
205        // The parent repartition procedure is responsible for rollback and recovery.
206        // Subprocedures are not recovered after metasrv restarts, so implementing rollback for them is meaningless.
207        Ok(())
208    }
209
210    #[tracing::instrument(skip_all, fields(
211        state = %self.state.name(),
212        table_id = self.context.persistent_ctx.table_id,
213        group_id = %self.context.persistent_ctx.group_id,
214    ))]
215    async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult<Status> {
216        let state = &mut self.state;
217        let state_name = state.name();
218        // Log state transition
219        common_telemetry::info!(
220            "Repartition group procedure executing state: {}, group id: {}, table id: {}",
221            state_name,
222            self.context.persistent_ctx.group_id,
223            self.context.persistent_ctx.table_id
224        );
225
226        match state.next(&mut self.context, _ctx).await {
227            Ok((next, status)) => {
228                *state = next;
229                Ok(status)
230            }
231            Err(e) => {
232                if e.is_retryable() {
233                    Err(ProcedureError::retry_later(e))
234                } else {
235                    error!(
236                        e;
237                        "Repartition group procedure failed, group id: {}, table id: {}",
238                        self.context.persistent_ctx.group_id,
239                        self.context.persistent_ctx.table_id,
240                    );
241                    Err(ProcedureError::external(e))
242                }
243            }
244        }
245    }
246
247    fn rollback_supported(&self) -> bool {
248        // Parent repartition owns rollback and recovery because subprocedures are
249        // not relied on as durable rollback units across metasrv restarts.
250        false
251    }
252
253    fn dump(&self) -> ProcedureResult<String> {
254        let data = RepartitionGroupData {
255            persistent_ctx: &self.context.persistent_ctx,
256            state: self.state.as_ref(),
257        };
258        serde_json::to_string(&data).context(ToJsonSnafu)
259    }
260
261    fn lock_key(&self) -> LockKey {
262        LockKey::new(self.context.persistent_ctx.lock_key())
263    }
264
265    fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn common_event_recorder::Event>> {
266        if !ctx.event_type_filter.allows(REPARTITION_GROUP_EVENT_TYPE) {
267            return None;
268        }
269
270        let event = if matches!(ctx.trigger, EventTrigger::Submitted) {
271            RepartitionGroupEvent::submitted(&self.context.persistent_ctx)
272        } else {
273            RepartitionGroupEvent::lifecycle(&self.context.persistent_ctx)
274        };
275        Some(Box::new(event))
276    }
277}
278
279pub struct Context {
280    pub persistent_ctx: PersistentContext,
281
282    pub cache_invalidator: CacheInvalidatorRef,
283
284    pub table_metadata_manager: TableMetadataManagerRef,
285
286    pub mailbox: MailboxRef,
287
288    pub server_addr: String,
289
290    pub start_time: Instant,
291
292    pub volatile_ctx: VolatileContext,
293}
294
295#[derive(Debug, Clone, Default)]
296pub struct VolatileContext {
297    pub metrics: Metrics,
298}
299
300impl Context {
301    pub fn new(
302        ddl_ctx: &DdlContext,
303        mailbox: MailboxRef,
304        server_addr: String,
305        persistent_ctx: PersistentContext,
306    ) -> Self {
307        Self {
308            persistent_ctx,
309            cache_invalidator: ddl_ctx.cache_invalidator.clone(),
310            table_metadata_manager: ddl_ctx.table_metadata_manager.clone(),
311            mailbox,
312            server_addr,
313            start_time: Instant::now(),
314            volatile_ctx: VolatileContext::default(),
315        }
316    }
317}
318
319/// The result of the group preparation phase, containing validated region routes.
320#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
321pub struct GroupPrepareResult {
322    /// The validated source region routes.
323    pub source_routes: Vec<RegionRoute>,
324    /// Validated target region routes used for metadata rollback (logical rollback).
325    pub target_routes: Vec<RegionRoute>,
326    /// The primary source region id (first source region), used for retrieving region options.
327    pub central_region: RegionId,
328    /// The peer where the primary source region is located.
329    pub central_region_datanode: Peer,
330}
331
332#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
333pub struct PersistentContext {
334    pub group_id: GroupId,
335    /// The parent repartition procedure id, when the group was created by a live parent.
336    #[serde(default)]
337    pub parent_procedure_id: Option<ProcedureId>,
338    /// The table id of the repartition group.
339    pub table_id: TableId,
340    /// The catalog name of the repartition group.
341    pub catalog_name: String,
342    /// The schema name of the repartition group.
343    pub schema_name: String,
344    /// The table name of the repartition group.
345    #[serde(default)]
346    pub table_name: Option<String>,
347    /// The source regions of the repartition group.
348    pub sources: Vec<SourceRegionDescriptor>,
349    /// The target regions of the repartition group.
350    pub targets: Vec<TargetRegionDescriptor>,
351    /// For each `source region`, the corresponding
352    /// `target regions` that overlap with it.
353    pub region_mapping: HashMap<RegionId, Vec<RegionId>>,
354    /// The result of group prepare.
355    /// The value will be set in [RepartitionStart](crate::procedure::repartition::group::repartition_start::RepartitionStart) state.
356    pub group_prepare_result: Option<GroupPrepareResult>,
357    /// The staging manifest paths of the repartition group.
358    /// The value will be set in [RemapManifest](crate::procedure::repartition::group::remap_manifest::RemapManifest) state.
359    pub staging_manifest_paths: HashMap<RegionId, String>,
360    /// Whether sync region is needed for this group.
361    pub sync_region: bool,
362    /// The region ids of the newly allocated regions.
363    pub allocated_region_ids: Vec<RegionId>,
364    /// The region ids of the regions that are pending deallocation.
365    pub pending_deallocate_region_ids: Vec<RegionId>,
366    /// The timeout for repartition operations.
367    #[serde(with = "humantime_serde")]
368    pub timeout: Duration,
369}
370
371impl PersistentContext {
372    #[allow(clippy::too_many_arguments)]
373    pub fn new(
374        group_id: GroupId,
375        parent_procedure_id: ProcedureId,
376        table_id: TableId,
377        catalog_name: String,
378        schema_name: String,
379        table_name: String,
380        sources: Vec<SourceRegionDescriptor>,
381        targets: Vec<TargetRegionDescriptor>,
382        region_mapping: HashMap<RegionId, Vec<RegionId>>,
383        sync_region: bool,
384        allocated_region_ids: Vec<RegionId>,
385        pending_deallocate_region_ids: Vec<RegionId>,
386        timeout: Duration,
387    ) -> Self {
388        Self {
389            group_id,
390            parent_procedure_id: Some(parent_procedure_id),
391            table_id,
392            catalog_name,
393            schema_name,
394            table_name: Some(table_name),
395            sources,
396            targets,
397            region_mapping,
398            group_prepare_result: None,
399            staging_manifest_paths: HashMap::new(),
400            sync_region,
401            allocated_region_ids,
402            pending_deallocate_region_ids,
403            timeout,
404        }
405    }
406
407    pub fn lock_key(&self) -> Vec<StringKey> {
408        let mut lock_keys = Vec::with_capacity(2 + self.sources.len());
409        lock_keys.extend([
410            CatalogLock::Read(&self.catalog_name).into(),
411            SchemaLock::read(&self.catalog_name, &self.schema_name).into(),
412        ]);
413        for source in &self.sources {
414            lock_keys.push(RegionLock::Write(source.region_id()).into());
415        }
416        lock_keys
417    }
418}
419
420impl Context {
421    /// Retrieves the table route value for the given table id.
422    ///
423    /// Retry:
424    /// - Failed to retrieve the metadata of table.
425    ///
426    /// Abort:
427    /// - Table route not found.
428    pub async fn get_table_route_value(
429        &self,
430    ) -> Result<DeserializedValueWithBytes<TableRouteValue>> {
431        let table_id = self.persistent_ctx.table_id;
432        let group_id = self.persistent_ctx.group_id;
433        let table_route_value = self
434            .table_metadata_manager
435            .table_route_manager()
436            .table_route_storage()
437            .get_with_raw_bytes(table_id)
438            .await
439            .map_err(BoxedError::new)
440            .with_context(|_| error::RetryLaterWithSourceSnafu {
441                reason: format!(
442                    "Failed to get table route for table: {}, repartition group: {}",
443                    table_id, group_id
444                ),
445            })?
446            .context(error::TableRouteNotFoundSnafu { table_id })?;
447
448        Ok(table_route_value)
449    }
450
451    /// Returns the `datanode_table_value`
452    ///
453    /// Retry:
454    /// - Failed to retrieve the metadata of datanode table.
455    pub async fn get_datanode_table_value(
456        &self,
457        table_id: TableId,
458        datanode_id: u64,
459    ) -> Result<DatanodeTableValue> {
460        get_datanode_table_value(&self.table_metadata_manager, table_id, datanode_id).await
461    }
462
463    /// Broadcasts the invalidate table cache message.
464    pub async fn invalidate_table_cache(&self) -> Result<()> {
465        let table_id = self.persistent_ctx.table_id;
466        let group_id = self.persistent_ctx.group_id;
467        let subject = format!(
468            "Invalidate table cache for repartition table, group: {}, table: {}",
469            group_id, table_id,
470        );
471        let ctx = common_meta::cache_invalidator::Context {
472            subject: Some(subject),
473        };
474        let _ = self
475            .cache_invalidator
476            .invalidate(&ctx, &[CacheIdent::TableId(table_id)])
477            .await;
478        Ok(())
479    }
480
481    /// Updates the table route.
482    ///
483    /// Retry:
484    /// - Failed to retrieve the metadata of datanode table.
485    ///
486    /// Abort:
487    /// - Table route not found.
488    /// - Failed to update the table route.
489    pub async fn update_table_route(
490        &self,
491        current_table_route_value: &DeserializedValueWithBytes<TableRouteValue>,
492        new_region_routes: Vec<RegionRoute>,
493    ) -> Result<()> {
494        let table_id = self.persistent_ctx.table_id;
495        let group_id = self.persistent_ctx.group_id;
496        // Safety: prepare result is set in [RepartitionStart] state.
497        let prepare_result = self.persistent_ctx.group_prepare_result.as_ref().unwrap();
498        let central_region_datanode_table_value = self
499            .get_datanode_table_value(table_id, prepare_result.central_region_datanode.id)
500            .await?;
501        let RegionInfo {
502            region_options,
503            region_wal_options,
504            ..
505        } = &central_region_datanode_table_value.region_info;
506
507        info!(
508            "Updating table route for table: {}, group_id: {}, new region routes: {:?}",
509            table_id, group_id, new_region_routes
510        );
511        self.table_metadata_manager
512            .update_table_route(
513                table_id,
514                central_region_datanode_table_value.region_info.clone(),
515                current_table_route_value,
516                new_region_routes,
517                region_options,
518                region_wal_options,
519            )
520            .await
521            .context(error::TableMetadataManagerSnafu)
522    }
523
524    /// Updates the table repart mapping.
525    pub async fn update_table_repart_mapping(&self) -> Result<()> {
526        info!(
527            "Updating table repart mapping for table: {}, group_id: {}, region mapping: {:?}",
528            self.persistent_ctx.table_id,
529            self.persistent_ctx.group_id,
530            self.persistent_ctx.region_mapping
531        );
532
533        self.table_metadata_manager
534            .table_repart_manager()
535            .update_mappings(
536                self.persistent_ctx.table_id,
537                &self.persistent_ctx.region_mapping,
538            )
539            .await
540            .context(error::TableMetadataManagerSnafu)
541    }
542
543    /// Returns the next operation timeout.
544    ///
545    /// If the next operation timeout is not set, it will return `None`.
546    pub fn next_operation_timeout(&self) -> Option<Duration> {
547        self.persistent_ctx
548            .timeout
549            .checked_sub(self.start_time.elapsed())
550    }
551
552    /// Updates the elapsed time of entering staging region.
553    pub fn update_enter_staging_region_elapsed(&mut self, elapsed: Duration) {
554        self.volatile_ctx
555            .metrics
556            .update_enter_staging_region_elapsed(elapsed);
557    }
558
559    /// Updates the elapsed time of flushing pending deallocate regions.
560    pub fn update_flush_pending_deallocate_regions_elapsed(&mut self, elapsed: Duration) {
561        self.volatile_ctx
562            .metrics
563            .update_flush_pending_deallocate_regions_elapsed(elapsed);
564    }
565
566    /// Updates the elapsed time of applying staging manifest.
567    pub fn update_apply_staging_manifest_elapsed(&mut self, elapsed: Duration) {
568        self.volatile_ctx
569            .metrics
570            .update_apply_staging_manifest_elapsed(elapsed);
571    }
572
573    /// Updates the elapsed time of remapping manifest.
574    pub fn update_remap_manifest_elapsed(&mut self, elapsed: Duration) {
575        self.volatile_ctx
576            .metrics
577            .update_remap_manifest_elapsed(elapsed);
578    }
579
580    /// Updates the elapsed time of updating metadata.
581    pub fn update_update_metadata_elapsed(&mut self, elapsed: Duration) {
582        self.volatile_ctx
583            .metrics
584            .update_update_metadata_elapsed(elapsed);
585    }
586}
587
588/// Returns the region routes of the given table route value.
589///
590/// Abort:
591/// - Table route value is not physical.
592pub fn region_routes(
593    table_id: TableId,
594    table_route_value: &TableRouteValue,
595) -> Result<&Vec<RegionRoute>> {
596    table_route_value
597        .region_routes()
598        .with_context(|_| error::UnexpectedLogicalRouteTableSnafu {
599            err_msg: format!(
600                "TableRoute({:?}) is a non-physical TableRouteValue.",
601                table_id
602            ),
603        })
604}
605
606#[async_trait::async_trait]
607#[typetag::serde(tag = "repartition_group_state")]
608pub(crate) trait State: Sync + Send + Debug {
609    fn name(&self) -> &'static str {
610        let type_name = std::any::type_name::<Self>();
611        // short name
612        type_name.split("::").last().unwrap_or(type_name)
613    }
614
615    /// Yields the next [State] and [Status].
616    async fn next(
617        &mut self,
618        ctx: &mut Context,
619        procedure_ctx: &ProcedureContext,
620    ) -> Result<(Box<dyn State>, Status)>;
621
622    fn as_any(&self) -> &dyn Any;
623}
624
625#[cfg(test)]
626mod tests {
627    use std::assert_matches;
628    use std::collections::HashSet;
629    use std::sync::Arc;
630
631    use common_event_recorder::EventTypeFilter;
632    use common_meta::key::TableMetadataManager;
633    use common_meta::kv_backend::test_util::MockKvBackendBuilder;
634    use common_procedure::{EventContext, EventTrigger, Procedure, ProcedureId, ProcedureState};
635
636    use crate::error::Error;
637    use crate::event::repartition::REPARTITION_GROUP_EVENT_TYPE;
638    use crate::procedure::repartition::group::repartition_start::RepartitionStart;
639    use crate::procedure::repartition::group::{PersistentContext, RepartitionGroupProcedure};
640    use crate::procedure::repartition::test_util::{TestingEnv, new_persistent_context};
641
642    #[tokio::test]
643    async fn test_get_table_route_value_not_found_error() {
644        let env = TestingEnv::new();
645        let persistent_context = new_persistent_context(1024, vec![], vec![]);
646        let ctx = env.create_context(persistent_context);
647        let err = ctx.get_table_route_value().await.unwrap_err();
648        assert_matches!(err, Error::TableRouteNotFound { .. });
649        assert!(!err.is_retryable());
650    }
651
652    #[tokio::test]
653    async fn test_get_table_route_value_retry_error() {
654        let kv = MockKvBackendBuilder::default()
655            .range_fn(Arc::new(|_| {
656                common_meta::error::UnexpectedSnafu {
657                    err_msg: "mock err",
658                }
659                .fail()
660            }))
661            .build()
662            .unwrap();
663        let mut env = TestingEnv::new();
664        env.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(kv)));
665        let persistent_context = new_persistent_context(1024, vec![], vec![]);
666        let ctx = env.create_context(persistent_context);
667        let err = ctx.get_table_route_value().await.unwrap_err();
668        assert!(err.is_retryable());
669    }
670
671    #[tokio::test]
672    async fn test_get_datanode_table_value_retry_error() {
673        let kv = MockKvBackendBuilder::default()
674            .range_fn(Arc::new(|_| {
675                common_meta::error::UnexpectedSnafu {
676                    err_msg: "mock err",
677                }
678                .fail()
679            }))
680            .build()
681            .unwrap();
682        let mut env = TestingEnv::new();
683        env.table_metadata_manager = Arc::new(TableMetadataManager::new(Arc::new(kv)));
684        let persistent_context = new_persistent_context(1024, vec![], vec![]);
685        let ctx = env.create_context(persistent_context);
686        let err = ctx.get_datanode_table_value(1024, 1).await.unwrap_err();
687        assert!(err.is_retryable());
688    }
689
690    #[test]
691    fn test_persistent_context_is_backward_compatible_without_event_fields() {
692        let persistent_context = new_persistent_context(1024, vec![], vec![]);
693        let mut serialized = serde_json::to_value(persistent_context).unwrap();
694        let object = serialized.as_object_mut().unwrap();
695        object.remove("parent_procedure_id");
696        object.remove("table_name");
697
698        let deserialized: PersistentContext = serde_json::from_value(serialized).unwrap();
699        assert_eq!(deserialized.parent_procedure_id, None);
700        assert_eq!(deserialized.table_name, None);
701    }
702
703    #[test]
704    fn test_repartition_group_event_filter() {
705        let env = TestingEnv::new();
706        let procedure = RepartitionGroupProcedure {
707            state: Box::new(RepartitionStart),
708            context: env.create_context(new_persistent_context(1024, vec![], vec![])),
709        };
710        let state = ProcedureState::Running;
711        let event_context = |event_type_filter| EventContext {
712            procedure_id: ProcedureId::random(),
713            lifecycle_state: &state,
714            trigger: EventTrigger::Submitted,
715            event_type_filter: Arc::new(event_type_filter),
716            event_context: None,
717        };
718
719        let allowed = procedure
720            .event(&event_context(EventTypeFilter::Only(HashSet::from([
721                REPARTITION_GROUP_EVENT_TYPE.to_string(),
722            ]))))
723            .unwrap();
724        assert_eq!(allowed.event_type(), REPARTITION_GROUP_EVENT_TYPE);
725
726        assert!(
727            procedure
728                .event(&event_context(EventTypeFilter::Only(HashSet::from([
729                    "another_event".to_string(),
730                ]))))
731                .is_none()
732        );
733        assert!(
734            procedure
735                .event(&event_context(EventTypeFilter::Only(HashSet::new())))
736                .is_none()
737        );
738    }
739}
meta_srv/procedure/repartition/group.rs

meta_srv/procedure/repartition/
group.rs
