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