1use std::future::ready;
16use std::mem::take;
17use std::sync::Arc;
18use std::time::Duration;
19
20use anyhow::anyhow;
21use futures::future::{Either, try_join_all};
22use futures::{Stream, TryFutureExt, TryStreamExt, pin_mut};
23use risingwave_common::array::StreamChunk;
24use risingwave_common::catalog::TableId;
25use risingwave_common::hash::VnodeBitmapExt;
26use risingwave_common::row::OwnedRow;
27use risingwave_common::util::epoch::EpochPair;
28use risingwave_common_rate_limit::{MonitoredRateLimiter, RateLimit, RateLimiter};
29use risingwave_hummock_sdk::HummockReadEpoch;
30use risingwave_pb::batch_plan::ScanRange;
31use risingwave_pb::common::PbThrottleType;
32use risingwave_storage::StateStore;
33use risingwave_storage::store::PrefetchOptions;
34use risingwave_storage::table::ChangeLogRow;
35use risingwave_storage::table::batch_table::{BatchTable, PkScanRange};
36use tokio::select;
37
38use crate::executor::backfill::snapshot_backfill::state::{
39 BackfillState, EpochBackfillProgress, VnodeBackfillProgress,
40};
41use crate::executor::backfill::snapshot_backfill::vnode_stream::VnodeStream;
42use crate::executor::backfill::utils::{
43 UpstreamStreamKeyUpdateNormalizer, create_builder, mapping_message,
44};
45use crate::executor::monitor::StreamingMetrics;
46use crate::executor::prelude::{StateTable, StreamExt, try_stream};
47use crate::executor::{
48 ActorContextRef, Barrier, BoxedMessageStream, Execute, Executor, MergeExecutorInput, Message,
49 Mutation, StreamExecutorError, StreamExecutorResult, UpstreamFragmentId, expect_first_barrier,
50};
51use crate::task::CreateMviewProgressReporter;
52
53pub struct SnapshotBackfillExecutor<S: StateStore> {
54 upstream_table: BatchTable<S>,
56
57 progress_state_table: StateTable<S>,
59
60 upstream: Executor,
62 initial_upstream_is_empty: bool,
63 upstream_fragment_id: UpstreamFragmentId,
64
65 output_indices: Vec<usize>,
67
68 stream_key: Vec<usize>,
70
71 progress: CreateMviewProgressReporter,
72
73 chunk_size: usize,
74 rate_limiter: MonitoredRateLimiter,
75
76 actor_ctx: ActorContextRef,
77 metrics: Arc<StreamingMetrics>,
78
79 snapshot_epoch: Option<u64>,
80 pk_scan_range: PkScanRange,
82}
83
84impl<S: StateStore> SnapshotBackfillExecutor<S> {
85 fn build_pk_scan_range(
86 pb_scan_range: Option<&ScanRange>,
87 upstream_table: &BatchTable<S>,
88 ) -> StreamExecutorResult<PkScanRange> {
89 match pb_scan_range {
90 Some(scan_range) => Ok(PkScanRange::new(
91 scan_range.clone(),
92 upstream_table.pk_serializer().get_data_types().to_vec(),
93 )?),
94 None => Ok(PkScanRange::full()),
95 }
96 }
97
98 #[expect(clippy::too_many_arguments)]
99 pub(crate) fn new(
100 upstream_table: BatchTable<S>,
101 progress_state_table: StateTable<S>,
102 upstream: MergeExecutorInput,
103 pb_pk_scan_range: Option<&ScanRange>,
104 output_indices: Vec<usize>,
105 stream_key: Vec<usize>,
106 actor_ctx: ActorContextRef,
107 progress: CreateMviewProgressReporter,
108 chunk_size: usize,
109 rate_limit: RateLimit,
110 metrics: Arc<StreamingMetrics>,
111 snapshot_epoch: Option<u64>,
112 ) -> StreamExecutorResult<Self> {
113 assert_eq!(&upstream.info.schema, upstream_table.schema());
114 let initial_upstream_is_empty = upstream.is_empty();
115 let upstream_fragment_id = upstream.upstream_fragment_id();
116 let upstream = upstream.into_executor();
117 if upstream_table.pk_in_output_indices().is_none() {
118 panic!(
119 "storage table should include all pk columns in output: pk_indices: {:?}, output_indices: {:?}, schema: {:?}",
120 upstream_table.pk_indices(),
121 upstream_table.output_indices(),
122 upstream_table.schema()
123 )
124 };
125 assert!(
126 stream_key.iter().all(|idx| *idx < output_indices.len()),
127 "stream key indices should refer to output schema: stream_key: {:?}, output_indices: {:?}",
128 stream_key,
129 output_indices
130 );
131 let pk_scan_range = Self::build_pk_scan_range(pb_pk_scan_range, &upstream_table)?;
132 if !matches!(rate_limit, RateLimit::Disabled) {
133 trace!(
134 ?rate_limit,
135 "create snapshot backfill executor with rate limit"
136 );
137 }
138 let rate_limiter = RateLimiter::new(rate_limit).monitored(upstream_table.table_id());
139 Ok(Self {
140 upstream_table,
141 progress_state_table,
142 upstream,
143 initial_upstream_is_empty,
144 upstream_fragment_id,
145 output_indices,
146 stream_key,
147 progress,
148 chunk_size,
149 rate_limiter,
150 actor_ctx,
151 metrics,
152 snapshot_epoch,
153 pk_scan_range,
154 })
155 }
156
157 #[try_stream(ok = Message, error = StreamExecutorError)]
158 async fn execute_inner(mut self) {
159 trace!("snapshot backfill executor start");
160 let current_stream_key_indices = self
161 .stream_key
162 .iter()
163 .map(|idx| self.output_indices[*idx])
164 .collect();
165 let update_normalizer = UpstreamStreamKeyUpdateNormalizer::new(
166 &self.upstream.info().stream_key,
167 current_stream_key_indices,
168 );
169 let mut upstream = self.upstream.execute();
170 let first_recv_barrier = expect_first_barrier(&mut upstream).await?;
171 trace!(epoch = ?first_recv_barrier.epoch, "get first barrier");
172 let should_snapshot_backfill: Option<u64> =
173 if let Some(snapshot_epoch) = self.snapshot_epoch {
174 if self.initial_upstream_is_empty {
175 Some(snapshot_epoch)
176 } else {
177 None
178 }
179 } else {
180 if cfg!(debug_assertions) {
182 panic!(
183 "snapshot epoch not set. first_recv_epoch: {:?}",
184 first_recv_barrier.epoch
185 );
186 } else {
187 if self.initial_upstream_is_empty {
188 Err(anyhow!("no upstream while snapshot epoch not set"))?;
189 }
190 warn!(first_recv_epoch = ?first_recv_barrier.epoch, "snapshot epoch not set");
191 None
192 }
193 };
194 let first_recv_barrier_epoch = first_recv_barrier.epoch;
195 let upstream_table_id = self.upstream_table.table_id();
196 let progress_table_id = self.progress_state_table.table_id();
197 let initial_backfill_paused =
198 first_recv_barrier.is_backfill_pause_on_startup(self.actor_ctx.fragment_id);
199 yield Message::Barrier(first_recv_barrier);
200 let mut backfill_state = BackfillState::new(
201 self.progress_state_table,
202 first_recv_barrier_epoch,
203 self.upstream_table.pk_serializer().clone(),
204 )
205 .await?;
206
207 let mut barrier_epoch = {
208 if let Some(snapshot_epoch) = should_snapshot_backfill {
209 let table_id_str = format!("{}", self.upstream_table.table_id());
210 let actor_id_str = format!("{}", self.actor_ctx.id);
211
212 let mut barrier_epoch = if first_recv_barrier_epoch.prev < snapshot_epoch {
216 trace!(
217 table_id = %self.upstream_table.table_id(),
218 snapshot_epoch,
219 barrier_epoch = ?first_recv_barrier_epoch,
220 "start consuming snapshot"
221 );
222 {
223 let consuming_snapshot_row_count = self
224 .metrics
225 .snapshot_backfill_consume_row_count
226 .with_guarded_label_values(&[
227 actor_id_str.as_str(),
228 table_id_str.as_str(),
229 "consuming_snapshot",
230 ]);
231 let snapshot_stream = make_consume_snapshot_stream(
232 &self.upstream_table,
233 snapshot_epoch,
234 self.chunk_size,
235 &self.rate_limiter,
236 &mut upstream,
237 &mut self.progress,
238 &mut backfill_state,
239 first_recv_barrier_epoch,
240 initial_backfill_paused,
241 &self.actor_ctx,
242 &self.pk_scan_range,
243 progress_table_id,
244 );
245
246 pin_mut!(snapshot_stream);
247
248 while let Some(message) = snapshot_stream.try_next().await? {
249 if let Message::Chunk(chunk) = &message {
250 consuming_snapshot_row_count.inc_by(chunk.cardinality() as _);
251 }
252 yield message;
253 }
254 }
255
256 let recv_barrier = receive_next_upstream_barrier(&mut upstream).await?;
257 let recv_barrier_epoch = recv_barrier.epoch;
258 assert_eq!(snapshot_epoch, recv_barrier_epoch.prev);
259 let post_commit = backfill_state.commit(recv_barrier.epoch).await?;
260 if recv_barrier.is_checkpoint() {
261 self.progress.report_table_change_log_progress(
262 recv_barrier.epoch,
263 upstream_table_id,
264 progress_table_id,
265 self.upstream_table.vnodes().clone(),
266 snapshot_epoch,
267 );
268 }
269 yield Message::Barrier(recv_barrier);
270 post_commit.post_yield_barrier(None).await?;
271 recv_barrier_epoch
272 } else {
273 trace!(
274 table_id = %self.upstream_table.table_id(),
275 snapshot_epoch,
276 barrier_epoch = ?first_recv_barrier_epoch,
277 "skip consuming snapshot"
278 );
279 first_recv_barrier_epoch
280 };
281
282 trace!(
284 ?barrier_epoch,
285 table_id = %self.upstream_table.table_id(),
286 "start consuming log store"
287 );
288
289 let consuming_log_store_row_count = self
290 .metrics
291 .snapshot_backfill_consume_row_count
292 .with_guarded_label_values(&[
293 actor_id_str.as_str(),
294 table_id_str.as_str(),
295 "consuming_log_store",
296 ]);
297 let mut pending_non_checkpoint_barrier: Vec<EpochPair> = vec![];
298 loop {
299 let barrier = receive_next_upstream_barrier(&mut upstream).await?;
300 assert_eq!(barrier_epoch.curr, barrier.epoch.prev);
301 barrier_epoch = barrier.epoch;
302 if barrier.kind.is_checkpoint() {
303 let pending_non_checkpoint_barrier =
304 take(&mut pending_non_checkpoint_barrier);
305 let end_epoch = barrier_epoch.prev;
306 let start_epoch = pending_non_checkpoint_barrier
307 .first()
308 .map(|epoch| epoch.prev)
309 .unwrap_or(end_epoch);
310 trace!(?barrier_epoch, kind = ?barrier.kind, ?pending_non_checkpoint_barrier, "start consume epoch change log");
311 let mut stream = make_log_stream(
312 &self.upstream_table,
313 start_epoch,
314 end_epoch,
315 None,
316 self.chunk_size,
317 )
318 .await?;
319 while let Some(chunk) = stream.try_next().await? {
320 trace!(
321 ?barrier_epoch,
322 size = chunk.cardinality(),
323 "consume change log yield chunk",
324 );
325 consuming_log_store_row_count.inc_by(chunk.cardinality() as _);
326 yield Message::Chunk(chunk);
327 }
328
329 trace!(?barrier_epoch, "after consume change log");
330
331 stream
332 .for_vnode_pk_progress(|vnode, row_count, progress| {
333 assert_eq!(progress, None);
334 backfill_state.finish_epoch(vnode, barrier.epoch.prev, row_count);
335 })
336 .await?;
337 } else {
338 pending_non_checkpoint_barrier.push(barrier.epoch);
339 }
340
341 let start_consuming_upstream = if barrier
344 .as_update_merge(self.actor_ctx.id, self.upstream_fragment_id)
345 .is_some()
346 {
347 assert!(
348 barrier.kind.is_checkpoint(),
349 "the first live-upstream barrier must be a checkpoint"
350 );
351 true
352 } else {
353 false
354 };
355 let post_commit = backfill_state.commit(barrier.epoch).await?;
356 if barrier.is_checkpoint() {
357 self.progress.report_table_change_log_progress(
358 barrier.epoch,
359 upstream_table_id,
360 progress_table_id,
361 self.upstream_table.vnodes().clone(),
362 barrier.epoch.prev,
363 );
364 }
365 let update_vnode_bitmap = barrier.as_update_vnode_bitmap(self.actor_ctx.id);
366 yield Message::Barrier(barrier);
367 post_commit.post_yield_barrier(None).await?;
368 if update_vnode_bitmap.is_some() {
369 return Err(anyhow!(
370 "should not update vnode bitmap during consuming log store"
371 )
372 .into());
373 }
374 if start_consuming_upstream {
375 break;
376 }
377 }
378 trace!(
379 ?barrier_epoch,
380 table_id = %self.upstream_table.table_id(),
381 "finish consuming log store"
382 );
383 barrier_epoch
384 } else {
385 backfill_state
386 .latest_progress()
387 .for_each(|(vnode, progress)| {
388 let progress = progress.expect("should not be empty");
389 assert_eq!(
394 progress.epoch, first_recv_barrier_epoch.prev,
395 "vnode: {:?}",
396 vnode
397 );
398 assert_eq!(
399 progress.progress,
400 EpochBackfillProgress::Consumed,
401 "vnode: {:?}",
402 vnode
403 );
404 });
405 trace!(
406 table_id = %self.upstream_table.table_id(),
407 "skip backfill"
408 );
409 first_recv_barrier_epoch
410 }
411 };
412 let mut epoch_row_count = 0;
413 while let Some(msg) = upstream.try_next().await? {
415 let Some(msg) = update_normalizer.normalize_message(msg) else {
416 continue;
417 };
418 match msg {
419 Message::Barrier(barrier) => {
420 assert_eq!(barrier.epoch.prev, barrier_epoch.curr);
421 self.upstream_table
422 .vnodes()
423 .iter_vnodes()
424 .for_each(|vnode| {
425 backfill_state.finish_epoch(vnode, barrier.epoch.prev, epoch_row_count);
428 });
429 epoch_row_count = 0;
430 let update_vnode_bitmap = barrier.as_update_vnode_bitmap(self.actor_ctx.id);
431 barrier_epoch = barrier.epoch;
432 let post_commit = backfill_state.commit(barrier.epoch).await?;
433 yield Message::Barrier(barrier);
434 if let Some(new_vnode_bitmap) =
435 post_commit.post_yield_barrier(update_vnode_bitmap).await?
436 {
437 let _prev_vnode_bitmap =
438 self.upstream_table.update_vnode_bitmap(new_vnode_bitmap);
439 backfill_state
440 .latest_progress()
441 .for_each(|(vnode, progress)| {
442 let progress = progress.expect("should not be empty");
443 assert_eq!(
444 progress.epoch, barrier_epoch.prev,
445 "vnode {:?} has unexpected progress epoch",
446 vnode
447 );
448 assert_eq!(
449 progress.progress,
450 EpochBackfillProgress::Consumed,
451 "vnode {:?} has unexpected progress",
452 vnode
453 );
454 });
455 }
456 }
457 msg => {
458 if let Message::Chunk(chunk) = &msg {
459 epoch_row_count += chunk.cardinality();
460 }
461 yield msg;
462 }
463 }
464 }
465 }
466}
467
468impl<S: StateStore> Execute for SnapshotBackfillExecutor<S> {
469 fn execute(self: Box<Self>) -> BoxedMessageStream {
470 let output_indices = self.output_indices.clone();
471 self.execute_inner()
472 .filter_map(move |result| {
473 ready({
474 match result {
475 Ok(message) => mapping_message(message, &output_indices).map(Ok),
476 Err(e) => Some(Err(e)),
477 }
478 })
479 })
480 .boxed()
481 }
482}
483
484async fn receive_next_upstream_barrier(
485 upstream: &mut BoxedMessageStream,
486) -> StreamExecutorResult<Barrier> {
487 let message = upstream
488 .try_next()
489 .await?
490 .ok_or_else(|| anyhow!("end of snapshot backfill merge input"))?;
491 match message {
492 Message::Barrier(barrier) => Ok(barrier),
493 message => Err(anyhow!(
494 "expected barrier from empty snapshot backfill merge input, got {message:?}"
495 )
496 .into()),
497 }
498}
499
500#[await_tree::instrument("make_log_stream: {start_epoch}-{end_epoch} table {}", upstream_table.table_id())]
501async fn make_log_stream(
502 upstream_table: &BatchTable<impl StateStore>,
503 start_epoch: u64,
504 end_epoch: u64,
505 start_pk: Option<OwnedRow>,
506 chunk_size: usize,
507) -> StreamExecutorResult<VnodeStream<impl super::vnode_stream::ChangeLogRowStream>> {
508 let data_types = upstream_table.schema().data_types();
509 let start_pk = start_pk.as_ref();
510 let vnode_streams = try_join_all(upstream_table.vnodes().iter_vnodes().map(move |vnode| {
512 upstream_table
513 .batch_iter_vnode_log(
514 start_epoch,
515 HummockReadEpoch::Committed(end_epoch),
516 start_pk,
517 vnode,
518 )
519 .map_ok(move |stream| {
520 let stream = stream.map_err(Into::into);
521 (vnode, stream, 0)
522 })
523 }))
524 .await?;
525 let builder = create_builder(RateLimit::Disabled, chunk_size, data_types.clone());
526 Ok(VnodeStream::new(
527 vnode_streams,
528 upstream_table.pk_in_output_indices().expect("should exist"),
529 builder,
530 ))
531}
532
533async fn make_snapshot_stream(
534 upstream_table: &BatchTable<impl StateStore>,
535 snapshot_epoch: u64,
536 backfill_state: &BackfillState<impl StateStore>,
537 rate_limit: RateLimit,
538 chunk_size: usize,
539 snapshot_rebuild_interval: Duration,
540 pk_scan_range: &PkScanRange,
541) -> StreamExecutorResult<VnodeStream<impl super::vnode_stream::ChangeLogRowStream>> {
542 let data_types = upstream_table.schema().data_types();
543 let vnode_streams = try_join_all(backfill_state.latest_progress().filter_map(
544 move |(vnode, progress)| {
545 let start_pk = match progress {
546 None => Some((None, 0)),
547 Some(VnodeBackfillProgress {
548 row_count,
549 progress: EpochBackfillProgress::Consuming { latest_pk },
550 ..
551 }) => Some((Some(latest_pk), *row_count)),
552 Some(VnodeBackfillProgress {
553 progress: EpochBackfillProgress::Consumed,
554 ..
555 }) => None,
556 };
557 start_pk.map(|(start_pk, row_count)| {
558 upstream_table
559 .batch_iter_vnode_with_pk_range(
560 HummockReadEpoch::Committed(snapshot_epoch),
561 start_pk,
562 &pk_scan_range.pk_prefix,
563 &pk_scan_range.range_bounds,
564 vnode,
565 PrefetchOptions::prefetch_for_large_range_scan(),
566 snapshot_rebuild_interval,
567 )
568 .map_ok(move |stream| {
569 let stream = stream.map_ok(ChangeLogRow::Insert).map_err(Into::into);
570 (vnode, stream, row_count)
571 })
572 })
573 },
574 ))
575 .await?;
576 let builder = create_builder(rate_limit, chunk_size, data_types.clone());
577 Ok(VnodeStream::new(
578 vnode_streams,
579 upstream_table.pk_in_output_indices().expect("should exist"),
580 builder,
581 ))
582}
583
584#[expect(clippy::too_many_arguments)]
585#[try_stream(ok = Message, error = StreamExecutorError)]
586async fn make_consume_snapshot_stream<'a, S: StateStore>(
587 upstream_table: &'a BatchTable<S>,
588 snapshot_epoch: u64,
589 chunk_size: usize,
590 rate_limiter: &'a MonitoredRateLimiter,
591 upstream: &'a mut BoxedMessageStream,
592 progress: &'a mut CreateMviewProgressReporter,
593 backfill_state: &'a mut BackfillState<S>,
594 first_recv_barrier_epoch: EpochPair,
595 initial_backfill_paused: bool,
596 actor_ctx: &'a ActorContextRef,
597 pk_scan_range: &'a PkScanRange,
598 progress_table_id: TableId,
599) {
600 let mut barrier_epoch = first_recv_barrier_epoch;
601
602 let mut snapshot_stream = make_snapshot_stream(
604 upstream_table,
605 snapshot_epoch,
606 &*backfill_state,
607 rate_limiter.rate_limit(),
608 chunk_size,
609 actor_ctx.config.developer.snapshot_iter_rebuild_interval(),
610 pk_scan_range,
611 )
612 .await?;
613
614 async fn select_barrier_and_snapshot_stream(
615 upstream: &mut BoxedMessageStream,
616 snapshot_stream: &mut (impl Stream<Item = StreamExecutorResult<StreamChunk>> + Unpin),
617 throttle_snapshot_stream: bool,
618 backfill_paused: bool,
619 ) -> StreamExecutorResult<Either<Barrier, Option<StreamChunk>>> {
620 select! {
621 biased;
622
623 result = receive_next_upstream_barrier(upstream) => {
624 Ok(Either::Left(result?))
625 },
626 result = snapshot_stream.try_next(), if !throttle_snapshot_stream && !backfill_paused => {
627 Ok(Either::Right(result?))
628 }
629 }
630 }
631
632 let mut backfill_paused = initial_backfill_paused;
633 loop {
634 let throttle_snapshot_stream = matches!(rate_limiter.rate_limit(), RateLimit::Pause);
635 match select_barrier_and_snapshot_stream(
636 upstream,
637 &mut snapshot_stream,
638 throttle_snapshot_stream,
639 backfill_paused,
640 )
641 .await?
642 {
643 Either::Left(barrier) => {
644 assert_eq!(barrier.epoch.prev, barrier_epoch.curr);
645 barrier_epoch = barrier.epoch;
646
647 if barrier_epoch.curr >= snapshot_epoch {
648 return Err(anyhow!("should not receive barrier with epoch {barrier_epoch:?} later than snapshot epoch {snapshot_epoch}").into());
649 }
650 if barrier.should_start_fragment_backfill(actor_ctx.fragment_id) {
651 backfill_paused = false;
652 }
653 if let Some(chunk) = snapshot_stream.consume_builder() {
654 rate_limiter.wait(chunk.cardinality() as _).await;
655 yield Message::Chunk(chunk);
656 }
657 snapshot_stream
658 .for_vnode_pk_progress(|vnode, row_count, pk_progress| {
659 if let Some(pk) = pk_progress {
660 backfill_state.update_epoch_progress(
661 vnode,
662 snapshot_epoch,
663 row_count,
664 pk,
665 );
666 } else {
667 backfill_state.finish_epoch(vnode, snapshot_epoch, row_count);
668 }
669 })
670 .await?;
671 let count = backfill_state.total_row_count();
672 let post_commit = backfill_state.commit(barrier.epoch).await?;
673 trace!(?barrier_epoch, count, "update progress");
674 progress.update(barrier_epoch, barrier_epoch.prev, count as _);
675 if barrier.is_checkpoint() {
676 progress.report_table_change_log_progress(
677 barrier.epoch,
678 upstream_table.table_id(),
679 progress_table_id,
680 upstream_table.vnodes().clone(),
681 snapshot_epoch,
682 );
683 }
684
685 let new_rate_limit = barrier.mutation.as_ref().and_then(|m| {
686 if let Mutation::Throttle(config) = &**m
687 && let Some(config) = config.get(&actor_ctx.fragment_id)
688 && config.throttle_type() == PbThrottleType::Backfill
689 {
690 Some(config.rate_limit)
691 } else {
692 None
693 }
694 });
695 yield Message::Barrier(barrier);
696 post_commit.post_yield_barrier(None).await?;
697
698 if let Some(new_rate_limit) = new_rate_limit {
699 let new_rate_limit = new_rate_limit.into();
700 rate_limiter.update(new_rate_limit);
701 snapshot_stream.update_rate_limiter(new_rate_limit, chunk_size);
702 }
703 }
704 Either::Right(Some(chunk)) => {
705 if backfill_paused {
706 return Err(
707 anyhow!("snapshot backfill paused, but received snapshot chunk").into(),
708 );
709 }
710 rate_limiter.wait(chunk.cardinality() as _).await;
711 yield Message::Chunk(chunk);
712 }
713 Either::Right(None) => {
714 break;
715 }
716 }
717 }
718
719 let barrier_to_report_finish = receive_next_upstream_barrier(upstream).await?;
721 assert_eq!(barrier_to_report_finish.epoch.prev, barrier_epoch.curr);
722 barrier_epoch = barrier_to_report_finish.epoch;
723 snapshot_stream
724 .for_vnode_pk_progress(|vnode, row_count, pk_progress| {
725 assert_eq!(pk_progress, None);
726 backfill_state.finish_epoch(vnode, snapshot_epoch, row_count);
727 })
728 .await?;
729 let count = backfill_state.total_row_count();
730 trace!(?barrier_epoch, count, "report finish");
731 let post_commit = backfill_state.commit(barrier_epoch).await?;
732 progress.finish(barrier_epoch, count as _);
733 if barrier_to_report_finish.is_checkpoint() {
734 progress.report_table_change_log_progress(
735 barrier_epoch,
736 upstream_table.table_id(),
737 progress_table_id,
738 upstream_table.vnodes().clone(),
739 snapshot_epoch,
740 );
741 }
742 yield Message::Barrier(barrier_to_report_finish);
743 post_commit.post_yield_barrier(None).await?;
744
745 loop {
747 let barrier = receive_next_upstream_barrier(upstream).await?;
748 assert_eq!(barrier.epoch.prev, barrier_epoch.curr);
749 barrier_epoch = barrier.epoch;
750 let post_commit = backfill_state.commit(barrier.epoch).await?;
751 if barrier.is_checkpoint() {
752 progress.report_table_change_log_progress(
753 barrier.epoch,
754 upstream_table.table_id(),
755 progress_table_id,
756 upstream_table.vnodes().clone(),
757 snapshot_epoch,
758 );
759 }
760 yield Message::Barrier(barrier);
761 post_commit.post_yield_barrier(None).await?;
762 if barrier_epoch.curr == snapshot_epoch {
763 break;
764 }
765 }
766 trace!(?barrier_epoch, "finish consuming snapshot");
767}
768
769#[cfg(test)]
770mod tests {
771 use std::collections::{HashMap, HashSet};
772 use std::sync::Arc;
773 use std::task::Poll;
774
775 use assert_matches::assert_matches;
776 use futures::poll;
777 use risingwave_common::array::StreamChunk;
778 use risingwave_common::catalog::{ColumnDesc, ColumnId, Field, Schema, TableId};
779 use risingwave_common::row::OwnedRow;
780 use risingwave_common::test_prelude::StreamChunkTestExt;
781 use risingwave_common::types::DataType;
782 use risingwave_common::util::epoch::{EpochPair, test_epoch};
783 use risingwave_common::util::sort_util::OrderType;
784 use risingwave_hummock_test::test_utils::{HummockTestEnv, prepare_hummock_test_env};
785 use risingwave_pb::stream_plan::update_mutation::MergeUpdate;
786 use risingwave_rpc_client::HummockMetaClient;
787 use risingwave_storage::hummock::HummockStorage;
788 use risingwave_storage::table::batch_table::BatchTable;
789 use tokio::time::{Duration, timeout};
790
791 use super::*;
792 use crate::common::table::state_table::{
793 StateTable, StateTableBuilder, StateTableOpConsistencyLevel,
794 };
795 use crate::common::table::test_utils::gen_pbtable_with_value_indices;
796 use crate::executor::exchange::input::{Input, LocalInput};
797 use crate::executor::exchange::permit::channel_for_test;
798 use crate::executor::{
799 ActorContext, DispatcherMessage, ExecutorInfo, MergeExecutor, MergeExecutorInput,
800 MergeExecutorUpstream, UpdateMutation,
801 };
802 use crate::task::barrier_test_utils::LocalBarrierTestEnv;
803 use crate::task::test_utils::helper_make_local_actor;
804 use crate::task::{LocalBarrierManager, NewOutputRequest};
805
806 const SOURCE_TABLE_ID: TableId = TableId::new(0x233);
807 const PROGRESS_TABLE_ID: TableId = TableId::new(0x234);
808
809 fn source_table_pb() -> risingwave_pb::catalog::PbTable {
810 gen_pbtable_with_value_indices(
811 SOURCE_TABLE_ID,
812 vec![ColumnDesc::unnamed(ColumnId::new(0), DataType::Int64)],
813 vec![OrderType::ascending()],
814 vec![0],
815 0,
816 vec![0],
817 )
818 }
819
820 fn progress_table_pb() -> risingwave_pb::catalog::PbTable {
821 gen_pbtable_with_value_indices(
822 PROGRESS_TABLE_ID,
823 vec![
824 ColumnDesc::unnamed(ColumnId::new(0), DataType::Int16),
825 ColumnDesc::unnamed(ColumnId::new(1), DataType::Int64),
826 ColumnDesc::unnamed(ColumnId::new(2), DataType::Int64),
827 ColumnDesc::unnamed(ColumnId::new(3), DataType::Boolean),
828 ColumnDesc::unnamed(ColumnId::new(4), DataType::Int64),
829 ],
830 vec![OrderType::ascending()],
831 vec![0],
832 1,
833 vec![1, 2, 3, 4],
834 )
835 }
836
837 fn source_batch_table(store: HummockStorage) -> BatchTable<HummockStorage> {
838 BatchTable::for_test(
839 store,
840 SOURCE_TABLE_ID,
841 vec![ColumnDesc::unnamed(ColumnId::new(0), DataType::Int64)],
842 vec![OrderType::ascending()],
843 vec![0],
844 vec![0],
845 )
846 }
847
848 async fn source_state_table(store: HummockStorage) -> StateTable<HummockStorage> {
849 StateTableBuilder::new(&source_table_pb(), store, None)
850 .with_op_consistency_level(StateTableOpConsistencyLevel::LogStoreEnabled)
851 .forbid_preload_all_rows()
852 .build()
853 .await
854 }
855
856 async fn progress_state_table(store: HummockStorage) -> StateTable<HummockStorage> {
857 StateTable::from_table_catalog(&progress_table_pb(), store, None).await
858 }
859
860 async fn commit_insert_epoch(
861 test_env: &HummockTestEnv,
862 source_state_table: &mut StateTable<HummockStorage>,
863 epoch: &mut EpochPair,
864 table_ids: HashSet<TableId>,
865 values: &[i64],
866 ) {
867 for value in values {
868 source_state_table.insert(OwnedRow::new(vec![Some((*value).into())]));
869 }
870 epoch.inc_for_test();
871 test_env.storage.start_epoch(epoch.curr, table_ids);
872 source_state_table.commit_for_test(*epoch).await.unwrap();
873 let res = test_env
874 .storage
875 .seal_and_sync_epoch(epoch.prev, HashSet::from_iter([SOURCE_TABLE_ID]))
876 .await
877 .unwrap();
878 test_env
879 .meta_client
880 .commit_epoch_with_change_log(epoch.prev, res, Some(vec![epoch.prev]))
881 .await
882 .unwrap();
883 test_env
884 .storage
885 .wait_version(test_env.manager.get_current_version().await)
886 .await;
887 }
888
889 async fn start_progress_epochs(test_env: &HummockTestEnv, max_epoch: u64) {
890 for epoch in 1..=max_epoch {
891 test_env
892 .storage
893 .start_epoch(test_epoch(epoch), HashSet::from_iter([PROGRESS_TABLE_ID]));
894 }
895 test_env.storage.flush_events_for_test().await;
896 }
897
898 async fn persist_finished_progress(
899 test_env: &HummockTestEnv,
900 write_epoch: EpochPair,
901 progress_epoch: u64,
902 ) -> StateTable<HummockStorage> {
903 let mut table = progress_state_table(test_env.storage.clone()).await;
904 table.init_epoch(write_epoch).await.unwrap();
905 let vnodes: Vec<_> = table.vnodes().iter_vnodes().collect();
906 for vnode in vnodes {
907 table.insert(OwnedRow::new(vec![
908 Some(vnode.to_scalar().into()),
909 Some((progress_epoch as i64).into()),
910 Some(0_i64.into()),
911 Some(true.into()),
912 None,
913 ]));
914 }
915
916 let mut commit_epoch = write_epoch;
917 commit_epoch.inc_for_test();
918 table.commit_for_test(commit_epoch).await.unwrap();
919 let result = test_env
920 .storage
921 .seal_and_sync_epoch(commit_epoch.prev, HashSet::from_iter([PROGRESS_TABLE_ID]))
922 .await
923 .unwrap();
924 test_env
925 .meta_client
926 .commit_epoch_with_change_log(commit_epoch.prev, result, None)
927 .await
928 .unwrap();
929 test_env
930 .storage
931 .wait_version(test_env.manager.get_current_version().await)
932 .await;
933
934 while commit_epoch.prev < progress_epoch {
935 commit_epoch.inc_for_test();
936 table.commit_for_test(commit_epoch).await.unwrap();
937 let result = test_env
938 .storage
939 .seal_and_sync_epoch(commit_epoch.prev, HashSet::from_iter([PROGRESS_TABLE_ID]))
940 .await
941 .unwrap();
942 test_env
943 .meta_client
944 .commit_epoch_with_change_log(commit_epoch.prev, result, None)
945 .await
946 .unwrap();
947 test_env
948 .storage
949 .wait_version(test_env.manager.get_current_version().await)
950 .await;
951 }
952
953 progress_state_table(test_env.storage.clone()).await
954 }
955
956 fn make_upstream_info() -> ExecutorInfo {
957 ExecutorInfo::for_test(
958 Schema::new(vec![Field::unnamed(DataType::Int64)]),
959 vec![0],
960 "SnapshotBackfillUpstream".to_owned(),
961 0,
962 )
963 }
964
965 fn make_upstream_input(
966 barrier_manager: LocalBarrierManager,
967 actor_ctx: ActorContextRef,
968 rx: Option<crate::executor::exchange::permit::Receiver>,
969 ) -> MergeExecutorInput {
970 let metrics = Arc::new(StreamingMetrics::unused());
971 let upstream = MergeExecutor::new_merge_upstream(
972 rx.into_iter()
973 .map(|rx| LocalInput::new(rx, 1001.into()).boxed_input())
974 .collect(),
975 &metrics,
976 &actor_ctx,
977 1024,
978 make_upstream_info().schema,
979 );
980 MergeExecutorInput::new(
981 MergeExecutorUpstream::Merge(upstream),
982 actor_ctx,
983 1919.into(),
984 barrier_manager,
985 metrics,
986 make_upstream_info(),
987 )
988 }
989
990 async fn expect_barrier_with_timeout(
991 executor: &mut BoxedMessageStream,
992 reason: &str,
993 ) -> Barrier {
994 let message = timeout(Duration::from_secs(10), executor.next())
995 .await
996 .unwrap_or_else(|_| panic!("timed out waiting for barrier: {reason}"))
997 .unwrap()
998 .unwrap();
999 match message {
1000 Message::Barrier(barrier) => barrier,
1001 other => panic!("expected barrier for {reason}, got {other:?}"),
1002 }
1003 }
1004
1005 async fn expect_chunk_with_timeout(
1006 executor: &mut BoxedMessageStream,
1007 reason: &str,
1008 ) -> StreamChunk {
1009 let message = timeout(Duration::from_secs(10), executor.next())
1010 .await
1011 .unwrap_or_else(|_| panic!("timed out waiting for chunk: {reason}"))
1012 .unwrap()
1013 .unwrap();
1014 match message {
1015 Message::Chunk(chunk) => chunk,
1016 other => panic!("expected chunk for {reason}, got {other:?}"),
1017 }
1018 }
1019
1020 async fn expect_pending_with_timeout(executor: &mut BoxedMessageStream, reason: &str) {
1021 assert!(
1022 timeout(Duration::from_millis(200), executor.next())
1023 .await
1024 .is_err(),
1025 "executor unexpectedly produced a message while waiting for {reason}"
1026 );
1027 }
1028
1029 #[tokio::test]
1030 async fn test_snapshot_backfill_without_upstream_on_hummock() {
1031 let source_env = prepare_hummock_test_env().await;
1032 source_env.register_table(source_table_pb()).await;
1033 let progress_env = prepare_hummock_test_env().await;
1034 progress_env.register_table(progress_table_pb()).await;
1035
1036 let mut source_state_table = source_state_table(source_env.storage.clone()).await;
1037 let source_table = source_batch_table(source_env.storage.clone());
1038 let progress_state_table = progress_state_table(progress_env.storage.clone()).await;
1039
1040 let mut epoch = EpochPair::new_test_epoch(test_epoch(1));
1041 source_env
1042 .storage
1043 .start_epoch(epoch.curr, HashSet::from_iter([SOURCE_TABLE_ID]));
1044 source_env.storage.flush_events_for_test().await;
1045 source_state_table.init_epoch(epoch).await.unwrap();
1046
1047 commit_insert_epoch(
1048 &source_env,
1049 &mut source_state_table,
1050 &mut epoch,
1051 HashSet::from_iter([SOURCE_TABLE_ID]),
1052 &[1],
1053 )
1054 .await;
1055 commit_insert_epoch(
1056 &source_env,
1057 &mut source_state_table,
1058 &mut epoch,
1059 HashSet::from_iter([SOURCE_TABLE_ID]),
1060 &[2],
1061 )
1062 .await;
1063 commit_insert_epoch(
1064 &source_env,
1065 &mut source_state_table,
1066 &mut epoch,
1067 HashSet::from_iter([SOURCE_TABLE_ID]),
1068 &[3],
1069 )
1070 .await;
1071 commit_insert_epoch(
1072 &source_env,
1073 &mut source_state_table,
1074 &mut epoch,
1075 HashSet::from_iter([SOURCE_TABLE_ID]),
1076 &[4],
1077 )
1078 .await;
1079 start_progress_epochs(&progress_env, 5).await;
1080
1081 let barrier_test_env = LocalBarrierTestEnv::for_test().await;
1082 let barrier_manager = barrier_test_env.local_barrier_manager.clone();
1083 let progress = CreateMviewProgressReporter::for_test(barrier_manager.clone());
1084 let actor_id = 1234.into();
1085 let actor_ctx = ActorContext::for_test(actor_id);
1086 let initial_barrier = Barrier::new_test_barrier(test_epoch(1));
1087 barrier_test_env.inject_barrier(&initial_barrier, [actor_id]);
1088 barrier_test_env.flush_all_events().await;
1089 let upstream = make_upstream_input(barrier_manager.clone(), actor_ctx.clone(), None);
1090
1091 let mut executor = SnapshotBackfillExecutor::new(
1092 source_table,
1093 progress_state_table,
1094 upstream,
1095 None,
1096 vec![0],
1097 vec![0],
1098 actor_ctx,
1099 progress,
1100 1024,
1101 RateLimit::Disabled,
1102 Arc::new(StreamingMetrics::unused()),
1103 Some(test_epoch(3)),
1104 )
1105 .expect("snapshot backfill executor should be created")
1106 .boxed()
1107 .execute();
1108
1109 assert_eq!(
1110 expect_barrier_with_timeout(&mut executor, "initial injected barrier")
1111 .await
1112 .epoch,
1113 Barrier::new_test_barrier(test_epoch(1)).epoch
1114 );
1115 assert_eq!(
1116 expect_chunk_with_timeout(&mut executor, "snapshot chunk without upstream").await,
1117 StreamChunk::from_pretty(
1118 " I
1119 + 1
1120 + 2
1121 + 3"
1122 )
1123 );
1124 expect_pending_with_timeout(&mut executor, "snapshot finish barrier 2").await;
1125
1126 let barrier = Barrier::new_test_barrier(test_epoch(2));
1127 barrier_test_env.inject_barrier(&barrier, [actor_id]);
1128 barrier_test_env.flush_all_events().await;
1129 assert_eq!(
1130 expect_barrier_with_timeout(&mut executor, "snapshot progress barrier 2")
1131 .await
1132 .epoch,
1133 Barrier::new_test_barrier(test_epoch(2)).epoch
1134 );
1135
1136 let barrier = Barrier::new_test_barrier(test_epoch(3));
1137 barrier_test_env.inject_barrier(&barrier, [actor_id]);
1138 barrier_test_env.flush_all_events().await;
1139 assert_eq!(
1140 expect_barrier_with_timeout(&mut executor, "snapshot progress barrier 3")
1141 .await
1142 .epoch,
1143 Barrier::new_test_barrier(test_epoch(3)).epoch
1144 );
1145
1146 let barrier = Barrier::new_test_barrier(test_epoch(4));
1147 barrier_test_env.inject_barrier(&barrier, [actor_id]);
1148 barrier_test_env.flush_all_events().await;
1149 assert_eq!(
1150 expect_barrier_with_timeout(&mut executor, "post-snapshot barrier 4")
1151 .await
1152 .epoch,
1153 Barrier::new_test_barrier(test_epoch(4)).epoch
1154 );
1155
1156 let upstream_actor_id = 1001.into();
1157 let upstream_fragment_id = 1919.into();
1158 let transition_barrier = Barrier::new_test_barrier(test_epoch(5)).with_mutation(
1159 Mutation::Update(UpdateMutation {
1160 merges: HashMap::from([(
1161 (actor_id, upstream_fragment_id),
1162 MergeUpdate {
1163 actor_id,
1164 upstream_fragment_id,
1165 new_upstream_fragment_id: None,
1166 added_upstream_actors: vec![helper_make_local_actor(upstream_actor_id)],
1167 removed_upstream_actor_id: vec![],
1168 },
1169 )]),
1170 ..Default::default()
1171 }),
1172 );
1173 barrier_test_env.inject_barrier(&transition_barrier, [actor_id]);
1174 barrier_test_env.flush_all_events().await;
1175 assert_matches!(poll!(executor.as_mut().next()), Poll::Pending);
1176 barrier_test_env.flush_all_events().await;
1177 let mut requests = barrier_test_env
1178 .take_pending_new_output_requests(upstream_actor_id)
1179 .await;
1180 assert_eq!(requests.len(), 1);
1181 let (downstream_actor_id, request) = requests.pop().unwrap();
1182 let NewOutputRequest::Local(upstream_tx) = request else {
1183 unreachable!()
1184 };
1185 assert_eq!(downstream_actor_id, actor_id);
1186 upstream_tx
1187 .send(DispatcherMessage::Barrier(transition_barrier.clone().into_dispatcher()).into())
1188 .await
1189 .unwrap();
1190 assert_eq!(
1191 expect_chunk_with_timeout(&mut executor, "log-store chunk without upstream").await,
1192 StreamChunk::from_pretty(" I\n + 4")
1193 );
1194 assert_eq!(
1195 expect_barrier_with_timeout(&mut executor, "steady-state barrier 5")
1196 .await
1197 .epoch,
1198 Barrier::new_test_barrier(test_epoch(5)).epoch
1199 );
1200
1201 upstream_tx
1202 .send(DispatcherMessage::Chunk(StreamChunk::from_pretty(" I\n + 5")).into())
1203 .await
1204 .unwrap();
1205 assert_eq!(
1206 expect_chunk_with_timeout(&mut executor, "live upstream chunk after handoff").await,
1207 StreamChunk::from_pretty(" I\n + 5")
1208 );
1209
1210 let next_barrier = Barrier::new_test_barrier(test_epoch(6));
1211 barrier_test_env.inject_barrier(&next_barrier, [actor_id]);
1212 barrier_test_env.flush_all_events().await;
1213 upstream_tx
1214 .send(DispatcherMessage::Barrier(next_barrier.clone().into_dispatcher()).into())
1215 .await
1216 .unwrap();
1217 assert_eq!(
1218 expect_barrier_with_timeout(&mut executor, "next upstream barrier")
1219 .await
1220 .epoch,
1221 next_barrier.epoch
1222 );
1223 }
1224
1225 #[tokio::test]
1226 async fn test_snapshot_backfill_restarted_with_upstream_on_hummock() {
1227 let source_env = prepare_hummock_test_env().await;
1228 source_env.register_table(source_table_pb()).await;
1229 let progress_env = prepare_hummock_test_env().await;
1230 progress_env.register_table(progress_table_pb()).await;
1231
1232 let mut source_state_table = source_state_table(source_env.storage.clone()).await;
1233 let source_table = source_batch_table(source_env.storage.clone());
1234
1235 let mut epoch = EpochPair::new_test_epoch(test_epoch(1));
1236 source_env
1237 .storage
1238 .start_epoch(epoch.curr, HashSet::from_iter([SOURCE_TABLE_ID]));
1239 source_env.storage.flush_events_for_test().await;
1240 source_state_table.init_epoch(epoch).await.unwrap();
1241
1242 commit_insert_epoch(
1243 &source_env,
1244 &mut source_state_table,
1245 &mut epoch,
1246 HashSet::from_iter([SOURCE_TABLE_ID]),
1247 &[],
1248 )
1249 .await;
1250 commit_insert_epoch(
1251 &source_env,
1252 &mut source_state_table,
1253 &mut epoch,
1254 HashSet::from_iter([SOURCE_TABLE_ID]),
1255 &[],
1256 )
1257 .await;
1258 commit_insert_epoch(
1259 &source_env,
1260 &mut source_state_table,
1261 &mut epoch,
1262 HashSet::from_iter([SOURCE_TABLE_ID]),
1263 &[],
1264 )
1265 .await;
1266 commit_insert_epoch(
1267 &source_env,
1268 &mut source_state_table,
1269 &mut epoch,
1270 HashSet::from_iter([SOURCE_TABLE_ID]),
1271 &[4],
1272 )
1273 .await;
1274 start_progress_epochs(&progress_env, 7).await;
1275 let initial_barrier = Barrier::new_test_barrier(test_epoch(6));
1276 let progress_state_table = persist_finished_progress(
1277 &progress_env,
1278 EpochPair::new_test_epoch(test_epoch(1)),
1279 initial_barrier.epoch.prev,
1280 )
1281 .await;
1282
1283 let barrier_test_env = LocalBarrierTestEnv::for_test().await;
1284 let barrier_manager = barrier_test_env.local_barrier_manager.clone();
1285 let progress = CreateMviewProgressReporter::for_test(barrier_manager.clone());
1286 let actor_id = 1235.into();
1287 let actor_ctx = ActorContext::for_test(actor_id);
1288 let (upstream_tx, upstream_rx) = channel_for_test();
1289
1290 upstream_tx
1291 .send(DispatcherMessage::Barrier(initial_barrier.clone().into_dispatcher()).into())
1292 .await
1293 .unwrap();
1294 barrier_test_env.inject_barrier(&initial_barrier, [actor_id]);
1295 barrier_test_env.flush_all_events().await;
1296 let upstream = make_upstream_input(
1297 barrier_manager.clone(),
1298 actor_ctx.clone(),
1299 Some(upstream_rx),
1300 );
1301
1302 let mut executor = SnapshotBackfillExecutor::new(
1303 source_table,
1304 progress_state_table,
1305 upstream,
1306 None,
1307 vec![0],
1308 vec![0],
1309 actor_ctx,
1310 progress,
1311 1024,
1312 RateLimit::Disabled,
1313 Arc::new(StreamingMetrics::unused()),
1314 Some(test_epoch(3)),
1315 )
1316 .expect("snapshot backfill executor should be created")
1317 .boxed()
1318 .execute();
1319
1320 assert_eq!(
1321 expect_barrier_with_timeout(&mut executor, "initial injected barrier")
1322 .await
1323 .epoch,
1324 initial_barrier.epoch
1325 );
1326
1327 upstream_tx
1328 .send(DispatcherMessage::Chunk(StreamChunk::from_pretty(" I\n + 5")).into())
1329 .await
1330 .unwrap();
1331 let next_barrier = Barrier::new_test_barrier(test_epoch(7));
1332 upstream_tx
1333 .send(DispatcherMessage::Barrier(next_barrier.clone().into_dispatcher()).into())
1334 .await
1335 .unwrap();
1336 barrier_test_env.inject_barrier(&next_barrier, [actor_id]);
1337 barrier_test_env.flush_all_events().await;
1338
1339 assert_eq!(
1340 expect_chunk_with_timeout(&mut executor, "live upstream chunk after handoff").await,
1341 StreamChunk::from_pretty(" I\n + 5")
1342 );
1343
1344 assert_eq!(
1345 expect_barrier_with_timeout(&mut executor, "next upstream barrier")
1346 .await
1347 .epoch,
1348 next_barrier.epoch
1349 );
1350 }
1351}