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::TableId;
19use risingwave_common::id::JobId;
20use risingwave_common::util::epoch::Epoch;
21use risingwave_pb::id::{FragmentId, PartialGraphId};
22use risingwave_pb::stream_plan::PbSubscriptionUpstreamInfo;
23use risingwave_pb::stream_plan::barrier::PbBarrierKind;
24
25pub(crate) mod batch_refresh_job;
26pub(crate) mod creating_job;
27
28pub(crate) use batch_refresh_job::{
29    BatchRefreshJobCheckpointControl, BatchRefreshJobTriggerContext, BatchRefreshLogicalFragments,
30    BatchRefreshRenderResult,
31};
32pub(crate) use creating_job::CreatingStreamingJobControl;
33
34use crate::barrier::info::BarrierInfo;
35use crate::barrier::notifier::{CollectionNotifier, NotifierStarter};
36use crate::barrier::partial_graph::{CollectedBarrier, PartialGraphManager};
37use crate::barrier::{BackfillProgress, BarrierKind, FragmentBackfillProgress, TracedEpoch};
38use crate::controller::fragment::InflightFragmentInfo;
39
40/// Build a fake `BarrierInfo` for independent partial-graph barriers.
41///
42/// Shared by both `CreatingStreamingJobControl` and `BatchRefreshJobCheckpointControl`.
43fn new_fake_barrier(
44    prev_epoch_fake_physical_time: &mut u64,
45    pending_non_checkpoint_barriers: &mut Vec<u64>,
46    kind: PbBarrierKind,
47) -> BarrierInfo {
48    let prev_epoch = TracedEpoch::new(Epoch::from_physical_time(*prev_epoch_fake_physical_time));
49    *prev_epoch_fake_physical_time += 1;
50    let curr_epoch = TracedEpoch::new(Epoch::from_physical_time(*prev_epoch_fake_physical_time));
51    let kind = match kind {
52        PbBarrierKind::Unspecified => unreachable!(),
53        PbBarrierKind::Initial => {
54            assert!(pending_non_checkpoint_barriers.is_empty());
55            BarrierKind::Initial
56        }
57        PbBarrierKind::Barrier => {
58            pending_non_checkpoint_barriers.push(prev_epoch.value().0);
59            BarrierKind::Barrier
60        }
61        PbBarrierKind::Checkpoint => {
62            pending_non_checkpoint_barriers.push(prev_epoch.value().0);
63            BarrierKind::Checkpoint(take(pending_non_checkpoint_barriers))
64        }
65    };
66    BarrierInfo {
67        prev_epoch,
68        curr_epoch,
69        kind,
70    }
71}
72
73// ── Enum unifying independent checkpoint job types ──────────────────────────
74
75/// The type-specific running state of a streaming job that checkpoints independently from the
76/// database's main graph.
77pub(crate) enum IndependentCheckpointJob {
78    CreatingStreamingJob(CreatingStreamingJobControl),
79    BatchRefresh(BatchRefreshJobCheckpointControl),
80}
81
82#[derive(Debug, PartialEq, Eq)]
83pub(crate) enum IndependentCheckpointJobStatus {
84    /// The initial barrier cannot be completed until the database has committed the snapshot
85    /// epoch.
86    Initial { snapshot_epoch: u64 },
87    /// The upstream database has committed the snapshot epoch, so barriers may complete.
88    Ready,
89}
90
91/// The lifecycle shared by all independent checkpoint jobs.
92pub(crate) enum IndependentCheckpointJobControl {
93    Running {
94        status: IndependentCheckpointJobStatus,
95        job_id: JobId,
96        partial_graph_id: PartialGraphId,
97        job: IndependentCheckpointJob,
98    },
99    Resetting {
100        /// Keep changelog pins until compute acknowledges the partial-graph reset. Before that
101        /// acknowledgment, snapshot-backfill executors may still be stopping and reading logs.
102        pinned_upstream_tables: HashSet<TableId>,
103        subscriptions_to_drop: Vec<PbSubscriptionUpstreamInfo>,
104        notifiers: Vec<CollectionNotifier>,
105    },
106}
107
108impl IndependentCheckpointJob {
109    fn can_drop_independently(&self) -> bool {
110        match self {
111            Self::CreatingStreamingJob(j) => j.can_drop_independently(),
112            Self::BatchRefresh(_) => true,
113        }
114    }
115
116    fn pinned_upstream_tables(&self) -> &HashSet<TableId> {
117        match self {
118            Self::CreatingStreamingJob(j) => j.pinned_upstream_tables(),
119            Self::BatchRefresh(j) => j.pinned_upstream_tables(),
120        }
121    }
122}
123
124impl IndependentCheckpointJobStatus {
125    fn on_upstream_database_ack_completed(&mut self, committed_epoch: u64) {
126        if let Self::Initial { snapshot_epoch } = self
127            && committed_epoch >= *snapshot_epoch
128        {
129            *self = Self::Ready;
130        }
131    }
132}
133
134impl IndependentCheckpointJobControl {
135    pub(crate) fn creating_streaming_job(
136        job_id: JobId,
137        partial_graph_id: PartialGraphId,
138        status: IndependentCheckpointJobStatus,
139        job: CreatingStreamingJobControl,
140    ) -> Self {
141        Self::Running {
142            status,
143            job_id,
144            partial_graph_id,
145            job: IndependentCheckpointJob::CreatingStreamingJob(job),
146        }
147    }
148
149    pub(crate) fn batch_refresh(
150        job_id: JobId,
151        partial_graph_id: PartialGraphId,
152        status: IndependentCheckpointJobStatus,
153        job: BatchRefreshJobCheckpointControl,
154    ) -> Self {
155        Self::Running {
156            status,
157            job_id,
158            partial_graph_id,
159            job: IndependentCheckpointJob::BatchRefresh(job),
160        }
161    }
162
163    pub(crate) fn running(&self) -> Option<&IndependentCheckpointJob> {
164        match self {
165            Self::Running { job, .. } => Some(job),
166            Self::Resetting { .. } => None,
167        }
168    }
169
170    pub(crate) fn running_mut(&mut self) -> Option<&mut IndependentCheckpointJob> {
171        match self {
172            Self::Running { job, .. } => Some(job),
173            Self::Resetting { .. } => None,
174        }
175    }
176
177    pub(crate) fn ready_mut(&mut self) -> Option<&mut IndependentCheckpointJob> {
178        match self {
179            Self::Running {
180                status: IndependentCheckpointJobStatus::Ready,
181                job,
182                ..
183            } => Some(job),
184            Self::Running {
185                status: IndependentCheckpointJobStatus::Initial { .. },
186                ..
187            }
188            | Self::Resetting { .. } => None,
189        }
190    }
191
192    pub(crate) fn on_upstream_database_ack_completed(&mut self, committed_epoch: u64) {
193        if let Self::Running { status, .. } = self {
194            status.on_upstream_database_ack_completed(committed_epoch);
195        }
196    }
197
198    pub(crate) fn gen_backfill_progress(&self) -> Option<BackfillProgress> {
199        match self.running()? {
200            IndependentCheckpointJob::CreatingStreamingJob(j) => Some(j.gen_backfill_progress()),
201            IndependentCheckpointJob::BatchRefresh(j) => j.gen_backfill_progress(),
202        }
203    }
204
205    /// Collect a barrier and return whether a checkpoint should be forced in the next barrier.
206    pub(crate) fn collect(&mut self, collected_barrier: CollectedBarrier<'_>) -> bool {
207        let job = self
208            .running_mut()
209            .expect("barriers should only be collected from a running partial graph");
210        match job {
211            IndependentCheckpointJob::CreatingStreamingJob(j) => j.collect(collected_barrier),
212            IndependentCheckpointJob::BatchRefresh(j) => j.collect(collected_barrier),
213        }
214    }
215
216    pub(crate) fn gen_fragment_backfill_progress(&self) -> Vec<FragmentBackfillProgress> {
217        match self.running() {
218            Some(IndependentCheckpointJob::CreatingStreamingJob(j)) => {
219                j.gen_fragment_backfill_progress()
220            }
221            Some(IndependentCheckpointJob::BatchRefresh(j)) => j.gen_fragment_backfill_progress(),
222            None => vec![],
223        }
224    }
225
226    pub(crate) fn pinned_upstream_tables(&self) -> &HashSet<TableId> {
227        match self {
228            Self::Running { job, .. } => job.pinned_upstream_tables(),
229            Self::Resetting {
230                pinned_upstream_tables,
231                ..
232            } => pinned_upstream_tables,
233        }
234    }
235
236    pub(crate) fn fragment_infos(&self) -> Option<&HashMap<FragmentId, InflightFragmentInfo>> {
237        match self.running()? {
238            IndependentCheckpointJob::CreatingStreamingJob(j) => j.fragment_infos(),
239            IndependentCheckpointJob::BatchRefresh(j) => j.fragment_infos(),
240        }
241    }
242
243    pub(crate) fn ack_completed(
244        &mut self,
245        partial_graph_manager: &mut PartialGraphManager,
246        epoch: u64,
247    ) {
248        match self {
249            Self::Running {
250                status: IndependentCheckpointJobStatus::Ready,
251                job: IndependentCheckpointJob::CreatingStreamingJob(j),
252                ..
253            } => j.ack_completed(partial_graph_manager, epoch),
254            Self::Running {
255                status: IndependentCheckpointJobStatus::Ready,
256                job: IndependentCheckpointJob::BatchRefresh(j),
257                ..
258            } => j.ack_completed(partial_graph_manager, epoch),
259            Self::Running {
260                status: IndependentCheckpointJobStatus::Initial { .. },
261                ..
262            } => {
263                panic!("an initial job should transition to ready before completing a barrier")
264            }
265            Self::Resetting { .. } => {
266                // The job was dropped while the completing task was running in the background.
267                // The partial graph has already been reset, so skip the ack.
268            }
269        }
270    }
271
272    pub(crate) fn on_partial_graph_reset(self) -> Vec<PbSubscriptionUpstreamInfo> {
273        match self {
274            Self::Resetting {
275                subscriptions_to_drop,
276                notifiers,
277                ..
278            } => {
279                for notifier in notifiers {
280                    notifier.notify_collected();
281                }
282                subscriptions_to_drop
283            }
284            Self::Running { .. } => {
285                panic!("should be resetting when receiving reset partial graph resp")
286            }
287        }
288    }
289
290    pub(crate) fn drop(
291        &mut self,
292        notifier: Option<&mut NotifierStarter>,
293        partial_graph_manager: &mut PartialGraphManager,
294    ) -> bool {
295        match self {
296            Self::Resetting { notifiers, .. } => {
297                notifiers.extend(notifier.map(NotifierStarter::add_notify));
298                true
299            }
300            Self::Running { job, .. } if !job.can_drop_independently() => false,
301            Self::Running {
302                job_id,
303                partial_graph_id,
304                job,
305                ..
306            } => {
307                let subscriptions_to_drop = job
308                    .pinned_upstream_tables()
309                    .iter()
310                    .map(|upstream_mv_table_id| PbSubscriptionUpstreamInfo {
311                        subscriber_id: job_id.as_subscriber_id(),
312                        upstream_mv_table_id: *upstream_mv_table_id,
313                    })
314                    .collect();
315                // Resetting must keep owning the pins after the concrete job is replaced.
316                let pinned_upstream_tables = job.pinned_upstream_tables().clone();
317                partial_graph_manager.reset_partial_graphs([*partial_graph_id]);
318                *self = Self::Resetting {
319                    pinned_upstream_tables,
320                    subscriptions_to_drop,
321                    notifiers: notifier
322                        .map(NotifierStarter::add_notify)
323                        .into_iter()
324                        .collect(),
325                };
326                true
327            }
328        }
329    }
330
331    /// Reset during database recovery.
332    ///
333    /// Returns `true` if the partial graph was already resetting (from a prior drop),
334    /// meaning caller should not issue a new reset request.
335    pub(crate) fn reset(self) -> bool {
336        match self {
337            // Running jobs have no reset state to drain. Recovery will issue the partial-graph
338            // reset after rebuilding its reset plan.
339            Self::Running { .. } => false,
340            Self::Resetting { notifiers, .. } => {
341                for notifier in notifiers {
342                    notifier.notify_collected();
343                }
344                true
345            }
346        }
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn test_resetting_has_shared_inactive_behavior_and_keeps_pin() {
356        let pinned_upstream_tables = HashSet::from([TableId::new(7)]);
357        let job = IndependentCheckpointJobControl::Resetting {
358            pinned_upstream_tables: pinned_upstream_tables.clone(),
359            subscriptions_to_drop: vec![],
360            notifiers: vec![],
361        };
362
363        assert!(job.running().is_none());
364        assert!(job.gen_backfill_progress().is_none());
365        assert!(job.gen_fragment_backfill_progress().is_empty());
366        assert_eq!(job.pinned_upstream_tables(), &pinned_upstream_tables);
367        assert!(job.fragment_infos().is_none());
368        assert!(job.reset());
369    }
370
371    #[test]
372    fn test_reset_completion_returns_subscriptions_to_drop() {
373        let subscriptions_to_drop = vec![PbSubscriptionUpstreamInfo {
374            subscriber_id: JobId::new(8).as_subscriber_id(),
375            upstream_mv_table_id: TableId::new(7),
376        }];
377        let job = IndependentCheckpointJobControl::Resetting {
378            pinned_upstream_tables: HashSet::from([TableId::new(7)]),
379            subscriptions_to_drop: subscriptions_to_drop.clone(),
380            notifiers: vec![],
381        };
382
383        assert_eq!(job.on_partial_graph_reset(), subscriptions_to_drop);
384    }
385
386    #[test]
387    fn test_initial_status_transitions_on_upstream_database_ack() {
388        let mut status = IndependentCheckpointJobStatus::Initial { snapshot_epoch: 10 };
389
390        status.on_upstream_database_ack_completed(9);
391        assert_eq!(
392            status,
393            IndependentCheckpointJobStatus::Initial { snapshot_epoch: 10 }
394        );
395        status.on_upstream_database_ack_completed(10);
396        assert_eq!(status, IndependentCheckpointJobStatus::Ready);
397        status.on_upstream_database_ack_completed(11);
398        assert_eq!(status, IndependentCheckpointJobStatus::Ready);
399    }
400}
risingwave_meta/barrier/checkpoint/independent_job/mod.rs

risingwave_meta/barrier/checkpoint/independent_job/
mod.rs
