Skip to main content

risingwave_meta/barrier/checkpoint/independent_job/
mod.rs

1// Copyright 2026 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, HashSet};
16use std::mem::take;
17
18use risingwave_common::catalog::{DatabaseId, TableId};
19use risingwave_common::id::JobId;
20use risingwave_common::util::epoch::Epoch;
21use risingwave_meta_model::WorkerId;
22use risingwave_pb::hummock::HummockVersionStats;
23use risingwave_pb::id::{ActorId, FragmentId, PartialGraphId};
24use risingwave_pb::source::ConnectorSplits;
25use risingwave_pb::stream_plan::barrier::PbBarrierKind;
26use risingwave_pb::stream_plan::barrier_mutation::Mutation;
27use risingwave_pb::stream_plan::{
28    AddMutation, PbSubscriptionUpstreamInfo, StartFragmentBackfillMutation,
29};
30use risingwave_pb::stream_service::barrier_complete_response::CreateMviewProgress;
31
32pub(crate) mod batch_refresh_job;
33pub(crate) mod creating_job;
34pub(crate) mod iceberg_v3_job;
35pub(crate) use batch_refresh_job::{
36    BatchRefreshJobCheckpointControl, BatchRefreshJobTriggerContext, BatchRefreshLogicalFragments,
37    BatchRefreshRenderResult, RenderedIndependentJobActors,
38};
39pub(crate) use creating_job::CreatingStreamingJobControl;
40pub(crate) use iceberg_v3_job::{
41    IcebergV3JobCheckpointControl, IcebergV3RenderResult, is_iceberg_v3_fragment_nodes,
42};
43
44use crate::MetaResult;
45use crate::barrier::backfill_order_control::get_nodes_with_backfill_dependencies;
46use crate::barrier::command::{CreateStreamingJobCommandInfo, ThrottleConfigMap};
47use crate::barrier::context::CreateIndependentStreamingJobCommandInfo;
48use crate::barrier::info::BarrierInfo;
49use crate::barrier::partial_graph::{
50    CollectedBarrier, PartialGraphAdder, PartialGraphBarrierInfo, PartialGraphManager,
51};
52use crate::barrier::progress::CreateMviewProgressTracker;
53use crate::barrier::rpc::to_partial_graph_id;
54use crate::barrier::{
55    BackfillOrderState, BackfillProgress, BarrierKind, FragmentBackfillProgress, TracedEpoch,
56};
57use crate::controller::fragment::InflightFragmentInfo;
58use crate::model::StreamJobActorsToCreate;
59use crate::notification::{CollectionNotifier, NotifierStarter};
60use crate::stream::ExtendedFragmentBackfillOrder;
61
62#[derive(Debug)]
63struct SnapshotPhaseControl {
64    /// Duplicated from `IndependentJobInfo` so snapshot-phase transitions are self-contained.
65    snapshot_epoch: u64,
66    prev_epoch_fake_physical_time: u64,
67    version_stats: HummockVersionStats,
68    create_mview_tracker: CreateMviewProgressTracker,
69    pending_non_checkpoint_barriers: Vec<u64>,
70}
71
72impl SnapshotPhaseControl {
73    fn for_new_job(
74        snapshot_epoch: u64,
75        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
76        create_info: &CreateStreamingJobCommandInfo,
77        version_stats: &HummockVersionStats,
78    ) -> (Self, BarrierInfo) {
79        let job_id = create_info.stream_job_fragments.stream_job_id();
80        let backfill_order_state = BackfillOrderState::new(
81            &create_info.fragment_backfill_ordering,
82            fragment_infos,
83            create_info.locality_fragment_state_table_mapping.clone(),
84        );
85        let create_mview_tracker = CreateMviewProgressTracker::recover(
86            job_id,
87            fragment_infos,
88            backfill_order_state,
89            version_stats,
90        );
91        let mut snapshot = Self {
92            snapshot_epoch,
93            prev_epoch_fake_physical_time: 0,
94            version_stats: version_stats.clone(),
95            create_mview_tracker,
96            pending_non_checkpoint_barriers: vec![],
97        };
98        let initial_barrier_info = new_fake_barrier(
99            &mut snapshot.prev_epoch_fake_physical_time,
100            &mut snapshot.pending_non_checkpoint_barriers,
101            PbBarrierKind::Checkpoint,
102        );
103        assert!(snapshot.pending_non_checkpoint_barriers.is_empty());
104        (snapshot, initial_barrier_info)
105    }
106
107    fn for_recovery(
108        job_id: JobId,
109        snapshot_epoch: u64,
110        committed_epoch: u64,
111        fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
112        backfill_order: &ExtendedFragmentBackfillOrder,
113        version_stats: &HummockVersionStats,
114    ) -> (Self, BarrierInfo) {
115        let backfill_order_state =
116            BackfillOrderState::recover_from_fragment_infos(backfill_order, fragment_infos);
117        let create_mview_tracker = CreateMviewProgressTracker::recover(
118            job_id,
119            fragment_infos,
120            backfill_order_state,
121            version_stats,
122        );
123        let mut snapshot = Self {
124            snapshot_epoch,
125            prev_epoch_fake_physical_time: Epoch(committed_epoch).physical_time(),
126            version_stats: version_stats.clone(),
127            create_mview_tracker,
128            pending_non_checkpoint_barriers: vec![],
129        };
130        let initial_barrier_info = new_fake_barrier(
131            &mut snapshot.prev_epoch_fake_physical_time,
132            &mut snapshot.pending_non_checkpoint_barriers,
133            PbBarrierKind::Initial,
134        );
135        assert!(snapshot.pending_non_checkpoint_barriers.is_empty());
136        (snapshot, initial_barrier_info)
137    }
138
139    fn apply_progress<'a>(
140        &mut self,
141        progress: impl IntoIterator<Item = &'a CreateMviewProgress>,
142    ) -> bool {
143        for progress in progress {
144            self.create_mview_tracker
145                .apply_progress(progress, &self.version_stats);
146        }
147        self.create_mview_tracker.is_finished()
148    }
149
150    fn take_start_backfill_mutation(&mut self) -> Option<Mutation> {
151        let fragment_ids = self
152            .create_mview_tracker
153            .take_pending_backfill_nodes()
154            .collect::<Vec<_>>();
155        (!fragment_ids.is_empty()).then_some(Mutation::StartFragmentBackfill(
156            StartFragmentBackfillMutation { fragment_ids },
157        ))
158    }
159
160    fn next_fake_barrier(&mut self, upstream_kind: &BarrierKind) -> BarrierInfo {
161        let kind = match upstream_kind {
162            BarrierKind::Barrier => PbBarrierKind::Barrier,
163            BarrierKind::Checkpoint(_) => PbBarrierKind::Checkpoint,
164            BarrierKind::Initial => unreachable!("upstream new epoch should not be initial"),
165        };
166        new_fake_barrier(
167            &mut self.prev_epoch_fake_physical_time,
168            &mut self.pending_non_checkpoint_barriers,
169            kind,
170        )
171    }
172
173    fn finish_snapshot_barrier(&mut self) -> BarrierInfo {
174        self.pending_non_checkpoint_barriers
175            .push(self.snapshot_epoch);
176        BarrierInfo {
177            curr_epoch: TracedEpoch::new(Epoch(self.snapshot_epoch)),
178            prev_epoch: TracedEpoch::new(Epoch::from_physical_time(
179                self.prev_epoch_fake_physical_time,
180            )),
181            kind: BarrierKind::Checkpoint(take(&mut self.pending_non_checkpoint_barriers)),
182        }
183    }
184}
185
186fn build_initial_add_mutation(
187    fragment_infos: &HashMap<FragmentId, InflightFragmentInfo>,
188    backfill_ordering: &ExtendedFragmentBackfillOrder,
189    actor_splits: HashMap<ActorId, ConnectorSplits>,
190) -> Mutation {
191    Mutation::Add(AddMutation {
192        actor_dispatchers: Default::default(),
193        added_actors: fragment_infos
194            .values()
195            .flat_map(|fragment| fragment.actors.keys().copied())
196            .collect(),
197        actor_splits,
198        pause: false,
199        subscriptions_to_add: Default::default(),
200        backfill_nodes_to_pause: get_nodes_with_backfill_dependencies(backfill_ordering)
201            .into_iter()
202            .collect(),
203        actor_cdc_table_snapshot_splits: None,
204        new_upstream_sinks: Default::default(),
205        dropped_actors: Default::default(),
206        sink_log_store_flush: Default::default(),
207    })
208}
209
210struct InitialPartialGraphRequest<'a> {
211    node_actors: &'a HashMap<WorkerId, HashSet<ActorId>>,
212    state_table_ids: &'a HashSet<TableId>,
213    barrier_info: BarrierInfo,
214    actors_to_create: StreamJobActorsToCreate,
215    mutation: Mutation,
216    notifier: Option<&'a mut NotifierStarter>,
217    create_info: CreateIndependentStreamingJobCommandInfo,
218}
219
220fn add_initial_partial_graph(
221    graph_adder: &mut PartialGraphAdder<'_>,
222    partial_graph_id: PartialGraphId,
223    request: InitialPartialGraphRequest<'_>,
224) -> MetaResult<()> {
225    let InitialPartialGraphRequest {
226        node_actors,
227        state_table_ids,
228        barrier_info,
229        actors_to_create,
230        mutation,
231        notifier,
232        create_info,
233    } = request;
234    graph_adder.manager().inject_barrier(
235        partial_graph_id,
236        Some(mutation),
237        node_actors,
238        state_table_ids.iter().copied(),
239        node_actors.keys().copied(),
240        Some(actors_to_create),
241        PartialGraphBarrierInfo::new(
242            create_info.into_post_collect(),
243            barrier_info,
244            notifier,
245            state_table_ids.clone(),
246        ),
247    )
248}
249
250/// The maximum number of upstream barriers a snapshot backfill partial graph may queue before
251/// the upstream is back-pressured.
252pub fn snapshot_backfill_max_pending_barrier_num(
253    in_flight_barrier_nums: usize,
254    snapshot_backfill_barrier_amplification_factor: usize,
255) -> usize {
256    in_flight_barrier_nums.saturating_mul(snapshot_backfill_barrier_amplification_factor.max(1))
257}
258
259#[derive(Debug)]
260pub(crate) struct IndependentJobInfo {
261    pub(crate) job_id: JobId,
262    pub(crate) partial_graph_id: PartialGraphId,
263    pub(crate) snapshot_backfill_upstream_tables: HashSet<TableId>,
264    pub(crate) snapshot_epoch: u64,
265    pub(crate) state_table_ids: HashSet<TableId>,
266}
267
268impl IndependentJobInfo {
269    pub(crate) fn from_fragment_infos<'a>(
270        database_id: DatabaseId,
271        job_id: JobId,
272        snapshot_epoch: u64,
273        snapshot_backfill_upstream_tables: HashSet<TableId>,
274        fragment_infos: impl IntoIterator<Item = &'a InflightFragmentInfo> + 'a,
275    ) -> Self {
276        Self {
277            job_id,
278            partial_graph_id: to_partial_graph_id(database_id, Some(job_id)),
279            snapshot_backfill_upstream_tables,
280            snapshot_epoch,
281            state_table_ids: InflightFragmentInfo::existing_table_ids(fragment_infos).collect(),
282        }
283    }
284}
285
286#[derive(Debug)]
287pub(crate) struct IndependentJobControl {
288    /// Immutable metadata shared by every phase of the independent job.
289    pub(crate) info: IndependentJobInfo,
290    /// Latest checkpoint persisted through the common complete-barrier path.
291    pub(crate) max_committed_epoch: Option<u64>,
292}
293
294impl IndependentJobControl {
295    pub(crate) fn new(info: IndependentJobInfo) -> Self {
296        Self {
297            info,
298            max_committed_epoch: None,
299        }
300    }
301
302    pub(crate) fn recovered(info: IndependentJobInfo, committed_epoch: u64) -> Self {
303        Self {
304            info,
305            max_committed_epoch: Some(committed_epoch),
306        }
307    }
308
309    pub(crate) fn ack_completed(&mut self, completed_epoch: u64) {
310        if let Some(previous_epoch) = self.max_committed_epoch.replace(completed_epoch) {
311            assert!(completed_epoch > previous_epoch);
312        }
313    }
314}
315
316/// Build a fake `BarrierInfo` for independent partial-graph barriers.
317///
318/// Shared by both `CreatingStreamingJobControl` and `BatchRefreshJobCheckpointControl`.
319fn new_fake_barrier(
320    prev_epoch_fake_physical_time: &mut u64,
321    pending_non_checkpoint_barriers: &mut Vec<u64>,
322    kind: PbBarrierKind,
323) -> BarrierInfo {
324    let prev_epoch = TracedEpoch::new(Epoch::from_physical_time(*prev_epoch_fake_physical_time));
325    *prev_epoch_fake_physical_time += 1;
326    let curr_epoch = TracedEpoch::new(Epoch::from_physical_time(*prev_epoch_fake_physical_time));
327    let kind = match kind {
328        PbBarrierKind::Unspecified => unreachable!(),
329        PbBarrierKind::Initial => {
330            assert!(pending_non_checkpoint_barriers.is_empty());
331            BarrierKind::Initial
332        }
333        PbBarrierKind::Barrier => {
334            pending_non_checkpoint_barriers.push(prev_epoch.value().0);
335            BarrierKind::Barrier
336        }
337        PbBarrierKind::Checkpoint => {
338            pending_non_checkpoint_barriers.push(prev_epoch.value().0);
339            BarrierKind::Checkpoint(take(pending_non_checkpoint_barriers))
340        }
341    };
342    BarrierInfo {
343        prev_epoch,
344        curr_epoch,
345        kind,
346    }
347}
348
349// ── Enum unifying independent checkpoint job types ──────────────────────────
350
351/// The type-specific running state of a streaming job that checkpoints independently from the
352/// database's main graph.
353pub(crate) enum IndependentCheckpointJob {
354    CreatingStreamingJob(CreatingStreamingJobControl),
355    BatchRefresh(BatchRefreshJobCheckpointControl),
356    IcebergV3(IcebergV3JobCheckpointControl),
357}
358
359#[derive(Debug, PartialEq, Eq)]
360pub(crate) enum IndependentCheckpointJobStatus {
361    /// The initial barrier cannot be completed until the database has committed the snapshot
362    /// epoch.
363    Initial { snapshot_epoch: u64 },
364    /// The upstream database has committed the snapshot epoch, so barriers may complete.
365    Ready,
366}
367
368/// The lifecycle shared by all independent checkpoint jobs.
369pub(crate) enum IndependentCheckpointJobControl {
370    Running {
371        status: IndependentCheckpointJobStatus,
372        job_id: JobId,
373        partial_graph_id: PartialGraphId,
374        job: IndependentCheckpointJob,
375    },
376    Resetting {
377        /// Keep changelog pins until compute acknowledges the partial-graph reset. Before that
378        /// acknowledgment, snapshot-backfill executors may still be stopping and reading logs.
379        pinned_upstream_tables: HashSet<TableId>,
380        subscriptions_to_drop: Vec<PbSubscriptionUpstreamInfo>,
381        notifiers: Vec<CollectionNotifier>,
382    },
383}
384
385impl IndependentCheckpointJob {
386    fn can_drop_independently(&self) -> bool {
387        match self {
388            Self::CreatingStreamingJob(j) => j.can_drop_independently(),
389            Self::BatchRefresh(_) => true,
390            Self::IcebergV3(_) => true,
391        }
392    }
393
394    fn can_drop(&self) -> bool {
395        match self {
396            Self::CreatingStreamingJob(j) => j.can_drop(),
397            Self::BatchRefresh(_) => true,
398            Self::IcebergV3(_) => true,
399        }
400    }
401
402    fn pinned_upstream_tables(&self) -> &HashSet<TableId> {
403        match self {
404            Self::CreatingStreamingJob(j) => j.pinned_upstream_tables(),
405            Self::BatchRefresh(j) => j.pinned_upstream_tables(),
406            Self::IcebergV3(j) => j.pinned_upstream_tables(),
407        }
408    }
409
410    pub(crate) fn pre_apply_throttle(
411        &mut self,
412        config: &mut ThrottleConfigMap,
413    ) -> Option<Mutation> {
414        match self {
415            Self::CreatingStreamingJob(job) => job.pre_apply_throttle(config),
416            Self::BatchRefresh(job) => job.pre_apply_throttle(config),
417            Self::IcebergV3(job) => job.pre_apply_throttle(config),
418        }
419    }
420
421    pub(crate) fn on_new_upstream_barrier(
422        &mut self,
423        partial_graph_manager: &mut PartialGraphManager,
424        barrier_info: &BarrierInfo,
425        mutation: Option<(Mutation, Option<&mut NotifierStarter>)>,
426    ) -> MetaResult<()> {
427        match self {
428            Self::CreatingStreamingJob(job) => {
429                job.on_new_upstream_barrier(partial_graph_manager, barrier_info, mutation)
430            }
431            Self::BatchRefresh(job) => {
432                job.on_new_upstream_barrier(partial_graph_manager, barrier_info, mutation)
433            }
434            Self::IcebergV3(job) => {
435                job.on_new_upstream_barrier(partial_graph_manager, barrier_info, mutation)
436            }
437        }
438    }
439}
440
441impl IndependentCheckpointJobStatus {
442    fn on_upstream_database_ack_completed(&mut self, committed_epoch: u64) {
443        if let Self::Initial { snapshot_epoch } = self
444            && committed_epoch >= *snapshot_epoch
445        {
446            *self = Self::Ready;
447        }
448    }
449}
450
451impl IndependentCheckpointJobControl {
452    pub(crate) fn resetting(pinned_upstream_tables: HashSet<TableId>) -> Self {
453        Self::Resetting {
454            pinned_upstream_tables,
455            subscriptions_to_drop: vec![],
456            notifiers: vec![],
457        }
458    }
459
460    pub(crate) fn creating_streaming_job(
461        job_id: JobId,
462        partial_graph_id: PartialGraphId,
463        status: IndependentCheckpointJobStatus,
464        job: CreatingStreamingJobControl,
465    ) -> Self {
466        Self::Running {
467            status,
468            job_id,
469            partial_graph_id,
470            job: IndependentCheckpointJob::CreatingStreamingJob(job),
471        }
472    }
473
474    pub(crate) fn batch_refresh(
475        job_id: JobId,
476        partial_graph_id: PartialGraphId,
477        status: IndependentCheckpointJobStatus,
478        job: BatchRefreshJobCheckpointControl,
479    ) -> Self {
480        Self::Running {
481            status,
482            job_id,
483            partial_graph_id,
484            job: IndependentCheckpointJob::BatchRefresh(job),
485        }
486    }
487
488    pub(crate) fn iceberg_v3(
489        job_id: JobId,
490        partial_graph_id: PartialGraphId,
491        status: IndependentCheckpointJobStatus,
492        job: IcebergV3JobCheckpointControl,
493    ) -> Self {
494        Self::Running {
495            status,
496            job_id,
497            partial_graph_id,
498            job: IndependentCheckpointJob::IcebergV3(job),
499        }
500    }
501
502    pub(crate) fn running(&self) -> Option<&IndependentCheckpointJob> {
503        match self {
504            Self::Running { job, .. } => Some(job),
505            Self::Resetting { .. } => None,
506        }
507    }
508
509    pub(crate) fn running_mut(&mut self) -> Option<&mut IndependentCheckpointJob> {
510        match self {
511            Self::Running { job, .. } => Some(job),
512            Self::Resetting { .. } => None,
513        }
514    }
515
516    pub(crate) fn ready_mut(&mut self) -> Option<&mut IndependentCheckpointJob> {
517        match self {
518            Self::Running {
519                status: IndependentCheckpointJobStatus::Ready,
520                job,
521                ..
522            } => Some(job),
523            Self::Running {
524                status: IndependentCheckpointJobStatus::Initial { .. },
525                ..
526            }
527            | Self::Resetting { .. } => None,
528        }
529    }
530
531    pub(crate) fn on_upstream_database_ack_completed(&mut self, committed_epoch: u64) {
532        if let Self::Running { status, .. } = self {
533            status.on_upstream_database_ack_completed(committed_epoch);
534        }
535    }
536
537    pub(crate) fn gen_backfill_progress(&self) -> Option<BackfillProgress> {
538        match self.running()? {
539            IndependentCheckpointJob::CreatingStreamingJob(j) => Some(j.gen_backfill_progress()),
540            IndependentCheckpointJob::BatchRefresh(j) => j.gen_backfill_progress(),
541            IndependentCheckpointJob::IcebergV3(j) => j.gen_backfill_progress(),
542        }
543    }
544
545    /// Collect a barrier and return whether a checkpoint should be forced in the next barrier.
546    pub(crate) fn collect(&mut self, collected_barrier: CollectedBarrier<'_>) -> bool {
547        let job = self
548            .running_mut()
549            .expect("barriers should only be collected from a running partial graph");
550        match job {
551            IndependentCheckpointJob::CreatingStreamingJob(j) => j.collect(collected_barrier),
552            IndependentCheckpointJob::BatchRefresh(j) => j.collect(collected_barrier),
553            IndependentCheckpointJob::IcebergV3(j) => j.collect(collected_barrier),
554        }
555    }
556
557    pub(crate) fn gen_fragment_backfill_progress(&self) -> Vec<FragmentBackfillProgress> {
558        match self.running() {
559            Some(IndependentCheckpointJob::CreatingStreamingJob(j)) => {
560                j.gen_fragment_backfill_progress()
561            }
562            Some(IndependentCheckpointJob::BatchRefresh(j)) => j.gen_fragment_backfill_progress(),
563            Some(IndependentCheckpointJob::IcebergV3(j)) => j.gen_fragment_backfill_progress(),
564            None => vec![],
565        }
566    }
567
568    pub(crate) fn pinned_upstream_tables(&self) -> &HashSet<TableId> {
569        match self {
570            Self::Running { job, .. } => job.pinned_upstream_tables(),
571            Self::Resetting {
572                pinned_upstream_tables,
573                ..
574            } => pinned_upstream_tables,
575        }
576    }
577
578    pub(crate) fn fragment_infos(&self) -> Option<&HashMap<FragmentId, InflightFragmentInfo>> {
579        match self.running()? {
580            IndependentCheckpointJob::CreatingStreamingJob(j) => j.fragment_infos(),
581            IndependentCheckpointJob::BatchRefresh(j) => j.fragment_infos(),
582            IndependentCheckpointJob::IcebergV3(j) => Some(j.fragment_infos()),
583        }
584    }
585
586    pub(crate) fn ack_completed(
587        &mut self,
588        partial_graph_manager: &mut PartialGraphManager,
589        epoch: u64,
590    ) -> MetaResult<()> {
591        match self {
592            Self::Running {
593                status: IndependentCheckpointJobStatus::Ready,
594                job: IndependentCheckpointJob::CreatingStreamingJob(j),
595                ..
596            } => {
597                j.ack_completed(partial_graph_manager, epoch);
598            }
599            Self::Running {
600                status: IndependentCheckpointJobStatus::Ready,
601                job: IndependentCheckpointJob::BatchRefresh(j),
602                ..
603            } => {
604                j.ack_completed(partial_graph_manager, epoch);
605            }
606            Self::Running {
607                status: IndependentCheckpointJobStatus::Ready,
608                job: IndependentCheckpointJob::IcebergV3(j),
609                ..
610            } => j.ack_completed(partial_graph_manager, epoch)?,
611            Self::Running {
612                status: IndependentCheckpointJobStatus::Initial { .. },
613                ..
614            } => {
615                panic!("an initial job should transition to ready before completing a barrier")
616            }
617            Self::Resetting { .. } => {
618                // The job was dropped while the completing task was running in the background.
619                // The partial graph has already been reset, so skip the ack.
620            }
621        }
622        Ok(())
623    }
624
625    pub(crate) fn on_partial_graph_reset(self) -> Vec<PbSubscriptionUpstreamInfo> {
626        match self {
627            Self::Resetting {
628                subscriptions_to_drop,
629                notifiers,
630                ..
631            } => {
632                for notifier in notifiers {
633                    notifier.notify_collected();
634                }
635                subscriptions_to_drop
636            }
637            Self::Running { .. } => {
638                panic!("should be resetting when receiving reset partial graph resp")
639            }
640        }
641    }
642
643    pub(crate) fn can_drop(&self) -> bool {
644        match self {
645            Self::Running { job, .. } => job.can_drop(),
646            Self::Resetting { .. } => true,
647        }
648    }
649
650    pub(crate) fn drop(
651        &mut self,
652        notifier: Option<&mut NotifierStarter>,
653        partial_graph_manager: &mut PartialGraphManager,
654    ) -> bool {
655        match self {
656            Self::Resetting { notifiers, .. } => {
657                notifiers.extend(notifier.map(NotifierStarter::add_notify));
658                true
659            }
660            Self::Running { job, .. } if !job.can_drop_independently() => false,
661            Self::Running {
662                job_id,
663                partial_graph_id,
664                job,
665                ..
666            } => {
667                let subscriptions_to_drop = job
668                    .pinned_upstream_tables()
669                    .iter()
670                    .map(|upstream_mv_table_id| PbSubscriptionUpstreamInfo {
671                        subscriber_id: job_id.as_subscriber_id(),
672                        upstream_mv_table_id: *upstream_mv_table_id,
673                    })
674                    .collect();
675                // Resetting must keep owning the pins after the concrete job is replaced.
676                let pinned_upstream_tables = job.pinned_upstream_tables().clone();
677                partial_graph_manager.reset_partial_graphs([*partial_graph_id]);
678                *self = Self::Resetting {
679                    pinned_upstream_tables,
680                    subscriptions_to_drop,
681                    notifiers: notifier
682                        .map(NotifierStarter::add_notify)
683                        .into_iter()
684                        .collect(),
685                };
686                true
687            }
688        }
689    }
690
691    /// Reset during database recovery.
692    ///
693    /// Returns `true` if the partial graph was already resetting (from a prior drop),
694    /// meaning caller should not issue a new reset request.
695    pub(crate) fn reset(self) -> bool {
696        match self {
697            // Running jobs have no reset state to drain. Recovery will issue the partial-graph
698            // reset after rebuilding its reset plan.
699            Self::Running { .. } => false,
700            Self::Resetting { notifiers, .. } => {
701                for notifier in notifiers {
702                    notifier.notify_collected();
703                }
704                true
705            }
706        }
707    }
708}
709
710#[cfg(test)]
711mod tests {
712    use super::*;
713
714    #[test]
715    fn test_resetting_has_shared_inactive_behavior_and_keeps_pin() {
716        let pinned_upstream_tables = HashSet::from([TableId::new(7)]);
717        let job = IndependentCheckpointJobControl::Resetting {
718            pinned_upstream_tables: pinned_upstream_tables.clone(),
719            subscriptions_to_drop: vec![],
720            notifiers: vec![],
721        };
722
723        assert!(job.running().is_none());
724        assert!(job.gen_backfill_progress().is_none());
725        assert!(job.gen_fragment_backfill_progress().is_empty());
726        assert_eq!(job.pinned_upstream_tables(), &pinned_upstream_tables);
727        assert!(job.fragment_infos().is_none());
728        assert!(job.reset());
729    }
730
731    #[test]
732    fn test_reset_completion_returns_subscriptions_to_drop() {
733        let subscriptions_to_drop = vec![PbSubscriptionUpstreamInfo {
734            subscriber_id: JobId::new(8).as_subscriber_id(),
735            upstream_mv_table_id: TableId::new(7),
736        }];
737        let job = IndependentCheckpointJobControl::Resetting {
738            pinned_upstream_tables: HashSet::from([TableId::new(7)]),
739            subscriptions_to_drop: subscriptions_to_drop.clone(),
740            notifiers: vec![],
741        };
742
743        assert_eq!(job.on_partial_graph_reset(), subscriptions_to_drop);
744    }
745
746    #[test]
747    fn test_initial_status_transitions_on_upstream_database_ack() {
748        let mut status = IndependentCheckpointJobStatus::Initial { snapshot_epoch: 10 };
749
750        status.on_upstream_database_ack_completed(9);
751        assert_eq!(
752            status,
753            IndependentCheckpointJobStatus::Initial { snapshot_epoch: 10 }
754        );
755        status.on_upstream_database_ack_completed(10);
756        assert_eq!(status, IndependentCheckpointJobStatus::Ready);
757        status.on_upstream_database_ack_completed(11);
758        assert_eq!(status, IndependentCheckpointJobStatus::Ready);
759    }
760}