1use 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 ConsumingSnapshot {
36 snapshot: SnapshotPhaseControl,
37 pending_upstream_barriers: Vec<BarrierInfo>,
38 info: CreatingJobInfo,
39 },
40 ConsumingLogStore {
44 tracking_job: TrackingJob,
45 info: CreatingJobInfo,
46 pending_barriers: VecDeque<BarrierInfo>,
47 },
48 ConsumingUpstream {
51 tracking_job: TrackingJob,
52 info: CreatingJobInfo,
53 },
54 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>, ) -> 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 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 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 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}