1mod upstream_table_ext {
16 use std::collections::HashMap;
17 use std::time::Duration;
18
19 use futures::future::{BoxFuture, try_join_all};
20 use futures::{TryFutureExt, TryStreamExt};
21 use risingwave_common::hash::VirtualNode;
22 use risingwave_common::row::OwnedRow;
23 use risingwave_common_rate_limit::RateLimit;
24 use risingwave_storage::table::ChangeLogRow;
25
26 use crate::executor::StreamExecutorResult;
27 use crate::executor::backfill::snapshot_backfill::consume_upstream::upstream_table_trait::UpstreamTable;
28 use crate::executor::backfill::snapshot_backfill::vnode_stream::{
29 ChangeLogRowStream, VnodeStream,
30 };
31 use crate::executor::backfill::utils::create_builder;
32
33 pub(super) type UpstreamTableSnapshotStream<T: UpstreamTable> =
34 VnodeStream<impl ChangeLogRowStream>;
35 pub(super) type UpstreamTableSnapshotStreamFuture<'a, T> =
36 BoxFuture<'a, StreamExecutorResult<UpstreamTableSnapshotStream<T>>>;
37
38 #[define_opaque(UpstreamTableSnapshotStream)]
39 pub(super) fn create_upstream_table_snapshot_stream<T: UpstreamTable>(
40 upstream_table: &T,
41 snapshot_epoch: u64,
42 rate_limit: RateLimit,
43 chunk_size: usize,
44 rebuild_interval: Duration,
45 vnode_progresses: HashMap<VirtualNode, (Option<OwnedRow>, usize)>,
46 ) -> UpstreamTableSnapshotStreamFuture<'_, T> {
47 Box::pin(async move {
48 let streams = try_join_all(vnode_progresses.into_iter().map(
49 |(vnode, (start_pk, row_count))| {
50 upstream_table
51 .snapshot_stream(vnode, snapshot_epoch, start_pk, rebuild_interval)
52 .map_ok(move |stream| {
53 (vnode, stream.map_ok(ChangeLogRow::Insert), row_count)
54 })
55 },
56 ))
57 .await?;
58 Ok(VnodeStream::new(
59 streams,
60 upstream_table.pk_in_output_indices(),
61 create_builder(rate_limit, chunk_size, upstream_table.output_data_types()),
62 ))
63 })
64 }
65
66 pub(super) type UpstreamTableChangeLogStream<T: UpstreamTable> =
67 VnodeStream<impl ChangeLogRowStream>;
68 pub(super) type UpstreamTableChangeLogStreamFuture<'a, T> =
69 BoxFuture<'a, StreamExecutorResult<UpstreamTableChangeLogStream<T>>>;
70
71 #[define_opaque(UpstreamTableChangeLogStreamFuture)]
72 pub(super) fn create_upstream_table_change_log_stream<T: UpstreamTable>(
73 upstream_table: &T,
74 epoch: u64,
75 rate_limit: RateLimit,
76 chunk_size: usize,
77 vnode_progresses: HashMap<VirtualNode, (Option<OwnedRow>, usize)>,
78 ) -> UpstreamTableChangeLogStreamFuture<'_, T> {
79 Box::pin(async move {
80 let streams = try_join_all(vnode_progresses.into_iter().map(
81 |(vnode, (start_pk, row_count))| {
82 upstream_table
83 .change_log_stream(vnode, epoch, start_pk)
84 .map_ok(move |stream| (vnode, stream, row_count))
85 },
86 ))
87 .await?;
88 Ok(VnodeStream::new(
89 streams,
90 upstream_table.pk_in_output_indices(),
91 create_builder(rate_limit, chunk_size, upstream_table.output_data_types()),
92 ))
93 })
94 }
95
96 pub(super) type NextEpochFuture<'a> = BoxFuture<'a, StreamExecutorResult<u64>>;
97 pub(super) fn next_epoch_future<T: UpstreamTable>(
98 upstream_table: &T,
99 epoch: u64,
100 ) -> NextEpochFuture<'_> {
101 Box::pin(async move { upstream_table.next_epoch(epoch).await })
102 }
103}
104
105use std::collections::{BTreeMap, HashMap};
106use std::mem::take;
107use std::pin::Pin;
108use std::task::{Context, Poll, ready};
109use std::time::Duration;
110
111use risingwave_common::array::StreamChunk;
112use risingwave_common::hash::VirtualNode;
113use risingwave_common::row::OwnedRow;
114use risingwave_common_rate_limit::RateLimit;
115use risingwave_storage::error::ErrorKind as StorageErrorKind;
116use risingwave_storage::hummock::HummockErrorInner;
117use thiserror_ext::AsReport;
118use upstream_table_ext::*;
119
120use crate::executor::backfill::snapshot_backfill::consume_upstream::upstream_table_trait::UpstreamTable;
121use crate::executor::backfill::snapshot_backfill::state::{
122 EpochBackfillProgress, VnodeBackfillProgress,
123};
124use crate::executor::error::ErrorKind as StreamExecutorErrorKind;
125use crate::executor::prelude::{Stream, StreamExt};
126use crate::executor::{StreamExecutorError, StreamExecutorResult};
127
128enum ConsumeUpstreamStreamState<'a, T: UpstreamTable> {
129 CreatingSnapshotStream {
130 future: UpstreamTableSnapshotStreamFuture<'a, T>,
131 snapshot_epoch: u64,
132 pre_finished_vnodes: HashMap<VirtualNode, usize>,
133 },
134 ConsumingSnapshotStream {
135 stream: UpstreamTableSnapshotStream<T>,
136 snapshot_epoch: u64,
137 pre_finished_vnodes: HashMap<VirtualNode, usize>,
138 },
139 CreatingChangeLogStream {
140 future: UpstreamTableChangeLogStreamFuture<'a, T>,
141 prev_epoch_finished_vnodes: Option<(u64, HashMap<VirtualNode, usize>)>,
142 epoch: u64,
143 pre_finished_vnodes: HashMap<VirtualNode, usize>,
144 },
145 ConsumingChangeLogStream {
146 stream: UpstreamTableChangeLogStream<T>,
147 epoch: u64,
148 pre_finished_vnodes: HashMap<VirtualNode, usize>,
149 },
150 ResolvingNextEpoch {
151 future: NextEpochFuture<'a>,
152 prev_epoch_finished_vnodes: Option<(u64, HashMap<VirtualNode, usize>)>,
153 },
154 StoppedOnRetentionMiss,
155 Err,
156}
157
158fn is_retention_or_snapshot_expired_error(error: &StreamExecutorError) -> bool {
159 let StreamExecutorErrorKind::Storage(storage_error) = error.inner() else {
160 return false;
161 };
162 let StorageErrorKind::Hummock(hummock_error) = storage_error.inner() else {
163 return false;
164 };
165
166 matches!(
167 hummock_error.inner(),
168 HummockErrorInner::ChangeLogRetentionMiss { .. }
169 | HummockErrorInner::TimeTravelVersionExpired { .. }
170 | HummockErrorInner::CommittedEpochMismatch { .. }
171 )
172}
173
174fn log_retention_or_snapshot_expired(error: &StreamExecutorError) {
175 tracing::warn!(
176 error = %error.as_report(),
177 "stop consuming cross-database upstream because upstream retention or snapshot availability was exceeded"
178 );
179}
180
181pub(super) struct ConsumeUpstreamStream<'a, T: UpstreamTable> {
182 upstream_table: &'a T,
183 pending_epoch_vnode_progress:
184 BTreeMap<u64, HashMap<VirtualNode, (EpochBackfillProgress, usize)>>,
185 state: ConsumeUpstreamStreamState<'a, T>,
186
187 chunk_size: usize,
188 rate_limit: RateLimit,
189}
190
191impl<T: UpstreamTable> ConsumeUpstreamStream<'_, T> {
192 pub(super) fn update_rate_limiter(&mut self, new_rate_limit: RateLimit) {
193 self.rate_limit = new_rate_limit;
194 let chunk_size = self.chunk_size;
195 match &mut self.state {
196 ConsumeUpstreamStreamState::ConsumingSnapshotStream { stream, .. } => {
197 stream.update_rate_limiter(new_rate_limit, chunk_size);
198 }
199 ConsumeUpstreamStreamState::ConsumingChangeLogStream { stream, .. } => {
200 stream.update_rate_limiter(new_rate_limit, chunk_size);
201 }
202 ConsumeUpstreamStreamState::ResolvingNextEpoch { .. }
203 | ConsumeUpstreamStreamState::CreatingChangeLogStream { .. }
204 | ConsumeUpstreamStreamState::CreatingSnapshotStream { .. } => {}
205 ConsumeUpstreamStreamState::StoppedOnRetentionMiss => {}
206 ConsumeUpstreamStreamState::Err => {
207 unreachable!("should not be accessed on Err")
208 }
209 }
210 }
211
212 pub(super) fn consume_builder(&mut self) -> Option<StreamChunk> {
213 match &mut self.state {
214 ConsumeUpstreamStreamState::ConsumingSnapshotStream { stream, .. } => {
215 stream.consume_builder()
216 }
217 ConsumeUpstreamStreamState::ConsumingChangeLogStream { stream, .. } => {
218 stream.consume_builder()
219 }
220 ConsumeUpstreamStreamState::ResolvingNextEpoch { .. }
221 | ConsumeUpstreamStreamState::CreatingChangeLogStream { .. }
222 | ConsumeUpstreamStreamState::CreatingSnapshotStream { .. } => None,
223 ConsumeUpstreamStreamState::StoppedOnRetentionMiss => None,
224 ConsumeUpstreamStreamState::Err => {
225 unreachable!("should not be accessed on Err")
226 }
227 }
228 }
229
230 pub(super) async fn for_vnode_pk_progress(
231 &mut self,
232 mut on_vnode_progress: impl FnMut(VirtualNode, u64, usize, Option<OwnedRow>),
233 ) -> StreamExecutorResult<()> {
234 match &mut self.state {
235 ConsumeUpstreamStreamState::CreatingSnapshotStream { .. } => {
236 }
238 ConsumeUpstreamStreamState::ConsumingSnapshotStream {
239 stream,
240 snapshot_epoch,
241 ..
242 } => {
243 stream
244 .for_vnode_pk_progress(|vnode, row_count, pk_progress| {
245 on_vnode_progress(vnode, *snapshot_epoch, row_count, pk_progress)
246 })
247 .await?;
248 }
249 &mut ConsumeUpstreamStreamState::ConsumingChangeLogStream {
250 ref mut stream,
251 ref epoch,
252 ..
253 } => {
254 stream
255 .for_vnode_pk_progress(|vnode, row_count, pk_progress| {
256 on_vnode_progress(vnode, *epoch, row_count, pk_progress)
257 })
258 .await?;
259 }
260 &mut ConsumeUpstreamStreamState::CreatingChangeLogStream {
261 ref prev_epoch_finished_vnodes,
262 ..
263 }
264 | &mut ConsumeUpstreamStreamState::ResolvingNextEpoch {
265 ref prev_epoch_finished_vnodes,
266 ..
267 } => {
268 if let Some((prev_epoch, prev_epoch_finished_vnodes)) = prev_epoch_finished_vnodes {
269 for (vnode, row_count) in prev_epoch_finished_vnodes {
270 on_vnode_progress(*vnode, *prev_epoch, *row_count, None);
271 }
272 }
273 }
274 ConsumeUpstreamStreamState::StoppedOnRetentionMiss => {}
275 ConsumeUpstreamStreamState::Err => {
276 unreachable!("should not be accessed on Err")
277 }
278 }
279 Ok(())
280 }
281}
282
283impl<T: UpstreamTable> Stream for ConsumeUpstreamStream<'_, T> {
284 type Item = StreamExecutorResult<StreamChunk>;
285
286 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
287 let result: Result<!, StreamExecutorError> = try {
288 loop {
289 let rate_limit = self.rate_limit;
290 let chunk_size = self.chunk_size;
291 match &mut self.state {
292 ConsumeUpstreamStreamState::CreatingSnapshotStream {
293 future,
294 snapshot_epoch,
295 pre_finished_vnodes,
296 } => {
297 let mut stream = ready!(future.as_mut().poll(cx))?;
298 stream.update_rate_limiter(rate_limit, chunk_size);
299 let snapshot_epoch = *snapshot_epoch;
300 let pre_finished_vnodes = take(pre_finished_vnodes);
301 self.state = ConsumeUpstreamStreamState::ConsumingSnapshotStream {
302 stream,
303 snapshot_epoch,
304 pre_finished_vnodes,
305 };
306 continue;
307 }
308 ConsumeUpstreamStreamState::ConsumingSnapshotStream {
309 stream,
310 snapshot_epoch,
311 pre_finished_vnodes,
312 } => match ready!(stream.poll_next_unpin(cx)).transpose()? {
313 None => {
314 let prev_epoch = *snapshot_epoch;
315 let mut prev_epoch_finished_vnodes = take(pre_finished_vnodes);
316 for (vnode, row_count) in stream.take_finished_vnodes() {
317 prev_epoch_finished_vnodes
318 .try_insert(vnode, row_count)
319 .expect("non-duplicate");
320 }
321 self.state = ConsumeUpstreamStreamState::ResolvingNextEpoch {
322 future: next_epoch_future(self.upstream_table, prev_epoch),
323 prev_epoch_finished_vnodes: Some((
324 prev_epoch,
325 prev_epoch_finished_vnodes,
326 )),
327 };
328 continue;
329 }
330 Some(chunk) => {
331 return Poll::Ready(Some(Ok(chunk)));
332 }
333 },
334 ConsumeUpstreamStreamState::CreatingChangeLogStream {
335 future,
336 epoch,
337 pre_finished_vnodes,
338 ..
339 } => {
340 let mut stream = ready!(future.as_mut().poll(cx))?;
341 stream.update_rate_limiter(rate_limit, chunk_size);
342 let epoch = *epoch;
343 let pre_finished_vnodes = take(pre_finished_vnodes);
344 self.state = ConsumeUpstreamStreamState::ConsumingChangeLogStream {
345 stream,
346 epoch,
347 pre_finished_vnodes,
348 };
349 continue;
350 }
351 ConsumeUpstreamStreamState::ConsumingChangeLogStream {
352 stream,
353 epoch,
354 pre_finished_vnodes,
355 } => {
356 match ready!(stream.poll_next_unpin(cx)).transpose()? {
357 None => {
358 let prev_epoch = *epoch;
359 let mut prev_epoch_finished_vnodes = take(pre_finished_vnodes);
360 for (vnode, row_count) in stream.take_finished_vnodes() {
361 prev_epoch_finished_vnodes
362 .try_insert(vnode, row_count)
363 .expect("non-duplicate");
364 }
365 self.state = ConsumeUpstreamStreamState::ResolvingNextEpoch {
366 future: next_epoch_future(self.upstream_table, prev_epoch),
367 prev_epoch_finished_vnodes: Some((
368 prev_epoch,
369 prev_epoch_finished_vnodes,
370 )),
371 };
372 continue;
373 }
374 Some(chunk) => {
375 return Poll::Ready(Some(Ok(chunk)));
376 }
377 };
378 }
379 ConsumeUpstreamStreamState::ResolvingNextEpoch {
380 future,
381 prev_epoch_finished_vnodes,
382 } => {
383 let epoch = ready!(future.as_mut().poll(cx))?;
384 let prev_epoch_finished_vnodes = take(prev_epoch_finished_vnodes);
385 let mut pre_finished_vnodes = HashMap::new();
386 let mut vnode_progresses = HashMap::new();
387 for prev_epoch_vnode in prev_epoch_finished_vnodes
388 .as_ref()
389 .map(|(_, vnodes)| vnodes.keys())
390 .into_iter()
391 .flatten()
392 {
393 vnode_progresses
394 .try_insert(*prev_epoch_vnode, (None, 0))
395 .expect("non-duplicate");
396 }
397 if let Some((pending_epoch, _)) =
398 self.pending_epoch_vnode_progress.first_key_value()
399 {
400 assert!(
402 epoch <= *pending_epoch,
403 "pending_epoch {} earlier than next epoch {}",
404 pending_epoch,
405 epoch
406 );
407 if epoch == *pending_epoch {
408 let (_, progress) = self
409 .pending_epoch_vnode_progress
410 .pop_first()
411 .expect("checked Some");
412 for (vnode, (progress, row_count)) in progress {
413 match progress {
414 EpochBackfillProgress::Consuming { latest_pk } => {
415 vnode_progresses
416 .try_insert(vnode, (Some(latest_pk), row_count))
417 .expect("non-duplicate");
418 }
419 EpochBackfillProgress::Consumed => {
420 pre_finished_vnodes
421 .try_insert(vnode, row_count)
422 .expect("non-duplicate");
423 }
424 }
425 }
426 }
427 }
428 self.state = ConsumeUpstreamStreamState::CreatingChangeLogStream {
429 future: create_upstream_table_change_log_stream(
430 self.upstream_table,
431 epoch,
432 self.rate_limit,
433 self.chunk_size,
434 vnode_progresses,
435 ),
436 prev_epoch_finished_vnodes,
437 epoch,
438 pre_finished_vnodes,
439 };
440 continue;
441 }
442 ConsumeUpstreamStreamState::StoppedOnRetentionMiss => {
443 return Poll::Pending;
444 }
445 ConsumeUpstreamStreamState::Err => {
446 unreachable!("should not be accessed on Err")
447 }
448 }
449 }
450 };
451 match result {
452 Err(error) if is_retention_or_snapshot_expired_error(&error) => {
453 log_retention_or_snapshot_expired(&error);
454 self.state = ConsumeUpstreamStreamState::StoppedOnRetentionMiss;
455 Poll::Pending
456 }
457 Err(error) => {
458 self.state = ConsumeUpstreamStreamState::Err;
459 Poll::Ready(Some(Err(error)))
460 }
461 Ok(unreachable) => match unreachable {},
462 }
463 }
464}
465
466impl<'a, T: UpstreamTable> ConsumeUpstreamStream<'a, T> {
467 pub(super) fn new<'p>(
468 initial_progress: impl Iterator<Item = (VirtualNode, Option<&'p VnodeBackfillProgress>)>,
469 upstream_table: &'a T,
470 snapshot_epoch: u64,
471 chunk_size: usize,
472 rate_limit: RateLimit,
473 snapshot_rebuild_interval: Duration,
474 ) -> Self {
475 let mut ongoing_snapshot_epoch_vnodes = HashMap::new();
476 let mut finished_snapshot_epoch_vnodes = HashMap::new();
477 let mut pending_epoch_vnode_progress: BTreeMap<_, HashMap<_, _>> = BTreeMap::new();
478 for (vnode, progress) in initial_progress {
479 match progress {
480 None => {
481 ongoing_snapshot_epoch_vnodes
482 .try_insert(vnode, (None, 0))
483 .expect("non-duplicate");
484 }
485 Some(progress) => {
486 let epoch = progress.epoch;
487 let row_count = progress.row_count;
488 if epoch == snapshot_epoch {
489 match &progress.progress {
490 EpochBackfillProgress::Consumed => {
491 finished_snapshot_epoch_vnodes
492 .try_insert(vnode, row_count)
493 .expect("non-duplicate");
494 }
495 EpochBackfillProgress::Consuming { latest_pk } => {
496 ongoing_snapshot_epoch_vnodes
497 .try_insert(vnode, (Some(latest_pk.clone()), row_count))
498 .expect("non-duplicate");
499 }
500 }
501 } else {
502 assert!(
503 epoch > snapshot_epoch,
504 "epoch {} earlier than snapshot_epoch {} on vnode {}",
505 epoch,
506 snapshot_epoch,
507 vnode
508 );
509 pending_epoch_vnode_progress
510 .entry(epoch)
511 .or_default()
512 .try_insert(vnode, (progress.progress.clone(), progress.row_count))
513 .expect("non-duplicate");
514 }
515 }
516 };
517 }
518 let (pending_epoch_vnode_progress, state) = {
519 if !ongoing_snapshot_epoch_vnodes.is_empty() {
520 (
522 pending_epoch_vnode_progress,
523 ConsumeUpstreamStreamState::CreatingSnapshotStream {
524 future: create_upstream_table_snapshot_stream(
525 upstream_table,
526 snapshot_epoch,
527 rate_limit,
528 chunk_size,
529 snapshot_rebuild_interval,
530 ongoing_snapshot_epoch_vnodes,
531 ),
532 snapshot_epoch,
533 pre_finished_vnodes: finished_snapshot_epoch_vnodes,
534 },
535 )
536 } else if !finished_snapshot_epoch_vnodes.is_empty() {
537 (
539 pending_epoch_vnode_progress,
540 ConsumeUpstreamStreamState::ResolvingNextEpoch {
541 future: next_epoch_future(upstream_table, snapshot_epoch),
542 prev_epoch_finished_vnodes: Some((
543 snapshot_epoch,
544 finished_snapshot_epoch_vnodes,
545 )),
546 },
547 )
548 } else {
549 let (first_epoch, first_vnodes) = pending_epoch_vnode_progress
551 .pop_first()
552 .expect("non-empty vnodes");
553 let mut ongoing_vnodes = HashMap::new();
554 let mut finished_vnodes = HashMap::new();
555 for (vnode, (progress, row_count)) in first_vnodes {
556 match progress {
557 EpochBackfillProgress::Consuming { latest_pk } => {
558 ongoing_vnodes
559 .try_insert(vnode, (Some(latest_pk), row_count))
560 .expect("non-duplicate");
561 }
562 EpochBackfillProgress::Consumed => {
563 finished_vnodes
564 .try_insert(vnode, row_count)
565 .expect("non-duplicate");
566 }
567 }
568 }
569 let state = if ongoing_vnodes.is_empty() {
570 ConsumeUpstreamStreamState::ResolvingNextEpoch {
572 future: next_epoch_future(upstream_table, first_epoch),
573 prev_epoch_finished_vnodes: Some((first_epoch, finished_vnodes)),
574 }
575 } else {
576 ConsumeUpstreamStreamState::CreatingChangeLogStream {
577 future: create_upstream_table_change_log_stream(
578 upstream_table,
579 first_epoch,
580 rate_limit,
581 chunk_size,
582 ongoing_vnodes,
583 ),
584 prev_epoch_finished_vnodes: None,
585 epoch: first_epoch,
586 pre_finished_vnodes: finished_vnodes,
587 }
588 };
589 (pending_epoch_vnode_progress, state)
590 }
591 };
592 Self {
593 upstream_table,
594 pending_epoch_vnode_progress,
595 state,
596 chunk_size,
597 rate_limit,
598 }
599 }
600}
601
602#[cfg(test)]
603mod tests {
604 use risingwave_common::catalog::TableId;
605 use risingwave_storage::hummock::HummockError;
606
607 use super::*;
608
609 fn stream_error_from_hummock(error: HummockError) -> StreamExecutorError {
610 let storage_error: risingwave_storage::error::StorageError = error.into();
611 storage_error.into()
612 }
613
614 #[test]
615 fn test_change_log_retention_miss_error() {
616 let retention_miss = stream_error_from_hummock(HummockError::change_log_retention_miss(
617 TableId::new(233),
618 10678350547714048,
619 ));
620 assert!(is_retention_or_snapshot_expired_error(&retention_miss));
621
622 let missing_time_travel_version = stream_error_from_hummock(
623 HummockError::time_travel_version_expired(TableId::new(233), 10678350547714048),
624 );
625 assert!(is_retention_or_snapshot_expired_error(
626 &missing_time_travel_version
627 ));
628
629 let committed_epoch_mismatch =
630 stream_error_from_hummock(HummockError::committed_epoch_mismatch(
631 TableId::new(233),
632 10682740646084608,
633 10678350547714048,
634 ));
635 assert!(is_retention_or_snapshot_expired_error(
636 &committed_epoch_mismatch
637 ));
638
639 let dropped_table =
640 stream_error_from_hummock(HummockError::next_epoch("table 233 has been dropped"));
641 assert!(!is_retention_or_snapshot_expired_error(&dropped_table));
642
643 let other_error = StreamExecutorError::from(anyhow::anyhow!("other error"));
644 assert!(!is_retention_or_snapshot_expired_error(&other_error));
645 }
646}