Skip to main content

risingwave_meta/barrier/
progress.rs

1// Copyright 2022 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::collections::HashMap;
16use std::mem::take;
17
18use risingwave_common::catalog::{FragmentTypeFlag, TableId};
19use risingwave_common::id::JobId;
20use risingwave_common::util::epoch::Epoch;
21use risingwave_pb::hummock::HummockVersionStats;
22use risingwave_pb::stream_plan::StreamNode;
23use risingwave_pb::stream_service::barrier_complete_response::CreateMviewProgress;
24
25use crate::MetaResult;
26use crate::barrier::backfill_order_control::BackfillOrderState;
27use crate::barrier::info::InflightStreamingJobInfo;
28use crate::barrier::{CreateStreamingJobCommandInfo, FragmentBackfillProgress};
29use crate::controller::fragment::InflightFragmentInfo;
30use crate::manager::MetadataManager;
31use crate::model::{ActorId, BackfillUpstreamType, FragmentId, StreamJobFragments};
32use crate::stream::{SourceChange, SourceManagerRef};
33
34type ConsumedRows = u64;
35type BufferedRows = u64;
36
37#[derive(Debug, Clone, Copy)]
38pub(crate) struct ActorBackfillProgress {
39    pub(crate) actor_id: ActorId,
40    pub(crate) upstream_type: BackfillUpstreamType,
41    pub(crate) consumed_rows: u64,
42    pub(crate) done: bool,
43}
44
45#[derive(Clone, Copy, Debug)]
46enum BackfillState {
47    Init,
48    ConsumingUpstream(#[expect(dead_code)] Epoch, ConsumedRows, BufferedRows),
49    Done(ConsumedRows, BufferedRows),
50}
51
52/// Represents the backfill nodes that need to be scheduled or cleaned up.
53#[derive(Debug, Default)]
54pub(super) struct PendingBackfillFragments {
55    /// Fragment IDs that should start backfilling in the next checkpoint
56    pub next_backfill_nodes: Vec<FragmentId>,
57    /// State tables of locality provider fragments that should be truncated
58    pub truncate_locality_provider_state_tables: Vec<TableId>,
59}
60
61/// Progress of all actors containing backfill executors while creating mview.
62#[derive(Debug)]
63pub(super) struct Progress {
64    job_id: JobId,
65    // `states` and `done_count` decides whether the progress is done. See `is_done`.
66    states: HashMap<ActorId, BackfillState>,
67    backfill_order_state: BackfillOrderState,
68    done_count: usize,
69
70    /// Tells whether the backfill is from source or mv.
71    backfill_upstream_types: HashMap<ActorId, BackfillUpstreamType>,
72
73    // The following row counts are used to calculate the progress. See `calculate_progress`.
74    /// Upstream mv count.
75    /// Keep track of how many times each upstream MV
76    /// appears in this stream job.
77    upstream_mv_count: HashMap<TableId, usize>,
78    /// Total key count of all the upstream materialized views
79    upstream_mvs_total_key_count: u64,
80    mv_backfill_consumed_rows: u64,
81    source_backfill_consumed_rows: u64,
82    /// Buffered rows (for locality backfill) that are yet to be consumed
83    /// This is used to calculate precise progress: consumed / (`upstream_total` + buffered)
84    mv_backfill_buffered_rows: u64,
85}
86
87impl Progress {
88    /// Create a [`Progress`] for some creating mview, with all `actors` containing the backfill executors.
89    fn new(
90        job_id: JobId,
91        actors: impl IntoIterator<Item = (ActorId, BackfillUpstreamType)>,
92        upstream_mv_count: HashMap<TableId, usize>,
93        upstream_total_key_count: u64,
94        backfill_order_state: BackfillOrderState,
95    ) -> Self {
96        let mut states = HashMap::new();
97        let mut backfill_upstream_types = HashMap::new();
98        for (actor, backfill_upstream_type) in actors {
99            states.insert(actor, BackfillState::Init);
100            backfill_upstream_types.insert(actor, backfill_upstream_type);
101        }
102        assert!(!states.is_empty());
103
104        Self {
105            job_id,
106            states,
107            backfill_upstream_types,
108            done_count: 0,
109            upstream_mv_count,
110            upstream_mvs_total_key_count: upstream_total_key_count,
111            mv_backfill_consumed_rows: 0,
112            source_backfill_consumed_rows: 0,
113            mv_backfill_buffered_rows: 0,
114            backfill_order_state,
115        }
116    }
117
118    /// Update the progress of `actor`.
119    /// Returns the backfill fragments that need to be scheduled or cleaned up.
120    fn update(
121        &mut self,
122        actor: ActorId,
123        new_state: BackfillState,
124        upstream_total_key_count: u64,
125    ) -> PendingBackfillFragments {
126        let mut result = PendingBackfillFragments::default();
127        self.upstream_mvs_total_key_count = upstream_total_key_count;
128        let total_actors = self.states.len();
129        let Some(backfill_upstream_type) = self.backfill_upstream_types.get(&actor) else {
130            tracing::warn!(%actor, "receive progress from unknown actor, likely removed after reschedule");
131            return result;
132        };
133
134        let mut old_consumed_row = 0;
135        let mut new_consumed_row = 0;
136        let mut old_buffered_row = 0;
137        let mut new_buffered_row = 0;
138        let Some(prev_state) = self.states.remove(&actor) else {
139            tracing::warn!(%actor, "receive progress for actor not in state map");
140            return result;
141        };
142        match prev_state {
143            BackfillState::Init => {}
144            BackfillState::ConsumingUpstream(_, consumed_rows, buffered_rows) => {
145                old_consumed_row = consumed_rows;
146                old_buffered_row = buffered_rows;
147            }
148            BackfillState::Done(_, _) => panic!("should not report done multiple times"),
149        };
150        match &new_state {
151            BackfillState::Init => {}
152            BackfillState::ConsumingUpstream(_, consumed_rows, buffered_rows) => {
153                new_consumed_row = *consumed_rows;
154                new_buffered_row = *buffered_rows;
155            }
156            BackfillState::Done(consumed_rows, buffered_rows) => {
157                tracing::debug!("actor {} done", actor);
158                new_consumed_row = *consumed_rows;
159                new_buffered_row = *buffered_rows;
160                self.done_count += 1;
161                let before_backfill_nodes = self
162                    .backfill_order_state
163                    .current_backfill_node_fragment_ids();
164                result.next_backfill_nodes = self.backfill_order_state.finish_actor(actor);
165                let after_backfill_nodes = self
166                    .backfill_order_state
167                    .current_backfill_node_fragment_ids();
168                // last_backfill_nodes = before_backfill_nodes - after_backfill_nodes
169                let last_backfill_nodes_iter = before_backfill_nodes
170                    .into_iter()
171                    .filter(|x| !after_backfill_nodes.contains(x));
172                result.truncate_locality_provider_state_tables = last_backfill_nodes_iter
173                    .filter_map(|fragment_id| {
174                        self.backfill_order_state
175                            .get_locality_fragment_state_table_mapping()
176                            .get(&fragment_id)
177                    })
178                    .flatten()
179                    .copied()
180                    .collect();
181                tracing::debug!(
182                    "{} actors out of {} complete",
183                    self.done_count,
184                    total_actors,
185                );
186            }
187        };
188        debug_assert!(
189            new_consumed_row >= old_consumed_row,
190            "backfill progress should not go backward"
191        );
192        debug_assert!(
193            new_buffered_row >= old_buffered_row,
194            "backfill progress should not go backward"
195        );
196        match backfill_upstream_type {
197            BackfillUpstreamType::MView => {
198                self.mv_backfill_consumed_rows += new_consumed_row - old_consumed_row;
199            }
200            BackfillUpstreamType::Source => {
201                self.source_backfill_consumed_rows += new_consumed_row - old_consumed_row;
202            }
203            BackfillUpstreamType::Values => {
204                // do not consider progress for values
205            }
206            BackfillUpstreamType::LocalityProvider => {
207                // Track LocalityProvider progress similar to MView
208                // Update buffered rows for precise progress calculation
209                self.mv_backfill_consumed_rows += new_consumed_row - old_consumed_row;
210                self.mv_backfill_buffered_rows += new_buffered_row - old_buffered_row;
211            }
212        }
213        self.states.insert(actor, new_state);
214        result
215    }
216
217    fn iter_actor_progress(&self) -> impl Iterator<Item = ActorBackfillProgress> + '_ {
218        self.states.iter().filter_map(|(actor_id, state)| {
219            let upstream_type = *self.backfill_upstream_types.get(actor_id)?;
220            let (consumed_rows, done) = match *state {
221                BackfillState::Init => (0, false),
222                BackfillState::ConsumingUpstream(_, consumed_rows, _) => (consumed_rows, false),
223                BackfillState::Done(consumed_rows, _) => (consumed_rows, true),
224            };
225            Some(ActorBackfillProgress {
226                actor_id: *actor_id,
227                upstream_type,
228                consumed_rows,
229                done,
230            })
231        })
232    }
233
234    /// Returns whether all backfill executors are done.
235    fn is_done(&self) -> bool {
236        tracing::trace!(
237            "Progress::is_done? {}, {}, {:?}",
238            self.done_count,
239            self.states.len(),
240            self.states
241        );
242        self.done_count == self.states.len()
243    }
244
245    /// `progress` = `consumed_rows` / `upstream_total_key_count`
246    fn calculate_progress(&self) -> String {
247        if self.is_done() || self.states.is_empty() {
248            return "100%".to_owned();
249        }
250        let mut mv_count = 0;
251        let mut source_count = 0;
252        for backfill_upstream_type in self.backfill_upstream_types.values() {
253            match backfill_upstream_type {
254                BackfillUpstreamType::MView => mv_count += 1,
255                BackfillUpstreamType::Source => source_count += 1,
256                BackfillUpstreamType::Values => (),
257                BackfillUpstreamType::LocalityProvider => mv_count += 1, /* Count LocalityProvider as an MView for progress */
258            }
259        }
260
261        let mv_progress = (mv_count > 0).then_some({
262            // Include buffered rows in total for precise progress calculation
263            // Progress = consumed / (upstream_total + buffered)
264            let total_rows_to_consume =
265                self.upstream_mvs_total_key_count + self.mv_backfill_buffered_rows;
266            if total_rows_to_consume == 0 {
267                "99.99%".to_owned()
268            } else {
269                let mut progress =
270                    self.mv_backfill_consumed_rows as f64 / (total_rows_to_consume as f64);
271                if progress > 1.0 {
272                    progress = 0.9999;
273                }
274                format!(
275                    "{:.2}% ({}/{})",
276                    progress * 100.0,
277                    self.mv_backfill_consumed_rows,
278                    total_rows_to_consume
279                )
280            }
281        });
282        let source_progress = (source_count > 0).then_some(format!(
283            "{} rows consumed",
284            self.source_backfill_consumed_rows
285        ));
286        match (mv_progress, source_progress) {
287            (Some(mv_progress), Some(source_progress)) => {
288                format!(
289                    "MView Backfill: {}, Source Backfill: {}",
290                    mv_progress, source_progress
291                )
292            }
293            (Some(mv_progress), None) => mv_progress,
294            (None, Some(source_progress)) => source_progress,
295            (None, None) => "Unknown".to_owned(),
296        }
297    }
298}
299
300/// There are two kinds of `TrackingJobs`:
301/// 1. if `is_recovered` is false, it is a "New" tracking job.
302///    It is instantiated and managed by the stream manager.
303///    On recovery, the stream manager will stop managing the job.
304/// 2. if `is_recovered` is true, it is a "Recovered" tracking job.
305///    On recovery, the barrier manager will recover and start managing the job.
306pub struct TrackingJob {
307    job_id: JobId,
308    is_recovered: bool,
309    source_change: Option<SourceChange>,
310}
311
312impl std::fmt::Display for TrackingJob {
313    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
314        write!(
315            f,
316            "{}{}",
317            self.job_id,
318            if self.is_recovered { "<recovered>" } else { "" }
319        )
320    }
321}
322
323impl TrackingJob {
324    /// Create a new tracking job.
325    pub(crate) fn new(stream_job_fragments: &StreamJobFragments) -> Self {
326        let finished_backfill_fragments = stream_job_fragments.source_backfill_fragments();
327        // `None` when empty, consistent with the recovered-job constructor.
328        let source_change = if finished_backfill_fragments.is_empty() {
329            None
330        } else {
331            Some(SourceChange::CreateJobFinished {
332                finished_backfill_fragments,
333            })
334        };
335        Self {
336            job_id: stream_job_fragments.stream_job_id,
337            is_recovered: false,
338            source_change,
339        }
340    }
341
342    /// Create a recovered tracking job.
343    pub(crate) fn recovered(
344        job_id: JobId,
345        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
346    ) -> Self {
347        Self::recovered_from_fragment_nodes(
348            job_id,
349            fragment_infos
350                .iter()
351                .map(|(fragment_id, fragment)| (*fragment_id, &fragment.nodes)),
352        )
353    }
354
355    pub(crate) fn recovered_from_fragment_nodes<'a>(
356        job_id: JobId,
357        fragment_nodes: impl Iterator<Item = (FragmentId, &'a StreamNode)>,
358    ) -> Self {
359        let source_backfill_fragments =
360            StreamJobFragments::source_backfill_fragments_impl(fragment_nodes);
361        let source_change = if source_backfill_fragments.is_empty() {
362            None
363        } else {
364            Some(SourceChange::CreateJobFinished {
365                finished_backfill_fragments: source_backfill_fragments,
366            })
367        };
368        Self {
369            job_id,
370            is_recovered: true,
371            source_change,
372        }
373    }
374
375    pub(crate) fn job_id(&self) -> JobId {
376        self.job_id
377    }
378
379    /// Notify the metadata manager that the job is finished.
380    pub(crate) async fn finish(
381        self,
382        metadata_manager: &MetadataManager,
383        source_manager: &SourceManagerRef,
384    ) -> MetaResult<()> {
385        metadata_manager
386            .catalog_controller
387            .finish_streaming_job(self.job_id)
388            .await?;
389        if let Some(source_change) = self.source_change {
390            source_manager.apply_source_change(source_change).await;
391        }
392        Ok(())
393    }
394}
395
396impl std::fmt::Debug for TrackingJob {
397    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
398        if !self.is_recovered {
399            write!(f, "TrackingJob::New({})", self.job_id)
400        } else {
401            write!(f, "TrackingJob::Recovered({})", self.job_id)
402        }
403    }
404}
405
406/// Information collected during barrier completion that needs to be committed.
407#[derive(Debug, Default)]
408pub(super) struct StagingCommitInfo {
409    /// Finished jobs that should be committed
410    pub finished_jobs: Vec<TrackingJob>,
411    /// Table IDs whose locality provider state tables need to be truncated
412    pub table_ids_to_truncate: Vec<TableId>,
413    pub finished_cdc_table_backfill: Vec<JobId>,
414}
415
416pub(super) enum UpdateProgressResult {
417    None,
418    /// The finished job, along with its pending backfill fragments for cleanup.
419    Finished {
420        truncate_locality_provider_state_tables: Vec<TableId>,
421    },
422    /// Backfill nodes have finished and new ones need to be scheduled.
423    BackfillNodeFinished(PendingBackfillFragments),
424}
425
426#[derive(Debug)]
427pub(super) struct CreateMviewProgressTracker {
428    tracking_job: TrackingJob,
429    status: CreateMviewStatus,
430}
431
432#[derive(Debug)]
433enum CreateMviewStatus {
434    Backfilling {
435        /// Progress of the create-mview DDL.
436        progress: Progress,
437
438        /// Stash of pending backfill nodes. They will start backfilling on checkpoint.
439        pending_backfill_nodes: Vec<FragmentId>,
440
441        /// Table IDs whose locality provider state tables need to be truncated
442        table_ids_to_truncate: Vec<TableId>,
443    },
444    CdcSourceInit,
445    Finished {
446        table_ids_to_truncate: Vec<TableId>,
447    },
448}
449
450impl CreateMviewProgressTracker {
451    pub fn recover(
452        creating_job_id: JobId,
453        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
454        backfill_order_state: BackfillOrderState,
455        version_stats: &HummockVersionStats,
456    ) -> Self {
457        let tracking_job = TrackingJob::recovered(creating_job_id, fragment_infos);
458        let actors = InflightStreamingJobInfo::tracking_progress_actor_ids(fragment_infos);
459        let status = if actors.is_empty() {
460            CreateMviewStatus::Finished {
461                table_ids_to_truncate: vec![],
462            }
463        } else {
464            let mut states = HashMap::new();
465            let mut backfill_upstream_types = HashMap::new();
466
467            for (actor, backfill_upstream_type) in actors {
468                states.insert(actor, BackfillState::ConsumingUpstream(Epoch(0), 0, 0));
469                backfill_upstream_types.insert(actor, backfill_upstream_type);
470            }
471
472            let progress = Self::recover_progress(
473                creating_job_id,
474                states,
475                backfill_upstream_types,
476                StreamJobFragments::upstream_table_counts_impl(
477                    fragment_infos.values().map(|fragment| &fragment.nodes),
478                ),
479                version_stats,
480                backfill_order_state,
481            );
482            let pending_backfill_nodes = progress
483                .backfill_order_state
484                .current_backfill_node_fragment_ids();
485            CreateMviewStatus::Backfilling {
486                progress,
487                pending_backfill_nodes,
488                table_ids_to_truncate: vec![],
489            }
490        };
491        Self {
492            tracking_job,
493            status,
494        }
495    }
496
497    /// ## How recovery works
498    ///
499    /// The progress (number of rows consumed) is persisted in state tables.
500    /// During recovery, the backfill executor will restore the number of rows consumed,
501    /// and then it will just report progress like newly created executors.
502    fn recover_progress(
503        job_id: JobId,
504        states: HashMap<ActorId, BackfillState>,
505        backfill_upstream_types: HashMap<ActorId, BackfillUpstreamType>,
506        upstream_mv_count: HashMap<TableId, usize>,
507        version_stats: &HummockVersionStats,
508        backfill_order_state: BackfillOrderState,
509    ) -> Progress {
510        let upstream_mvs_total_key_count =
511            calculate_total_key_count(&upstream_mv_count, version_stats);
512        Progress {
513            job_id,
514            states,
515            backfill_order_state,
516            backfill_upstream_types,
517            done_count: 0, // Fill only after first barrier pass
518            upstream_mv_count,
519            upstream_mvs_total_key_count,
520            mv_backfill_consumed_rows: 0, // Fill only after first barrier pass
521            source_backfill_consumed_rows: 0, // Fill only after first barrier pass
522            mv_backfill_buffered_rows: 0, // Fill only after first barrier pass
523        }
524    }
525
526    pub fn gen_backfill_progress(&self) -> String {
527        match &self.status {
528            CreateMviewStatus::Backfilling { progress, .. } => progress.calculate_progress(),
529            CreateMviewStatus::CdcSourceInit => "Initializing CDC source...".to_owned(),
530            CreateMviewStatus::Finished { .. } => "100%".to_owned(),
531        }
532    }
533
534    pub(crate) fn actor_progresses(&self) -> Vec<ActorBackfillProgress> {
535        match &self.status {
536            CreateMviewStatus::Backfilling { progress, .. } => {
537                progress.iter_actor_progress().collect()
538            }
539            CreateMviewStatus::CdcSourceInit | CreateMviewStatus::Finished { .. } => vec![],
540        }
541    }
542
543    /// Update the progress of tracked jobs, and add a new job to track if `info` is `Some`.
544    /// Return the table ids whose locality provider state tables need to be truncated.
545    pub(super) fn apply_progress(
546        &mut self,
547        create_mview_progress: &CreateMviewProgress,
548        version_stats: &HummockVersionStats,
549    ) {
550        let CreateMviewStatus::Backfilling {
551            progress,
552            pending_backfill_nodes,
553            table_ids_to_truncate,
554        } = &mut self.status
555        else {
556            tracing::warn!(
557                "update the progress of an backfill finished streaming job: {create_mview_progress:?}"
558            );
559            return;
560        };
561        {
562            // Update the progress of all commands.
563            {
564                // Those with actors complete can be finished immediately.
565                match progress.apply(create_mview_progress, version_stats) {
566                    UpdateProgressResult::None => {
567                        tracing::trace!(?progress, "update progress");
568                    }
569                    UpdateProgressResult::Finished {
570                        truncate_locality_provider_state_tables,
571                    } => {
572                        let mut table_ids_to_truncate = take(table_ids_to_truncate);
573                        table_ids_to_truncate.extend(truncate_locality_provider_state_tables);
574                        tracing::trace!(?progress, "finish progress");
575                        self.status = CreateMviewStatus::Finished {
576                            table_ids_to_truncate,
577                        };
578                    }
579                    UpdateProgressResult::BackfillNodeFinished(pending) => {
580                        table_ids_to_truncate
581                            .extend(pending.truncate_locality_provider_state_tables.clone());
582                        tracing::trace!(
583                            ?progress,
584                            next_backfill_nodes = ?pending.next_backfill_nodes,
585                            "start next backfill node"
586                        );
587                        pending_backfill_nodes.extend(pending.next_backfill_nodes);
588                    }
589                }
590            }
591        }
592    }
593
594    /// Refresh tracker state after reschedule so new actors can report progress correctly.
595    pub fn refresh_after_reschedule(
596        &mut self,
597        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
598        version_stats: &HummockVersionStats,
599    ) {
600        let CreateMviewStatus::Backfilling {
601            progress,
602            pending_backfill_nodes,
603            ..
604        } = &mut self.status
605        else {
606            return;
607        };
608
609        let new_tracking_actors = StreamJobFragments::tracking_progress_actor_ids_impl(
610            fragment_infos.values().map(|fragment| {
611                (
612                    fragment.fragment_type_mask,
613                    &fragment.nodes,
614                    fragment.actors.keys().copied(),
615                )
616            }),
617        );
618
619        #[cfg(debug_assertions)]
620        {
621            use std::collections::HashSet;
622            let old_actor_ids: HashSet<_> = progress.states.keys().copied().collect();
623            let new_actor_ids: HashSet<_> = new_tracking_actors
624                .iter()
625                .map(|(actor_id, _)| *actor_id)
626                .collect();
627            debug_assert!(
628                old_actor_ids.is_disjoint(&new_actor_ids),
629                "reschedule should rebuild backfill actors; old={old_actor_ids:?}, new={new_actor_ids:?}"
630            );
631        }
632
633        let mut new_states = HashMap::new();
634        let mut new_backfill_types = HashMap::new();
635        for (actor_id, upstream_type) in new_tracking_actors {
636            new_states.insert(actor_id, BackfillState::Init);
637            new_backfill_types.insert(actor_id, upstream_type);
638        }
639
640        let fragment_actors: HashMap<_, _> = fragment_infos
641            .iter()
642            .map(|(fragment_id, info)| (*fragment_id, info.actors.keys().copied().collect()))
643            .collect();
644
645        let newly_scheduled = progress
646            .backfill_order_state
647            .refresh_actors(&fragment_actors);
648
649        progress.backfill_upstream_types = new_backfill_types;
650        progress.states = new_states;
651        progress.done_count = 0;
652
653        progress.upstream_mv_count = StreamJobFragments::upstream_table_counts_impl(
654            fragment_infos.values().map(|fragment| &fragment.nodes),
655        );
656        progress.upstream_mvs_total_key_count =
657            calculate_total_key_count(&progress.upstream_mv_count, version_stats);
658
659        progress.mv_backfill_consumed_rows = 0;
660        progress.source_backfill_consumed_rows = 0;
661        progress.mv_backfill_buffered_rows = 0;
662
663        let mut pending = progress
664            .backfill_order_state
665            .current_backfill_node_fragment_ids();
666        pending.extend(newly_scheduled);
667        pending.sort_unstable();
668        pending.dedup();
669        *pending_backfill_nodes = pending;
670    }
671
672    pub(super) fn take_pending_backfill_nodes(&mut self) -> impl Iterator<Item = FragmentId> + '_ {
673        match &mut self.status {
674            CreateMviewStatus::Backfilling {
675                pending_backfill_nodes,
676                ..
677            } => Some(pending_backfill_nodes.drain(..)),
678            CreateMviewStatus::CdcSourceInit => None,
679            CreateMviewStatus::Finished { .. } => None,
680        }
681        .into_iter()
682        .flatten()
683    }
684
685    pub(super) fn collect_staging_commit_info(
686        &mut self,
687    ) -> (bool, Box<dyn Iterator<Item = TableId> + '_>) {
688        match &mut self.status {
689            CreateMviewStatus::Backfilling {
690                table_ids_to_truncate,
691                ..
692            } => (false, Box::new(table_ids_to_truncate.drain(..))),
693            CreateMviewStatus::CdcSourceInit => (false, Box::new(std::iter::empty())),
694            CreateMviewStatus::Finished {
695                table_ids_to_truncate,
696                ..
697            } => (true, Box::new(table_ids_to_truncate.drain(..))),
698        }
699    }
700
701    pub(super) fn is_finished(&self) -> bool {
702        matches!(self.status, CreateMviewStatus::Finished { .. })
703    }
704
705    /// Mark CDC source as finished when offset is updated.
706    pub(super) fn mark_cdc_source_finished(&mut self) {
707        if matches!(self.status, CreateMviewStatus::CdcSourceInit) {
708            self.status = CreateMviewStatus::Finished {
709                table_ids_to_truncate: vec![],
710            };
711        }
712    }
713
714    pub(super) fn into_tracking_job(self) -> TrackingJob {
715        let CreateMviewStatus::Finished { .. } = self.status else {
716            panic!("should be called when finished");
717        };
718        self.tracking_job
719    }
720
721    pub(crate) fn job_id(&self) -> JobId {
722        self.tracking_job.job_id
723    }
724
725    pub(crate) fn collect_fragment_progress(
726        &self,
727        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
728        mark_done_when_empty: bool,
729    ) -> Vec<FragmentBackfillProgress> {
730        let actor_progresses = self.actor_progresses();
731        if actor_progresses.is_empty() {
732            if mark_done_when_empty && self.is_finished() {
733                return collect_done_fragments(self.job_id(), fragment_infos);
734            }
735            return vec![];
736        }
737        collect_fragment_progress_from_actors(self.job_id(), fragment_infos, &actor_progresses)
738    }
739
740    /// Add a new create-mview DDL command to track.
741    ///
742    /// If the actors to track are empty, return the given command as it can be finished immediately.
743    /// For CDC sources, mark as `CdcSourceInit` instead of Finished.
744    pub fn new(
745        info: &CreateStreamingJobCommandInfo,
746        version_stats: &HummockVersionStats,
747        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
748    ) -> Self {
749        tracing::trace!(?info, "add job to track");
750        let CreateStreamingJobCommandInfo {
751            stream_job_fragments,
752            fragment_backfill_ordering,
753            locality_fragment_state_table_mapping,
754            streaming_job,
755            ..
756        } = info;
757        let job_id = stream_job_fragments.stream_job_id();
758        let actors = InflightStreamingJobInfo::tracking_progress_actor_ids(fragment_infos);
759        let tracking_job = TrackingJob::new(&info.stream_job_fragments);
760        if actors.is_empty() {
761            // NOTE: This CDC source detection uses hardcoded property checks and should be replaced
762            // with a more reliable identification method in the future.
763            let is_cdc_source = matches!(
764                streaming_job,
765                crate::manager::StreamingJob::Source(source)
766                    if source.info.as_ref().map(|info| info.is_shared()).unwrap_or(false) && source
767                    .get_with_properties()
768                    .get("connector")
769                    .map(|connector| connector.to_lowercase().contains("-cdc"))
770                    .unwrap_or(false)
771            );
772            if is_cdc_source {
773                // Mark CDC source as CdcSourceInit, will be finished when offset is updated
774                return Self {
775                    tracking_job,
776                    status: CreateMviewStatus::CdcSourceInit,
777                };
778            }
779            // The command can be finished immediately.
780            return Self {
781                tracking_job,
782                status: CreateMviewStatus::Finished {
783                    table_ids_to_truncate: vec![],
784                },
785            };
786        }
787
788        let upstream_mv_count = stream_job_fragments.upstream_table_counts();
789        let upstream_total_key_count: u64 =
790            calculate_total_key_count(&upstream_mv_count, version_stats);
791
792        let backfill_order_state = BackfillOrderState::new(
793            fragment_backfill_ordering,
794            fragment_infos,
795            locality_fragment_state_table_mapping.clone(),
796        );
797        let progress = Progress::new(
798            job_id,
799            actors,
800            upstream_mv_count,
801            upstream_total_key_count,
802            backfill_order_state,
803        );
804        let pending_backfill_nodes = progress
805            .backfill_order_state
806            .current_backfill_node_fragment_ids();
807        Self {
808            tracking_job,
809            status: CreateMviewStatus::Backfilling {
810                progress,
811                pending_backfill_nodes,
812                table_ids_to_truncate: vec![],
813            },
814        }
815    }
816}
817
818impl Progress {
819    /// Update the progress of `actor` according to the Pb struct.
820    ///
821    /// If all actors in this MV have finished, return the command.
822    fn apply(
823        &mut self,
824        progress: &CreateMviewProgress,
825        version_stats: &HummockVersionStats,
826    ) -> UpdateProgressResult {
827        tracing::trace!(?progress, "update progress");
828        let actor = progress.backfill_actor_id;
829        let job_id = self.job_id;
830
831        let new_state = if progress.done {
832            BackfillState::Done(progress.consumed_rows, progress.buffered_rows)
833        } else {
834            BackfillState::ConsumingUpstream(
835                progress.consumed_epoch.into(),
836                progress.consumed_rows,
837                progress.buffered_rows,
838            )
839        };
840
841        {
842            {
843                let progress_state = self;
844
845                let upstream_total_key_count: u64 =
846                    calculate_total_key_count(&progress_state.upstream_mv_count, version_stats);
847
848                tracing::trace!(%job_id, "updating progress for table");
849                let pending = progress_state.update(actor, new_state, upstream_total_key_count);
850
851                if progress_state.is_done() {
852                    tracing::debug!(
853                        %job_id,
854                        "all actors done for creating mview!",
855                    );
856
857                    let PendingBackfillFragments {
858                        next_backfill_nodes,
859                        truncate_locality_provider_state_tables,
860                    } = pending;
861
862                    assert!(next_backfill_nodes.is_empty());
863                    UpdateProgressResult::Finished {
864                        truncate_locality_provider_state_tables,
865                    }
866                } else if !pending.next_backfill_nodes.is_empty()
867                    || !pending.truncate_locality_provider_state_tables.is_empty()
868                {
869                    UpdateProgressResult::BackfillNodeFinished(pending)
870                } else {
871                    UpdateProgressResult::None
872                }
873            }
874        }
875    }
876}
877
878fn calculate_total_key_count(
879    table_count: &HashMap<TableId, usize>,
880    version_stats: &HummockVersionStats,
881) -> u64 {
882    table_count
883        .iter()
884        .map(|(table_id, count)| {
885            assert_ne!(*count, 0);
886            *count as u64
887                * version_stats
888                    .table_stats
889                    .get(table_id)
890                    .map_or(0, |stat| stat.total_key_count as u64)
891        })
892        .sum()
893}
894
895pub(crate) fn collect_fragment_progress_from_actors(
896    job_id: JobId,
897    fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
898    actor_progresses: &[ActorBackfillProgress],
899) -> Vec<FragmentBackfillProgress> {
900    let mut actor_to_fragment = HashMap::new();
901    for (fragment_id, info) in fragment_infos {
902        for actor_id in info.actors.keys() {
903            actor_to_fragment.insert(*actor_id, *fragment_id);
904        }
905    }
906
907    let mut per_fragment: HashMap<FragmentId, (u64, usize, usize, BackfillUpstreamType)> =
908        HashMap::new();
909    for progress in actor_progresses {
910        let Some(fragment_id) = actor_to_fragment.get(&progress.actor_id) else {
911            continue;
912        };
913        let entry = per_fragment
914            .entry(*fragment_id)
915            .or_insert((0, 0, 0, progress.upstream_type));
916        entry.0 = entry.0.saturating_add(progress.consumed_rows);
917        entry.1 += progress.done as usize;
918        entry.2 += 1;
919    }
920
921    per_fragment
922        .into_iter()
923        .map(
924            |(fragment_id, (consumed_rows, done_cnt, total_cnt, upstream_type))| {
925                FragmentBackfillProgress {
926                    job_id,
927                    fragment_id,
928                    consumed_rows,
929                    done: total_cnt > 0 && done_cnt == total_cnt,
930                    upstream_type,
931                }
932            },
933        )
934        .collect()
935}
936
937pub(crate) fn collect_done_fragments(
938    job_id: JobId,
939    fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
940) -> Vec<FragmentBackfillProgress> {
941    fragment_infos
942        .iter()
943        .filter(|(_, fragment)| {
944            fragment.fragment_type_mask.contains_any([
945                FragmentTypeFlag::StreamScan,
946                FragmentTypeFlag::SourceScan,
947                FragmentTypeFlag::LocalityProvider,
948            ])
949        })
950        .map(|(fragment_id, fragment)| FragmentBackfillProgress {
951            job_id,
952            fragment_id: *fragment_id,
953            consumed_rows: 0,
954            done: true,
955            upstream_type: BackfillUpstreamType::from_fragment_type_mask(
956                fragment.fragment_type_mask,
957            ),
958        })
959        .collect()
960}
961
962#[cfg(test)]
963mod tests {
964    use std::collections::HashSet;
965
966    use risingwave_common::catalog::{FragmentTypeFlag, FragmentTypeMask};
967    use risingwave_common::id::WorkerId;
968    use risingwave_meta_model::fragment::DistributionType;
969    use risingwave_pb::stream_plan::StreamNode as PbStreamNode;
970
971    use super::*;
972    use crate::controller::fragment::InflightActorInfo;
973
974    fn sample_inflight_fragment(
975        fragment_id: FragmentId,
976        actor_ids: &[ActorId],
977        flag: FragmentTypeFlag,
978    ) -> InflightFragmentInfo {
979        let mut fragment_type_mask = FragmentTypeMask::empty();
980        fragment_type_mask.add(flag);
981        InflightFragmentInfo {
982            fragment_id,
983            distribution_type: DistributionType::Single,
984            fragment_type_mask,
985            vnode_count: 0,
986            nodes: PbStreamNode::default(),
987            actors: actor_ids
988                .iter()
989                .map(|actor_id| {
990                    (
991                        *actor_id,
992                        InflightActorInfo {
993                            worker_id: WorkerId::new(1),
994                            vnode_bitmap: None,
995                            splits: vec![],
996                        },
997                    )
998                })
999                .collect(),
1000            state_table_ids: HashSet::new(),
1001        }
1002    }
1003
1004    #[test]
1005    fn test_recover_legacy_cdc_progress_in_mixed_job() {
1006        use risingwave_pb::stream_plan::stream_node::NodeBody;
1007
1008        for modern_mask in [false, true] {
1009            let mut cdc = sample_inflight_fragment(
1010                FragmentId::new(10),
1011                &[ActorId::new(1)],
1012                FragmentTypeFlag::StreamScan,
1013            );
1014            if modern_mask {
1015                cdc.fragment_type_mask.add(FragmentTypeFlag::StreamCdcScan);
1016            }
1017            // The CDC scan can be nested under a materialize/project node. Old plans may
1018            // also have no scan options; neither affects progress-reporter ownership.
1019            cdc.nodes = PbStreamNode {
1020                node_body: Some(NodeBody::Materialize(Default::default())),
1021                input: vec![PbStreamNode {
1022                    node_body: Some(NodeBody::StreamCdcScan(Default::default())),
1023                    ..Default::default()
1024                }],
1025                ..Default::default()
1026            };
1027            let mut fragments = HashMap::from([(cdc.fragment_id, cdc)]);
1028            let recovered = CreateMviewProgressTracker::recover(
1029                JobId::new(1),
1030                &fragments,
1031                Default::default(),
1032                &HummockVersionStats::default(),
1033            );
1034            assert!(matches!(
1035                recovered.status,
1036                CreateMviewStatus::Finished { .. }
1037            ));
1038
1039            let mv = sample_inflight_fragment(
1040                FragmentId::new(20),
1041                &[ActorId::new(2)],
1042                FragmentTypeFlag::StreamScan,
1043            );
1044            fragments.insert(mv.fragment_id, mv);
1045            let source = sample_inflight_fragment(
1046                FragmentId::new(30),
1047                &[ActorId::new(3)],
1048                FragmentTypeFlag::SourceScan,
1049            );
1050            fragments.insert(source.fragment_id, source);
1051            let mut recovered = CreateMviewProgressTracker::recover(
1052                JobId::new(1),
1053                &fragments,
1054                Default::default(),
1055                &HummockVersionStats::default(),
1056            );
1057            let CreateMviewStatus::Backfilling { progress, .. } = &recovered.status else {
1058                panic!("ordinary backfill fragments must remain tracked");
1059            };
1060            assert_eq!(progress.states.len(), 2);
1061            assert!(!progress.states.contains_key(&ActorId::new(1)));
1062            assert_eq!(
1063                progress.backfill_upstream_types[&ActorId::new(2)],
1064                BackfillUpstreamType::MView
1065            );
1066            assert_eq!(
1067                progress.backfill_upstream_types[&ActorId::new(3)],
1068                BackfillUpstreamType::Source
1069            );
1070
1071            let mv = sample_inflight_fragment(
1072                FragmentId::new(20),
1073                &[ActorId::new(4)],
1074                FragmentTypeFlag::StreamScan,
1075            );
1076            fragments.insert(mv.fragment_id, mv);
1077            let source = sample_inflight_fragment(
1078                FragmentId::new(30),
1079                &[ActorId::new(5)],
1080                FragmentTypeFlag::SourceScan,
1081            );
1082            fragments.insert(source.fragment_id, source);
1083            recovered.refresh_after_reschedule(&fragments, &HummockVersionStats::default());
1084            let CreateMviewStatus::Backfilling { progress, .. } = recovered.status else {
1085                panic!("ordinary backfill fragments must remain tracked after reschedule");
1086            };
1087            assert_eq!(
1088                progress.states.keys().copied().collect::<HashSet<_>>(),
1089                HashSet::from([ActorId::new(4), ActorId::new(5)])
1090            );
1091        }
1092    }
1093
1094    fn sample_progress(actor_id: ActorId) -> Progress {
1095        Progress {
1096            job_id: JobId::new(1),
1097            states: HashMap::from([(actor_id, BackfillState::Init)]),
1098            backfill_order_state: BackfillOrderState::default(),
1099            done_count: 0,
1100            backfill_upstream_types: HashMap::from([(actor_id, BackfillUpstreamType::MView)]),
1101            upstream_mv_count: HashMap::new(),
1102            upstream_mvs_total_key_count: 0,
1103            mv_backfill_consumed_rows: 0,
1104            source_backfill_consumed_rows: 0,
1105            mv_backfill_buffered_rows: 0,
1106        }
1107    }
1108
1109    #[test]
1110    fn update_ignores_unknown_actor() {
1111        let actor_known = ActorId::new(1);
1112        let actor_unknown = ActorId::new(2);
1113        let mut progress = sample_progress(actor_known);
1114
1115        let pending = progress.update(
1116            actor_unknown,
1117            BackfillState::Done(0, 0),
1118            progress.upstream_mvs_total_key_count,
1119        );
1120
1121        assert!(pending.next_backfill_nodes.is_empty());
1122        assert_eq!(progress.states.len(), 1);
1123        assert!(progress.states.contains_key(&actor_known));
1124    }
1125
1126    #[test]
1127    fn refresh_rebuilds_tracking_after_reschedule() {
1128        let actor_old = ActorId::new(1);
1129        let actor_new = ActorId::new(2);
1130
1131        let progress = Progress {
1132            job_id: JobId::new(1),
1133            states: HashMap::from([(actor_old, BackfillState::Done(5, 0))]),
1134            backfill_order_state: BackfillOrderState::default(),
1135            done_count: 1,
1136            backfill_upstream_types: HashMap::from([(actor_old, BackfillUpstreamType::MView)]),
1137            upstream_mv_count: HashMap::new(),
1138            upstream_mvs_total_key_count: 0,
1139            mv_backfill_consumed_rows: 5,
1140            source_backfill_consumed_rows: 0,
1141            mv_backfill_buffered_rows: 0,
1142        };
1143
1144        let mut tracker = CreateMviewProgressTracker {
1145            tracking_job: TrackingJob {
1146                job_id: JobId::new(1),
1147                is_recovered: false,
1148                source_change: None,
1149            },
1150            status: CreateMviewStatus::Backfilling {
1151                progress,
1152                pending_backfill_nodes: vec![],
1153                table_ids_to_truncate: vec![],
1154            },
1155        };
1156
1157        let fragment_infos = HashMap::from([(
1158            FragmentId::new(10),
1159            sample_inflight_fragment(
1160                FragmentId::new(10),
1161                &[actor_new],
1162                FragmentTypeFlag::StreamScan,
1163            ),
1164        )]);
1165
1166        tracker.refresh_after_reschedule(&fragment_infos, &HummockVersionStats::default());
1167
1168        let CreateMviewStatus::Backfilling { progress, .. } = tracker.status else {
1169            panic!("expected backfilling status");
1170        };
1171        assert!(progress.states.contains_key(&actor_new));
1172        assert!(!progress.states.contains_key(&actor_old));
1173        assert_eq!(progress.done_count, 0);
1174        assert_eq!(progress.mv_backfill_consumed_rows, 0);
1175        assert_eq!(progress.source_backfill_consumed_rows, 0);
1176    }
1177
1178    // CDC sources should be initialized as CdcSourceInit
1179    #[test]
1180    fn test_cdc_source_initialized_as_cdc_source_init() {
1181        use std::collections::BTreeMap;
1182
1183        use risingwave_meta_model::streaming_job;
1184        use risingwave_pb::catalog::{CreateType, PbSource, StreamSourceInfo};
1185
1186        use crate::barrier::command::CreateStreamingJobCommandInfo;
1187        use crate::manager::{StreamingJob, StreamingJobType};
1188        use crate::model::StreamJobFragmentsToCreate;
1189
1190        // Create a CDC source with cdc_source_job = true
1191        let source_info = StreamSourceInfo {
1192            cdc_source_job: true,
1193            ..Default::default()
1194        };
1195
1196        let source = PbSource {
1197            id: risingwave_common::id::SourceId::new(100),
1198            info: Some(source_info),
1199            with_properties: BTreeMap::from([("connector".to_owned(), "fake-cdc".to_owned())]),
1200            ..Default::default()
1201        };
1202
1203        // Create empty fragments (no actors to track)
1204        let fragments = StreamJobFragments::for_test(JobId::new(100), BTreeMap::new());
1205        let stream_job_fragments = StreamJobFragmentsToCreate {
1206            inner: fragments,
1207            downstreams: Default::default(),
1208        };
1209
1210        let info = CreateStreamingJobCommandInfo {
1211            stream_job_fragments,
1212            upstream_fragment_downstreams: Default::default(),
1213            init_split_assignment: Default::default(),
1214            definition: "CREATE SOURCE ...".to_owned(),
1215            job_type: StreamingJobType::Source,
1216            create_type: CreateType::Foreground,
1217            streaming_job: StreamingJob::Source(source),
1218            database_resource_group: risingwave_common::util::worker_util::DEFAULT_RESOURCE_GROUP
1219                .to_owned(),
1220            fragment_backfill_ordering: Default::default(),
1221            cdc_table_snapshot_splits: None,
1222            locality_fragment_state_table_mapping: Default::default(),
1223            is_serverless: false,
1224            replace_sink: None,
1225            refresh_interval_sec: None,
1226            streaming_job_model: streaming_job::Model {
1227                job_id: JobId::new(100),
1228                job_status: risingwave_meta_model::JobStatus::Creating,
1229                create_type: risingwave_meta_model::CreateType::Foreground,
1230                timezone: None,
1231                config_override: None,
1232                adaptive_parallelism_strategy: None,
1233                parallelism: risingwave_meta_model::StreamingParallelism::Adaptive,
1234                backfill_parallelism: None,
1235                backfill_adaptive_parallelism_strategy: None,
1236                backfill_orders: None,
1237                max_parallelism: 256,
1238                specific_resource_group: None,
1239                is_serverless_backfill: false,
1240                refresh_interval_sec: None,
1241            },
1242        };
1243
1244        let tracker = CreateMviewProgressTracker::new(
1245            &info,
1246            &HummockVersionStats::default(),
1247            &HashMap::new(),
1248        );
1249
1250        // CDC source should be in CdcSourceInit state
1251        assert!(matches!(tracker.status, CreateMviewStatus::CdcSourceInit));
1252        assert!(!tracker.is_finished());
1253    }
1254
1255    // CDC source should transition from CdcSourceInit to Finished when offset is updated
1256    #[test]
1257    fn test_cdc_source_transitions_to_finished_on_offset_update() {
1258        let mut tracker = CreateMviewProgressTracker {
1259            tracking_job: TrackingJob {
1260                job_id: JobId::new(300),
1261                is_recovered: false,
1262                source_change: None,
1263            },
1264            status: CreateMviewStatus::CdcSourceInit,
1265        };
1266
1267        // Initially in CdcSourceInit state
1268        assert!(matches!(tracker.status, CreateMviewStatus::CdcSourceInit));
1269        assert!(!tracker.is_finished());
1270
1271        // Mark as finished when offset is updated
1272        tracker.mark_cdc_source_finished();
1273
1274        // Should now be in Finished state
1275        assert!(matches!(tracker.status, CreateMviewStatus::Finished { .. }));
1276        assert!(tracker.is_finished());
1277    }
1278
1279    #[test]
1280    fn tracking_job_new_without_source_backfill_has_no_source_change() {
1281        use std::collections::BTreeMap;
1282
1283        let fragments = StreamJobFragments::for_test(JobId::new(1), BTreeMap::new());
1284        let job = TrackingJob::new(&fragments);
1285        assert!(job.source_change.is_none());
1286    }
1287
1288    #[test]
1289    fn tracking_job_new_with_source_backfill_tracks_finished_fragments() {
1290        use std::collections::{BTreeMap, BTreeSet};
1291
1292        use risingwave_common::id::SourceId;
1293        use risingwave_pb::stream_plan::stream_node::NodeBody;
1294        use risingwave_pb::stream_plan::{MergeNode, SourceBackfillNode};
1295
1296        use crate::model::Fragment;
1297
1298        let source_id = SourceId::new(42);
1299        let backfill_fragment_id = FragmentId::new(2);
1300        let upstream_source_fragment_id = FragmentId::new(1);
1301
1302        let nodes = PbStreamNode {
1303            node_body: Some(NodeBody::SourceBackfill(Box::new(SourceBackfillNode {
1304                upstream_source_id: source_id,
1305                ..Default::default()
1306            }))),
1307            input: vec![PbStreamNode {
1308                node_body: Some(NodeBody::Merge(Box::new(MergeNode {
1309                    upstream_fragment_id: upstream_source_fragment_id,
1310                    ..Default::default()
1311                }))),
1312                ..Default::default()
1313            }],
1314            ..Default::default()
1315        };
1316        let fragment = Fragment {
1317            fragment_id: backfill_fragment_id,
1318            nodes,
1319            ..Default::default()
1320        };
1321        let fragments = StreamJobFragments::for_test(
1322            JobId::new(1),
1323            BTreeMap::from([(backfill_fragment_id, fragment)]),
1324        );
1325
1326        let job = TrackingJob::new(&fragments);
1327        let Some(SourceChange::CreateJobFinished {
1328            finished_backfill_fragments,
1329        }) = job.source_change
1330        else {
1331            panic!("expected CreateJobFinished");
1332        };
1333        assert_eq!(
1334            finished_backfill_fragments,
1335            HashMap::from([(
1336                source_id,
1337                BTreeSet::from([(backfill_fragment_id, upstream_source_fragment_id)]),
1338            )])
1339        );
1340    }
1341}