Remove revision var, update usage in FragmentManagerCore, and add new revisions.

Author: shanicky

src/meta/src/barrier/command.rs

@@ -40,7 +40,7 @@ use crate::manager::{FragmentManagerRef, WorkerId};
 use crate::model::{ActorId, DispatcherId, FragmentId, TableFragments};
 use crate::storage::MetaStore;
 use crate::stream::{
-    build_actor_connector_splits, RescheduleRevision, SourceManagerRef, SplitAssignment,
+    build_actor_connector_splits, SourceManagerRef, SplitAssignment, TableRevision,
 };
 use crate::MetaResult;
 
@@ -120,7 +120,6 @@ pub enum Command {
     /// very similar to `Create` and `Drop` commands, for added and removed actors, respectively.
     RescheduleFragment {
         reschedules: HashMap<FragmentId, Reschedule>,
-        revision: RescheduleRevision,
     },
 
     /// `ReplaceTable` command generates a `Update` barrier with the given `merge_updates`. This is
@@ -620,7 +619,6 @@ where
 
             Command::RescheduleFragment {
                 reschedules,
-                revision,
             } => {
                 let mut node_dropped_actors = HashMap::new();
                 for table_fragments in self.fragment_manager.list_table_fragments().await? {
@@ -647,7 +645,7 @@ where
 
                 // Update fragment info after rescheduling in meta store.
                 self.fragment_manager
-                    .post_apply_reschedules(reschedules.clone(), *revision)
+                    .post_apply_reschedules(reschedules.clone())
                     .await?;
 
                 let mut stream_source_actor_splits = HashMap::new();

src/meta/src/manager/catalog/fragment.rs

@@ -42,11 +42,12 @@ use crate::model::{
     ValTransaction,
 };
 use crate::storage::{MetaStore, Transaction};
-use crate::stream::{RescheduleRevision, SplitAssignment};
+use crate::stream::{SplitAssignment, TableRevision};
 use crate::MetaResult;
 
 pub struct FragmentManagerCore {
     table_fragments: BTreeMap<TableId, TableFragments>,
+    table_revision: TableRevision,
 }
 
 impl FragmentManagerCore {
@@ -102,9 +103,14 @@ where
             .map(|tf| (tf.table_id(), tf))
             .collect();
 
+        let table_revision = TableRevision::get(env.meta_store()).await?;
+
         Ok(Self {
             env,
-            core: RwLock::new(FragmentManagerCore { table_fragments }),
+            core: RwLock::new(FragmentManagerCore {
+                table_fragments,
+                table_revision,
+            }),
         })
     }
 
@@ -118,6 +124,10 @@ where
         Ok(map.values().cloned().collect())
     }
 
+    pub async fn get_revision(&self) -> TableRevision {
+        self.core.read().await.table_revision
+    }
+
     pub async fn has_any_table_fragments(&self) -> bool {
         !self.core.read().await.table_fragments.is_empty()
     }
@@ -347,7 +357,10 @@ where
     /// Drop table fragments info and remove downstream actor infos in fragments from its dependent
     /// tables.
     pub async fn drop_table_fragments_vec(&self, table_ids: &HashSet<TableId>) -> MetaResult<()> {
-        let map = &mut self.core.write().await.table_fragments;
+        let mut guard = self.core.write().await;
+        let current_revision = guard.table_revision;
+
+        let map = &mut guard.table_fragments;
         let to_delete_table_fragments = table_ids
             .iter()
             .filter_map(|table_id| map.get(table_id).cloned())
@@ -385,7 +398,19 @@ where
                     });
             }
         }
-        commit_meta!(self, table_fragments)?;
+
+        // new empty transaction
+        let mut trx = Transaction::default();
+
+        // save next revision
+        let next_revision = current_revision.next();
+        next_revision.store(&mut trx);
+
+        // commit
+        commit_meta_with_trx!(self, trx, table_fragments)?;
+
+        // update revision in memory
+        guard.table_revision = next_revision;
 
         for table_fragments in to_delete_table_fragments {
             if table_fragments.state() != State::Initial {
@@ -605,9 +630,11 @@ where
     pub async fn post_apply_reschedules(
         &self,
         mut reschedules: HashMap<FragmentId, Reschedule>,
-        revision: RescheduleRevision,
     ) -> MetaResult<()> {
-        let map = &mut self.core.write().await.table_fragments;
+        let mut guard = self.core.write().await;
+        let current_version = guard.table_revision;
+
+        let map = &mut guard.table_fragments;
 
         fn update_actors(
             actors: &mut Vec<ActorId>,
@@ -822,10 +849,20 @@ where
         }
 
         assert!(reschedules.is_empty(), "all reschedules must be applied");
+
+        // new empty transaction
         let mut trx = Transaction::default();
-        revision.store(&mut trx);
+
+        // save next revision
+        let next_revision = current_version.next();
+        next_revision.store(&mut trx);
+
+        // commit
         commit_meta_with_trx!(self, trx, table_fragments)?;
 
+        // update revision in memory
+        guard.table_revision = next_revision;
+
         for mapping in fragment_mapping_to_notify {
             self.env
                 .notification_manager()

src/meta/src/rpc/service/scale_service.rs

@@ -139,17 +139,16 @@ where
 
         let _guard = self.stream_manager.reschedule_revision_lock.lock().await;
 
-        let current_revision = self.stream_manager.reschedule_revision().await?;
+        let current_revision = self.fragment_manager.get_revision().await;
 
         if req.revision != current_revision.inner() {
             return Ok(Response::new(RescheduleResponse {
                 success: false,
-                revision: current_revision.into(),
+                revision: current_revision.inner(),
             }));
         }
 
-        let next_revision = self
-            .stream_manager
+        self.stream_manager
             .reschedule_actors(
                 req.reschedules
                     .into_iter()
@@ -176,10 +175,11 @@ where
                         )
                     })
                     .collect(),
-                current_revision,
             )
             .await?;
 
+        let next_revision = self.fragment_manager.get_revision().await;
+
         Ok(Response::new(RescheduleResponse {
             success: true,
             revision: next_revision.into(),

src/meta/src/stream/scale.rs

@@ -44,17 +44,17 @@ use crate::stream::GlobalStreamManager;
 use crate::{MetaError, MetaResult};
 
 #[derive(Copy, Clone, Debug)]
-pub struct RescheduleRevision(u64);
+pub struct TableRevision(u64);
 
 const RESCHEDULE_VERSION_KEY: &[u8] = b"reschedule_version";
 
-impl From<RescheduleRevision> for u64 {
-    fn from(value: RescheduleRevision) -> Self {
+impl From<TableRevision> for u64 {
+    fn from(value: TableRevision) -> Self {
         value.0
     }
 }
 
-impl RescheduleRevision {
+impl TableRevision {
     pub async fn get<S>(store: &S) -> MetaResult<Self>
     where
         S: MetaStore,
@@ -71,8 +71,8 @@ impl RescheduleRevision {
         Ok(Self(version))
     }
 
-    pub fn increase(&self) -> Self {
-        RescheduleRevision(self.0 + 1)
+    pub fn next(&self) -> Self {
+        TableRevision(self.0 + 1)
     }
 
     pub fn store(&self, txn: &mut Transaction) {
@@ -591,29 +591,26 @@ where
     pub async fn reschedule_actors(
         &self,
         reschedules: HashMap<FragmentId, ParallelUnitReschedule>,
-        current_revision: RescheduleRevision,
-    ) -> MetaResult<RescheduleRevision> {
+    ) -> MetaResult<()> {
         let mut revert_funcs = vec![];
-        match self
-            .reschedule_actors_impl(&mut revert_funcs, reschedules, current_revision)
+        if let Err(e) = self
+            .reschedule_actors_impl(&mut revert_funcs, reschedules)
             .await
         {
-            Err(e) => {
-                for revert_func in revert_funcs.into_iter().rev() {
-                    revert_func.await;
-                }
-                Err(e)
+            for revert_func in revert_funcs.into_iter().rev() {
+                revert_func.await;
             }
-            Ok(revision) => Ok(revision),
+            return Err(e);
         }
+
+        Ok(())
     }
 
     async fn reschedule_actors_impl(
         &self,
         revert_funcs: &mut Vec<BoxFuture<'_, ()>>,
         mut reschedules: HashMap<FragmentId, ParallelUnitReschedule>,
-        current_revision: RescheduleRevision,
-    ) -> MetaResult<RescheduleRevision> {
+    ) -> MetaResult<()> {
         let ctx = self.build_reschedule_context(&mut reschedules).await?;
         // Index of actors to create/remove
         // Fragment Id => ( Actor Id => Parallel Unit Id )
@@ -1175,16 +1172,13 @@ where
 
         tracing::debug!("reschedule plan: {:#?}", reschedule_fragment);
 
-        let next_revision = current_revision.increase();
-
         self.barrier_scheduler
             .run_command_with_paused(Command::RescheduleFragment {
                 reschedules: reschedule_fragment,
-                revision: next_revision,
             })
             .await?;
 
-        Ok(next_revision)
+        Ok(())
     }
 
     async fn create_actors_on_compute_node(

src/meta/src/stream/stream_manager.rs

@@ -34,7 +34,7 @@ use crate::hummock::HummockManagerRef;
 use crate::manager::{ClusterManagerRef, FragmentManagerRef, MetaSrvEnv};
 use crate::model::{ActorId, TableFragments};
 use crate::storage::MetaStore;
-use crate::stream::{RescheduleRevision, SourceManagerRef};
+use crate::stream::{TableRevision, SourceManagerRef};
 use crate::{MetaError, MetaResult};
 
 pub type GlobalStreamManagerRef<S> = Arc<GlobalStreamManager<S>>;
@@ -540,8 +540,8 @@ where
         self.creating_job_info.cancel_jobs(table_ids).await;
     }
 
-    pub async fn reschedule_revision(&self) -> MetaResult<RescheduleRevision> {
-        RescheduleRevision::get(self.env.meta_store()).await
+    pub async fn reschedule_revision(&self) -> MetaResult<TableRevision> {
+        TableRevision::get(self.env.meta_store()).await
     }
 }