Skip to main content

risingwave_meta/barrier/checkpoint/
state.rs

1// Copyright 2024 RisingWave Labs
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
15use std::assert_matches;
16use std::collections::hash_map::Entry;
17use std::collections::{HashMap, HashSet};
18use std::mem::take;
19use std::sync::atomic::AtomicU32;
20
21use risingwave_common::bail;
22use risingwave_common::bitmap::Bitmap;
23use risingwave_common::catalog::TableId;
24use risingwave_common::hash::VnodeCountCompat;
25use risingwave_common::id::JobId;
26use risingwave_common::util::epoch::Epoch;
27use risingwave_meta_model::fragment::DistributionType;
28use risingwave_meta_model::{DispatcherType, WorkerId, streaming_job};
29use risingwave_pb::common::WorkerNode;
30use risingwave_pb::hummock::HummockVersionStats;
31use risingwave_pb::source::{ConnectorSplit, ConnectorSplits};
32use risingwave_pb::stream_plan::barrier_mutation::{Mutation, PbMutation};
33use risingwave_pb::stream_plan::throttle_mutation::ThrottleConfig;
34use risingwave_pb::stream_plan::update_mutation::PbDispatcherUpdate;
35use risingwave_pb::stream_plan::{
36    AddMutation, PbStartFragmentBackfillMutation, PbStreamNode, PbSubscriptionUpstreamInfo,
37    PbUpdateMutation, PbUpstreamSinkInfo, ThrottleMutation,
38};
39use tracing::warn;
40
41use crate::barrier::cdc_progress::CdcTableBackfillTracker;
42use crate::barrier::checkpoint::{
43    BatchRefreshJobCheckpointControl, BatchRefreshLogicalFragments, CreatingStreamingJobControl,
44    DatabaseCheckpointControl, IndependentCheckpointJobControl,
45};
46use crate::barrier::command::{CreateStreamingJobCommandInfo, PostCollectCommand, ReschedulePlan};
47use crate::barrier::context::CreateSnapshotBackfillJobCommandInfo;
48use crate::barrier::edge_builder::{EdgeBuilderFragmentInfo, FragmentEdgeBuilder};
49use crate::barrier::info::{
50    BarrierInfo, CreateStreamingJobStatus, InflightDatabaseInfo, InflightStreamingJobInfo,
51    SubscriberType,
52};
53use crate::barrier::notifier::Notifier;
54use crate::barrier::partial_graph::{PartialGraphBarrierInfo, PartialGraphManager};
55use crate::barrier::rpc::to_partial_graph_id;
56use crate::barrier::{BarrierKind, Command, CreateStreamingJobType, TracedEpoch};
57use crate::controller::fragment::{InflightActorInfo, InflightFragmentInfo};
58use crate::controller::scale::{
59    ComponentFragmentAligner, EnsembleActorTemplate, LoadedFragment, NoShuffleEnsemble,
60    build_no_shuffle_fragment_graph_edges, find_no_shuffle_graphs,
61};
62use crate::model::{
63    ActorId, ActorNewNoShuffle, FragmentDownstreamRelation, FragmentId, StreamActor, StreamContext,
64    StreamJobActorsToCreate, StreamJobFragmentsToCreate,
65};
66use crate::stream::cdc::parallel_cdc_table_backfill_fragment;
67use crate::stream::{
68    GlobalActorIdGen, ReplaceJobSplitPlan, SourceManager, SplitAssignment,
69    fill_snapshot_backfill_epoch,
70};
71use crate::{MetaError, MetaResult};
72
73type ThrottleForCreatingJobs = (
74    HashSet<JobId>,
75    HashMap<FragmentId, (ThrottleConfig, PbStreamNode)>,
76);
77
78/// The latest state of `GlobalBarrierWorker` after injecting the latest barrier.
79pub(in crate::barrier) struct BarrierWorkerState {
80    /// The last sent `prev_epoch`
81    ///
82    /// There's no need to persist this field. On recovery, we will restore this from the latest
83    /// committed snapshot in `HummockManager`.
84    in_flight_prev_epoch: TracedEpoch,
85
86    /// The `prev_epoch` of pending non checkpoint barriers
87    pending_non_checkpoint_barriers: Vec<u64>,
88
89    /// Whether the cluster is paused.
90    is_paused: bool,
91}
92
93impl BarrierWorkerState {
94    pub(super) fn new() -> Self {
95        Self {
96            in_flight_prev_epoch: TracedEpoch::new(Epoch::now()),
97            pending_non_checkpoint_barriers: vec![],
98            is_paused: false,
99        }
100    }
101
102    pub fn recovery(in_flight_prev_epoch: TracedEpoch, is_paused: bool) -> Self {
103        Self {
104            in_flight_prev_epoch,
105            pending_non_checkpoint_barriers: vec![],
106            is_paused,
107        }
108    }
109
110    pub fn is_paused(&self) -> bool {
111        self.is_paused
112    }
113
114    fn set_is_paused(&mut self, is_paused: bool) {
115        if self.is_paused != is_paused {
116            tracing::info!(
117                currently_paused = self.is_paused,
118                newly_paused = is_paused,
119                "update paused state"
120            );
121            self.is_paused = is_paused;
122        }
123    }
124
125    pub fn in_flight_prev_epoch(&self) -> &TracedEpoch {
126        &self.in_flight_prev_epoch
127    }
128
129    /// Returns the `BarrierInfo` for the next barrier, and updates the state.
130    pub fn next_barrier_info(
131        &mut self,
132        is_checkpoint: bool,
133        curr_epoch: TracedEpoch,
134    ) -> BarrierInfo {
135        assert!(
136            self.in_flight_prev_epoch.value() < curr_epoch.value(),
137            "curr epoch regress. {} > {}",
138            self.in_flight_prev_epoch.value(),
139            curr_epoch.value()
140        );
141        let prev_epoch = self.in_flight_prev_epoch.clone();
142        self.in_flight_prev_epoch = curr_epoch.clone();
143        self.pending_non_checkpoint_barriers
144            .push(prev_epoch.value().0);
145        let kind = if is_checkpoint {
146            let epochs = take(&mut self.pending_non_checkpoint_barriers);
147            BarrierKind::Checkpoint(epochs)
148        } else {
149            BarrierKind::Barrier
150        };
151        BarrierInfo {
152            prev_epoch,
153            curr_epoch,
154            kind,
155        }
156    }
157}
158
159pub(super) struct ApplyCommandInfo {
160    pub jobs_to_wait: HashSet<JobId>,
161}
162
163/// Result tuple of `apply_command`: mutation, table IDs to commit, actors to create,
164/// node actors, and post-collect command.
165type ApplyCommandResult = (
166    Option<Mutation>,
167    HashSet<TableId>,
168    Option<StreamJobActorsToCreate>,
169    HashMap<WorkerId, HashSet<ActorId>>,
170    PostCollectCommand,
171);
172
173/// Result of actor rendering for a create/replace streaming job.
174pub(crate) struct RenderResult {
175    /// Rendered actors grouped by fragment.
176    pub stream_actors: HashMap<FragmentId, Vec<StreamActor>>,
177    /// Worker placement for each actor.
178    pub actor_location: HashMap<ActorId, WorkerId>,
179}
180
181/// Derive `NoShuffle` edges from fragment downstream relations and resolve ensembles.
182///
183/// This scans both the internal downstream relations (`fragments.downstreams`) and
184/// the cross-boundary upstream-to-new-fragment relations (`upstream_fragment_downstreams`)
185/// to find all `NoShuffle` edges. It then runs BFS to find connected components (ensembles)
186/// and categorizes them into:
187/// - Ensembles whose entry fragments include existing (non-new) fragments
188/// - Ensembles whose entry fragments are all newly created
189pub(crate) fn resolve_no_shuffle_ensembles(
190    fragments: &StreamJobFragmentsToCreate,
191    upstream_fragment_downstreams: &FragmentDownstreamRelation,
192) -> MetaResult<Vec<NoShuffleEnsemble>> {
193    // Derive FragmentNewNoShuffle from the two downstream relation maps.
194    let mut new_no_shuffle: HashMap<_, HashSet<_>> = HashMap::new();
195
196    // Internal edges (new → new) and edges from new → existing downstream (replace job).
197    for (upstream_fid, relations) in &fragments.downstreams {
198        for rel in relations {
199            if rel.dispatcher_type == DispatcherType::NoShuffle {
200                new_no_shuffle
201                    .entry(*upstream_fid)
202                    .or_default()
203                    .insert(rel.downstream_fragment_id);
204            }
205        }
206    }
207
208    // Cross-boundary edges: existing upstream → new downstream.
209    for (upstream_fid, relations) in upstream_fragment_downstreams {
210        for rel in relations {
211            if rel.dispatcher_type == DispatcherType::NoShuffle {
212                new_no_shuffle
213                    .entry(*upstream_fid)
214                    .or_default()
215                    .insert(rel.downstream_fragment_id);
216            }
217        }
218    }
219
220    let mut ensembles = if new_no_shuffle.is_empty() {
221        Vec::new()
222    } else {
223        // Flatten into directed edge pairs for BFS.
224        let no_shuffle_edges: Vec<(FragmentId, FragmentId)> = new_no_shuffle
225            .iter()
226            .flat_map(|(upstream_fid, downstream_fids)| {
227                downstream_fids
228                    .iter()
229                    .map(move |downstream_fid| (*upstream_fid, *downstream_fid))
230            })
231            .collect();
232
233        let all_fragment_ids: Vec<FragmentId> = no_shuffle_edges
234            .iter()
235            .flat_map(|(u, d)| [*u, *d])
236            .collect::<HashSet<_>>()
237            .into_iter()
238            .collect();
239
240        let (fwd, bwd) = build_no_shuffle_fragment_graph_edges(no_shuffle_edges);
241        find_no_shuffle_graphs(&all_fragment_ids, &fwd, &bwd)?
242    };
243
244    // Add standalone fragments (not covered by any ensemble) as single-fragment ensembles.
245    let covered: HashSet<FragmentId> = ensembles
246        .iter()
247        .flat_map(|e| e.component_fragments())
248        .collect();
249    for fragment_id in fragments.inner.fragments.keys() {
250        if !covered.contains(fragment_id) {
251            ensembles.push(NoShuffleEnsemble::singleton(*fragment_id));
252        }
253    }
254
255    Ok(ensembles)
256}
257
258/// Render actors for a create or replace streaming job.
259///
260/// This determines the parallelism for each no-shuffle ensemble (either from an existing
261/// inflight upstream or computed fresh), and produces `StreamActor` instances with worker
262/// placements and actor-level no-shuffle mappings.
263///
264/// The process follows three steps:
265/// 1. For each ensemble, resolve `EnsembleActorTemplate` (from existing or fresh).
266/// 2. For each new component fragment, allocate actor IDs and compute worker/vnode assignments.
267/// 3. Expand the simple assignments into full `StreamActor` structures.
268pub(super) fn render_actors(
269    fragments: &StreamJobFragmentsToCreate,
270    database_info: &InflightDatabaseInfo,
271    definition: &str,
272    ctx: &StreamContext,
273    streaming_job_model: &streaming_job::Model,
274    actor_id_counter: &AtomicU32,
275    worker_map: &HashMap<WorkerId, WorkerNode>,
276    ensembles: &[NoShuffleEnsemble],
277    database_resource_group: &str,
278) -> MetaResult<RenderResult> {
279    // Step 2: Render actors for each ensemble.
280    // For each new fragment, produce a simple assignment: actor_id -> (worker_id, vnode_bitmap).
281    let mut actor_assignments: HashMap<FragmentId, HashMap<ActorId, (WorkerId, Option<Bitmap>)>> =
282        HashMap::new();
283
284    for ensemble in ensembles {
285        // Determine the EnsembleActorTemplate for this ensemble.
286        //
287        // Check if any component fragment in the ensemble already exists (i.e. is inflight).
288        // If so, derive the actor assignment from an existing fragment. Otherwise render fresh.
289        let existing_fragment_ids: Vec<FragmentId> = ensemble
290            .component_fragments()
291            .filter(|fragment_id| !fragments.inner.fragments.contains_key(fragment_id))
292            .collect();
293
294        let actor_template = if let Some(&first_existing) = existing_fragment_ids.first() {
295            let template = EnsembleActorTemplate::from_existing_inflight_fragment(
296                database_info.fragment(first_existing),
297            );
298
299            // Sanity check: all existing fragments in the same ensemble must be aligned —
300            // same actor count and same worker placement per vnode.
301            for &other_fragment_id in &existing_fragment_ids[1..] {
302                let other = EnsembleActorTemplate::from_existing_inflight_fragment(
303                    database_info.fragment(other_fragment_id),
304                );
305                template.assert_aligned_with(&other, first_existing, other_fragment_id);
306            }
307
308            template
309        } else {
310            // All fragments are new — render from scratch.
311            let first_component = ensemble
312                .component_fragments()
313                .next()
314                .expect("ensemble must have at least one component");
315            let fragment = &fragments.inner.fragments[&first_component];
316            let distribution_type: DistributionType = fragment.distribution_type.into();
317            let vnode_count = fragment.vnode_count();
318
319            // Assert all component fragments in this ensemble share the same vnode count.
320            for fragment_id in ensemble.component_fragments() {
321                let f = &fragments.inner.fragments[&fragment_id];
322                assert_eq!(
323                    vnode_count,
324                    f.vnode_count(),
325                    "component fragments {} and {} in the same no-shuffle ensemble have \
326                     different vnode counts: {} vs {}",
327                    first_component,
328                    fragment_id,
329                    vnode_count,
330                    f.vnode_count(),
331                );
332            }
333
334            EnsembleActorTemplate::render_new(
335                streaming_job_model,
336                worker_map,
337                None,
338                database_resource_group.to_owned(),
339                distribution_type,
340                vnode_count,
341            )?
342        };
343
344        // Render each new component fragment in this ensemble.
345        for fragment_id in ensemble.component_fragments() {
346            if !fragments.inner.fragments.contains_key(&fragment_id) {
347                continue; // Skip existing fragments.
348            }
349            let fragment = &fragments.inner.fragments[&fragment_id];
350            let distribution_type: DistributionType = fragment.distribution_type.into();
351            let aligner =
352                ComponentFragmentAligner::new_persistent(&actor_template, actor_id_counter);
353            let assignments = aligner.align_component_actor(distribution_type);
354            actor_assignments.insert(fragment_id, assignments);
355        }
356    }
357
358    // Step 3: Expand simple assignments into full StreamActor structures.
359    let mut result_stream_actors: HashMap<FragmentId, Vec<StreamActor>> = HashMap::new();
360    let mut result_actor_location: HashMap<ActorId, WorkerId> = HashMap::new();
361
362    for (fragment_id, assignments) in &actor_assignments {
363        let mut actors = Vec::with_capacity(assignments.len());
364        for (&actor_id, (worker_id, vnode_bitmap)) in assignments {
365            result_actor_location.insert(actor_id, *worker_id);
366            actors.push(StreamActor {
367                actor_id,
368                fragment_id: *fragment_id,
369                vnode_bitmap: vnode_bitmap.clone(),
370                mview_definition: definition.to_owned(),
371                expr_context: Some(ctx.to_expr_context()),
372                config_override: ctx.config_override.clone(),
373            });
374        }
375        result_stream_actors.insert(*fragment_id, actors);
376    }
377
378    Ok(RenderResult {
379        stream_actors: result_stream_actors,
380        actor_location: result_actor_location,
381    })
382}
383impl DatabaseCheckpointControl {
384    /// Collect table IDs to commit and actor IDs to collect from current fragment infos.
385    fn collect_base_info(&self) -> (HashSet<TableId>, HashMap<WorkerId, HashSet<ActorId>>) {
386        let table_ids_to_commit = self.database_info.existing_table_ids().collect();
387        let node_actors =
388            InflightFragmentInfo::actor_ids_to_collect(self.database_info.fragment_infos());
389        (table_ids_to_commit, node_actors)
390    }
391
392    /// Helper for the simplest command variants: those that only need a
393    /// pre-computed mutation and a command name, with no actors to create
394    /// and no additional side effects on `self`.
395    fn apply_simple_command(
396        &self,
397        mutation: Option<Mutation>,
398        command_name: &'static str,
399    ) -> ApplyCommandResult {
400        let (table_ids, node_actors) = self.collect_base_info();
401        (
402            mutation,
403            table_ids,
404            None,
405            node_actors,
406            PostCollectCommand::Command(command_name.to_owned()),
407        )
408    }
409
410    /// Returns the inflight actor infos that have included the newly added actors in the given command. The dropped actors
411    /// will be removed from the state after the info get resolved.
412    pub(super) fn apply_command(
413        &mut self,
414        command: Option<Command>,
415        notifiers: &mut Vec<Notifier>,
416        barrier_info: BarrierInfo,
417        partial_graph_manager: &mut PartialGraphManager,
418        hummock_version_stats: &HummockVersionStats,
419        worker_nodes: &HashMap<WorkerId, WorkerNode>,
420    ) -> MetaResult<ApplyCommandInfo> {
421        debug_assert!(
422            !matches!(
423                command,
424                Some(Command::RescheduleIntent {
425                    reschedule_plan: None,
426                    ..
427                })
428            ),
429            "reschedule intent must be resolved before apply"
430        );
431        if matches!(
432            command,
433            Some(Command::RescheduleIntent {
434                reschedule_plan: None,
435                ..
436            })
437        ) {
438            bail!("reschedule intent must be resolved before apply");
439        }
440
441        /// Resolve source splits for a create streaming job command.
442        ///
443        /// Combines source fragment split resolution and backfill split alignment
444        /// into one step, looking up existing upstream actor splits from the inflight database info.
445        fn resolve_source_splits(
446            info: &CreateStreamingJobCommandInfo,
447            render_result: &RenderResult,
448            actor_no_shuffle: &ActorNewNoShuffle,
449            database_info: &InflightDatabaseInfo,
450        ) -> MetaResult<SplitAssignment> {
451            let fragment_actor_ids: HashMap<FragmentId, Vec<ActorId>> = render_result
452                .stream_actors
453                .iter()
454                .map(|(fragment_id, actors)| {
455                    (
456                        *fragment_id,
457                        actors.iter().map(|a| a.actor_id).collect::<Vec<_>>(),
458                    )
459                })
460                .collect();
461            let mut resolved = SourceManager::resolve_fragment_to_actor_splits(
462                &info.stream_job_fragments,
463                &info.init_split_assignment,
464                &fragment_actor_ids,
465            )?;
466            resolved.extend(SourceManager::resolve_backfill_splits(
467                &info.stream_job_fragments,
468                actor_no_shuffle,
469                |fragment_id, actor_id| {
470                    database_info
471                        .fragment(fragment_id)
472                        .actors
473                        .get(&actor_id)
474                        .map(|info| info.splits.clone())
475                },
476            )?);
477            Ok(resolved)
478        }
479
480        // Throttle data for creating jobs (set only in the Throttle arm). Creating jobs live
481        // outside `database_info`, so they need both the compute mutation and the stream-node
482        // pre-apply.
483        let mut throttle_for_creating_jobs: Option<ThrottleForCreatingJobs> = None;
484
485        // Each variant handles its own pre-apply, edge building, mutation generation,
486        // collect base info, and post-apply. The match produces values consumed by the
487        // common snapshot-backfill-merging code that follows.
488        let (
489            mutation,
490            mut table_ids_to_commit,
491            mut actors_to_create,
492            mut node_actors,
493            post_collect_command,
494        ) = match command {
495            None => self.apply_simple_command(None, "barrier"),
496            Some(Command::CreateStreamingJob {
497                mut info,
498                job_type:
499                    CreateStreamingJobType::SnapshotBackfill {
500                        mut snapshot_backfill_info,
501                        since_epoch,
502                    },
503                cross_db_snapshot_backfill_info,
504            }) => {
505                let ensembles = resolve_no_shuffle_ensembles(
506                    &info.stream_job_fragments,
507                    &info.upstream_fragment_downstreams,
508                )?;
509                let actors = render_actors(
510                    &info.stream_job_fragments,
511                    &self.database_info,
512                    &info.definition,
513                    &info.stream_job_fragments.inner.ctx,
514                    &info.streaming_job_model,
515                    partial_graph_manager
516                        .control_stream_manager()
517                        .env
518                        .actor_id_generator(),
519                    worker_nodes,
520                    &ensembles,
521                    &info.database_resource_group,
522                )?;
523                {
524                    assert!(!self.state.is_paused());
525                    let (snapshot_epoch, since_timestamp_upstream_log_epochs) =
526                        if let Some(since_epoch) = &since_epoch {
527                            let (snapshot_epoch, log_epochs) =
528                                since_epoch.resolved.as_ref().ok_or_else(|| {
529                            MetaError::from(anyhow::anyhow!(
530                                "since_timestamp epoch has not been resolved for snapshot backfill"
531                            ))
532                        })?;
533                            (
534                                *snapshot_epoch,
535                                Some((
536                                    log_epochs,
537                                    to_partial_graph_id(self.database_id, None),
538                                    barrier_info.prev_epoch(),
539                                )),
540                            )
541                        } else {
542                            (barrier_info.prev_epoch(), None)
543                        };
544                    // set snapshot epoch of upstream table for snapshot backfill
545                    for snapshot_backfill_epoch in snapshot_backfill_info
546                        .upstream_mv_table_id_to_backfill_epoch
547                        .values_mut()
548                    {
549                        assert_eq!(
550                            snapshot_backfill_epoch.replace(snapshot_epoch),
551                            None,
552                            "must not set previously"
553                        );
554                    }
555                    for fragment in info.stream_job_fragments.inner.fragments.values_mut() {
556                        fill_snapshot_backfill_epoch(
557                            &mut fragment.nodes,
558                            Some(&snapshot_backfill_info),
559                            &cross_db_snapshot_backfill_info,
560                        )?;
561                    }
562                    let job_id = info.stream_job_fragments.stream_job_id();
563                    let snapshot_backfill_upstream_tables = snapshot_backfill_info
564                        .upstream_mv_table_id_to_backfill_epoch
565                        .keys()
566                        .cloned()
567                        .collect();
568                    // Build edges first (needed for no-shuffle mapping used in split resolution)
569                    let mut edges = self.database_info.build_edge(
570                        Some((&info, true)),
571                        None,
572                        None,
573                        partial_graph_manager.control_stream_manager(),
574                        &actors.stream_actors,
575                        &actors.actor_location,
576                    );
577                    // Phase 2: Resolve source-level DiscoveredSplits to actor-level SplitAssignment
578                    let resolved_split_assignment = resolve_source_splits(
579                        &info,
580                        &actors,
581                        edges.actor_new_no_shuffle(),
582                        &self.database_info,
583                    )?;
584
585                    let Entry::Vacant(entry) =
586                        self.independent_checkpoint_job_controls.entry(job_id)
587                    else {
588                        panic!("duplicated creating snapshot backfill job {job_id}");
589                    };
590
591                    let job = CreatingStreamingJobControl::new(
592                        entry,
593                        CreateSnapshotBackfillJobCommandInfo {
594                            info: info.clone(),
595                            snapshot_backfill_info: snapshot_backfill_info.clone(),
596                            cross_db_snapshot_backfill_info,
597                            resolved_split_assignment: resolved_split_assignment.clone(),
598                            refresh_interval_sec: None,
599                        },
600                        take(notifiers),
601                        snapshot_backfill_upstream_tables,
602                        snapshot_epoch,
603                        since_timestamp_upstream_log_epochs,
604                        hummock_version_stats,
605                        partial_graph_manager,
606                        &mut edges,
607                        &resolved_split_assignment,
608                        &actors,
609                    )?;
610
611                    if let Some(fragment_infos) = job.fragment_infos() {
612                        self.database_info.shared_actor_infos.upsert(
613                            self.database_id,
614                            fragment_infos.values().map(|f| (f, job_id)),
615                        );
616                    }
617
618                    for upstream_mv_table_id in snapshot_backfill_info
619                        .upstream_mv_table_id_to_backfill_epoch
620                        .keys()
621                    {
622                        self.database_info.register_subscriber(
623                            upstream_mv_table_id.as_job_id(),
624                            info.streaming_job.id().as_subscriber_id(),
625                            SubscriberType::SnapshotBackfill,
626                        );
627                    }
628
629                    let mutation = Command::create_streaming_job_to_mutation(
630                        &info,
631                        &CreateStreamingJobType::SnapshotBackfill {
632                            snapshot_backfill_info,
633                            since_epoch,
634                        },
635                        [],
636                        self.state.is_paused(),
637                        &mut edges,
638                        partial_graph_manager.control_stream_manager(),
639                        None,
640                        &resolved_split_assignment,
641                        &actors.stream_actors,
642                        &actors.actor_location,
643                    )?;
644
645                    let (table_ids, node_actors) = self.collect_base_info();
646                    (
647                        Some(mutation),
648                        table_ids,
649                        None,
650                        node_actors,
651                        PostCollectCommand::barrier(),
652                    )
653                }
654            }
655            Some(Command::CreateStreamingJob {
656                mut info,
657                job_type: CreateStreamingJobType::BatchRefresh(mut batch_refresh_info),
658                cross_db_snapshot_backfill_info,
659            }) => {
660                {
661                    if self.state.is_paused() {
662                        bail!("cannot create batch refresh job while database barrier is paused");
663                    }
664                    let snapshot_epoch = barrier_info.prev_epoch();
665                    let job_id = info.stream_job_fragments.stream_job_id();
666                    let database_id = info.streaming_job.database_id();
667
668                    // 1. Fill snapshot backfill epochs.
669                    let snapshot_backfill_info = &mut batch_refresh_info.snapshot_backfill_info;
670                    for snapshot_backfill_epoch in snapshot_backfill_info
671                        .upstream_mv_table_id_to_backfill_epoch
672                        .values_mut()
673                    {
674                        assert_eq!(
675                            snapshot_backfill_epoch.replace(snapshot_epoch),
676                            None,
677                            "must not set previously"
678                        );
679                    }
680                    for fragment in info.stream_job_fragments.inner.fragments.values_mut() {
681                        fill_snapshot_backfill_epoch(
682                            &mut fragment.nodes,
683                            Some(snapshot_backfill_info),
684                            &cross_db_snapshot_backfill_info,
685                        )?;
686                    }
687                    let snapshot_backfill_upstream_tables: HashSet<TableId> =
688                        snapshot_backfill_info
689                            .upstream_mv_table_id_to_backfill_epoch
690                            .keys()
691                            .cloned()
692                            .collect();
693
694                    // 2. Build BatchRefreshLogicalFragments (after epoch filling).
695                    let logical = BatchRefreshLogicalFragments {
696                        fragments: info
697                            .stream_job_fragments
698                            .inner
699                            .fragments
700                            .iter()
701                            .map(|(&fid, fragment)| {
702                                (
703                                    fid,
704                                    LoadedFragment {
705                                        fragment_id: fid,
706                                        job_id,
707                                        fragment_type_mask: fragment.fragment_type_mask,
708                                        distribution_type: fragment.distribution_type.into(),
709                                        vnode_count: fragment.vnode_count(),
710                                        nodes: fragment.nodes.clone(),
711                                        state_table_ids: fragment
712                                            .state_table_ids
713                                            .iter()
714                                            .cloned()
715                                            .collect(),
716                                        parallelism: None,
717                                    },
718                                )
719                            })
720                            .collect(),
721                        downstreams: info.stream_job_fragments.downstreams.clone(),
722                    };
723
724                    // 3. Create BatchRefreshJobCheckpointControl. `new()` handles actor
725                    //    rendering, the partial-graph initial barrier, and produces the
726                    //    database-graph mutation for the main barrier.
727                    assert!(
728                        !self
729                            .independent_checkpoint_job_controls
730                            .contains_key(&job_id),
731                        "duplicated creating batch refresh job {job_id}"
732                    );
733
734                    let snapshot_backfill_info_clone =
735                        batch_refresh_info.snapshot_backfill_info.clone();
736                    let refresh_interval_sec = batch_refresh_info.refresh_interval_sec;
737
738                    // Database-graph `Add` mutation: batch refresh has no actors in the
739                    // database graph; it only needs to register snapshot-backfill
740                    // subscribers on the upstream MV tables.
741                    let subscriber_id =
742                        info.stream_job_fragments.stream_job_id().as_subscriber_id();
743                    let mutation = Mutation::Add(AddMutation {
744                        actor_dispatchers: Default::default(),
745                        added_actors: Default::default(),
746                        actor_splits: Default::default(),
747                        pause: false,
748                        subscriptions_to_add: snapshot_backfill_info_clone
749                            .upstream_mv_table_id_to_backfill_epoch
750                            .keys()
751                            .map(|table_id| PbSubscriptionUpstreamInfo {
752                                subscriber_id,
753                                upstream_mv_table_id: *table_id,
754                            })
755                            .collect(),
756                        backfill_nodes_to_pause: Default::default(),
757                        actor_cdc_table_snapshot_splits: None,
758                        new_upstream_sinks: Default::default(),
759                        dropped_actors: Default::default(),
760                        sink_log_store_flush: Default::default(),
761                    });
762
763                    let job = BatchRefreshJobCheckpointControl::new(
764                        database_id,
765                        job_id,
766                        CreateSnapshotBackfillJobCommandInfo {
767                            info: info.clone(),
768                            snapshot_backfill_info: snapshot_backfill_info_clone.clone(),
769                            cross_db_snapshot_backfill_info,
770                            resolved_split_assignment: Default::default(),
771                            refresh_interval_sec: Some(refresh_interval_sec),
772                        },
773                        take(notifiers),
774                        snapshot_backfill_upstream_tables,
775                        snapshot_epoch,
776                        hummock_version_stats,
777                        partial_graph_manager,
778                        &logical,
779                        worker_nodes,
780                        refresh_interval_sec,
781                    )?;
782
783                    if let Some(fragment_infos) = job.fragment_infos() {
784                        self.database_info.shared_actor_infos.upsert(
785                            self.database_id,
786                            fragment_infos.values().map(|f| (f, job_id)),
787                        );
788                    }
789
790                    self.independent_checkpoint_job_controls
791                        .insert(job_id, IndependentCheckpointJobControl::BatchRefresh(job));
792
793                    // Register permanent subscriber (never unregistered until MV is dropped)
794                    for upstream_mv_table_id in snapshot_backfill_info_clone
795                        .upstream_mv_table_id_to_backfill_epoch
796                        .keys()
797                    {
798                        self.database_info.register_subscriber(
799                            upstream_mv_table_id.as_job_id(),
800                            info.streaming_job.id().as_subscriber_id(),
801                            SubscriberType::SnapshotBackfill,
802                        );
803                    }
804
805                    let (table_ids, node_actors) = self.collect_base_info();
806                    (
807                        Some(mutation),
808                        table_ids,
809                        None,
810                        node_actors,
811                        PostCollectCommand::barrier(),
812                    )
813                }
814            }
815            Some(Command::CreateStreamingJob {
816                mut info,
817                job_type,
818                cross_db_snapshot_backfill_info,
819            }) => {
820                let ensembles = resolve_no_shuffle_ensembles(
821                    &info.stream_job_fragments,
822                    &info.upstream_fragment_downstreams,
823                )?;
824                let actors = render_actors(
825                    &info.stream_job_fragments,
826                    &self.database_info,
827                    &info.definition,
828                    &info.stream_job_fragments.inner.ctx,
829                    &info.streaming_job_model,
830                    partial_graph_manager
831                        .control_stream_manager()
832                        .env
833                        .actor_id_generator(),
834                    worker_nodes,
835                    &ensembles,
836                    &info.database_resource_group,
837                )?;
838                for fragment in info.stream_job_fragments.inner.fragments.values_mut() {
839                    fill_snapshot_backfill_epoch(
840                        &mut fragment.nodes,
841                        None,
842                        &cross_db_snapshot_backfill_info,
843                    )?;
844                }
845
846                // Build edges
847                let new_upstream_sink =
848                    if let CreateStreamingJobType::SinkIntoTable(ref ctx) = job_type {
849                        Some(ctx)
850                    } else {
851                        None
852                    };
853
854                let mut edges = self.database_info.build_edge(
855                    Some((&info, false)),
856                    None,
857                    new_upstream_sink,
858                    partial_graph_manager.control_stream_manager(),
859                    &actors.stream_actors,
860                    &actors.actor_location,
861                );
862                // Phase 2: Resolve source-level DiscoveredSplits to actor-level SplitAssignment
863                let resolved_split_assignment = resolve_source_splits(
864                    &info,
865                    &actors,
866                    edges.actor_new_no_shuffle(),
867                    &self.database_info,
868                )?;
869
870                let old_sink_job_id = info
871                    .replace_sink
872                    .as_ref()
873                    .map(|old_sink_id| old_sink_id.as_job_id());
874                if old_sink_job_id.is_some()
875                    && matches!(
876                        job_type,
877                        CreateStreamingJobType::SnapshotBackfill { .. }
878                            | CreateStreamingJobType::BatchRefresh(_)
879                    )
880                {
881                    bail!("replace sink must not use snapshot backfill");
882                }
883
884                // Pre-apply: add new job and fragments
885                let cdc_tracker = if let Some(splits) = &info.cdc_table_snapshot_splits {
886                    let (fragment, _) =
887                        parallel_cdc_table_backfill_fragment(info.stream_job_fragments.fragments())
888                            .expect("should have parallel cdc fragment");
889                    Some(CdcTableBackfillTracker::new(
890                        fragment.fragment_id,
891                        splits.clone(),
892                    ))
893                } else {
894                    None
895                };
896                self.database_info
897                    .pre_apply_new_job(info.streaming_job.id(), cdc_tracker);
898                self.database_info.pre_apply_new_fragments(
899                    info.stream_job_fragments
900                        .new_fragment_info(
901                            &actors.stream_actors,
902                            &actors.actor_location,
903                            &resolved_split_assignment,
904                        )
905                        .map(|(fragment_id, fragment_infos)| {
906                            (fragment_id, info.streaming_job.id(), fragment_infos)
907                        }),
908                );
909                if let CreateStreamingJobType::SinkIntoTable(ref ctx) = job_type {
910                    let downstream_fragment_id = ctx.new_sink_downstream.downstream_fragment_id;
911                    self.database_info.pre_apply_add_node_upstream(
912                        downstream_fragment_id,
913                        &PbUpstreamSinkInfo {
914                            upstream_fragment_id: ctx.sink_fragment_id,
915                            sink_output_schema: ctx.sink_output_fields.clone(),
916                            project_exprs: ctx.project_exprs.clone(),
917                        },
918                    );
919                }
920
921                let (table_ids, node_actors) = self.collect_base_info();
922                let dropped_actors = if let Some(old_sink_job_id) = old_sink_job_id {
923                    let Some(job) = self.database_info.post_apply_remove_job(old_sink_job_id)
924                    else {
925                        bail!(
926                            "old sink job {} not found in barrier state",
927                            old_sink_job_id
928                        );
929                    };
930                    job.fragment_infos
931                        .values()
932                        .flat_map(|fragment| fragment.actors.keys().copied())
933                        .collect()
934                } else {
935                    vec![]
936                };
937
938                // Actors to create
939                let actors_to_create = Some(Command::create_streaming_job_actors_to_create(
940                    &info,
941                    &mut edges,
942                    &actors.stream_actors,
943                    &actors.actor_location,
944                ));
945
946                // CDC table snapshot splits
947                let actor_cdc_table_snapshot_splits = self
948                    .database_info
949                    .assign_cdc_backfill_splits(info.stream_job_fragments.stream_job_id())?;
950
951                // Mutation
952                let is_currently_paused = self.state.is_paused();
953                let mutation = Command::create_streaming_job_to_mutation(
954                    &info,
955                    &job_type,
956                    dropped_actors,
957                    is_currently_paused,
958                    &mut edges,
959                    partial_graph_manager.control_stream_manager(),
960                    actor_cdc_table_snapshot_splits,
961                    &resolved_split_assignment,
962                    &actors.stream_actors,
963                    &actors.actor_location,
964                )?;
965
966                (
967                    Some(mutation),
968                    table_ids,
969                    actors_to_create,
970                    node_actors,
971                    PostCollectCommand::CreateStreamingJob {
972                        info,
973                        job_type,
974                        cross_db_snapshot_backfill_info,
975                        resolved_split_assignment,
976                    },
977                )
978            }
979
980            Some(Command::Flush) => self.apply_simple_command(None, "Flush"),
981
982            Some(Command::Pause) => {
983                let prev_is_paused = self.state.is_paused();
984                self.state.set_is_paused(true);
985                let mutation = Command::pause_to_mutation(prev_is_paused);
986                let (table_ids, node_actors) = self.collect_base_info();
987                (
988                    mutation,
989                    table_ids,
990                    None,
991                    node_actors,
992                    PostCollectCommand::Command("Pause".to_owned()),
993                )
994            }
995
996            Some(Command::Resume) => {
997                let prev_is_paused = self.state.is_paused();
998                self.state.set_is_paused(false);
999                let mutation = Command::resume_to_mutation(prev_is_paused);
1000                let (table_ids, node_actors) = self.collect_base_info();
1001                (
1002                    mutation,
1003                    table_ids,
1004                    None,
1005                    node_actors,
1006                    PostCollectCommand::Command("Resume".to_owned()),
1007                )
1008            }
1009
1010            Some(Command::Throttle { jobs, config }) => {
1011                let mutation = Some(Command::throttle_to_mutation(&config));
1012                for (fragment_id, (_, stream_node)) in &config {
1013                    self.database_info
1014                        .pre_apply_throttle(*fragment_id, stream_node);
1015                }
1016                throttle_for_creating_jobs = Some((jobs, config));
1017                self.apply_simple_command(mutation, "Throttle")
1018            }
1019
1020            Some(Command::DropStreamingJobs {
1021                streaming_job_ids,
1022                unregistered_state_table_ids: _,
1023                dropped_sink_fragment_by_targets,
1024            }) => {
1025                // pre_apply: drop node upstream for sink targets
1026                for (target_fragment, sink_fragments) in &dropped_sink_fragment_by_targets {
1027                    self.database_info
1028                        .pre_apply_drop_node_upstream(*target_fragment, sink_fragments);
1029                }
1030
1031                let (table_ids, node_actors) = self.collect_base_info();
1032
1033                let mut actors = Vec::new();
1034                for job_id in streaming_job_ids {
1035                    let Some(job) = self.database_info.post_apply_remove_job(job_id) else {
1036                        warn!(
1037                            %job_id,
1038                            "skip drop payload for streaming job that has already been removed from barrier worker"
1039                        );
1040                        continue;
1041                    };
1042
1043                    for fragment in job.fragment_infos.values() {
1044                        actors.extend(fragment.actors.keys().copied());
1045                    }
1046                }
1047
1048                let mutation = Some(Command::drop_streaming_jobs_to_mutation(
1049                    &actors,
1050                    &dropped_sink_fragment_by_targets,
1051                ));
1052                (
1053                    mutation,
1054                    table_ids,
1055                    None,
1056                    node_actors,
1057                    PostCollectCommand::DropStreamingJobs,
1058                )
1059            }
1060
1061            Some(Command::RescheduleIntent {
1062                reschedule_plan, ..
1063            }) => {
1064                let ReschedulePlan {
1065                    reschedules,
1066                    fragment_actors,
1067                } = reschedule_plan
1068                    .as_ref()
1069                    .expect("reschedule intent should be resolved in global barrier worker");
1070
1071                // Pre-apply: reschedule fragments
1072                for (fragment_id, reschedule) in reschedules {
1073                    self.database_info.pre_apply_reschedule(
1074                        *fragment_id,
1075                        reschedule
1076                            .added_actors
1077                            .iter()
1078                            .flat_map(|(node_id, actors): (&WorkerId, &Vec<ActorId>)| {
1079                                actors.iter().map(|actor_id| {
1080                                    (
1081                                        *actor_id,
1082                                        InflightActorInfo {
1083                                            worker_id: *node_id,
1084                                            vnode_bitmap: reschedule
1085                                                .newly_created_actors
1086                                                .get(actor_id)
1087                                                .expect("should exist")
1088                                                .0
1089                                                .0
1090                                                .vnode_bitmap
1091                                                .clone(),
1092                                            splits: reschedule
1093                                                .actor_splits
1094                                                .get(actor_id)
1095                                                .cloned()
1096                                                .unwrap_or_default(),
1097                                        },
1098                                    )
1099                                })
1100                            })
1101                            .collect(),
1102                        reschedule
1103                            .vnode_bitmap_updates
1104                            .iter()
1105                            .filter(|(actor_id, _)| {
1106                                !reschedule.newly_created_actors.contains_key(*actor_id)
1107                            })
1108                            .map(|(actor_id, bitmap)| (*actor_id, bitmap.clone()))
1109                            .collect(),
1110                        reschedule.actor_splits.clone(),
1111                    );
1112                }
1113
1114                let (table_ids, node_actors) = self.collect_base_info();
1115
1116                // Actors to create
1117                let actors_to_create = Some(Command::reschedule_actors_to_create(
1118                    reschedules,
1119                    fragment_actors,
1120                    &self.database_info,
1121                    partial_graph_manager.control_stream_manager(),
1122                ));
1123
1124                // Post-apply: remove old actors
1125                self.database_info
1126                    .post_apply_reschedules(reschedules.iter().map(|(fragment_id, reschedule)| {
1127                        (
1128                            *fragment_id,
1129                            reschedule.removed_actors.iter().cloned().collect(),
1130                        )
1131                    }));
1132
1133                // Mutation
1134                let mutation = Command::reschedule_to_mutation(
1135                    reschedules,
1136                    fragment_actors,
1137                    partial_graph_manager.control_stream_manager(),
1138                    &mut self.database_info,
1139                )?;
1140
1141                let reschedules = reschedule_plan
1142                    .expect("reschedule intent should be resolved in global barrier worker")
1143                    .reschedules;
1144                (
1145                    mutation,
1146                    table_ids,
1147                    actors_to_create,
1148                    node_actors,
1149                    PostCollectCommand::Reschedule { reschedules },
1150                )
1151            }
1152
1153            Some(Command::ReplaceStreamJob(plan)) => {
1154                let ensembles = resolve_no_shuffle_ensembles(
1155                    &plan.new_fragments,
1156                    &plan.upstream_fragment_downstreams,
1157                )?;
1158                let mut render_result = render_actors(
1159                    &plan.new_fragments,
1160                    &self.database_info,
1161                    "", // replace jobs don't need mview definition
1162                    &plan.new_fragments.inner.ctx,
1163                    &plan.streaming_job_model,
1164                    partial_graph_manager
1165                        .control_stream_manager()
1166                        .env
1167                        .actor_id_generator(),
1168                    worker_nodes,
1169                    &ensembles,
1170                    &plan.database_resource_group,
1171                )?;
1172
1173                // Render actors for auto_refresh_schema_sinks.
1174                // Each sink's new_fragment inherits parallelism from its original_fragment.
1175                if let Some(sinks) = &plan.auto_refresh_schema_sinks {
1176                    let actor_id_counter = partial_graph_manager
1177                        .control_stream_manager()
1178                        .env
1179                        .actor_id_generator();
1180                    for sink_ctx in sinks {
1181                        let original_fragment_id = sink_ctx.original_fragment.fragment_id;
1182                        let original_frag_info = self.database_info.fragment(original_fragment_id);
1183                        let actor_template = EnsembleActorTemplate::from_existing_inflight_fragment(
1184                            original_frag_info,
1185                        );
1186                        let new_aligner = ComponentFragmentAligner::new_persistent(
1187                            &actor_template,
1188                            actor_id_counter,
1189                        );
1190                        let distribution_type: DistributionType =
1191                            sink_ctx.new_fragment.distribution_type.into();
1192                        let actor_assignments =
1193                            new_aligner.align_component_actor(distribution_type);
1194                        let new_fragment_id = sink_ctx.new_fragment.fragment_id;
1195                        let mut actors = Vec::with_capacity(actor_assignments.len());
1196                        for (&actor_id, (worker_id, vnode_bitmap)) in &actor_assignments {
1197                            render_result.actor_location.insert(actor_id, *worker_id);
1198                            actors.push(StreamActor {
1199                                actor_id,
1200                                fragment_id: new_fragment_id,
1201                                vnode_bitmap: vnode_bitmap.clone(),
1202                                mview_definition: String::new(),
1203                                expr_context: Some(sink_ctx.ctx.to_expr_context()),
1204                                config_override: sink_ctx.ctx.config_override.clone(),
1205                            });
1206                        }
1207                        render_result.stream_actors.insert(new_fragment_id, actors);
1208                    }
1209                }
1210
1211                // Build edges first (needed for no-shuffle mapping used in split resolution)
1212                let mut edges = self.database_info.build_edge(
1213                    None,
1214                    Some(&plan),
1215                    None,
1216                    partial_graph_manager.control_stream_manager(),
1217                    &render_result.stream_actors,
1218                    &render_result.actor_location,
1219                );
1220
1221                // Phase 2: Resolve splits to actor-level assignment.
1222                let fragment_actor_ids: HashMap<FragmentId, Vec<ActorId>> = render_result
1223                    .stream_actors
1224                    .iter()
1225                    .map(|(fragment_id, actors)| {
1226                        (
1227                            *fragment_id,
1228                            actors.iter().map(|a| a.actor_id).collect::<Vec<_>>(),
1229                        )
1230                    })
1231                    .collect();
1232                let resolved_split_assignment = match &plan.split_plan {
1233                    ReplaceJobSplitPlan::Discovered(discovered) => {
1234                        SourceManager::resolve_fragment_to_actor_splits(
1235                            &plan.new_fragments,
1236                            discovered,
1237                            &fragment_actor_ids,
1238                        )?
1239                    }
1240                    ReplaceJobSplitPlan::AlignFromPrevious => {
1241                        SourceManager::resolve_replace_source_splits(
1242                            &plan.new_fragments,
1243                            &plan.replace_upstream,
1244                            edges.actor_new_no_shuffle(),
1245                            |_fragment_id, actor_id| {
1246                                self.database_info.fragment_infos().find_map(|fragment| {
1247                                    fragment
1248                                        .actors
1249                                        .get(&actor_id)
1250                                        .map(|info| info.splits.clone())
1251                                })
1252                            },
1253                        )?
1254                    }
1255                };
1256
1257                // Pre-apply: add new fragments and replace upstream
1258                self.database_info.pre_apply_new_fragments(
1259                    plan.new_fragments
1260                        .new_fragment_info(
1261                            &render_result.stream_actors,
1262                            &render_result.actor_location,
1263                            &resolved_split_assignment,
1264                        )
1265                        .map(|(fragment_id, new_fragment)| {
1266                            (fragment_id, plan.streaming_job.id(), new_fragment)
1267                        }),
1268                );
1269                for (fragment_id, replace_map) in &plan.replace_upstream {
1270                    self.database_info
1271                        .pre_apply_replace_node_upstream(*fragment_id, replace_map);
1272                }
1273                if let Some(sinks) = &plan.auto_refresh_schema_sinks {
1274                    self.database_info
1275                        .pre_apply_new_fragments(sinks.iter().map(|sink| {
1276                            (
1277                                sink.new_fragment.fragment_id,
1278                                sink.original_sink.id.as_job_id(),
1279                                sink.new_fragment_info(
1280                                    &render_result.stream_actors,
1281                                    &render_result.actor_location,
1282                                ),
1283                            )
1284                        }));
1285                }
1286
1287                let (table_ids, node_actors) = self.collect_base_info();
1288
1289                // Actors to create
1290                let actors_to_create = Some(Command::replace_stream_job_actors_to_create(
1291                    &plan,
1292                    &mut edges,
1293                    &self.database_info,
1294                    &render_result.stream_actors,
1295                    &render_result.actor_location,
1296                ));
1297
1298                // Mutation (must be generated before removing old fragments,
1299                // because it reads actor info from database_info)
1300                let mutation = Command::replace_stream_job_to_mutation(
1301                    &plan,
1302                    &mut edges,
1303                    &mut self.database_info,
1304                    &resolved_split_assignment,
1305                )?;
1306
1307                // Post-apply: remove old fragments
1308                {
1309                    let mut fragment_ids_to_remove: Vec<_> = plan
1310                        .old_fragments
1311                        .fragments
1312                        .values()
1313                        .map(|f| f.fragment_id)
1314                        .collect();
1315                    if let Some(sinks) = &plan.auto_refresh_schema_sinks {
1316                        fragment_ids_to_remove
1317                            .extend(sinks.iter().map(|sink| sink.original_fragment.fragment_id));
1318                    }
1319                    self.database_info
1320                        .post_apply_remove_fragments(fragment_ids_to_remove);
1321                }
1322
1323                (
1324                    mutation,
1325                    table_ids,
1326                    actors_to_create,
1327                    node_actors,
1328                    PostCollectCommand::ReplaceStreamJob {
1329                        plan,
1330                        resolved_split_assignment,
1331                    },
1332                )
1333            }
1334
1335            Some(Command::SourceChangeSplit(split_state)) => {
1336                // Pre-apply: split assignments
1337                self.database_info.pre_apply_split_assignments(
1338                    split_state
1339                        .split_assignment
1340                        .iter()
1341                        .map(|(&fragment_id, splits)| (fragment_id, splits.clone())),
1342                );
1343
1344                let mutation = Some(Command::source_change_split_to_mutation(
1345                    &split_state.split_assignment,
1346                ));
1347                let (table_ids, node_actors) = self.collect_base_info();
1348                (
1349                    mutation,
1350                    table_ids,
1351                    None,
1352                    node_actors,
1353                    PostCollectCommand::SourceChangeSplit {
1354                        split_assignment: split_state.split_assignment,
1355                    },
1356                )
1357            }
1358
1359            Some(Command::CreateSubscription {
1360                subscription_id,
1361                upstream_mv_table_id,
1362                retention_second,
1363            }) => {
1364                self.database_info.register_subscriber(
1365                    upstream_mv_table_id.as_job_id(),
1366                    subscription_id.as_subscriber_id(),
1367                    SubscriberType::Subscription(retention_second),
1368                );
1369                let mutation = Some(Command::create_subscription_to_mutation(
1370                    upstream_mv_table_id,
1371                    subscription_id,
1372                ));
1373                let (table_ids, node_actors) = self.collect_base_info();
1374                (
1375                    mutation,
1376                    table_ids,
1377                    None,
1378                    node_actors,
1379                    PostCollectCommand::CreateSubscription { subscription_id },
1380                )
1381            }
1382
1383            Some(Command::DropSubscription {
1384                subscription_id,
1385                upstream_mv_table_id,
1386            }) => {
1387                if self
1388                    .database_info
1389                    .unregister_subscriber(
1390                        upstream_mv_table_id.as_job_id(),
1391                        subscription_id.as_subscriber_id(),
1392                    )
1393                    .is_none()
1394                {
1395                    warn!(%subscription_id, %upstream_mv_table_id, "no subscription to drop");
1396                }
1397                let mutation = Some(Command::drop_subscription_to_mutation(
1398                    upstream_mv_table_id,
1399                    subscription_id,
1400                ));
1401                let (table_ids, node_actors) = self.collect_base_info();
1402                (
1403                    mutation,
1404                    table_ids,
1405                    None,
1406                    node_actors,
1407                    PostCollectCommand::Command("DropSubscription".to_owned()),
1408                )
1409            }
1410
1411            Some(Command::AlterSubscriptionRetention {
1412                subscription_id,
1413                upstream_mv_table_id,
1414                retention_second,
1415            }) => {
1416                self.database_info.update_subscription_retention(
1417                    upstream_mv_table_id.as_job_id(),
1418                    subscription_id.as_subscriber_id(),
1419                    retention_second,
1420                );
1421                self.apply_simple_command(None, "AlterSubscriptionRetention")
1422            }
1423
1424            Some(Command::ConnectorPropsChange(config)) => {
1425                let mutation = Some(Command::connector_props_change_to_mutation(&config));
1426                let (table_ids, node_actors) = self.collect_base_info();
1427                (
1428                    mutation,
1429                    table_ids,
1430                    None,
1431                    node_actors,
1432                    PostCollectCommand::ConnectorPropsChange(config),
1433                )
1434            }
1435
1436            Some(Command::Refresh {
1437                table_id,
1438                associated_source_id,
1439            }) => {
1440                let mutation = Some(Command::refresh_to_mutation(table_id, associated_source_id));
1441                self.apply_simple_command(mutation, "Refresh")
1442            }
1443
1444            Some(Command::ListFinish {
1445                table_id: _,
1446                associated_source_id,
1447            }) => {
1448                let mutation = Some(Command::list_finish_to_mutation(associated_source_id));
1449                self.apply_simple_command(mutation, "ListFinish")
1450            }
1451
1452            Some(Command::LoadFinish {
1453                table_id: _,
1454                associated_source_id,
1455            }) => {
1456                let mutation = Some(Command::load_finish_to_mutation(associated_source_id));
1457                self.apply_simple_command(mutation, "LoadFinish")
1458            }
1459
1460            Some(Command::ResetSource { source_id }) => {
1461                let mutation = Some(Command::reset_source_to_mutation(source_id));
1462                self.apply_simple_command(mutation, "ResetSource")
1463            }
1464
1465            Some(Command::ResumeBackfill { target }) => {
1466                let mutation = Command::resume_backfill_to_mutation(&target, &self.database_info)?;
1467                let (table_ids, node_actors) = self.collect_base_info();
1468                (
1469                    mutation,
1470                    table_ids,
1471                    None,
1472                    node_actors,
1473                    PostCollectCommand::ResumeBackfill { target },
1474                )
1475            }
1476
1477            Some(Command::InjectSourceOffsets {
1478                source_id,
1479                split_offsets,
1480            }) => {
1481                let mutation = Some(Command::inject_source_offsets_to_mutation(
1482                    source_id,
1483                    &split_offsets,
1484                ));
1485                self.apply_simple_command(mutation, "InjectSourceOffsets")
1486            }
1487        };
1488
1489        let mut finished_snapshot_backfill_jobs = HashSet::new();
1490        let mutation = match mutation {
1491            Some(mutation) => Some(mutation),
1492            None => {
1493                let mut finished_snapshot_backfill_job_info = HashMap::new();
1494                if barrier_info.kind.is_checkpoint() {
1495                    for (&job_id, job) in &mut self.independent_checkpoint_job_controls {
1496                        if let IndependentCheckpointJobControl::CreatingStreamingJob(creating_job) =
1497                            job
1498                            && creating_job.should_merge_to_upstream(partial_graph_manager)
1499                        {
1500                            let info = creating_job
1501                                .start_consume_upstream(partial_graph_manager, &barrier_info)?;
1502                            finished_snapshot_backfill_job_info
1503                                .try_insert(job_id, info)
1504                                .expect("non-duplicated");
1505                        }
1506                    }
1507                }
1508
1509                if !finished_snapshot_backfill_job_info.is_empty() {
1510                    let actors_to_create = actors_to_create.get_or_insert_default();
1511                    let mut subscriptions_to_drop = vec![];
1512                    let mut dispatcher_update = vec![];
1513                    let mut actor_splits = HashMap::new();
1514                    for (job_id, info) in finished_snapshot_backfill_job_info {
1515                        finished_snapshot_backfill_jobs.insert(job_id);
1516                        subscriptions_to_drop.extend(
1517                            info.snapshot_backfill_upstream_tables.iter().map(
1518                                |upstream_table_id| PbSubscriptionUpstreamInfo {
1519                                    subscriber_id: job_id.as_subscriber_id(),
1520                                    upstream_mv_table_id: *upstream_table_id,
1521                                },
1522                            ),
1523                        );
1524                        for upstream_mv_table_id in &info.snapshot_backfill_upstream_tables {
1525                            assert_matches!(
1526                                self.database_info.unregister_subscriber(
1527                                    upstream_mv_table_id.as_job_id(),
1528                                    job_id.as_subscriber_id()
1529                                ),
1530                                Some(SubscriberType::SnapshotBackfill)
1531                            );
1532                        }
1533
1534                        table_ids_to_commit.extend(
1535                            info.fragment_infos
1536                                .values()
1537                                .flat_map(|fragment| fragment.state_table_ids.iter())
1538                                .copied(),
1539                        );
1540
1541                        let actor_len = info
1542                            .fragment_infos
1543                            .values()
1544                            .map(|fragment| fragment.actors.len() as u64)
1545                            .sum();
1546                        let id_gen = GlobalActorIdGen::new(
1547                            partial_graph_manager
1548                                .control_stream_manager()
1549                                .env
1550                                .actor_id_generator(),
1551                            actor_len,
1552                        );
1553                        let mut next_local_actor_id = 0;
1554                        // mapping from old_actor_id to new_actor_id
1555                        let actor_mapping: HashMap<_, _> = info
1556                            .fragment_infos
1557                            .values()
1558                            .flat_map(|fragment| fragment.actors.keys())
1559                            .map(|old_actor_id| {
1560                                let new_actor_id = id_gen.to_global_id(next_local_actor_id);
1561                                next_local_actor_id += 1;
1562                                (*old_actor_id, new_actor_id.as_global_id())
1563                            })
1564                            .collect();
1565                        let actor_mapping = &actor_mapping;
1566                        let new_stream_actors: HashMap<_, _> = info
1567                            .stream_actors
1568                            .into_iter()
1569                            .map(|(old_actor_id, mut actor)| {
1570                                let new_actor_id = actor_mapping[&old_actor_id];
1571                                actor.actor_id = new_actor_id;
1572                                (new_actor_id, actor)
1573                            })
1574                            .collect();
1575                        let new_fragment_info: HashMap<_, _> = info
1576                            .fragment_infos
1577                            .into_iter()
1578                            .map(|(fragment_id, mut fragment)| {
1579                                let actors = take(&mut fragment.actors);
1580                                fragment.actors = actors
1581                                    .into_iter()
1582                                    .map(|(old_actor_id, actor)| {
1583                                        let new_actor_id = actor_mapping[&old_actor_id];
1584                                        (new_actor_id, actor)
1585                                    })
1586                                    .collect();
1587                                (fragment_id, fragment)
1588                            })
1589                            .collect();
1590                        actor_splits.extend(
1591                            new_fragment_info
1592                                .values()
1593                                .flat_map(|fragment| &fragment.actors)
1594                                .map(|(actor_id, actor)| {
1595                                    (
1596                                        *actor_id,
1597                                        ConnectorSplits {
1598                                            splits: actor
1599                                                .splits
1600                                                .iter()
1601                                                .map(ConnectorSplit::from)
1602                                                .collect(),
1603                                        },
1604                                    )
1605                                }),
1606                        );
1607                        // new actors belong to the database partial graph
1608                        let partial_graph_id = to_partial_graph_id(self.database_id, None);
1609                        let mut edge_builder = FragmentEdgeBuilder::new(
1610                            info.upstream_fragment_downstreams
1611                                .keys()
1612                                .map(|upstream_fragment_id| {
1613                                    self.database_info.fragment(*upstream_fragment_id)
1614                                })
1615                                .chain(new_fragment_info.values())
1616                                .map(|fragment| {
1617                                    (
1618                                        fragment.fragment_id,
1619                                        EdgeBuilderFragmentInfo::from_inflight(
1620                                            fragment,
1621                                            partial_graph_id,
1622                                            partial_graph_manager.control_stream_manager(),
1623                                        ),
1624                                    )
1625                                }),
1626                        );
1627                        edge_builder.add_relations(&info.upstream_fragment_downstreams);
1628                        edge_builder.add_relations(&info.downstreams);
1629                        let mut edges = edge_builder.build();
1630                        let new_actors_to_create = edges.collect_actors_to_create(
1631                            new_fragment_info.values().map(|fragment| {
1632                                (
1633                                    fragment.fragment_id,
1634                                    &fragment.nodes,
1635                                    fragment.actors.iter().map(|(actor_id, actor)| {
1636                                        (&new_stream_actors[actor_id], actor.worker_id)
1637                                    }),
1638                                    [], // no initial subscriber for backfilling job
1639                                )
1640                            }),
1641                        );
1642                        dispatcher_update.extend(
1643                            info.upstream_fragment_downstreams.keys().flat_map(
1644                                |upstream_fragment_id| {
1645                                    let new_actor_dispatchers = edges
1646                                        .dispatchers
1647                                        .remove(upstream_fragment_id)
1648                                        .expect("should exist");
1649                                    new_actor_dispatchers.into_iter().flat_map(
1650                                        |(upstream_actor_id, dispatchers)| {
1651                                            dispatchers.into_iter().map(move |dispatcher| {
1652                                                PbDispatcherUpdate {
1653                                                    actor_id: upstream_actor_id,
1654                                                    dispatcher_id: dispatcher.dispatcher_id,
1655                                                    hash_mapping: dispatcher.hash_mapping,
1656                                                    removed_downstream_actor_id: dispatcher
1657                                                        .downstream_actor_id
1658                                                        .iter()
1659                                                        .map(|new_downstream_actor_id| {
1660                                                            actor_mapping
1661                                                            .iter()
1662                                                            .find_map(
1663                                                                |(old_actor_id, new_actor_id)| {
1664                                                                    (new_downstream_actor_id
1665                                                                        == new_actor_id)
1666                                                                        .then_some(*old_actor_id)
1667                                                                },
1668                                                            )
1669                                                            .expect("should exist")
1670                                                        })
1671                                                        .collect(),
1672                                                    added_downstream_actor_id: dispatcher
1673                                                        .downstream_actor_id,
1674                                                }
1675                                            })
1676                                        },
1677                                    )
1678                                },
1679                            ),
1680                        );
1681                        assert!(edges.is_empty(), "remaining edges: {:?}", edges);
1682                        for (worker_id, worker_actors) in new_actors_to_create {
1683                            node_actors.entry(worker_id).or_default().extend(
1684                                worker_actors.values().flat_map(|(_, actors, _)| {
1685                                    actors.iter().map(|(actor, _, _)| actor.actor_id)
1686                                }),
1687                            );
1688                            actors_to_create
1689                                .entry(worker_id)
1690                                .or_default()
1691                                .extend(worker_actors);
1692                        }
1693                        self.database_info.add_existing(InflightStreamingJobInfo {
1694                            job_id,
1695                            fragment_infos: new_fragment_info,
1696                            subscribers: Default::default(), // no initial subscribers for newly created snapshot backfill
1697                            status: CreateStreamingJobStatus::Created,
1698                            cdc_table_backfill_tracker: None, // no cdc table backfill for snapshot backfill
1699                        });
1700                    }
1701
1702                    Some(PbMutation::Update(PbUpdateMutation {
1703                        dispatcher_update,
1704                        merge_update: vec![], // no upstream update on existing actors
1705                        actor_vnode_bitmap_update: Default::default(), /* no in place update vnode bitmap happened */
1706                        dropped_actors: vec![], /* no actors to drop in the partial graph of database */
1707                        actor_splits,
1708                        actor_new_dispatchers: Default::default(), // no new dispatcher
1709                        actor_cdc_table_snapshot_splits: None, /* no cdc table backfill in snapshot backfill */
1710                        sink_schema_change: Default::default(), /* no sink auto schema change happened here */
1711                        subscriptions_to_drop,
1712                    }))
1713                } else {
1714                    let fragment_ids = self.database_info.take_pending_backfill_nodes();
1715                    if fragment_ids.is_empty() {
1716                        None
1717                    } else {
1718                        Some(PbMutation::StartFragmentBackfill(
1719                            PbStartFragmentBackfillMutation { fragment_ids },
1720                        ))
1721                    }
1722                }
1723            }
1724        };
1725
1726        // Forward barrier to independent job controls
1727        for (job_id, job) in &mut self.independent_checkpoint_job_controls {
1728            match job {
1729                IndependentCheckpointJobControl::CreatingStreamingJob(creating_job) => {
1730                    if !finished_snapshot_backfill_jobs.contains(job_id) {
1731                        let throttle_mutation = if let Some((ref jobs, ref config)) =
1732                            throttle_for_creating_jobs
1733                            && jobs.contains(job_id)
1734                        {
1735                            assert_eq!(
1736                                jobs.len(),
1737                                1,
1738                                "should not alter rate limit of snapshot backfill job with other jobs"
1739                            );
1740                            creating_job.pre_apply_throttle(config.iter().map(
1741                                |(fragment_id, (_, stream_node))| (*fragment_id, stream_node),
1742                            ));
1743                            Some((
1744                                Mutation::Throttle(ThrottleMutation {
1745                                    fragment_throttle: config
1746                                        .iter()
1747                                        .map(|(fragment_id, (throttle_config, _))| {
1748                                            (*fragment_id, *throttle_config)
1749                                        })
1750                                        .collect(),
1751                                }),
1752                                take(notifiers),
1753                            ))
1754                        } else {
1755                            None
1756                        };
1757                        creating_job.on_new_upstream_barrier(
1758                            partial_graph_manager,
1759                            &barrier_info,
1760                            throttle_mutation,
1761                        )?;
1762                    }
1763                }
1764                IndependentCheckpointJobControl::BatchRefresh(batch_refresh_job) => {
1765                    batch_refresh_job.on_new_upstream_barrier(
1766                        partial_graph_manager,
1767                        &barrier_info,
1768                        None, // no throttle mutation for batch refresh jobs
1769                    )?;
1770                }
1771            }
1772        }
1773
1774        partial_graph_manager.inject_barrier(
1775            to_partial_graph_id(self.database_id, None),
1776            mutation,
1777            &node_actors,
1778            InflightFragmentInfo::existing_table_ids(self.database_info.fragment_infos()),
1779            InflightFragmentInfo::workers(self.database_info.fragment_infos()),
1780            actors_to_create,
1781            PartialGraphBarrierInfo::new(
1782                post_collect_command,
1783                barrier_info,
1784                take(notifiers),
1785                table_ids_to_commit,
1786            ),
1787        )?;
1788
1789        Ok(ApplyCommandInfo {
1790            jobs_to_wait: finished_snapshot_backfill_jobs,
1791        })
1792    }
1793}