risingwave_meta/barrier/checkpoint/

control.rs

// Copyright 2024 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::hash_map::Entry;
use std::collections::{HashMap, HashSet};
use std::future::{Future, poll_fn};
use std::ops::Bound::{Excluded, Unbounded};
use std::task::Poll;

use anyhow::anyhow;
use fail::fail_point;
use itertools::Itertools;
use risingwave_common::catalog::{DatabaseId, TableId};
use risingwave_common::id::JobId;
use risingwave_common::metrics::{LabelGuardedHistogram, LabelGuardedIntGauge};
use risingwave_common::system_param::AdaptiveParallelismStrategy;
use risingwave_common::util::epoch::EpochPair;
use risingwave_common::util::stream_graph_visitor::visit_stream_node_cont;
use risingwave_meta_model::WorkerId;
use risingwave_pb::common::WorkerNode;
use risingwave_pb::hummock::HummockVersionStats;
use risingwave_pb::id::{FragmentId, PartialGraphId};
use risingwave_pb::stream_plan::DispatcherType as PbDispatcherType;
use risingwave_pb::stream_plan::stream_node::NodeBody;
use risingwave_pb::stream_service::BarrierCompleteResponse;
use risingwave_pb::stream_service::streaming_control_stream_response::ResetPartialGraphResponse;
use tracing::{debug, warn};

use crate::barrier::cdc_progress::CdcProgress;
use crate::barrier::checkpoint::creating_job::CreatingStreamingJobControl;
use crate::barrier::checkpoint::recovery::{
    DatabaseRecoveringState, DatabaseStatusAction, EnterInitializing, EnterRunning,
    RecoveringStateAction,
};
use crate::barrier::checkpoint::state::{ApplyCommandInfo, BarrierWorkerState};
use crate::barrier::complete_task::{BarrierCompleteOutput, CompleteBarrierTask};
use crate::barrier::info::{InflightDatabaseInfo, SharedActorInfos};
use crate::barrier::notifier::Notifier;
use crate::barrier::partial_graph::{CollectedBarrier, PartialGraphManager, PartialGraphStat};
use crate::barrier::progress::TrackingJob;
use crate::barrier::rpc::{from_partial_graph_id, to_partial_graph_id};
use crate::barrier::schedule::{NewBarrier, PeriodicBarriers};
use crate::barrier::utils::{BarrierItemCollector, collect_creating_job_commit_epoch_info};
use crate::barrier::{
    BackfillProgress, Command, CreateStreamingJobType, FragmentBackfillProgress, Reschedule,
};
use crate::controller::scale::{build_no_shuffle_fragment_graph_edges, find_no_shuffle_graphs};
use crate::manager::MetaSrvEnv;
use crate::model::ActorId;
use crate::rpc::metrics::GLOBAL_META_METRICS;
use crate::{MetaError, MetaResult};

pub(crate) struct CheckpointControl {
    pub(crate) env: MetaSrvEnv,
    pub(super) databases: HashMap<DatabaseId, DatabaseCheckpointControlStatus>,
    pub(super) hummock_version_stats: HummockVersionStats,
    /// The max barrier nums in flight
    pub(crate) in_flight_barrier_nums: usize,
}

impl CheckpointControl {
    pub fn new(env: MetaSrvEnv) -> Self {
        Self {
            in_flight_barrier_nums: env.opts.in_flight_barrier_nums,
            env,
            databases: Default::default(),
            hummock_version_stats: Default::default(),
        }
    }

    pub(crate) fn recover(
        databases: HashMap<DatabaseId, DatabaseCheckpointControl>,
        failed_databases: HashMap<DatabaseId, HashSet<PartialGraphId>>, /* `database_id` -> set of resetting partial graph ids */
        hummock_version_stats: HummockVersionStats,
        env: MetaSrvEnv,
    ) -> Self {
        env.shared_actor_infos()
            .retain_databases(databases.keys().chain(failed_databases.keys()).cloned());
        Self {
            in_flight_barrier_nums: env.opts.in_flight_barrier_nums,
            env,
            databases: databases
                .into_iter()
                .map(|(database_id, control)| {
                    (
                        database_id,
                        DatabaseCheckpointControlStatus::Running(control),
                    )
                })
                .chain(failed_databases.into_iter().map(
                    |(database_id, resetting_partial_graphs)| {
                        (
                            database_id,
                            DatabaseCheckpointControlStatus::Recovering(
                                DatabaseRecoveringState::new_resetting(
                                    database_id,
                                    resetting_partial_graphs,
                                ),
                            ),
                        )
                    },
                ))
                .collect(),
            hummock_version_stats,
        }
    }

    pub(crate) fn ack_completed(
        &mut self,
        partial_graph_manager: &mut PartialGraphManager,
        output: BarrierCompleteOutput,
    ) {
        self.hummock_version_stats = output.hummock_version_stats;
        for (database_id, (command_prev_epoch, creating_job_epochs)) in output.epochs_to_ack {
            self.databases
                .get_mut(&database_id)
                .expect("should exist")
                .expect_running("should have wait for completing command before enter recovery")
                .ack_completed(
                    partial_graph_manager,
                    command_prev_epoch,
                    creating_job_epochs,
                );
        }
    }

    pub(crate) fn next_complete_barrier_task(
        &mut self,
        periodic_barriers: &mut PeriodicBarriers,
        partial_graph_manager: &mut PartialGraphManager,
    ) -> Option<CompleteBarrierTask> {
        let mut task = None;
        for database in self.databases.values_mut() {
            let Some(database) = database.running_state_mut() else {
                continue;
            };
            database.next_complete_barrier_task(
                periodic_barriers,
                partial_graph_manager,
                &mut task,
                &self.hummock_version_stats,
            );
        }
        task
    }

    pub(crate) fn barrier_collected(
        &mut self,
        partial_graph_id: PartialGraphId,
        collected_barrier: CollectedBarrier<'_>,
        periodic_barriers: &mut PeriodicBarriers,
    ) -> MetaResult<()> {
        let (database_id, _) = from_partial_graph_id(partial_graph_id);
        let database_status = self.databases.get_mut(&database_id).expect("should exist");
        match database_status {
            DatabaseCheckpointControlStatus::Running(database) => {
                database.barrier_collected(partial_graph_id, collected_barrier, periodic_barriers)
            }
            DatabaseCheckpointControlStatus::Recovering(_) => {
                if cfg!(debug_assertions) {
                    panic!(
                        "receive collected barrier {:?} on recovering database {} from partial graph {}",
                        collected_barrier, database_id, partial_graph_id
                    );
                } else {
                    warn!(?collected_barrier, %partial_graph_id, "ignore collected barrier on recovering database");
                }
                Ok(())
            }
        }
    }

    pub(crate) fn recovering_databases(&self) -> impl Iterator<Item = DatabaseId> + '_ {
        self.databases.iter().filter_map(|(database_id, database)| {
            database.running_state().is_none().then_some(*database_id)
        })
    }

    pub(crate) fn running_databases(&self) -> impl Iterator<Item = DatabaseId> + '_ {
        self.databases.iter().filter_map(|(database_id, database)| {
            database.running_state().is_some().then_some(*database_id)
        })
    }

    pub(crate) fn database_info(&self, database_id: DatabaseId) -> Option<&InflightDatabaseInfo> {
        self.databases
            .get(&database_id)
            .and_then(|database| database.running_state())
            .map(|database| &database.database_info)
    }

    pub(crate) fn may_have_snapshot_backfilling_jobs(&self) -> bool {
        self.databases
            .values()
            .any(|database| database.may_have_snapshot_backfilling_jobs())
    }

    /// return Some(failed `database_id` -> `err`)
    pub(crate) fn handle_new_barrier(
        &mut self,
        new_barrier: NewBarrier,
        partial_graph_manager: &mut PartialGraphManager,
        adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
        worker_nodes: &HashMap<WorkerId, WorkerNode>,
    ) -> MetaResult<()> {
        let NewBarrier {
            database_id,
            command,
            span,
            checkpoint,
        } = new_barrier;

        if let Some((mut command, notifiers)) = command {
            if let &mut Command::CreateStreamingJob {
                ref mut cross_db_snapshot_backfill_info,
                ref info,
                ..
            } = &mut command
            {
                for (table_id, snapshot_epoch) in
                    &mut cross_db_snapshot_backfill_info.upstream_mv_table_id_to_backfill_epoch
                {
                    for database in self.databases.values() {
                        if let Some(database) = database.running_state()
                            && database.database_info.contains_job(table_id.as_job_id())
                        {
                            if let Some(committed_epoch) = database.committed_epoch {
                                *snapshot_epoch = Some(committed_epoch);
                            }
                            break;
                        }
                    }
                    if snapshot_epoch.is_none() {
                        let table_id = *table_id;
                        warn!(
                            ?cross_db_snapshot_backfill_info,
                            ?table_id,
                            ?info,
                            "database of cross db upstream table not found"
                        );
                        let err: MetaError =
                            anyhow!("database of cross db upstream table {} not found", table_id)
                                .into();
                        for notifier in notifiers {
                            notifier.notify_start_failed(err.clone());
                        }

                        return Ok(());
                    }
                }
            }

            let database = match self.databases.entry(database_id) {
                Entry::Occupied(entry) => entry
                    .into_mut()
                    .expect_running("should not have command when not running"),
                Entry::Vacant(entry) => match &command {
                    Command::CreateStreamingJob { info, job_type, .. } => {
                        let CreateStreamingJobType::Normal = job_type else {
                            if cfg!(debug_assertions) {
                                panic!(
                                    "unexpected first job of type {job_type:?} with info {info:?}"
                                );
                            } else {
                                for notifier in notifiers {
                                    notifier.notify_start_failed(anyhow!("unexpected job_type {job_type:?} for first job {} in database {database_id}", info.streaming_job.id()).into());
                                }
                                return Ok(());
                            }
                        };
                        let new_database = DatabaseCheckpointControl::new(
                            database_id,
                            self.env.shared_actor_infos().clone(),
                        );
                        let adder = partial_graph_manager.add_partial_graph(
                            to_partial_graph_id(database_id, None),
                            DatabaseCheckpointControlMetrics::new(database_id),
                        );
                        adder.added();
                        entry
                            .insert(DatabaseCheckpointControlStatus::Running(new_database))
                            .expect_running("just initialized as running")
                    }
                    Command::Flush
                    | Command::Pause
                    | Command::Resume
                    | Command::DropStreamingJobs { .. }
                    | Command::DropSubscription { .. } => {
                        for mut notifier in notifiers {
                            notifier.notify_started();
                            notifier.notify_collected();
                        }
                        warn!(?command, "skip command for empty database");
                        return Ok(());
                    }
                    Command::RescheduleIntent { .. }
                    | Command::ReplaceStreamJob(_)
                    | Command::SourceChangeSplit(_)
                    | Command::Throttle { .. }
                    | Command::CreateSubscription { .. }
                    | Command::AlterSubscriptionRetention { .. }
                    | Command::ConnectorPropsChange(_)
                    | Command::Refresh { .. }
                    | Command::ListFinish { .. }
                    | Command::LoadFinish { .. }
                    | Command::ResetSource { .. }
                    | Command::ResumeBackfill { .. }
                    | Command::InjectSourceOffsets { .. } => {
                        if cfg!(debug_assertions) {
                            panic!(
                                "new database graph info can only be created for normal creating streaming job, but get command: {} {:?}",
                                database_id, command
                            )
                        } else {
                            warn!(%database_id, ?command, "database not exist when handling command");
                            for notifier in notifiers {
                                notifier.notify_start_failed(anyhow!("database {database_id} not exist when handling command {command:?}").into());
                            }
                            return Ok(());
                        }
                    }
                },
            };

            database.handle_new_barrier(
                Some((command, notifiers)),
                checkpoint,
                span,
                partial_graph_manager,
                &self.hummock_version_stats,
                adaptive_parallelism_strategy,
                worker_nodes,
            )
        } else {
            let database = match self.databases.entry(database_id) {
                Entry::Occupied(entry) => entry.into_mut(),
                Entry::Vacant(_) => {
                    // If it does not exist in the HashMap yet, it means that the first streaming
                    // job has not been created, and we do not need to send a barrier.
                    return Ok(());
                }
            };
            let Some(database) = database.running_state_mut() else {
                // Skip new barrier for database which is not running.
                return Ok(());
            };
            if partial_graph_manager.inflight_barrier_num(database.partial_graph_id)
                >= self.in_flight_barrier_nums
            {
                // Skip new barrier with no explicit command when the database should pause inject additional barrier
                return Ok(());
            }
            database.handle_new_barrier(
                None,
                checkpoint,
                span,
                partial_graph_manager,
                &self.hummock_version_stats,
                adaptive_parallelism_strategy,
                worker_nodes,
            )
        }
    }

    pub(crate) fn gen_backfill_progress(&self) -> HashMap<JobId, BackfillProgress> {
        let mut progress = HashMap::new();
        for status in self.databases.values() {
            let Some(database_checkpoint_control) = status.running_state() else {
                continue;
            };
            // Progress of normal backfill
            progress.extend(
                database_checkpoint_control
                    .database_info
                    .gen_backfill_progress(),
            );
            // Progress of snapshot backfill
            for (job_id, creating_job) in
                &database_checkpoint_control.creating_streaming_job_controls
            {
                progress.extend([(*job_id, creating_job.gen_backfill_progress())]);
            }
        }
        progress
    }

    pub(crate) fn gen_fragment_backfill_progress(&self) -> Vec<FragmentBackfillProgress> {
        let mut progress = Vec::new();
        for status in self.databases.values() {
            let Some(database_checkpoint_control) = status.running_state() else {
                continue;
            };
            progress.extend(
                database_checkpoint_control
                    .database_info
                    .gen_fragment_backfill_progress(),
            );
            for creating_job in database_checkpoint_control
                .creating_streaming_job_controls
                .values()
            {
                progress.extend(creating_job.gen_fragment_backfill_progress());
            }
        }
        progress
    }

    pub(crate) fn gen_cdc_progress(&self) -> HashMap<JobId, CdcProgress> {
        let mut progress = HashMap::new();
        for status in self.databases.values() {
            let Some(database_checkpoint_control) = status.running_state() else {
                continue;
            };
            // Progress of normal backfill
            progress.extend(database_checkpoint_control.database_info.gen_cdc_progress());
        }
        progress
    }

    pub(crate) fn databases_failed_at_worker_err(
        &mut self,
        worker_id: WorkerId,
    ) -> impl Iterator<Item = DatabaseId> + '_ {
        self.databases
            .iter_mut()
            .filter_map(
                move |(database_id, database_status)| match database_status {
                    DatabaseCheckpointControlStatus::Running(control) => {
                        if !control.is_valid_after_worker_err(worker_id) {
                            Some(*database_id)
                        } else {
                            None
                        }
                    }
                    DatabaseCheckpointControlStatus::Recovering(state) => {
                        if !state.is_valid_after_worker_err(worker_id) {
                            Some(*database_id)
                        } else {
                            None
                        }
                    }
                },
            )
    }
}

pub(crate) enum CheckpointControlEvent<'a> {
    EnteringInitializing(DatabaseStatusAction<'a, EnterInitializing>),
    EnteringRunning(DatabaseStatusAction<'a, EnterRunning>),
}

impl CheckpointControl {
    pub(crate) fn on_partial_graph_reset(
        &mut self,
        partial_graph_id: PartialGraphId,
        reset_resps: HashMap<WorkerId, ResetPartialGraphResponse>,
    ) {
        let (database_id, creating_job) = from_partial_graph_id(partial_graph_id);
        match self.databases.get_mut(&database_id).expect("should exist") {
            DatabaseCheckpointControlStatus::Running(database) => {
                if let Some(creating_job_id) = creating_job {
                    let Some(job) = database
                        .creating_streaming_job_controls
                        .remove(&creating_job_id)
                    else {
                        if cfg!(debug_assertions) {
                            panic!(
                                "receive reset partial graph resp on non-existing creating job {creating_job_id} in database {database_id}"
                            )
                        }
                        warn!(
                            %database_id,
                            %creating_job_id,
                            "ignore reset partial graph resp on non-existing creating job on running database"
                        );
                        return;
                    };
                    job.on_partial_graph_reset();
                } else {
                    unreachable!("should not receive reset database resp when database running")
                }
            }
            DatabaseCheckpointControlStatus::Recovering(state) => {
                state.on_partial_graph_reset(partial_graph_id, reset_resps);
            }
        }
    }

    pub(crate) fn on_partial_graph_initialized(&mut self, partial_graph_id: PartialGraphId) {
        let (database_id, _) = from_partial_graph_id(partial_graph_id);
        match self.databases.get_mut(&database_id).expect("should exist") {
            DatabaseCheckpointControlStatus::Running(_) => {
                unreachable!("should not have partial graph initialized when running")
            }
            DatabaseCheckpointControlStatus::Recovering(state) => {
                state.partial_graph_initialized(partial_graph_id);
            }
        }
    }

    pub(crate) fn next_event(
        &mut self,
    ) -> impl Future<Output = CheckpointControlEvent<'_>> + Send + '_ {
        let mut this = Some(self);
        poll_fn(move |cx| {
            let Some(this_mut) = this.as_mut() else {
                unreachable!("should not be polled after poll ready")
            };
            for (&database_id, database_status) in &mut this_mut.databases {
                match database_status {
                    DatabaseCheckpointControlStatus::Running(_) => {}
                    DatabaseCheckpointControlStatus::Recovering(state) => {
                        let poll_result = state.poll_next_event(cx);
                        if let Poll::Ready(action) = poll_result {
                            let this = this.take().expect("checked Some");
                            return Poll::Ready(match action {
                                RecoveringStateAction::EnterInitializing(reset_workers) => {
                                    CheckpointControlEvent::EnteringInitializing(
                                        this.new_database_status_action(
                                            database_id,
                                            EnterInitializing(reset_workers),
                                        ),
                                    )
                                }
                                RecoveringStateAction::EnterRunning => {
                                    CheckpointControlEvent::EnteringRunning(
                                        this.new_database_status_action(database_id, EnterRunning),
                                    )
                                }
                            });
                        }
                    }
                }
            }
            Poll::Pending
        })
    }
}

pub(crate) enum DatabaseCheckpointControlStatus {
    Running(DatabaseCheckpointControl),
    Recovering(DatabaseRecoveringState),
}

impl DatabaseCheckpointControlStatus {
    fn running_state(&self) -> Option<&DatabaseCheckpointControl> {
        match self {
            DatabaseCheckpointControlStatus::Running(state) => Some(state),
            DatabaseCheckpointControlStatus::Recovering(_) => None,
        }
    }

    fn running_state_mut(&mut self) -> Option<&mut DatabaseCheckpointControl> {
        match self {
            DatabaseCheckpointControlStatus::Running(state) => Some(state),
            DatabaseCheckpointControlStatus::Recovering(_) => None,
        }
    }

    fn may_have_snapshot_backfilling_jobs(&self) -> bool {
        self.running_state()
            .map(|database| !database.creating_streaming_job_controls.is_empty())
            .unwrap_or(true) // there can be snapshot backfilling jobs when the database is recovering.
    }

    fn expect_running(&mut self, reason: &'static str) -> &mut DatabaseCheckpointControl {
        match self {
            DatabaseCheckpointControlStatus::Running(state) => state,
            DatabaseCheckpointControlStatus::Recovering(_) => {
                panic!("should be at running: {}", reason)
            }
        }
    }
}

pub(in crate::barrier) struct DatabaseCheckpointControlMetrics {
    barrier_latency: LabelGuardedHistogram,
    in_flight_barrier_nums: LabelGuardedIntGauge,
    all_barrier_nums: LabelGuardedIntGauge,
}

impl DatabaseCheckpointControlMetrics {
    pub(in crate::barrier) fn new(database_id: DatabaseId) -> Self {
        let database_id_str = database_id.to_string();
        let barrier_latency = GLOBAL_META_METRICS
            .barrier_latency
            .with_guarded_label_values(&[&database_id_str]);
        let in_flight_barrier_nums = GLOBAL_META_METRICS
            .in_flight_barrier_nums
            .with_guarded_label_values(&[&database_id_str]);
        let all_barrier_nums = GLOBAL_META_METRICS
            .all_barrier_nums
            .with_guarded_label_values(&[&database_id_str]);
        Self {
            barrier_latency,
            in_flight_barrier_nums,
            all_barrier_nums,
        }
    }
}

impl PartialGraphStat for DatabaseCheckpointControlMetrics {
    fn observe_barrier_latency(&self, _epoch: EpochPair, barrier_latency_secs: f64) {
        self.barrier_latency.observe(barrier_latency_secs);
    }

    fn observe_barrier_num(&self, inflight_barrier_num: usize, collected_barrier_num: usize) {
        self.in_flight_barrier_nums.set(inflight_barrier_num as _);
        self.all_barrier_nums
            .set((inflight_barrier_num + collected_barrier_num) as _);
    }
}

/// Controls the concurrent execution of commands.
pub(in crate::barrier) struct DatabaseCheckpointControl {
    pub(super) database_id: DatabaseId,
    partial_graph_id: PartialGraphId,
    pub(super) state: BarrierWorkerState,

    finishing_jobs_collector:
        BarrierItemCollector<JobId, (Vec<BarrierCompleteResponse>, TrackingJob), ()>,
    /// The barrier that are completing.
    /// Some(`prev_epoch`)
    completing_barrier: Option<u64>,

    committed_epoch: Option<u64>,

    pub(super) database_info: InflightDatabaseInfo,
    pub creating_streaming_job_controls: HashMap<JobId, CreatingStreamingJobControl>,
}

impl DatabaseCheckpointControl {
    fn new(database_id: DatabaseId, shared_actor_infos: SharedActorInfos) -> Self {
        Self {
            database_id,
            partial_graph_id: to_partial_graph_id(database_id, None),
            state: BarrierWorkerState::new(),
            finishing_jobs_collector: BarrierItemCollector::new(),
            completing_barrier: None,
            committed_epoch: None,
            database_info: InflightDatabaseInfo::empty(database_id, shared_actor_infos),
            creating_streaming_job_controls: Default::default(),
        }
    }

    pub(crate) fn recovery(
        database_id: DatabaseId,
        state: BarrierWorkerState,
        committed_epoch: u64,
        database_info: InflightDatabaseInfo,
        creating_streaming_job_controls: HashMap<JobId, CreatingStreamingJobControl>,
    ) -> Self {
        Self {
            database_id,
            partial_graph_id: to_partial_graph_id(database_id, None),
            state,
            finishing_jobs_collector: BarrierItemCollector::new(),
            completing_barrier: None,
            committed_epoch: Some(committed_epoch),
            database_info,
            creating_streaming_job_controls,
        }
    }

    pub(crate) fn is_valid_after_worker_err(&self, worker_id: WorkerId) -> bool {
        !self.database_info.contains_worker(worker_id as _)
            && self
                .creating_streaming_job_controls
                .values()
                .all(|job| job.is_valid_after_worker_err(worker_id))
    }

    /// Enqueue a barrier command
    fn enqueue_command(&mut self, epoch: EpochPair, creating_jobs_to_wait: HashSet<JobId>) {
        let prev_epoch = epoch.prev;
        tracing::trace!(prev_epoch, ?creating_jobs_to_wait, "enqueue command");
        if !creating_jobs_to_wait.is_empty() {
            self.finishing_jobs_collector
                .enqueue(epoch, creating_jobs_to_wait, ());
        }
    }

    /// Change the state of this `prev_epoch` to `Completed`. Return continuous nodes
    /// with `Completed` starting from first node [`Completed`..`InFlight`) and remove them.
    fn barrier_collected(
        &mut self,
        partial_graph_id: PartialGraphId,
        collected_barrier: CollectedBarrier<'_>,
        periodic_barriers: &mut PeriodicBarriers,
    ) -> MetaResult<()> {
        let prev_epoch = collected_barrier.epoch.prev;
        tracing::trace!(
            prev_epoch,
            partial_graph_id = %partial_graph_id,
            "barrier collected"
        );
        let (database_id, creating_job_id) = from_partial_graph_id(partial_graph_id);
        assert_eq!(self.database_id, database_id);
        if let Some(creating_job_id) = creating_job_id {
            let should_merge_to_upstream = self
                .creating_streaming_job_controls
                .get_mut(&creating_job_id)
                .expect("should exist")
                .collect(collected_barrier);
            if should_merge_to_upstream {
                periodic_barriers.force_checkpoint_in_next_barrier(self.database_id);
            }
        }
        Ok(())
    }
}

impl DatabaseCheckpointControl {
    /// return creating job table fragment id -> (backfill progress epoch , {`upstream_mv_table_id`})
    fn collect_backfill_pinned_upstream_log_epoch(
        &self,
    ) -> HashMap<JobId, (u64, HashSet<TableId>)> {
        self.creating_streaming_job_controls
            .iter()
            .map(|(job_id, creating_job)| (*job_id, creating_job.pinned_upstream_log_epoch()))
            .collect()
    }

    fn collect_no_shuffle_fragment_relations_for_reschedule_check(
        &self,
    ) -> Vec<(FragmentId, FragmentId)> {
        let mut no_shuffle_relations = Vec::new();
        for fragment in self.database_info.fragment_infos() {
            let downstream_fragment_id = fragment.fragment_id;
            visit_stream_node_cont(&fragment.nodes, |node| {
                if let Some(NodeBody::Merge(merge)) = node.node_body.as_ref()
                    && merge.upstream_dispatcher_type == PbDispatcherType::NoShuffle as i32
                {
                    no_shuffle_relations.push((merge.upstream_fragment_id, downstream_fragment_id));
                }
                true
            });
        }

        for creating_job in self.creating_streaming_job_controls.values() {
            creating_job.collect_no_shuffle_fragment_relations(&mut no_shuffle_relations);
        }
        no_shuffle_relations
    }

    fn collect_reschedule_blocked_jobs_for_creating_jobs_inflight(
        &self,
    ) -> MetaResult<HashSet<JobId>> {
        let mut initial_blocked_fragment_ids = HashSet::new();
        for creating_job in self.creating_streaming_job_controls.values() {
            creating_job.collect_reschedule_blocked_fragment_ids(&mut initial_blocked_fragment_ids);
        }

        let mut blocked_fragment_ids = initial_blocked_fragment_ids.clone();
        if !initial_blocked_fragment_ids.is_empty() {
            let no_shuffle_relations =
                self.collect_no_shuffle_fragment_relations_for_reschedule_check();
            let (forward_edges, backward_edges) =
                build_no_shuffle_fragment_graph_edges(no_shuffle_relations);
            let initial_blocked_fragment_ids: Vec<_> =
                initial_blocked_fragment_ids.iter().copied().collect();
            for ensemble in find_no_shuffle_graphs(
                &initial_blocked_fragment_ids,
                &forward_edges,
                &backward_edges,
            )? {
                blocked_fragment_ids.extend(ensemble.fragments());
            }
        }

        let mut blocked_job_ids = HashSet::new();
        blocked_job_ids.extend(
            blocked_fragment_ids
                .into_iter()
                .filter_map(|fragment_id| self.database_info.job_id_by_fragment(fragment_id)),
        );
        Ok(blocked_job_ids)
    }

    fn collect_reschedule_blocked_job_ids(
        &self,
        reschedules: &HashMap<FragmentId, Reschedule>,
        fragment_actors: &HashMap<FragmentId, HashSet<ActorId>>,
        blocked_job_ids: &HashSet<JobId>,
    ) -> HashSet<JobId> {
        let mut affected_fragment_ids: HashSet<FragmentId> = reschedules.keys().copied().collect();
        affected_fragment_ids.extend(fragment_actors.keys().copied());
        for reschedule in reschedules.values() {
            affected_fragment_ids.extend(reschedule.downstream_fragment_ids.iter().copied());
            affected_fragment_ids.extend(
                reschedule
                    .upstream_fragment_dispatcher_ids
                    .iter()
                    .map(|(fragment_id, _)| *fragment_id),
            );
        }

        affected_fragment_ids
            .into_iter()
            .filter_map(|fragment_id| self.database_info.job_id_by_fragment(fragment_id))
            .filter(|job_id| blocked_job_ids.contains(job_id))
            .collect()
    }

    fn next_complete_barrier_task(
        &mut self,
        periodic_barriers: &mut PeriodicBarriers,
        partial_graph_manager: &mut PartialGraphManager,
        task: &mut Option<CompleteBarrierTask>,
        hummock_version_stats: &HummockVersionStats,
    ) {
        // `Vec::new` is a const fn, and do not have memory allocation, and therefore is lightweight enough
        let mut creating_jobs_task = vec![];
        if let Some(committed_epoch) = self.committed_epoch {
            // `Vec::new` is a const fn, and do not have memory allocation, and therefore is lightweight enough
            let mut finished_jobs = Vec::new();
            let min_upstream_inflight_barrier = partial_graph_manager
                .first_inflight_barrier(self.partial_graph_id)
                .map(|epoch| epoch.prev);
            for (job_id, job) in &mut self.creating_streaming_job_controls {
                if let Some((epoch, resps, info, is_finish_epoch)) = job.start_completing(
                    partial_graph_manager,
                    min_upstream_inflight_barrier,
                    committed_epoch,
                ) {
                    let resps = resps.into_values().collect_vec();
                    if is_finish_epoch {
                        assert!(info.notifiers.is_empty());
                        finished_jobs.push((*job_id, epoch, resps));
                        continue;
                    };
                    creating_jobs_task.push((*job_id, epoch, resps, info));
                }
            }

            if !finished_jobs.is_empty() {
                partial_graph_manager.remove_partial_graphs(
                    finished_jobs
                        .iter()
                        .map(|(job_id, ..)| to_partial_graph_id(self.database_id, Some(*job_id)))
                        .collect(),
                );
            }
            for (job_id, epoch, resps) in finished_jobs {
                debug!(epoch, %job_id, "finish creating job");
                // It's safe to remove the creating job, because on CompleteJobType::Finished,
                // all previous barriers have been collected and completed.
                let creating_streaming_job = self
                    .creating_streaming_job_controls
                    .remove(&job_id)
                    .expect("should exist");
                let tracking_job = creating_streaming_job.into_tracking_job();
                self.finishing_jobs_collector
                    .collect(epoch, job_id, (resps, tracking_job));
            }
        }
        let mut observed_non_checkpoint = false;
        self.finishing_jobs_collector.advance_collected();
        let epoch_end_bound = self
            .finishing_jobs_collector
            .first_inflight_epoch()
            .map_or(Unbounded, |epoch| Excluded(epoch.prev));
        if let Some((epoch, resps, info)) = partial_graph_manager.start_completing(
            self.partial_graph_id,
            epoch_end_bound,
            |_, resps, post_collect_command| {
                observed_non_checkpoint = true;
                self.handle_refresh_table_info(task, &resps);
                self.database_info.apply_collected_command(
                    &post_collect_command,
                    &resps,
                    hummock_version_stats,
                );
            },
        ) {
            self.handle_refresh_table_info(task, &resps);
            self.database_info.apply_collected_command(
                &info.post_collect_command,
                &resps,
                hummock_version_stats,
            );
            let mut resps_to_commit = resps.into_values().collect_vec();
            let mut staging_commit_info = self.database_info.take_staging_commit_info();
            if let Some((_, finished_jobs, _)) =
                self.finishing_jobs_collector
                    .take_collected_if(|collected_epoch| {
                        assert!(epoch <= collected_epoch.prev);
                        epoch == collected_epoch.prev
                    })
            {
                finished_jobs
                    .into_iter()
                    .for_each(|(_, (resps, tracking_job))| {
                        resps_to_commit.extend(resps);
                        staging_commit_info.finished_jobs.push(tracking_job);
                    });
            }
            {
                let task = task.get_or_insert_default();
                Command::collect_commit_epoch_info(
                    &self.database_info,
                    &info,
                    &mut task.commit_info,
                    resps_to_commit,
                    self.collect_backfill_pinned_upstream_log_epoch(),
                );
                self.completing_barrier = Some(info.barrier_info.prev_epoch());
                task.finished_jobs.extend(staging_commit_info.finished_jobs);
                task.finished_cdc_table_backfill
                    .extend(staging_commit_info.finished_cdc_table_backfill);
                task.epoch_infos
                    .try_insert(self.partial_graph_id, info)
                    .expect("non duplicate");
                task.commit_info
                    .truncate_tables
                    .extend(staging_commit_info.table_ids_to_truncate);
            }
        } else if observed_non_checkpoint
            && self.database_info.has_pending_finished_jobs()
            && !partial_graph_manager.has_pending_checkpoint_barrier(self.partial_graph_id)
        {
            periodic_barriers.force_checkpoint_in_next_barrier(self.database_id);
        }
        if !creating_jobs_task.is_empty() {
            let task = task.get_or_insert_default();
            for (job_id, epoch, resps, info) in creating_jobs_task {
                collect_creating_job_commit_epoch_info(&mut task.commit_info, epoch, resps, &info);
                task.epoch_infos
                    .try_insert(to_partial_graph_id(self.database_id, Some(job_id)), info)
                    .expect("non duplicate");
            }
        }
    }

    fn ack_completed(
        &mut self,
        partial_graph_manager: &mut PartialGraphManager,
        command_prev_epoch: Option<u64>,
        creating_job_epochs: Vec<(JobId, u64)>,
    ) {
        {
            if let Some(prev_epoch) = self.completing_barrier.take() {
                assert_eq!(command_prev_epoch, Some(prev_epoch));
                self.committed_epoch = Some(prev_epoch);
                partial_graph_manager.ack_completed(self.partial_graph_id, prev_epoch);
            } else {
                assert_eq!(command_prev_epoch, None);
            };
            for (job_id, epoch) in creating_job_epochs {
                if let Some(job) = self.creating_streaming_job_controls.get_mut(&job_id) {
                    job.ack_completed(partial_graph_manager, epoch);
                }
                // If the job is not found, it was dropped and already removed
                // by `on_partial_graph_reset` while the completing task was running.
            }
        }
    }

    fn handle_refresh_table_info(
        &self,
        task: &mut Option<CompleteBarrierTask>,
        resps: &HashMap<WorkerId, BarrierCompleteResponse>,
    ) {
        let list_finished_info = resps
            .values()
            .flat_map(|resp| resp.list_finished_sources.clone())
            .collect::<Vec<_>>();
        if !list_finished_info.is_empty() {
            let task = task.get_or_insert_default();
            task.list_finished_source_ids.extend(list_finished_info);
        }

        let load_finished_info = resps
            .values()
            .flat_map(|resp| resp.load_finished_sources.clone())
            .collect::<Vec<_>>();
        if !load_finished_info.is_empty() {
            let task = task.get_or_insert_default();
            task.load_finished_source_ids.extend(load_finished_info);
        }

        let refresh_finished_table_ids: Vec<JobId> = resps
            .values()
            .flat_map(|resp| {
                resp.refresh_finished_tables
                    .iter()
                    .map(|table_id| table_id.as_job_id())
            })
            .collect::<Vec<_>>();
        if !refresh_finished_table_ids.is_empty() {
            let task = task.get_or_insert_default();
            task.refresh_finished_table_job_ids
                .extend(refresh_finished_table_ids);
        }
    }
}

impl DatabaseCheckpointControl {
    /// Handle the new barrier from the scheduled queue and inject it.
    fn handle_new_barrier(
        &mut self,
        command: Option<(Command, Vec<Notifier>)>,
        checkpoint: bool,
        span: tracing::Span,
        partial_graph_manager: &mut PartialGraphManager,
        hummock_version_stats: &HummockVersionStats,
        adaptive_parallelism_strategy: AdaptiveParallelismStrategy,
        worker_nodes: &HashMap<WorkerId, WorkerNode>,
    ) -> MetaResult<()> {
        let curr_epoch = self.state.in_flight_prev_epoch().next();

        let (command, mut notifiers) = if let Some((command, notifiers)) = command {
            (Some(command), notifiers)
        } else {
            (None, vec![])
        };

        debug_assert!(
            !matches!(
                &command,
                Some(Command::RescheduleIntent {
                    reschedule_plan: None,
                    ..
                })
            ),
            "reschedule intent should be resolved before injection"
        );

        if let Some(Command::DropStreamingJobs {
            streaming_job_ids, ..
        }) = &command
        {
            if streaming_job_ids.len() > 1 {
                for job_to_cancel in streaming_job_ids {
                    if self
                        .creating_streaming_job_controls
                        .contains_key(job_to_cancel)
                    {
                        warn!(
                            job_id = %job_to_cancel,
                            "ignore multi-job cancel command on creating snapshot backfill streaming job"
                        );
                        for notifier in notifiers {
                            notifier
                                .notify_start_failed(anyhow!("cannot cancel creating snapshot backfill streaming job with other jobs, \
                                the job will continue creating until created or recovery. Please cancel the snapshot backfilling job in a single DDL ").into());
                        }
                        return Ok(());
                    }
                }
            } else if let Some(job_to_drop) = streaming_job_ids.iter().next()
                && let Some(creating_job) =
                    self.creating_streaming_job_controls.get_mut(job_to_drop)
                && creating_job.drop(&mut notifiers, partial_graph_manager)
            {
                return Ok(());
            }
        }

        if let Some(Command::Throttle { jobs, .. }) = &command
            && jobs.len() > 1
            && let Some(creating_job_id) = jobs
                .iter()
                .find(|job| self.creating_streaming_job_controls.contains_key(*job))
        {
            warn!(
                job_id = %creating_job_id,
                "ignore multi-job throttle command on creating snapshot backfill streaming job"
            );
            for notifier in notifiers {
                notifier
                    .notify_start_failed(anyhow!("cannot alter rate limit for snapshot backfill streaming job with other jobs, \
                                the original rate limit will be kept recovery.").into());
            }
            return Ok(());
        };

        if let Some(Command::RescheduleIntent {
            reschedule_plan: Some(reschedule_plan),
            ..
        }) = &command
            && !self.creating_streaming_job_controls.is_empty()
        {
            let blocked_job_ids =
                self.collect_reschedule_blocked_jobs_for_creating_jobs_inflight()?;
            let blocked_reschedule_job_ids = self.collect_reschedule_blocked_job_ids(
                &reschedule_plan.reschedules,
                &reschedule_plan.fragment_actors,
                &blocked_job_ids,
            );
            if !blocked_reschedule_job_ids.is_empty() {
                warn!(
                    blocked_reschedule_job_ids = ?blocked_reschedule_job_ids,
                    "reject reschedule fragments related to creating unreschedulable backfill jobs"
                );
                for notifier in notifiers {
                    notifier.notify_start_failed(
                        anyhow!(
                            "cannot reschedule jobs {:?} when creating jobs with unreschedulable backfill fragments",
                            blocked_reschedule_job_ids
                        )
                            .into(),
                    );
                }
                return Ok(());
            }
        }

        if !matches!(&command, Some(Command::CreateStreamingJob { .. }))
            && self.database_info.is_empty()
        {
            assert!(
                self.creating_streaming_job_controls.is_empty(),
                "should not have snapshot backfill job when there is no normal job in database"
            );
            // skip the command when there is nothing to do with the barrier
            for mut notifier in notifiers {
                notifier.notify_started();
                notifier.notify_collected();
            }
            return Ok(());
        };

        if let Some(Command::CreateStreamingJob {
            job_type: CreateStreamingJobType::SnapshotBackfill(_),
            ..
        }) = &command
            && self.state.is_paused()
        {
            warn!("cannot create streaming job with snapshot backfill when paused");
            for notifier in notifiers {
                notifier.notify_start_failed(
                    anyhow!("cannot create streaming job with snapshot backfill when paused",)
                        .into(),
                );
            }
            return Ok(());
        }

        let barrier_info = self.state.next_barrier_info(checkpoint, curr_epoch);
        // Tracing related stuff
        barrier_info.prev_epoch.span().in_scope(|| {
            tracing::info!(target: "rw_tracing", epoch = barrier_info.curr_epoch(), "new barrier enqueued");
        });
        span.record("epoch", barrier_info.curr_epoch());

        let epoch = barrier_info.epoch();
        let ApplyCommandInfo { jobs_to_wait } = match self.apply_command(
            command,
            &mut notifiers,
            barrier_info,
            partial_graph_manager,
            hummock_version_stats,
            adaptive_parallelism_strategy,
            worker_nodes,
        ) {
            Ok(info) => {
                assert!(notifiers.is_empty());
                info
            }
            Err(err) => {
                for notifier in notifiers {
                    notifier.notify_start_failed(err.clone());
                }
                fail_point!("inject_barrier_err_success");
                return Err(err);
            }
        };

        // Record the in-flight barrier.
        self.enqueue_command(epoch, jobs_to_wait);

        Ok(())
    }
}