feat: add update metadata step for downgrading leader region (#2771)

src/common/meta/src/key.rs

@@ -631,7 +631,7 @@ impl TableMetadataManager {
     pub async fn update_leader_region_status<F>(
         &self,
         table_id: TableId,
-        current_table_route_value: DeserializedValueWithBytes<TableRouteValue>,
+        current_table_route_value: &DeserializedValueWithBytes<TableRouteValue>,
         next_region_route_status: F,
     ) -> Result<()>
     where
@@ -658,7 +658,7 @@ impl TableMetadataManager {
 
         let (update_table_route_txn, on_update_table_route_failure) = self
             .table_route_manager()
-            .build_update_txn(table_id, &current_table_route_value, &new_table_route_value)?;
+            .build_update_txn(table_id, current_table_route_value, &new_table_route_value)?;
 
         let r = self.kv_backend.txn(update_table_route_txn).await?;
 
@@ -1094,7 +1094,7 @@ mod tests {
             .unwrap();
 
         table_metadata_manager
-            .update_leader_region_status(table_id, current_table_route_value, |region_route| {
+            .update_leader_region_status(table_id, &current_table_route_value, |region_route| {
                 if region_route.leader_status.is_some() {
                     None
                 } else {

src/meta-srv/src/procedure/region_failover/deactivate_region.rs

@@ -58,7 +58,7 @@ impl DeactivateRegion {
             .context(error::TableRouteNotFoundSnafu { table_id })?;
 
         ctx.table_metadata_manager
-            .update_leader_region_status(table_id, table_route_value, |region| {
+            .update_leader_region_status(table_id, &table_route_value, |region| {
                 if region.region.id.region_number() == failed_region.region_number {
                     Some(Some(RegionStatus::Downgraded))
                 } else {

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

@@ -18,23 +18,25 @@ pub(crate) mod migration_start;
 pub(crate) mod open_candidate_region;
 #[cfg(test)]
 pub(crate) mod test_util;
+pub(crate) mod update_metadata;
 
 use std::any::Any;
 use std::fmt::Debug;
 
-use common_meta::key::TableMetadataManagerRef;
+use common_meta::key::table_route::TableRouteValue;
+use common_meta::key::{DeserializedValueWithBytes, TableMetadataManagerRef};
 use common_meta::peer::Peer;
 use common_meta::ClusterId;
 use common_procedure::error::{
     Error as ProcedureError, FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu,
 };
 use common_procedure::{Context as ProcedureContext, LockKey, Procedure, Status};
 use serde::{Deserialize, Serialize};
-use snafu::ResultExt;
+use snafu::{location, Location, OptionExt, ResultExt};
 use store_api::storage::RegionId;
 
 use self::migration_start::RegionMigrationStart;
-use crate::error::{Error, Result};
+use crate::error::{self, Error, Result};
 use crate::procedure::utils::region_lock_key;
 use crate::region::lease_keeper::{OpeningRegionGuard, OpeningRegionKeeperRef};
 use crate::service::mailbox::MailboxRef;
@@ -74,8 +76,10 @@ pub struct VolatileContext {
     ///
     /// `opening_region_guard` should be consumed after
     /// the corresponding [RegionRoute](common_meta::rpc::router::RegionRoute) of the opening region
-    /// was written into [TableRouteValue](common_meta::key::table_route::TableRouteValue) .
+    /// was written into [TableRouteValue](common_meta::key::table_route::TableRouteValue).
     opening_region_guard: Option<OpeningRegionGuard>,
+    /// `table_route_info` is stored via previous steps for future use.
+    table_route_info: Option<DeserializedValueWithBytes<TableRouteValue>>,
 }
 
 /// Used to generate new [Context].
@@ -122,6 +126,47 @@ impl Context {
     pub fn server_addr(&self) -> &str {
         &self.server_addr
     }
+
+    /// Returns the `table_route_value` of [VolatileContext] if any.
+    /// Otherwise, returns the value retrieved from remote.
+    ///
+    /// Retry:
+    /// - Failed to retrieve the metadata of table.
+    pub async fn get_table_route_value(
+        &mut self,
+    ) -> Result<&DeserializedValueWithBytes<TableRouteValue>> {
+        let table_route_value = &mut self.volatile_ctx.table_route_info;
+
+        if table_route_value.is_none() {
+            let table_id = self.persistent_ctx.region_id.table_id();
+            let table_route = self
+                .table_metadata_manager
+                .table_route_manager()
+                .get(table_id)
+                .await
+                .context(error::TableMetadataManagerSnafu)
+                .map_err(|e| error::Error::RetryLater {
+                    reason: e.to_string(),
+                    location: location!(),
+                })?
+                .context(error::TableRouteNotFoundSnafu { table_id })?;
+
+            *table_route_value = Some(table_route);
+        }
+
+        Ok(table_route_value.as_ref().unwrap())
+    }
+
+    /// Removes the `table_route_value` of [VolatileContext], returns true if any.
+    pub fn remove_table_route_value(&mut self) -> bool {
+        let value = self.volatile_ctx.table_route_info.take();
+        value.is_some()
+    }
+
+    /// Returns the [RegionId].
+    pub fn region_id(&self) -> RegionId {
+        self.persistent_ctx.region_id
+    }
 }
 
 #[async_trait::async_trait]

src/meta-srv/src/procedure/region_migration/migration_start.rs

@@ -17,7 +17,7 @@ use std::any::Any;
 use common_meta::peer::Peer;
 use common_meta::rpc::router::RegionRoute;
 use serde::{Deserialize, Serialize};
-use snafu::{location, Location, OptionExt, ResultExt};
+use snafu::OptionExt;
 use store_api::storage::RegionId;
 
 use super::downgrade_leader_region::DowngradeLeaderRegion;
@@ -41,9 +41,8 @@ impl State for RegionMigrationStart {
     /// Otherwise go to the OpenCandidateRegion state.
     async fn next(&mut self, ctx: &mut Context) -> Result<Box<dyn State>> {
         let region_id = ctx.persistent_ctx.region_id;
-        let to_peer = &ctx.persistent_ctx.to_peer;
-
         let region_route = self.retrieve_region_route(ctx, region_id).await?;
+        let to_peer = &ctx.persistent_ctx.to_peer;
 
         if self.check_leader_region_on_peer(&region_route, to_peer)? {
             Ok(Box::new(RegionMigrationEnd))
@@ -70,21 +69,11 @@ impl RegionMigrationStart {
     /// - Failed to retrieve the metadata of table.
     async fn retrieve_region_route(
         &self,
-        ctx: &Context,
+        ctx: &mut Context,
         region_id: RegionId,
     ) -> Result<RegionRoute> {
         let table_id = region_id.table_id();
-        let table_route = ctx
-            .table_metadata_manager
-            .table_route_manager()
-            .get(table_id)
-            .await
-            .context(error::TableMetadataManagerSnafu)
-            .map_err(|e| error::Error::RetryLater {
-                reason: e.to_string(),
-                location: location!(),
-            })?
-            .context(error::TableRouteNotFoundSnafu { table_id })?;
+        let table_route = ctx.get_table_route_value().await?;
 
         let region_route = table_route
             .region_routes
@@ -165,10 +154,10 @@ mod tests {
         let state = RegionMigrationStart;
         let env = TestingEnv::new();
         let persistent_context = new_persistent_context();
-        let ctx = env.context_factory().new_context(persistent_context);
+        let mut ctx = env.context_factory().new_context(persistent_context);
 
         let err = state
-            .retrieve_region_route(&ctx, RegionId::new(1024, 1))
+            .retrieve_region_route(&mut ctx, RegionId::new(1024, 1))
             .await
             .unwrap_err();
 
@@ -184,7 +173,7 @@ mod tests {
         let from_peer = persistent_context.from_peer.clone();
 
         let env = TestingEnv::new();
-        let ctx = env.context_factory().new_context(persistent_context);
+        let mut ctx = env.context_factory().new_context(persistent_context);
 
         let table_info = new_test_table_info(1024, vec![1]).into();
         let region_route = RegionRoute {
@@ -199,7 +188,7 @@ mod tests {
             .unwrap();
 
         let err = state
-            .retrieve_region_route(&ctx, RegionId::new(1024, 3))
+            .retrieve_region_route(&mut ctx, RegionId::new(1024, 3))
             .await
             .unwrap_err();