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