Skip to main content

risingwave_meta/barrier/checkpoint/independent_job/creating_job/
status.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, VecDeque};
16use std::mem::replace;
17
18use risingwave_pb::id::{FragmentId, PartialGraphId};
19use risingwave_pb::stream_plan::barrier_mutation::Mutation;
20use risingwave_pb::stream_service::barrier_complete_response::CreateMviewProgress;
21
22use crate::barrier::BarrierInfo;
23use crate::barrier::checkpoint::independent_job::SnapshotPhaseControl;
24use crate::barrier::checkpoint::independent_job::creating_job::CreatingJobInfo;
25use crate::barrier::command::{ThrottleConfigMap, extract_throttle_config};
26use crate::barrier::partial_graph::PartialGraphManager;
27use crate::barrier::progress::TrackingJob;
28use crate::controller::fragment::InflightFragmentInfo;
29
30#[derive(Debug)]
31pub(super) enum CreatingStreamingJobStatus {
32    /// The creating job is consuming upstream snapshot.
33    /// Will transit to `ConsumingLogStore` on `update_progress` when
34    /// the snapshot has been fully consumed after `update_progress`.
35    ConsumingSnapshot {
36        snapshot: SnapshotPhaseControl,
37        pending_upstream_barriers: Vec<BarrierInfo>,
38        info: CreatingJobInfo,
39    },
40    /// The creating job is consuming log store.
41    ///
42    /// Will transit to `ConsumingUpstream` after the delayed runtime edge is installed.
43    ConsumingLogStore {
44        tracking_job: TrackingJob,
45        info: CreatingJobInfo,
46        pending_barriers: VecDeque<BarrierInfo>,
47    },
48    /// The delayed runtime edge has been installed and the existing actors consume the live
49    /// upstream while remaining in their independent partial graph.
50    ConsumingUpstream {
51        tracking_job: TrackingJob,
52        info: CreatingJobInfo,
53    },
54    /// All backfill actors have started consuming upstream, and the job
55    /// will be finished when all previously injected barriers have been collected
56    /// Store the `prev_epoch` that will finish at.
57    Finishing(u64, TrackingJob),
58    PlaceHolder,
59}
60
61impl CreatingStreamingJobStatus {
62    pub(super) fn update_progress(
63        &mut self,
64        create_mview_progress: impl IntoIterator<Item = &CreateMviewProgress>,
65    ) {
66        match self {
67            &mut Self::ConsumingSnapshot {
68                ref mut snapshot,
69                ref mut pending_upstream_barriers,
70                ..
71            } => {
72                if snapshot.apply_progress(create_mview_progress) {
73                    let pending_barriers: VecDeque<_> = [snapshot.finish_snapshot_barrier()]
74                        .into_iter()
75                        .chain(pending_upstream_barriers.drain(..))
76                        .collect();
77
78                    let CreatingStreamingJobStatus::ConsumingSnapshot { snapshot, info, .. } =
79                        replace(self, CreatingStreamingJobStatus::PlaceHolder)
80                    else {
81                        unreachable!()
82                    };
83
84                    let tracking_job = snapshot.create_mview_tracker.into_tracking_job();
85
86                    *self = CreatingStreamingJobStatus::ConsumingLogStore {
87                        tracking_job,
88                        info,
89                        pending_barriers,
90                    };
91                }
92            }
93            CreatingStreamingJobStatus::ConsumingLogStore { .. }
94            | CreatingStreamingJobStatus::ConsumingUpstream { .. } => {}
95            CreatingStreamingJobStatus::Finishing(..) => {}
96            CreatingStreamingJobStatus::PlaceHolder => {
97                unreachable!()
98            }
99        }
100    }
101
102    pub(super) fn start_consume_upstream(
103        &mut self,
104        barrier_info: &BarrierInfo,
105    ) -> &CreatingJobInfo {
106        match self {
107            CreatingStreamingJobStatus::ConsumingSnapshot { .. } => {
108                unreachable!(
109                    "should not start consuming upstream for a job that are consuming snapshot"
110                )
111            }
112            CreatingStreamingJobStatus::ConsumingLogStore {
113                pending_barriers, ..
114            } => {
115                assert!(pending_barriers.is_empty());
116                assert!(barrier_info.kind.is_checkpoint());
117                let CreatingStreamingJobStatus::ConsumingLogStore {
118                    info, tracking_job, ..
119                } = replace(self, CreatingStreamingJobStatus::PlaceHolder)
120                else {
121                    unreachable!()
122                };
123                *self = CreatingStreamingJobStatus::ConsumingUpstream { tracking_job, info };
124            }
125            CreatingStreamingJobStatus::ConsumingUpstream { .. } => {
126                unreachable!("should not start consuming upstream for a job again")
127            }
128            CreatingStreamingJobStatus::Finishing { .. } => {
129                unreachable!("should not start consuming upstream for a job again")
130            }
131            CreatingStreamingJobStatus::PlaceHolder => {
132                unreachable!()
133            }
134        }
135        let CreatingStreamingJobStatus::ConsumingUpstream { info, .. } = self else {
136            unreachable!("should be consuming upstream")
137        };
138        info
139    }
140
141    pub(super) fn start_finishing(&mut self, barrier_info: &BarrierInfo) -> CreatingJobInfo {
142        let CreatingStreamingJobStatus::ConsumingUpstream { tracking_job, info } =
143            replace(self, CreatingStreamingJobStatus::PlaceHolder)
144        else {
145            unreachable!("should only finish after starting to consume upstream")
146        };
147        *self = CreatingStreamingJobStatus::Finishing(barrier_info.prev_epoch(), tracking_job);
148        info
149    }
150
151    pub(super) fn on_new_upstream_epoch(
152        &mut self,
153        partial_graph_manager: &PartialGraphManager,
154        partial_graph_id: PartialGraphId,
155        max_pending_barrier_num: usize,
156        barrier_info: &BarrierInfo,
157        mutation: Option<Mutation>, // mutation to be set for the first barrier to inject
158    ) -> Vec<(BarrierInfo, Option<Mutation>)> {
159        let resolve_initial_barrier_num_to_inject = || {
160            max_pending_barrier_num
161                .saturating_sub(partial_graph_manager.pending_barrier_num(partial_graph_id))
162        };
163        match self {
164            CreatingStreamingJobStatus::ConsumingSnapshot {
165                pending_upstream_barriers,
166                snapshot,
167                ..
168            } => {
169                let mutation = mutation.or_else(|| snapshot.take_start_backfill_mutation());
170                let barrier_num_to_inject = resolve_initial_barrier_num_to_inject();
171                pending_upstream_barriers.push(barrier_info.clone());
172                // Mutation barriers must be forwarded even when the partial graph has reached the
173                // configured pending-barrier limit.
174                if barrier_num_to_inject == 0 && mutation.is_none() {
175                    return vec![];
176                }
177                vec![(snapshot.next_fake_barrier(&barrier_info.kind), mutation)]
178            }
179            CreatingStreamingJobStatus::ConsumingLogStore {
180                pending_barriers, ..
181            } => {
182                // Throttle has no effect on the snapshot executor after it starts consuming the
183                // log store. The updated fragment plan is kept for the actors created on merge,
184                // so the mutation does not need to be forwarded in this phase.
185                drain_pending_barriers(
186                    pending_barriers,
187                    barrier_info.clone(),
188                    resolve_initial_barrier_num_to_inject(),
189                )
190                .into_iter()
191                .map(|barrier_info| (barrier_info, None))
192                .collect()
193            }
194            CreatingStreamingJobStatus::ConsumingUpstream { .. } => {
195                vec![(barrier_info.clone(), mutation)]
196            }
197            CreatingStreamingJobStatus::Finishing { .. } => vec![],
198            CreatingStreamingJobStatus::PlaceHolder => {
199                unreachable!()
200            }
201        }
202    }
203
204    pub(super) fn fragment_infos(&self) -> Option<&HashMap<FragmentId, InflightFragmentInfo>> {
205        match self {
206            CreatingStreamingJobStatus::ConsumingSnapshot { info, .. }
207            | CreatingStreamingJobStatus::ConsumingLogStore { info, .. }
208            | CreatingStreamingJobStatus::ConsumingUpstream { info, .. } => {
209                Some(&info.fragment_infos)
210            }
211            CreatingStreamingJobStatus::Finishing(..) => None,
212            CreatingStreamingJobStatus::PlaceHolder => {
213                unreachable!()
214            }
215        }
216    }
217
218    /// Dropping a job while its delayed upstream edge is being activated can leave an upstream
219    /// actor waiting for an output request from an actor that has already been stopped.
220    pub(super) fn can_drop(&self) -> bool {
221        match self {
222            CreatingStreamingJobStatus::ConsumingSnapshot { .. }
223            | CreatingStreamingJobStatus::ConsumingLogStore { .. }
224            | CreatingStreamingJobStatus::Finishing(_, _) => true,
225            CreatingStreamingJobStatus::ConsumingUpstream { .. } => false,
226            CreatingStreamingJobStatus::PlaceHolder => {
227                unreachable!()
228            }
229        }
230    }
231
232    pub(super) fn pre_apply_throttle(
233        &mut self,
234        config: &mut ThrottleConfigMap,
235    ) -> Option<Mutation> {
236        let fragment_infos = match self {
237            CreatingStreamingJobStatus::ConsumingSnapshot { info, .. }
238            | CreatingStreamingJobStatus::ConsumingLogStore { info, .. }
239            | CreatingStreamingJobStatus::ConsumingUpstream { info, .. } => {
240                &mut info.fragment_infos
241            }
242            CreatingStreamingJobStatus::Finishing(..) => return None,
243            CreatingStreamingJobStatus::PlaceHolder => {
244                unreachable!()
245            }
246        };
247
248        extract_throttle_config(config, |fragment_id, stream_node| {
249            if let Some(fragment_info) = fragment_infos.get_mut(&fragment_id) {
250                fragment_info.nodes = stream_node.clone();
251                true
252            } else {
253                false
254            }
255        })
256    }
257}
258
259fn drain_pending_barriers(
260    pending_barriers: &mut VecDeque<BarrierInfo>,
261    new_upstream_barrier: BarrierInfo,
262    barrier_num_to_inject: usize,
263) -> Vec<BarrierInfo> {
264    pending_barriers.push_back(new_upstream_barrier);
265    let barrier_count = pending_barriers.len().min(barrier_num_to_inject);
266    pending_barriers.drain(..barrier_count).collect()
267}
268
269#[cfg(test)]
270mod tests {
271    use risingwave_common::util::epoch::Epoch;
272    use risingwave_pb::stream_plan::PbStreamNode;
273
274    use super::*;
275    use crate::barrier::{BarrierKind, TracedEpoch};
276
277    fn barrier(prev_epoch: u64, curr_epoch: u64) -> BarrierInfo {
278        BarrierInfo {
279            prev_epoch: TracedEpoch::new(Epoch(prev_epoch)),
280            curr_epoch: TracedEpoch::new(Epoch(curr_epoch)),
281            kind: BarrierKind::Barrier,
282        }
283    }
284
285    fn epochs(barriers: &[BarrierInfo]) -> Vec<(u64, u64)> {
286        barriers
287            .iter()
288            .map(|barrier| (barrier.prev_epoch(), barrier.curr_epoch()))
289            .collect()
290    }
291
292    #[test]
293    fn test_drain_pending_barriers_with_available_capacity() {
294        let mut pending_barriers = VecDeque::from([barrier(1, 2), barrier(2, 3), barrier(3, 4)]);
295
296        let injected = drain_pending_barriers(&mut pending_barriers, barrier(4, 5), 0);
297        assert!(injected.is_empty());
298        assert_eq!(
299            epochs(pending_barriers.make_contiguous()),
300            vec![(1, 2), (2, 3), (3, 4), (4, 5)]
301        );
302
303        let injected = drain_pending_barriers(&mut pending_barriers, barrier(5, 6), 2);
304        assert_eq!(epochs(&injected), vec![(1, 2), (2, 3)]);
305        assert_eq!(
306            epochs(pending_barriers.make_contiguous()),
307            vec![(3, 4), (4, 5), (5, 6)]
308        );
309
310        let injected = drain_pending_barriers(&mut pending_barriers, barrier(6, 7), 2);
311        assert_eq!(epochs(&injected), vec![(3, 4), (4, 5)]);
312        assert_eq!(
313            epochs(pending_barriers.make_contiguous()),
314            vec![(5, 6), (6, 7)]
315        );
316
317        let injected = drain_pending_barriers(&mut pending_barriers, barrier(7, 8), 2);
318        assert_eq!(epochs(&injected), vec![(5, 6), (6, 7)]);
319        assert_eq!(epochs(pending_barriers.make_contiguous()), vec![(7, 8)]);
320    }
321
322    #[test]
323    fn test_drain_pending_barriers_without_backlog() {
324        let mut pending_barriers = VecDeque::new();
325
326        let injected = drain_pending_barriers(&mut pending_barriers, barrier(1, 2), 100);
327
328        assert_eq!(epochs(&injected), vec![(1, 2)]);
329        assert!(pending_barriers.is_empty());
330    }
331
332    #[test]
333    fn test_drop_admission_during_upstream_handoff() {
334        let job_id = risingwave_common::id::JobId::new(1);
335        let fragment_infos = HashMap::new();
336        let mut status = CreatingStreamingJobStatus::ConsumingLogStore {
337            tracking_job: TrackingJob::recovered(job_id, &fragment_infos),
338            info: CreatingJobInfo {
339                fragment_infos,
340                upstream_fragment_downstreams: Default::default(),
341                downstreams: Default::default(),
342                snapshot_backfill_upstream_tables: Default::default(),
343                stream_actors: Default::default(),
344            },
345            pending_barriers: Default::default(),
346        };
347        let transition_barrier = BarrierInfo {
348            prev_epoch: TracedEpoch::new(Epoch(1)),
349            curr_epoch: TracedEpoch::new(Epoch(2)),
350            kind: BarrierKind::Checkpoint(vec![1]),
351        };
352
353        assert!(status.can_drop());
354        status.start_consume_upstream(&transition_barrier);
355        assert!(!status.can_drop());
356        status.start_finishing(&BarrierInfo {
357            prev_epoch: TracedEpoch::new(Epoch(2)),
358            curr_epoch: TracedEpoch::new(Epoch(3)),
359            kind: BarrierKind::Checkpoint(vec![2]),
360        });
361        assert!(status.can_drop());
362    }
363
364    #[test]
365    fn test_pre_apply_throttle_before_merge() {
366        let job_id = risingwave_common::id::JobId::new(1);
367        let fragment_id = FragmentId::new(1);
368        let old_node = PbStreamNode {
369            identity: "old".to_owned(),
370            ..Default::default()
371        };
372        let new_node = PbStreamNode {
373            identity: "new".to_owned(),
374            ..Default::default()
375        };
376        let fragment_infos = HashMap::from([(
377            fragment_id,
378            InflightFragmentInfo {
379                fragment_id,
380                distribution_type: risingwave_meta_model::fragment::DistributionType::Single,
381                fragment_type_mask: Default::default(),
382                vnode_count: 1,
383                nodes: old_node,
384                actors: Default::default(),
385                state_table_ids: Default::default(),
386            },
387        )]);
388        let mut status = CreatingStreamingJobStatus::ConsumingLogStore {
389            tracking_job: TrackingJob::recovered(job_id, &fragment_infos),
390            info: CreatingJobInfo {
391                fragment_infos,
392                upstream_fragment_downstreams: Default::default(),
393                downstreams: Default::default(),
394                snapshot_backfill_upstream_tables: Default::default(),
395                stream_actors: Default::default(),
396            },
397            pending_barriers: Default::default(),
398        };
399        let mut config = HashMap::from([(
400            fragment_id,
401            (
402                risingwave_pb::stream_plan::throttle_mutation::ThrottleConfig {
403                    rate_limit: Some(1_000),
404                    throttle_type: Default::default(),
405                },
406                new_node.clone(),
407            ),
408        )]);
409
410        assert!(status.pre_apply_throttle(&mut config).is_some());
411        assert!(config.is_empty());
412
413        let transition_barrier = BarrierInfo {
414            prev_epoch: TracedEpoch::new(Epoch(1)),
415            curr_epoch: TracedEpoch::new(Epoch(2)),
416            kind: BarrierKind::Checkpoint(vec![1]),
417        };
418        let info = status.start_consume_upstream(&transition_barrier);
419        assert_eq!(info.fragment_infos[&fragment_id].nodes, new_node);
420        assert!(matches!(
421            status,
422            CreatingStreamingJobStatus::ConsumingUpstream { .. }
423        ));
424        let info = status.start_finishing(&BarrierInfo {
425            prev_epoch: TracedEpoch::new(Epoch(2)),
426            curr_epoch: TracedEpoch::new(Epoch(3)),
427            kind: BarrierKind::Checkpoint(vec![2]),
428        });
429        assert_eq!(info.fragment_infos[&fragment_id].nodes, new_node);
430        assert!(matches!(
431            status,
432            CreatingStreamingJobStatus::Finishing(2, _)
433        ));
434    }
435}