risingwave_stream/executor/backfill/snapshot_backfill/consume_upstream/
executor.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::future::ready;
16
17use anyhow::anyhow;
18use futures::future::{Either, select};
19use futures::{FutureExt, TryStreamExt};
20use futures_async_stream::try_stream;
21use risingwave_common::catalog::TableId;
22use risingwave_common::metrics::LabelGuardedIntGauge;
23use risingwave_common_rate_limit::{MonitoredRateLimiter, RateLimit, RateLimiter};
24use risingwave_pb::common::ThrottleType;
25use risingwave_storage::StateStore;
26use rw_futures_util::drop_either_future;
27use tokio::sync::mpsc::UnboundedReceiver;
28
29use crate::executor::backfill::snapshot_backfill::consume_upstream::stream::ConsumeUpstreamStream;
30use crate::executor::backfill::snapshot_backfill::consume_upstream::upstream_table_trait::UpstreamTable;
31use crate::executor::backfill::snapshot_backfill::receive_next_barrier;
32use crate::executor::backfill::snapshot_backfill::state::{BackfillState, EpochBackfillProgress};
33use crate::executor::backfill::utils::mapping_message;
34use crate::executor::monitor::BackfillMetrics;
35use crate::executor::prelude::{StateTable, *};
36use crate::executor::{Barrier, Message, Mutation, StreamExecutorError};
37use crate::task::CreateMviewProgressReporter;
38
39pub struct UpstreamTableExecutor<T: UpstreamTable, S: StateStore> {
40    upstream_table_id: TableId,
41    upstream_table: T,
42    progress_state_table: StateTable<S>,
43    snapshot_epoch: u64,
44    output_indices: Vec<usize>,
45
46    chunk_size: usize,
47    rate_limiter: MonitoredRateLimiter,
48    actor_ctx: ActorContextRef,
49    barrier_rx: UnboundedReceiver<Barrier>,
50    progress: CreateMviewProgressReporter,
51    crossdb_last_consumed_min_epoch: LabelGuardedIntGauge,
52    metrics: BackfillMetrics,
53}
54
55impl<T: UpstreamTable, S: StateStore> UpstreamTableExecutor<T, S> {
56    #[expect(clippy::too_many_arguments)]
57    pub fn new(
58        upstream_table_id: TableId,
59        upstream_table: T,
60        progress_state_table: StateTable<S>,
61        snapshot_epoch: u64,
62        output_indices: Vec<usize>,
63
64        chunk_size: usize,
65        rate_limit: RateLimit,
66        actor_ctx: ActorContextRef,
67        barrier_rx: UnboundedReceiver<Barrier>,
68        progress: CreateMviewProgressReporter,
69    ) -> Self {
70        let rate_limiter = RateLimiter::new(rate_limit).monitored(upstream_table_id);
71        let table_id_label = upstream_table_id.to_string();
72        let actor_id_label = actor_ctx.id.to_string();
73        let fragment_id_label = actor_ctx.fragment_id.to_string();
74        let crossdb_last_consumed_min_epoch = actor_ctx
75            .streaming_metrics
76            .crossdb_last_consumed_min_epoch
77            .with_guarded_label_values(&[
78                table_id_label.as_str(),
79                actor_id_label.as_str(),
80                fragment_id_label.as_str(),
81            ]);
82        let metrics = actor_ctx
83            .streaming_metrics
84            .new_backfill_metrics(upstream_table_id, actor_ctx.id);
85        Self {
86            upstream_table_id,
87            upstream_table,
88            progress_state_table,
89            snapshot_epoch,
90            output_indices,
91            chunk_size,
92            rate_limiter,
93            actor_ctx,
94            barrier_rx,
95            progress,
96            crossdb_last_consumed_min_epoch,
97            metrics,
98        }
99    }
100
101    fn extract_last_consumed_min_epoch(progress_state: &BackfillState<S>) -> u64 {
102        let mut min_epoch = u64::MAX;
103        for (_, progress) in progress_state.latest_progress() {
104            let Some(progress) = progress else {
105                // If any vnode has no progress yet, report `0` explicitly to indicate the
106                // progress is not fully ready, instead of hiding this state.
107                return 0;
108            };
109            min_epoch = min_epoch.min(progress.epoch);
110        }
111        if min_epoch == u64::MAX { 0 } else { min_epoch }
112    }
113
114    #[try_stream(ok = Message, error = StreamExecutorError)]
115    pub async fn into_stream(mut self) {
116        self.upstream_table
117            .check_initial_vnode_bitmap(self.progress_state_table.vnodes())?;
118        let first_barrier = receive_next_barrier(&mut self.barrier_rx).await?;
119        let first_barrier_epoch = first_barrier.epoch;
120        yield Message::Barrier(first_barrier);
121        let mut progress_state = BackfillState::new(
122            self.progress_state_table,
123            first_barrier_epoch,
124            self.upstream_table.pk_serde(),
125        )
126        .await?;
127        let mut finish_reported = false;
128        let mut prev_reported_row_count = 0;
129        let mut upstream_table = self.upstream_table;
130        let snapshot_rebuild_interval = self
131            .actor_ctx
132            .config
133            .developer
134            .snapshot_iter_rebuild_interval();
135        let mut stream = ConsumeUpstreamStream::new(
136            progress_state.latest_progress(),
137            &upstream_table,
138            self.snapshot_epoch,
139            self.chunk_size,
140            self.rate_limiter.rate_limit(),
141            snapshot_rebuild_interval,
142        );
143
144        'on_new_stream: loop {
145            loop {
146                let barrier = {
147                    loop {
148                        if self.rate_limiter.rate_limit().is_paused() {
149                            break receive_next_barrier(&mut self.barrier_rx).await?;
150                        }
151                        let future1 = receive_next_barrier(&mut self.barrier_rx);
152                        let future2 = stream.try_next().map(|result| {
153                            result
154                                .and_then(|opt| opt.ok_or_else(|| anyhow!("end of stream").into()))
155                        });
156                        pin_mut!(future1);
157                        pin_mut!(future2);
158                        match drop_either_future(select(future1, future2).await) {
159                            Either::Left(Ok(barrier)) => {
160                                break barrier;
161                            }
162                            Either::Right(Ok(chunk)) => {
163                                assert!(!self.rate_limiter.rate_limit().is_paused());
164                                self.rate_limiter.wait(chunk.cardinality() as _).await;
165                                yield Message::Chunk(chunk);
166                            }
167                            Either::Left(Err(e)) | Either::Right(Err(e)) => {
168                                return Err(e);
169                            }
170                        }
171                    }
172                };
173
174                if let Some(chunk) = stream.consume_builder() {
175                    self.rate_limiter.wait(chunk.cardinality() as _).await;
176                    yield Message::Chunk(chunk);
177                }
178                stream
179                    .for_vnode_pk_progress(|vnode, epoch, row_count, progress| {
180                        if let Some(progress) = progress {
181                            progress_state.update_epoch_progress(vnode, epoch, row_count, progress);
182                        } else {
183                            progress_state.finish_epoch(vnode, epoch, row_count);
184                        }
185                    })
186                    .await?;
187
188                let last_consumed_min_epoch =
189                    Self::extract_last_consumed_min_epoch(&progress_state);
190                self.crossdb_last_consumed_min_epoch
191                    .set(last_consumed_min_epoch as i64);
192
193                if !finish_reported {
194                    let mut row_count = 0;
195                    let mut is_finished = true;
196                    for (_, progress) in progress_state.latest_progress() {
197                        if let Some(progress) = progress {
198                            if progress.epoch == self.snapshot_epoch {
199                                if let EpochBackfillProgress::Consuming { .. } = &progress.progress
200                                {
201                                    is_finished = false;
202                                }
203                                row_count += progress.row_count;
204                            }
205                        } else {
206                            is_finished = false;
207                        }
208                    }
209                    // ensure that the reported row count is non-decreasing.
210                    let row_count_to_report = std::cmp::max(prev_reported_row_count, row_count);
211                    self.metrics
212                        .backfill_snapshot_read_row_count
213                        .inc_by(row_count_to_report.saturating_sub(prev_reported_row_count) as _);
214                    prev_reported_row_count = row_count_to_report;
215
216                    if is_finished {
217                        self.progress
218                            .finish(barrier.epoch, row_count_to_report as _);
219                        finish_reported = true;
220                    } else {
221                        self.progress.update(
222                            barrier.epoch,
223                            self.snapshot_epoch,
224                            row_count_to_report as _,
225                        );
226                    }
227                }
228
229                let post_commit = progress_state.commit(barrier.epoch).await?;
230                let update_vnode_bitmap = barrier.as_update_vnode_bitmap(self.actor_ctx.id);
231                if let Some(new_rate_limit) = barrier.mutation.as_ref().and_then(|mutation| {
232                    if let Mutation::Throttle(config) = &**mutation
233                        && let Some(config) = config.get(&self.actor_ctx.fragment_id)
234                        && config.throttle_type() == ThrottleType::Backfill
235                    {
236                        Some(config.rate_limit)
237                    } else {
238                        None
239                    }
240                }) {
241                    let new_rate_limit = new_rate_limit.into();
242                    let old_rate_limit = self.rate_limiter.update(new_rate_limit);
243                    if old_rate_limit != new_rate_limit {
244                        stream.update_rate_limiter(new_rate_limit);
245                        tracing::info!(
246                            old_rate_limit = ?old_rate_limit,
247                            new_rate_limit = ?new_rate_limit,
248                            upstream_table_id = %self.upstream_table_id,
249                            actor_id = %self.actor_ctx.id,
250                            "cross-db backfill rate limit changed",
251                        );
252                    }
253                }
254                yield Message::Barrier(barrier);
255                if let Some(new_vnode_bitmap) =
256                    post_commit.post_yield_barrier(update_vnode_bitmap).await?
257                {
258                    drop(stream);
259                    upstream_table.update_vnode_bitmap(new_vnode_bitmap);
260                    // recreate the stream on update vnode bitmap
261                    stream = ConsumeUpstreamStream::new(
262                        progress_state.latest_progress(),
263                        &upstream_table,
264                        self.snapshot_epoch,
265                        self.chunk_size,
266                        self.rate_limiter.rate_limit(),
267                        snapshot_rebuild_interval,
268                    );
269                    continue 'on_new_stream;
270                }
271            }
272        }
273    }
274}
275
276impl<T: UpstreamTable, S: StateStore> Execute for UpstreamTableExecutor<T, S> {
277    fn execute(self: Box<Self>) -> BoxedMessageStream {
278        let output_indices = self.output_indices.clone();
279        self.into_stream()
280            .filter_map(move |result| {
281                ready({
282                    match result {
283                        Ok(message) => mapping_message(message, &output_indices).map(Ok),
284                        Err(e) => Some(Err(e)),
285                    }
286                })
287            })
288            .boxed()
289    }
290}
risingwave_stream/executor/backfill/snapshot_backfill/consume_upstream/executor.rs

risingwave_stream/executor/backfill/snapshot_backfill/consume_upstream/
executor.rs
