risingwave_meta/barrier/

progress.rs

// Copyright 2025 RisingWave Labs
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::HashMap;
use std::mem::take;

use risingwave_common::catalog::TableId;
use risingwave_common::id::JobId;
use risingwave_common::util::epoch::Epoch;
use risingwave_meta_model::CreateType;
use risingwave_pb::ddl_service::DdlProgress;
use risingwave_pb::hummock::HummockVersionStats;
use risingwave_pb::stream_service::barrier_complete_response::CreateMviewProgress;

use crate::MetaResult;
use crate::barrier::CreateStreamingJobCommandInfo;
use crate::barrier::backfill_order_control::BackfillOrderState;
use crate::barrier::info::InflightStreamingJobInfo;
use crate::controller::fragment::InflightFragmentInfo;
use crate::manager::MetadataManager;
use crate::model::{ActorId, BackfillUpstreamType, FragmentId, StreamJobFragments};
use crate::stream::{SourceChange, SourceManagerRef};

type ConsumedRows = u64;
type BufferedRows = u64;

#[derive(Clone, Copy, Debug)]
enum BackfillState {
    Init,
    ConsumingUpstream(#[expect(dead_code)] Epoch, ConsumedRows, BufferedRows),
    Done(ConsumedRows, BufferedRows),
}

/// Represents the backfill nodes that need to be scheduled or cleaned up.
#[derive(Debug, Default)]
pub(super) struct PendingBackfillFragments {
    /// Fragment IDs that should start backfilling in the next checkpoint
    pub next_backfill_nodes: Vec<FragmentId>,
    /// State tables of locality provider fragments that should be truncated
    pub truncate_locality_provider_state_tables: Vec<TableId>,
}

/// Progress of all actors containing backfill executors while creating mview.
#[derive(Debug, Clone)]
pub(super) struct Progress {
    job_id: JobId,
    // `states` and `done_count` decides whether the progress is done. See `is_done`.
    states: HashMap<ActorId, BackfillState>,
    backfill_order_state: BackfillOrderState,
    done_count: usize,

    /// Tells whether the backfill is from source or mv.
    backfill_upstream_types: HashMap<ActorId, BackfillUpstreamType>,

    // The following row counts are used to calculate the progress. See `calculate_progress`.
    /// Upstream mv count.
    /// Keep track of how many times each upstream MV
    /// appears in this stream job.
    upstream_mv_count: HashMap<TableId, usize>,
    /// Total key count of all the upstream materialized views
    upstream_mvs_total_key_count: u64,
    mv_backfill_consumed_rows: u64,
    source_backfill_consumed_rows: u64,
    /// Buffered rows (for locality backfill) that are yet to be consumed
    /// This is used to calculate precise progress: consumed / (`upstream_total` + buffered)
    mv_backfill_buffered_rows: u64,
}

impl Progress {
    /// Create a [`Progress`] for some creating mview, with all `actors` containing the backfill executors.
    fn new(
        job_id: JobId,
        actors: impl IntoIterator<Item = (ActorId, BackfillUpstreamType)>,
        upstream_mv_count: HashMap<TableId, usize>,
        upstream_total_key_count: u64,
        backfill_order_state: BackfillOrderState,
    ) -> Self {
        let mut states = HashMap::new();
        let mut backfill_upstream_types = HashMap::new();
        for (actor, backfill_upstream_type) in actors {
            states.insert(actor, BackfillState::Init);
            backfill_upstream_types.insert(actor, backfill_upstream_type);
        }
        assert!(!states.is_empty());

        Self {
            job_id,
            states,
            backfill_upstream_types,
            done_count: 0,
            upstream_mv_count,
            upstream_mvs_total_key_count: upstream_total_key_count,
            mv_backfill_consumed_rows: 0,
            source_backfill_consumed_rows: 0,
            mv_backfill_buffered_rows: 0,
            backfill_order_state,
        }
    }

    /// Update the progress of `actor`.
    /// Returns the backfill fragments that need to be scheduled or cleaned up.
    fn update(
        &mut self,
        actor: ActorId,
        new_state: BackfillState,
        upstream_total_key_count: u64,
    ) -> PendingBackfillFragments {
        let mut result = PendingBackfillFragments::default();
        self.upstream_mvs_total_key_count = upstream_total_key_count;
        let total_actors = self.states.len();
        let backfill_upstream_type = self.backfill_upstream_types.get(&actor).unwrap();

        let mut old_consumed_row = 0;
        let mut new_consumed_row = 0;
        let mut old_buffered_row = 0;
        let mut new_buffered_row = 0;
        match self.states.remove(&actor).unwrap() {
            BackfillState::Init => {}
            BackfillState::ConsumingUpstream(_, consumed_rows, buffered_rows) => {
                old_consumed_row = consumed_rows;
                old_buffered_row = buffered_rows;
            }
            BackfillState::Done(_, _) => panic!("should not report done multiple times"),
        };
        match &new_state {
            BackfillState::Init => {}
            BackfillState::ConsumingUpstream(_, consumed_rows, buffered_rows) => {
                new_consumed_row = *consumed_rows;
                new_buffered_row = *buffered_rows;
            }
            BackfillState::Done(consumed_rows, buffered_rows) => {
                tracing::debug!("actor {} done", actor);
                new_consumed_row = *consumed_rows;
                new_buffered_row = *buffered_rows;
                self.done_count += 1;
                let before_backfill_nodes = self
                    .backfill_order_state
                    .current_backfill_node_fragment_ids();
                result.next_backfill_nodes = self.backfill_order_state.finish_actor(actor);
                let after_backfill_nodes = self
                    .backfill_order_state
                    .current_backfill_node_fragment_ids();
                // last_backfill_nodes = before_backfill_nodes - after_backfill_nodes
                let last_backfill_nodes_iter = before_backfill_nodes
                    .into_iter()
                    .filter(|x| !after_backfill_nodes.contains(x));
                result.truncate_locality_provider_state_tables = last_backfill_nodes_iter
                    .filter_map(|fragment_id| {
                        self.backfill_order_state
                            .get_locality_fragment_state_table_mapping()
                            .get(&fragment_id)
                    })
                    .flatten()
                    .copied()
                    .collect();
                tracing::debug!(
                    "{} actors out of {} complete",
                    self.done_count,
                    total_actors,
                );
            }
        };
        debug_assert!(
            new_consumed_row >= old_consumed_row,
            "backfill progress should not go backward"
        );
        debug_assert!(
            new_buffered_row >= old_buffered_row,
            "backfill progress should not go backward"
        );
        match backfill_upstream_type {
            BackfillUpstreamType::MView => {
                self.mv_backfill_consumed_rows += new_consumed_row - old_consumed_row;
            }
            BackfillUpstreamType::Source => {
                self.source_backfill_consumed_rows += new_consumed_row - old_consumed_row;
            }
            BackfillUpstreamType::Values => {
                // do not consider progress for values
            }
            BackfillUpstreamType::LocalityProvider => {
                // Track LocalityProvider progress similar to MView
                // Update buffered rows for precise progress calculation
                self.mv_backfill_consumed_rows += new_consumed_row - old_consumed_row;
                self.mv_backfill_buffered_rows += new_buffered_row - old_buffered_row;
            }
        }
        self.states.insert(actor, new_state);
        result
    }

    /// Returns whether all backfill executors are done.
    fn is_done(&self) -> bool {
        tracing::trace!(
            "Progress::is_done? {}, {}, {:?}",
            self.done_count,
            self.states.len(),
            self.states
        );
        self.done_count == self.states.len()
    }

    /// `progress` = `consumed_rows` / `upstream_total_key_count`
    fn calculate_progress(&self) -> String {
        if self.is_done() || self.states.is_empty() {
            return "100%".to_owned();
        }
        let mut mv_count = 0;
        let mut source_count = 0;
        for backfill_upstream_type in self.backfill_upstream_types.values() {
            match backfill_upstream_type {
                BackfillUpstreamType::MView => mv_count += 1,
                BackfillUpstreamType::Source => source_count += 1,
                BackfillUpstreamType::Values => (),
                BackfillUpstreamType::LocalityProvider => mv_count += 1, /* Count LocalityProvider as an MView for progress */
            }
        }

        let mv_progress = (mv_count > 0).then_some({
            // Include buffered rows in total for precise progress calculation
            // Progress = consumed / (upstream_total + buffered)
            let total_rows_to_consume =
                self.upstream_mvs_total_key_count + self.mv_backfill_buffered_rows;
            if total_rows_to_consume == 0 {
                "99.99%".to_owned()
            } else {
                let mut progress =
                    self.mv_backfill_consumed_rows as f64 / (total_rows_to_consume as f64);
                if progress > 1.0 {
                    progress = 0.9999;
                }
                format!(
                    "{:.2}% ({}/{})",
                    progress * 100.0,
                    self.mv_backfill_consumed_rows,
                    total_rows_to_consume
                )
            }
        });
        let source_progress = (source_count > 0).then_some(format!(
            "{} rows consumed",
            self.source_backfill_consumed_rows
        ));
        match (mv_progress, source_progress) {
            (Some(mv_progress), Some(source_progress)) => {
                format!(
                    "MView Backfill: {}, Source Backfill: {}",
                    mv_progress, source_progress
                )
            }
            (Some(mv_progress), None) => mv_progress,
            (None, Some(source_progress)) => source_progress,
            (None, None) => "Unknown".to_owned(),
        }
    }
}

/// There are two kinds of `TrackingJobs`:
/// 1. if `is_recovered` is false, it is a "New" tracking job.
///    It is instantiated and managed by the stream manager.
///    On recovery, the stream manager will stop managing the job.
/// 2. if `is_recovered` is true, it is a "Recovered" tracking job.
///    On recovery, the barrier manager will recover and start managing the job.
#[derive(Clone)]
pub struct TrackingJob {
    job_id: JobId,
    is_recovered: bool,
    source_change: Option<SourceChange>,
}

impl std::fmt::Display for TrackingJob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{}{}",
            self.job_id,
            if self.is_recovered { "<recovered>" } else { "" }
        )
    }
}

impl TrackingJob {
    /// Create a new tracking job.
    pub(crate) fn new(stream_job_fragments: &StreamJobFragments) -> Self {
        Self {
            job_id: stream_job_fragments.stream_job_id,
            is_recovered: false,
            source_change: Some(SourceChange::CreateJobFinished {
                finished_backfill_fragments: stream_job_fragments.source_backfill_fragments(),
            }),
        }
    }

    /// Create a recovered tracking job.
    pub(crate) fn recovered(
        job_id: JobId,
        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
    ) -> Self {
        let source_backfill_fragments = StreamJobFragments::source_backfill_fragments_impl(
            fragment_infos
                .iter()
                .map(|(fragment_id, fragment)| (*fragment_id, &fragment.nodes)),
        );
        let source_change = if source_backfill_fragments.is_empty() {
            None
        } else {
            Some(SourceChange::CreateJobFinished {
                finished_backfill_fragments: source_backfill_fragments,
            })
        };
        Self {
            job_id,
            is_recovered: true,
            source_change,
        }
    }

    /// Notify the metadata manager that the job is finished.
    pub(crate) async fn finish(
        self,
        metadata_manager: &MetadataManager,
        source_manager: &SourceManagerRef,
    ) -> MetaResult<()> {
        metadata_manager
            .catalog_controller
            .finish_streaming_job(self.job_id)
            .await?;
        if let Some(source_change) = self.source_change {
            source_manager.apply_source_change(source_change).await;
        }
        Ok(())
    }
}

impl std::fmt::Debug for TrackingJob {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if !self.is_recovered {
            write!(f, "TrackingJob::New({})", self.job_id)
        } else {
            write!(f, "TrackingJob::Recovered({})", self.job_id)
        }
    }
}

/// Information collected during barrier completion that needs to be committed.
#[derive(Debug, Default)]
pub(super) struct StagingCommitInfo {
    /// Finished jobs that should be committed
    pub finished_jobs: Vec<TrackingJob>,
    /// Table IDs whose locality provider state tables need to be truncated
    pub table_ids_to_truncate: Vec<TableId>,
}

pub(super) enum UpdateProgressResult {
    None,
    /// The finished job, along with its pending backfill fragments for cleanup.
    Finished {
        truncate_locality_provider_state_tables: Vec<TableId>,
    },
    /// Backfill nodes have finished and new ones need to be scheduled.
    BackfillNodeFinished(PendingBackfillFragments),
}

#[derive(Debug, Clone)]
pub(super) struct CreateMviewProgressTracker {
    job_id: JobId,
    definition: String,
    create_type: CreateType,
    tracking_job: TrackingJob,
    status: CreateMviewStatus,
}

#[derive(Debug, Clone)]
enum CreateMviewStatus {
    Backfilling {
        /// Progress of the create-mview DDL.
        progress: Progress,

        /// Stash of pending backfill nodes. They will start backfilling on checkpoint.
        pending_backfill_nodes: Vec<FragmentId>,

        /// Table IDs whose locality provider state tables need to be truncated
        table_ids_to_truncate: Vec<TableId>,
    },
    Finished {
        table_ids_to_truncate: Vec<TableId>,
    },
}

impl CreateMviewProgressTracker {
    pub fn recover(
        creating_job_id: JobId,
        definition: String,
        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
        backfill_order_state: BackfillOrderState,
        version_stats: &HummockVersionStats,
    ) -> Self {
        {
            let create_type = CreateType::Background;
            let tracking_job = TrackingJob::recovered(creating_job_id, fragment_infos);
            let actors = InflightStreamingJobInfo::tracking_progress_actor_ids(fragment_infos);
            let status = if actors.is_empty() {
                CreateMviewStatus::Finished {
                    table_ids_to_truncate: vec![],
                }
            } else {
                let mut states = HashMap::new();
                let mut backfill_upstream_types = HashMap::new();

                for (actor, backfill_upstream_type) in actors {
                    states.insert(actor, BackfillState::ConsumingUpstream(Epoch(0), 0, 0));
                    backfill_upstream_types.insert(actor, backfill_upstream_type);
                }

                let progress = Self::recover_progress(
                    creating_job_id,
                    states,
                    backfill_upstream_types,
                    StreamJobFragments::upstream_table_counts_impl(
                        fragment_infos.values().map(|fragment| &fragment.nodes),
                    ),
                    version_stats,
                    backfill_order_state,
                );
                CreateMviewStatus::Backfilling {
                    progress,
                    pending_backfill_nodes: vec![],
                    table_ids_to_truncate: vec![],
                }
            };
            Self {
                job_id: creating_job_id,
                definition,
                create_type,
                tracking_job,
                status,
            }
        }
    }

    /// ## How recovery works
    ///
    /// The progress (number of rows consumed) is persisted in state tables.
    /// During recovery, the backfill executor will restore the number of rows consumed,
    /// and then it will just report progress like newly created executors.
    fn recover_progress(
        job_id: JobId,
        states: HashMap<ActorId, BackfillState>,
        backfill_upstream_types: HashMap<ActorId, BackfillUpstreamType>,
        upstream_mv_count: HashMap<TableId, usize>,
        version_stats: &HummockVersionStats,
        backfill_order_state: BackfillOrderState,
    ) -> Progress {
        let upstream_mvs_total_key_count =
            calculate_total_key_count(&upstream_mv_count, version_stats);
        Progress {
            job_id,
            states,
            backfill_order_state,
            backfill_upstream_types,
            done_count: 0, // Fill only after first barrier pass
            upstream_mv_count,
            upstream_mvs_total_key_count,
            mv_backfill_consumed_rows: 0, // Fill only after first barrier pass
            source_backfill_consumed_rows: 0, // Fill only after first barrier pass
            mv_backfill_buffered_rows: 0, // Fill only after first barrier pass
        }
    }

    pub fn gen_ddl_progress(&self) -> DdlProgress {
        let progress = match &self.status {
            CreateMviewStatus::Backfilling { progress, .. } => progress.calculate_progress(),
            CreateMviewStatus::Finished { .. } => "100%".to_owned(),
        };
        DdlProgress {
            id: self.job_id.as_raw_id() as u64,
            statement: self.definition.clone(),
            create_type: self.create_type.as_str().to_owned(),
            progress,
        }
    }

    /// Update the progress of tracked jobs, and add a new job to track if `info` is `Some`.
    /// Return the table ids whose locality provider state tables need to be truncated.
    pub(super) fn apply_progress(
        &mut self,
        create_mview_progress: &CreateMviewProgress,
        version_stats: &HummockVersionStats,
    ) {
        let CreateMviewStatus::Backfilling {
            progress,
            pending_backfill_nodes,
            table_ids_to_truncate,
        } = &mut self.status
        else {
            tracing::warn!(
                "update the progress of an backfill finished streaming job: {create_mview_progress:?}"
            );
            return;
        };
        {
            // Update the progress of all commands.
            {
                // Those with actors complete can be finished immediately.
                match progress.apply(create_mview_progress, version_stats) {
                    UpdateProgressResult::None => {
                        tracing::trace!(?progress, "update progress");
                    }
                    UpdateProgressResult::Finished {
                        truncate_locality_provider_state_tables,
                    } => {
                        let mut table_ids_to_truncate = take(table_ids_to_truncate);
                        table_ids_to_truncate.extend(truncate_locality_provider_state_tables);
                        tracing::trace!(?progress, "finish progress");
                        self.status = CreateMviewStatus::Finished {
                            table_ids_to_truncate,
                        };
                    }
                    UpdateProgressResult::BackfillNodeFinished(pending) => {
                        table_ids_to_truncate
                            .extend(pending.truncate_locality_provider_state_tables.clone());
                        tracing::trace!(
                            ?progress,
                            next_backfill_nodes = ?pending.next_backfill_nodes,
                            "start next backfill node"
                        );
                        pending_backfill_nodes.extend(pending.next_backfill_nodes);
                    }
                }
            }
        }
    }

    pub(super) fn take_pending_backfill_nodes(&mut self) -> impl Iterator<Item = FragmentId> + '_ {
        match &mut self.status {
            CreateMviewStatus::Backfilling {
                pending_backfill_nodes,
                ..
            } => Some(pending_backfill_nodes.drain(..)),
            CreateMviewStatus::Finished { .. } => None,
        }
        .into_iter()
        .flatten()
    }

    pub(super) fn collect_staging_commit_info(
        &mut self,
    ) -> (bool, impl Iterator<Item = TableId> + '_) {
        let (is_finished, table_ids) = match &mut self.status {
            CreateMviewStatus::Backfilling {
                table_ids_to_truncate,
                ..
            } => (false, table_ids_to_truncate),
            CreateMviewStatus::Finished {
                table_ids_to_truncate,
                ..
            } => (true, table_ids_to_truncate),
        };
        (is_finished, table_ids.drain(..))
    }

    pub(super) fn is_finished(&self) -> bool {
        matches!(self.status, CreateMviewStatus::Finished { .. })
    }

    pub(super) fn into_tracking_job(self) -> TrackingJob {
        let CreateMviewStatus::Finished { .. } = self.status else {
            panic!("should be called when finished");
        };
        self.tracking_job
    }

    /// Add a new create-mview DDL command to track.
    ///
    /// If the actors to track are empty, return the given command as it can be finished immediately.
    pub fn new(info: &CreateStreamingJobCommandInfo, version_stats: &HummockVersionStats) -> Self {
        tracing::trace!(?info, "add job to track");
        let CreateStreamingJobCommandInfo {
            stream_job_fragments,
            definition,
            create_type,
            fragment_backfill_ordering,
            locality_fragment_state_table_mapping,
            ..
        } = info;
        let job_id = stream_job_fragments.stream_job_id();
        let definition = definition.clone();
        let create_type = (*create_type).into();
        let actors = stream_job_fragments.tracking_progress_actor_ids();
        let tracking_job = TrackingJob::new(&info.stream_job_fragments);
        if actors.is_empty() {
            // The command can be finished immediately.
            return Self {
                job_id,
                definition,
                create_type,
                tracking_job,
                status: CreateMviewStatus::Finished {
                    table_ids_to_truncate: vec![],
                },
            };
        }

        let upstream_mv_count = stream_job_fragments.upstream_table_counts();
        let upstream_total_key_count: u64 =
            calculate_total_key_count(&upstream_mv_count, version_stats);

        let backfill_order_state = BackfillOrderState::new(
            fragment_backfill_ordering,
            stream_job_fragments,
            locality_fragment_state_table_mapping.clone(),
        );
        let progress = Progress::new(
            job_id,
            actors,
            upstream_mv_count,
            upstream_total_key_count,
            backfill_order_state,
        );
        Self {
            job_id,
            definition,
            create_type,
            tracking_job,
            status: CreateMviewStatus::Backfilling {
                progress,
                pending_backfill_nodes: vec![],
                table_ids_to_truncate: vec![],
            },
        }
    }
}

impl Progress {
    /// Update the progress of `actor` according to the Pb struct.
    ///
    /// If all actors in this MV have finished, return the command.
    fn apply(
        &mut self,
        progress: &CreateMviewProgress,
        version_stats: &HummockVersionStats,
    ) -> UpdateProgressResult {
        tracing::trace!(?progress, "update progress");
        let actor = progress.backfill_actor_id;
        let job_id = self.job_id;

        let new_state = if progress.done {
            BackfillState::Done(progress.consumed_rows, progress.buffered_rows)
        } else {
            BackfillState::ConsumingUpstream(
                progress.consumed_epoch.into(),
                progress.consumed_rows,
                progress.buffered_rows,
            )
        };

        {
            {
                let progress_state = self;

                let upstream_total_key_count: u64 =
                    calculate_total_key_count(&progress_state.upstream_mv_count, version_stats);

                tracing::debug!(%job_id, "updating progress for table");
                let pending = progress_state.update(actor, new_state, upstream_total_key_count);

                if progress_state.is_done() {
                    tracing::debug!(
                        %job_id,
                        "all actors done for creating mview!",
                    );

                    let PendingBackfillFragments {
                        next_backfill_nodes,
                        truncate_locality_provider_state_tables,
                    } = pending;

                    assert!(next_backfill_nodes.is_empty());
                    UpdateProgressResult::Finished {
                        truncate_locality_provider_state_tables,
                    }
                } else if !pending.next_backfill_nodes.is_empty()
                    || !pending.truncate_locality_provider_state_tables.is_empty()
                {
                    UpdateProgressResult::BackfillNodeFinished(pending)
                } else {
                    UpdateProgressResult::None
                }
            }
        }
    }
}

fn calculate_total_key_count(
    table_count: &HashMap<TableId, usize>,
    version_stats: &HummockVersionStats,
) -> u64 {
    table_count
        .iter()
        .map(|(table_id, count)| {
            assert_ne!(*count, 0);
            *count as u64
                * version_stats
                    .table_stats
                    .get(table_id)
                    .map_or(0, |stat| stat.total_key_count as u64)
        })
        .sum()
}