refactor(meta): unify code for different kinds of source change (#19991)

Author: xxchan

src/meta/src/barrier/command.rs

@@ -307,7 +307,7 @@ pub enum Command {
 
     /// `ReplaceStreamJob` command generates a `Update` barrier with the given `merge_updates`. This is
     /// essentially switching the downstream of the old job fragments to the new ones, and
-    /// dropping the old job fragments. Used for table schema change.
+    /// dropping the old job fragments. Used for schema change.
     ///
     /// This can be treated as a special case of `RescheduleFragment`, while the upstream fragment
     /// of the Merge executors are changed additionally.

src/meta/src/barrier/context/context_impl.rs

@@ -33,6 +33,7 @@ use crate::barrier::{
     Scheduled,
 };
 use crate::hummock::CommitEpochInfo;
+use crate::stream::SourceChange;
 use crate::{MetaError, MetaResult};
 
 impl GlobalBarrierWorkerContext for GlobalBarrierWorkerContextImpl {
@@ -166,7 +167,7 @@ impl CommandContext {
                     .await?;
                 barrier_manager_context
                     .source_manager
-                    .apply_source_change(None, None, Some(split_assignment.clone()), None, None)
+                    .apply_source_change(SourceChange::SplitChange(split_assignment.clone()))
                     .await;
             }
 
@@ -180,17 +181,18 @@ impl CommandContext {
                     .await?;
             }
             Command::CreateStreamingJob { info, job_type } => {
-                let mut fragment_replacements = None;
-                let mut dropped_actors = None;
+                let mut is_sink_into_table = false;
                 match job_type {
                     CreateStreamingJobType::SinkIntoTable(
                         replace_plan @ ReplaceStreamJobPlan {
+                            old_fragments,
                             new_fragments,
                             dispatchers,
                             init_split_assignment,
                             ..
                         },
                     ) => {
+                        is_sink_into_table = true;
                         barrier_manager_context
                             .metadata_manager
                             .catalog_controller
@@ -201,8 +203,15 @@ impl CommandContext {
                                 init_split_assignment,
                             )
                             .await?;
-                        fragment_replacements = Some(replace_plan.fragment_replacements());
-                        dropped_actors = Some(replace_plan.dropped_actors());
+                        barrier_manager_context
+                            .source_manager
+                            .handle_replace_job(
+                                old_fragments,
+                                new_fragments.stream_source_fragments(),
+                                init_split_assignment.clone(),
+                                replace_plan.fragment_replacements(),
+                            )
+                            .await;
                     }
                     CreateStreamingJobType::Normal => {}
                     CreateStreamingJobType::SnapshotBackfill(snapshot_backfill_info) => {
@@ -247,19 +256,17 @@ impl CommandContext {
                     )
                     .await?;
 
-                // Extract the fragments that include source operators.
-                let source_fragments = stream_job_fragments.stream_source_fragments();
-                let backfill_fragments = stream_job_fragments.source_backfill_fragments()?;
-                barrier_manager_context
-                    .source_manager
-                    .apply_source_change(
-                        Some(source_fragments),
-                        Some(backfill_fragments),
-                        Some(init_split_assignment.clone()),
-                        dropped_actors,
-                        fragment_replacements,
-                    )
-                    .await;
+                if !is_sink_into_table {
+                    barrier_manager_context
+                        .source_manager
+                        .apply_source_change(SourceChange::CreateJob {
+                            added_source_fragments: stream_job_fragments.stream_source_fragments(),
+                            added_backfill_fragments: stream_job_fragments
+                                .source_backfill_fragments()?,
+                            split_assignment: init_split_assignment.clone(),
+                        })
+                        .await;
+                }
             }
             Command::RescheduleFragment {
                 reschedules,
@@ -296,19 +303,11 @@ impl CommandContext {
                 // Apply the split changes in source manager.
                 barrier_manager_context
                     .source_manager
-                    .drop_source_fragments_vec(std::slice::from_ref(old_fragments))
-                    .await;
-                let source_fragments = new_fragments.stream_source_fragments();
-                // XXX: is it possible to have backfill fragments here?
-                let backfill_fragments = new_fragments.source_backfill_fragments()?;
-                barrier_manager_context
-                    .source_manager
-                    .apply_source_change(
-                        Some(source_fragments),
-                        Some(backfill_fragments),
-                        Some(init_split_assignment.clone()),
-                        Some(replace_plan.dropped_actors()),
-                        Some(replace_plan.fragment_replacements()),
+                    .handle_replace_job(
+                        old_fragments,
+                        new_fragments.stream_source_fragments(),
+                        init_split_assignment.clone(),
+                        replace_plan.fragment_replacements(),
                     )
                     .await;
             }

src/meta/src/barrier/context/recovery.rs

@@ -35,7 +35,7 @@ use crate::barrier::{DatabaseRuntimeInfoSnapshot, InflightSubscriptionInfo};
 use crate::controller::fragment::InflightFragmentInfo;
 use crate::manager::ActiveStreamingWorkerNodes;
 use crate::model::{ActorId, StreamJobFragments, TableParallelism};
-use crate::stream::{RescheduleOptions, TableResizePolicy};
+use crate::stream::{RescheduleOptions, SourceChange, TableResizePolicy};
 use crate::{model, MetaResult};
 
 impl GlobalBarrierWorkerContextImpl {
@@ -46,14 +46,18 @@ impl GlobalBarrierWorkerContextImpl {
             .catalog_controller
             .clean_dirty_subscription(database_id)
             .await?;
-        let source_ids = self
+        let dirty_associated_source_ids = self
             .metadata_manager
             .catalog_controller
             .clean_dirty_creating_jobs(database_id)
             .await?;
 
         // unregister cleaned sources.
-        self.source_manager.unregister_sources(source_ids).await;
+        self.source_manager
+            .apply_source_change(SourceChange::DropSource {
+                dropped_source_ids: dirty_associated_source_ids,
+            })
+            .await;
 
         Ok(())
     }

src/meta/src/controller/catalog/mod.rs

@@ -364,7 +364,9 @@ impl CatalogController {
             })
             .collect_vec();
 
-        let associated_source_ids: Vec<SourceId> = Table::find()
+        // The source ids for dirty tables with connector.
+        // FIXME: we should also clean dirty sources.
+        let dirty_associated_source_ids: Vec<SourceId> = Table::find()
             .select_only()
             .column(table::Column::OptionalAssociatedSourceId)
             .filter(
@@ -402,7 +404,7 @@ impl CatalogController {
             .clone()
             .into_iter()
             .chain(dirty_state_table_ids.into_iter())
-            .chain(associated_source_ids.clone().into_iter())
+            .chain(dirty_associated_source_ids.clone().into_iter())
             .collect();
 
         let res = Object::delete_many()
@@ -424,7 +426,7 @@ impl CatalogController {
             .notify_frontend(NotificationOperation::Delete, relation_group)
             .await;
 
-        Ok(associated_source_ids)
+        Ok(dirty_associated_source_ids)
     }
 
     pub async fn comment_on(&self, comment: PbComment) -> MetaResult<NotificationVersion> {

src/meta/src/rpc/ddl_controller.rs

@@ -77,7 +77,8 @@ use crate::model::{StreamContext, StreamJobFragments, TableParallelism};
 use crate::stream::{
     create_source_worker_handle, validate_sink, ActorGraphBuildResult, ActorGraphBuilder,
     CompleteStreamFragmentGraph, CreateStreamingJobContext, CreateStreamingJobOption,
-    GlobalStreamManagerRef, ReplaceStreamJobContext, SourceManagerRef, StreamFragmentGraph,
+    GlobalStreamManagerRef, ReplaceStreamJobContext, SourceChange, SourceManagerRef,
+    StreamFragmentGraph,
 };
 use crate::{MetaError, MetaResult};
 
@@ -974,7 +975,9 @@ impl DdlController {
                     tracing::warn!(id = job_id, "aborted streaming job");
                     if let Some(source_id) = source_id {
                         self.source_manager
-                            .unregister_sources(vec![source_id as SourceId])
+                            .apply_source_change(SourceChange::DropSource {
+                                dropped_source_ids: vec![source_id as SourceId],
+                            })
                             .await;
                     }
                 }
@@ -1263,23 +1266,28 @@ impl DdlController {
 
         // unregister sources.
         self.source_manager
-            .unregister_sources(removed_source_ids.into_iter().map(|id| id as _).collect())
+            .apply_source_change(SourceChange::DropSource {
+                dropped_source_ids: removed_source_ids.into_iter().map(|id| id as _).collect(),
+            })
             .await;
 
         // unregister fragments and actors from source manager.
+        // FIXME: need also unregister source backfill fragments.
+        let dropped_source_fragments = removed_source_fragments
+            .into_iter()
+            .map(|(source_id, fragments)| {
+                (
+                    source_id,
+                    fragments.into_iter().map(|id| id as u32).collect(),
+                )
+            })
+            .collect();
+        let dropped_actors = removed_actors.iter().map(|id| *id as _).collect();
         self.source_manager
-            .drop_source_fragments(
-                removed_source_fragments
-                    .into_iter()
-                    .map(|(source_id, fragments)| {
-                        (
-                            source_id,
-                            fragments.into_iter().map(|id| id as u32).collect(),
-                        )
-                    })
-                    .collect(),
-                removed_actors.iter().map(|id| *id as _).collect(),
-            )
+            .apply_source_change(SourceChange::DropMv {
+                dropped_source_fragments,
+                dropped_actors,
+            })
             .await;
 
         // drop streaming jobs.

src/meta/src/stream/scale.rs

@@ -140,6 +140,7 @@ impl CustomFragmentInfo {
 
 use educe::Educe;
 
+use super::SourceChange;
 use crate::controller::id::IdCategory;
 
 // The debug implementation is arbitrary. Just used in debug logs.
@@ -1769,13 +1770,10 @@ impl ScaleController {
 
         if !stream_source_actor_splits.is_empty() {
             self.source_manager
-                .apply_source_change(
-                    None,
-                    None,
-                    Some(stream_source_actor_splits),
-                    Some(stream_source_dropped_actors),
-                    None,
-                )
+                .apply_source_change(SourceChange::Reschedule {
+                    split_assignment: stream_source_actor_splits,
+                    dropped_actors: stream_source_dropped_actors,
+                })
                 .await;
         }