risingwave_stream/executor/source/
reader_stream.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::collections::HashMap;
16
17use futures::StreamExt;
18use itertools::Itertools;
19use risingwave_common::array::StreamChunk;
20use risingwave_common::catalog::ColumnId;
21use risingwave_common::id::SourceId;
22use risingwave_common::metrics::GLOBAL_ERROR_METRICS;
23use risingwave_connector::parser::schema_change::SchemaChangeEnvelope;
24use risingwave_connector::source::reader::desc::SourceDesc;
25use risingwave_connector::source::{
26    BoxSourceReaderEventStream, CdcAutoSchemaChangeFailCallback, ConnectorState,
27    CreateSplitReaderResult, SourceContext, SourceCtrlOpts, SourceReaderEvent, SplitId, SplitImpl,
28    SplitMetaData, StreamChunkWithState,
29};
30use thiserror_ext::AsReport;
31use tokio::sync::{mpsc, oneshot};
32
33use super::{
34    apply_rate_limit_to_source_reader_event, get_infinite_backoff_strategy,
35    get_split_offset_col_idx,
36};
37use crate::common::rate_limit::limited_chunk_size;
38use crate::executor::prelude::*;
39
40type AutoSchemaChangeSetup = (
41    Option<mpsc::Sender<(SchemaChangeEnvelope, oneshot::Sender<()>)>>,
42    Option<CdcAutoSchemaChangeFailCallback>,
43);
44
45#[derive(Debug)]
46pub(crate) enum SourceReaderEventWithState {
47    Data(StreamChunkWithState),
48    Progress(HashMap<SplitId, SplitImpl>),
49}
50
51pub(crate) struct StreamReaderBuilder {
52    pub source_desc: SourceDesc,
53    pub rate_limit: Option<u32>,
54    pub source_id: SourceId,
55    pub source_name: String,
56    pub reader_stream: Option<BoxSourceReaderEventStream>,
57
58    // cdc related
59    pub is_auto_schema_change_enable: bool,
60    pub actor_ctx: ActorContextRef,
61}
62
63impl StreamReaderBuilder {
64    fn update_offsets_from_chunk(
65        latest_splits_info: &mut HashMap<SplitId, SplitImpl>,
66        chunk: &StreamChunk,
67        split_idx: usize,
68        offset_idx: usize,
69    ) {
70        // All rows (including those visible or invisible) will be used to update the source offset.
71        for i in 0..chunk.capacity() {
72            let (_, row, _) = chunk.row_at(i);
73            let split = row.datum_at(split_idx).unwrap().into_utf8();
74            let offset = row.datum_at(offset_idx).unwrap().into_utf8();
75            latest_splits_info
76                .get_mut(&Arc::from(split.to_owned()))
77                .map(|split_impl| split_impl.update_in_place(offset.to_owned()));
78        }
79    }
80
81    fn apply_split_progress(
82        latest_splits_info: &mut HashMap<SplitId, SplitImpl>,
83        split_progress: HashMap<SplitId, String>,
84    ) {
85        for (split_id, offset) in split_progress {
86            if let Some(split_impl) = latest_splits_info.get_mut(&split_id) {
87                if let Err(e) = split_impl.update_in_place(offset) {
88                    tracing::warn!(
89                        error = %e.as_report(),
90                        split_id = %split_id,
91                        "failed to apply split progress update",
92                    );
93                }
94            } else {
95                tracing::warn!(
96                    split_id = %split_id,
97                    "ignore progress for unknown split",
98                );
99            }
100        }
101    }
102
103    fn setup_auto_schema_change(&self) -> AutoSchemaChangeSetup {
104        if self.is_auto_schema_change_enable {
105            let (schema_change_tx, mut schema_change_rx) =
106                mpsc::channel::<(SchemaChangeEnvelope, oneshot::Sender<()>)>(16);
107            let meta_client = self.actor_ctx.meta_client.clone();
108            // spawn a task to handle schema change event from source parser
109            let _join_handle = tokio::task::spawn(async move {
110                while let Some((schema_change, finish_tx)) = schema_change_rx.recv().await {
111                    let table_ids = schema_change.table_ids();
112                    tracing::info!(
113                        target: "auto_schema_change",
114                        "recv a schema change event for tables: {:?}", table_ids);
115                    // TODO: retry on rpc error
116                    if let Some(ref meta_client) = meta_client {
117                        match meta_client
118                            .auto_schema_change(schema_change.to_protobuf())
119                            .await
120                        {
121                            Ok(_) => {
122                                tracing::info!(
123                                    target: "auto_schema_change",
124                                    "schema change success for tables: {:?}", table_ids);
125                                let _ = finish_tx.send(());
126                            }
127                            Err(e) => {
128                                tracing::error!(
129                                    target: "auto_schema_change",
130                                    error = %e.as_report(), "schema change error");
131
132                                let _ = finish_tx.send(());
133                            }
134                        }
135                    }
136                }
137            });
138
139            // Create callback function for reporting CDC auto schema change fail events
140            let on_cdc_auto_schema_change_failure = if let Some(ref meta_client) =
141                self.actor_ctx.meta_client
142            {
143                let meta_client = meta_client.clone();
144                Some(CdcAutoSchemaChangeFailCallback::new(
145                    move |source_id: SourceId,
146                          table_name: String,
147                          cdc_table_id: String,
148                          upstream_ddl: String,
149                          fail_info: String| {
150                        let meta_client = meta_client.clone();
151                        tokio::spawn(async move {
152                            if let Err(e) = meta_client
153                                .add_cdc_auto_schema_change_fail_event(
154                                    source_id,
155                                    table_name,
156                                    cdc_table_id,
157                                    upstream_ddl,
158                                    fail_info,
159                                )
160                                .await
161                            {
162                                tracing::warn!(
163                                    error = %e.as_report(),
164                                    %source_id,
165                                    "Failed to add CDC auto schema change fail event to event log."
166                                );
167                            }
168                        });
169                    },
170                ))
171            } else {
172                None
173            };
174
175            (Some(schema_change_tx), on_cdc_auto_schema_change_failure)
176        } else {
177            info!("auto schema change is disabled in config");
178            (None, None)
179        }
180    }
181
182    fn prepare_source_stream_build(&self) -> (Vec<ColumnId>, SourceContext) {
183        let column_ids = self
184            .source_desc
185            .columns
186            .iter()
187            .map(|column_desc| column_desc.column_id)
188            .collect_vec();
189        debug_assert!(column_ids.iter().all_unique(), "column_ids must be unique");
190
191        let (schema_change_tx, on_cdc_auto_schema_change_failure) = self.setup_auto_schema_change();
192
193        let source_ctx = SourceContext::new_with_auto_schema_change_callback(
194            self.actor_ctx.id,
195            self.source_id,
196            self.actor_ctx.fragment_id,
197            self.source_name.clone(),
198            self.source_desc.metrics.clone(),
199            SourceCtrlOpts {
200                chunk_size: limited_chunk_size(self.rate_limit),
201                split_txn: self.rate_limit.is_some(), // when rate limiting, we may split txn
202            },
203            self.source_desc.source.config.clone(),
204            schema_change_tx,
205            on_cdc_auto_schema_change_failure,
206        );
207
208        (column_ids, source_ctx)
209    }
210
211    pub(crate) async fn fetch_latest_splits(
212        &mut self,
213        state: ConnectorState,
214        seek_to_latest: bool,
215    ) -> StreamExecutorResult<CreateSplitReaderResult> {
216        let (column_ids, source_ctx) = self.prepare_source_stream_build();
217        let source_ctx_ref = Arc::new(source_ctx);
218        let (stream, res) = self
219            .source_desc
220            .source
221            .build_stream(
222                state.clone(),
223                column_ids.clone(),
224                source_ctx_ref.clone(),
225                seek_to_latest,
226            )
227            .await
228            .map_err(StreamExecutorError::connector_error)?;
229        self.reader_stream = Some(stream);
230        Ok(res)
231    }
232
233    #[try_stream(ok = SourceReaderEventWithState, error = StreamExecutorError)]
234    pub(crate) async fn into_retry_stream(mut self, state: ConnectorState, is_initial_build: bool) {
235        let (column_ids, source_ctx) = self.prepare_source_stream_build();
236        let source_ctx_ref = Arc::new(source_ctx);
237
238        let mut latest_splits_info = {
239            if let Some(splits) = state.as_ref() {
240                splits
241                    .iter()
242                    .map(|split| (split.id(), split.clone()))
243                    .collect::<HashMap<_, _>>()
244            } else {
245                HashMap::new()
246            }
247        };
248
249        let (Some(split_idx), Some(offset_idx), _) =
250            get_split_offset_col_idx(&self.source_desc.columns)
251        else {
252            unreachable!("Partition and offset columns must be set.");
253        };
254
255        let mut backoff = get_infinite_backoff_strategy();
256
257        'build_consume_loop: loop {
258            let bootstrap_state = if latest_splits_info.is_empty() {
259                None
260            } else {
261                Some(latest_splits_info.values().cloned().collect_vec())
262            };
263            tracing::debug!(
264                "build stream source reader with state: {:?}",
265                bootstrap_state
266            );
267            let build_stream_result = if let Some(exist_stream) = self.reader_stream.take() {
268                Ok((exist_stream, CreateSplitReaderResult::default()))
269            } else {
270                self.source_desc
271                    .source
272                    .build_stream(
273                        bootstrap_state,
274                        column_ids.clone(),
275                        source_ctx_ref.clone(),
276                        // just `seek_to_latest` for initial build
277                        is_initial_build,
278                    )
279                    .await
280            };
281            if let Err(e) = build_stream_result {
282                if is_initial_build {
283                    return Err(StreamExecutorError::connector_error(e));
284                } else {
285                    let delay = backoff.next().unwrap();
286                    tracing::error!(
287                        error = %e.as_report(),
288                        source_name = self.source_name,
289                        source_id = %self.source_id,
290                        actor_id = %self.actor_ctx.id,
291                        "build stream source reader error, retry in {delay:?}"
292                    );
293                    GLOBAL_ERROR_METRICS.user_source_error.report([
294                        e.variant_name().to_owned(),
295                        self.source_id.to_string(),
296                        self.source_name.clone(),
297                        self.actor_ctx.fragment_id.to_string(),
298                    ]);
299                    tokio::time::sleep(delay).await;
300                    continue 'build_consume_loop;
301                }
302            }
303
304            let (stream, _) = build_stream_result.unwrap();
305            let stream = apply_rate_limit_to_source_reader_event(stream, self.rate_limit).boxed();
306            let mut is_error = false;
307            let mut has_progress = false;
308            #[for_await]
309            'consume: for event in stream {
310                let event = match event {
311                    Ok(event) => {
312                        has_progress = true;
313                        event
314                    }
315                    Err(e) => {
316                        tracing::error!(
317                            error = %e.as_report(),
318                            source_name = self.source_name,
319                            source_id = %self.source_id,
320                            actor_id = %self.actor_ctx.id,
321                            "stream source reader error"
322                        );
323                        GLOBAL_ERROR_METRICS.user_source_error.report([
324                            e.variant_name().to_owned(),
325                            self.source_id.to_string(),
326                            self.source_name.clone(),
327                            self.actor_ctx.fragment_id.to_string(),
328                        ]);
329                        is_error = true;
330                        break 'consume;
331                    }
332                };
333
334                match event {
335                    SourceReaderEvent::DataChunk(chunk) => {
336                        Self::update_offsets_from_chunk(
337                            &mut latest_splits_info,
338                            &chunk,
339                            split_idx,
340                            offset_idx,
341                        );
342                        yield SourceReaderEventWithState::Data((chunk, latest_splits_info.clone()));
343                    }
344                    SourceReaderEvent::SplitProgress(split_progress) => {
345                        Self::apply_split_progress(&mut latest_splits_info, split_progress);
346                        yield SourceReaderEventWithState::Progress(latest_splits_info.clone());
347                    }
348                }
349            }
350            if !is_error {
351                tracing::info!("stream source reader consume finished");
352                latest_splits_info.values_mut().for_each(|split_impl| {
353                    if let Some(mut batch_split) = split_impl.clone().into_batch_split() {
354                        batch_split.finish();
355                        *split_impl = batch_split.into();
356                    }
357                });
358                yield SourceReaderEventWithState::Data((
359                    StreamChunk::empty(
360                        self.source_desc
361                            .columns
362                            .iter()
363                            .map(|c| c.data_type.clone())
364                            .collect_vec()
365                            .as_slice(),
366                    ),
367                    latest_splits_info.clone(),
368                ));
369                break 'build_consume_loop;
370            }
371            if has_progress {
372                backoff = get_infinite_backoff_strategy();
373            }
374            let delay = backoff.next().unwrap();
375            tracing::info!("stream source reader error, retry in {delay:?}");
376            tokio::time::sleep(delay).await;
377        }
378    }
379}
risingwave_stream/executor/source/reader_stream.rs

risingwave_stream/executor/source/
reader_stream.rs
