feat(meta-srv): add repartition procedure skeleton (#7487)

Signed-off-by: WenyXu <[email protected]>

Author: WenyXu

src/meta-srv/src/error.rs

@@ -17,6 +17,7 @@ use common_error::ext::{BoxedError, ErrorExt};
 use common_error::status_code::StatusCode;
 use common_macro::stack_trace_debug;
 use common_meta::DatanodeId;
+use common_procedure::ProcedureId;
 use common_runtime::JoinError;
 use snafu::{Location, Snafu};
 use store_api::storage::RegionId;
@@ -768,6 +769,35 @@ pub enum Error {
         location: Location,
     },
 
+    #[snafu(display("Failed to create repartition subtasks"))]
+    RepartitionCreateSubtasks {
+        source: partition::error::Error,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display(
+        "Source partition expression '{}' does not match any existing region",
+        expr
+    ))]
+    RepartitionSourceExprMismatch {
+        expr: String,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
+    #[snafu(display(
+        "Failed to get the state receiver for repartition subprocedure {}",
+        procedure_id
+    ))]
+    RepartitionSubprocedureStateReceiver {
+        procedure_id: ProcedureId,
+        #[snafu(source)]
+        source: common_procedure::Error,
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display("Unsupported operation {}", operation))]
     Unsupported {
         operation: String,
@@ -1113,7 +1143,8 @@ impl ErrorExt for Error {
             | Error::LeaderPeerChanged { .. }
             | Error::RepartitionSourceRegionMissing { .. }
             | Error::RepartitionTargetRegionMissing { .. }
-            | Error::PartitionExprMismatch { .. } => StatusCode::InvalidArguments,
+            | Error::PartitionExprMismatch { .. }
+            | Error::RepartitionSourceExprMismatch { .. } => StatusCode::InvalidArguments,
             Error::LeaseKeyFromUtf8 { .. }
             | Error::LeaseValueFromUtf8 { .. }
             | Error::InvalidRegionKeyFromUtf8 { .. }
@@ -1173,6 +1204,8 @@ impl ErrorExt for Error {
 
             Error::BuildTlsOptions { source, .. } => source.status_code(),
             Error::Other { source, .. } => source.status_code(),
+            Error::RepartitionCreateSubtasks { source, .. } => source.status_code(),
+            Error::RepartitionSubprocedureStateReceiver { source, .. } => source.status_code(),
             Error::NoEnoughAvailableNode { .. } => StatusCode::RuntimeResourcesExhausted,
 
             #[cfg(feature = "pg_kvbackend")]

src/meta-srv/src/procedure/repartition.rs

@@ -12,8 +12,63 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+pub mod allocate_region;
+pub mod collect;
+pub mod deallocate_region;
+pub mod dispatch;
 pub mod group;
 pub mod plan;
+pub mod repartition_end;
+pub mod repartition_start;
+
+use std::any::Any;
+use std::fmt::Debug;
+
+use common_meta::cache_invalidator::CacheInvalidatorRef;
+use common_meta::key::TableMetadataManagerRef;
+use common_procedure::{Context as ProcedureContext, Status};
+use serde::{Deserialize, Serialize};
+use store_api::storage::TableId;
+
+use crate::error::Result;
+use crate::procedure::repartition::plan::RepartitionPlanEntry;
+use crate::service::mailbox::MailboxRef;
 
 #[cfg(test)]
 pub mod test_util;
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
+pub struct PersistentContext {
+    pub catalog_name: String,
+    pub schema_name: String,
+    pub table_name: String,
+    pub table_id: TableId,
+    pub plans: Vec<RepartitionPlanEntry>,
+}
+
+pub struct Context {
+    pub persistent_ctx: PersistentContext,
+    pub table_metadata_manager: TableMetadataManagerRef,
+    pub mailbox: MailboxRef,
+    pub server_addr: String,
+    pub cache_invalidator: CacheInvalidatorRef,
+}
+
+#[async_trait::async_trait]
+#[typetag::serde(tag = "repartition_state")]
+pub(crate) trait State: Sync + Send + Debug {
+    fn name(&self) -> &'static str {
+        let type_name = std::any::type_name::<Self>();
+        // short name
+        type_name.split("::").last().unwrap_or(type_name)
+    }
+
+    /// Yields the next [State] and [Status].
+    async fn next(
+        &mut self,
+        ctx: &mut Context,
+        procedure_ctx: &ProcedureContext,
+    ) -> Result<(Box<dyn State>, Status)>;
+
+    fn as_any(&self) -> &dyn Any;
+}

src/meta-srv/src/procedure/repartition/allocate_region.rs

@@ -0,0 +1,67 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::any::Any;
+
+use common_procedure::{Context as ProcedureContext, Status};
+use serde::{Deserialize, Serialize};
+
+use crate::error::Result;
+use crate::procedure::repartition::dispatch::Dispatch;
+use crate::procedure::repartition::plan::{AllocationPlanEntry, RepartitionPlanEntry};
+use crate::procedure::repartition::{Context, State};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct AllocateRegion {
+    plan_entries: Vec<AllocationPlanEntry>,
+}
+
+impl AllocateRegion {
+    pub fn new(plan_entries: Vec<AllocationPlanEntry>) -> Self {
+        Self { plan_entries }
+    }
+}
+
+#[async_trait::async_trait]
+#[typetag::serde]
+impl State for AllocateRegion {
+    async fn next(
+        &mut self,
+        ctx: &mut Context,
+        _procedure_ctx: &ProcedureContext,
+    ) -> Result<(Box<dyn State>, Status)> {
+        let region_to_allocate = self
+            .plan_entries
+            .iter()
+            .map(|p| p.regions_to_allocate)
+            .sum::<usize>();
+
+        if region_to_allocate == 0 {
+            let repartition_plan_entries = self
+                .plan_entries
+                .iter()
+                .map(RepartitionPlanEntry::from_allocation_plan_entry)
+                .collect::<Vec<_>>();
+            ctx.persistent_ctx.plans = repartition_plan_entries;
+            return Ok((Box::new(Dispatch), Status::executing(true)));
+        }
+
+        // TODO(weny): allocate regions.
+        todo!()
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}

src/meta-srv/src/procedure/repartition/collect.rs

@@ -0,0 +1,106 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::any::Any;
+
+use common_procedure::{Context as ProcedureContext, ProcedureId, Status, watcher};
+use common_telemetry::error;
+use serde::{Deserialize, Serialize};
+use snafu::ResultExt;
+
+use crate::error::{RepartitionSubprocedureStateReceiverSnafu, Result};
+use crate::procedure::repartition::deallocate_region::DeallocateRegion;
+use crate::procedure::repartition::group::GroupId;
+use crate::procedure::repartition::{Context, State};
+
+/// Metadata for tracking a dispatched sub-procedure.
+#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
+pub struct ProcedureMeta {
+    /// The index of the plan entry in the parent procedure's plan list.
+    pub plan_index: usize,
+    /// The group id of the repartition group.
+    pub group_id: GroupId,
+    /// The procedure id of the sub-procedure.
+    pub procedure_id: ProcedureId,
+}
+
+/// State for collecting results from dispatched sub-procedures.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Collect {
+    /// Sub-procedures that are currently in-flight.
+    pub inflight_procedures: Vec<ProcedureMeta>,
+    /// Sub-procedures that have completed successfully.
+    pub succeeded_procedures: Vec<ProcedureMeta>,
+    /// Sub-procedures that have failed.
+    pub failed_procedures: Vec<ProcedureMeta>,
+    /// Sub-procedures whose state could not be determined.
+    pub unknown_procedures: Vec<ProcedureMeta>,
+}
+
+impl Collect {
+    pub fn new(inflight_procedures: Vec<ProcedureMeta>) -> Self {
+        Self {
+            inflight_procedures,
+            succeeded_procedures: Vec::new(),
+            failed_procedures: Vec::new(),
+            unknown_procedures: Vec::new(),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+#[typetag::serde]
+impl State for Collect {
+    async fn next(
+        &mut self,
+        _ctx: &mut Context,
+        procedure_ctx: &ProcedureContext,
+    ) -> Result<(Box<dyn State>, Status)> {
+        for procedure_meta in self.inflight_procedures.iter() {
+            let procedure_id = procedure_meta.procedure_id;
+            let group_id = procedure_meta.group_id;
+            let Some(mut receiver) = procedure_ctx
+                .provider
+                .procedure_state_receiver(procedure_id)
+                .await
+                .context(RepartitionSubprocedureStateReceiverSnafu { procedure_id })?
+            else {
+                error!(
+                    "failed to get procedure state receiver, procedure_id: {}, group_id: {}",
+                    procedure_id, group_id
+                );
+                self.unknown_procedures.push(*procedure_meta);
+                continue;
+            };
+
+            match watcher::wait(&mut receiver).await {
+                Ok(_) => self.succeeded_procedures.push(*procedure_meta),
+                Err(e) => {
+                    error!(e; "failed to wait for repartition subprocedure, procedure_id: {}, group_id: {}", procedure_id, group_id);
+                    self.failed_procedures.push(*procedure_meta);
+                }
+            }
+        }
+
+        if !self.failed_procedures.is_empty() || !self.unknown_procedures.is_empty() {
+            // TODO(weny): retry the failed or unknown procedures.
+        }
+
+        Ok((Box::new(DeallocateRegion), Status::executing(true)))
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}

src/meta-srv/src/procedure/repartition/deallocate_region.rs

@@ -0,0 +1,52 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::any::Any;
+
+use common_procedure::{Context as ProcedureContext, Status};
+use serde::{Deserialize, Serialize};
+
+use crate::error::Result;
+use crate::procedure::repartition::repartition_end::RepartitionEnd;
+use crate::procedure::repartition::{Context, State};
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DeallocateRegion;
+
+#[async_trait::async_trait]
+#[typetag::serde]
+impl State for DeallocateRegion {
+    async fn next(
+        &mut self,
+        ctx: &mut Context,
+        _procedure_ctx: &ProcedureContext,
+    ) -> Result<(Box<dyn State>, Status)> {
+        let region_to_deallocate = ctx
+            .persistent_ctx
+            .plans
+            .iter()
+            .map(|p| p.pending_deallocate_region_ids.len())
+            .sum::<usize>();
+        if region_to_deallocate == 0 {
+            return Ok((Box::new(RepartitionEnd), Status::done()));
+        }
+
+        // TODO(weny): deallocate regions.
+        todo!()
+    }
+
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+}