Skip to main content

risingwave_meta/stream/
cdc.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, HashSet};
16use std::iter;
17use std::time::Duration;
18
19use anyhow::Context;
20use futures::pin_mut;
21use futures_async_stream::for_await;
22use itertools::Itertools;
23use risingwave_common::bail;
24use risingwave_common::catalog::{FragmentTypeFlag, FragmentTypeMask, Schema};
25use risingwave_common::id::{ActorId, JobId};
26use risingwave_common::row::{OwnedRow, Row};
27use risingwave_common::util::iter_util::ZipEqDebug;
28use risingwave_common::util::stream_graph_visitor::visit_stream_node_cont;
29use risingwave_connector::source::cdc::external::{
30    CdcTableSnapshotSplitOption, ExternalCdcTableType, ExternalTableConfig, ExternalTableReader,
31    SchemaTableName,
32};
33use risingwave_connector::source::cdc::{CdcScanOptions, build_cdc_table_snapshot_split};
34use risingwave_connector::source::{CdcTableSnapshotSplit, CdcTableSnapshotSplitRaw};
35use risingwave_meta_model::cdc_table_snapshot_split::Relation::Object;
36use risingwave_meta_model::{cdc_table_snapshot_split, object};
37use risingwave_meta_model_migration::JoinType;
38use risingwave_pb::id::{DatabaseId, TableId};
39use risingwave_pb::plan_common::ExternalTableDesc;
40use risingwave_pb::source::PbCdcTableSnapshotSplits;
41use risingwave_pb::stream_plan::stream_node::NodeBody;
42use risingwave_pb::stream_plan::{StreamCdcScanNode, StreamCdcScanOptions, StreamNode};
43use sea_orm::{
44    ColumnTrait, ConnectionTrait, EntityTrait, QueryFilter, QuerySelect, RelationTrait, Set,
45    TransactionTrait,
46};
47
48use crate::MetaResult;
49use crate::controller::SqlMetaStore;
50use crate::model::Fragment;
51
52/// A CDC table snapshot splits can only be successfully initialized once.
53/// Subsequent attempts to write to the metastore with the same primary key will be rejected.
54pub(crate) async fn try_init_parallel_cdc_table_snapshot_splits(
55    table_id: TableId,
56    table_desc: &ExternalTableDesc,
57    meta_store: &SqlMetaStore,
58    per_table_options: &StreamCdcScanOptions,
59    insert_batch_size: u64,
60    sleep_split_interval: u64,
61    sleep_duration_millis: u64,
62) -> MetaResult<Vec<CdcTableSnapshotSplitRaw>> {
63    let split_options = CdcTableSnapshotSplitOption {
64        backfill_num_rows_per_split: per_table_options.backfill_num_rows_per_split,
65        backfill_as_even_splits: per_table_options.backfill_as_even_splits,
66        backfill_split_pk_column_index: per_table_options.backfill_split_pk_column_index,
67    };
68    let table_type = ExternalCdcTableType::from_properties(&table_desc.connect_properties);
69    // Filter out additional columns to construct the external table schema
70    let table_schema: Schema = table_desc
71        .columns
72        .iter()
73        .filter(|col| {
74            col.additional_column
75                .as_ref()
76                .is_none_or(|a_col| a_col.column_type.is_none())
77        })
78        .map(Into::into)
79        .collect();
80    let table_pk_indices = if let Some(pk) = &table_desc.pk_ordering {
81        pk.columns
82            .iter()
83            .map(|column| column.pk_col_idx as usize)
84            .collect_vec()
85    } else {
86        table_desc
87            .pk
88            .iter()
89            .map(|column| column.column_index as usize)
90            .collect_vec()
91    };
92    let table_config = ExternalTableConfig::try_from_btreemap(
93        table_desc.connect_properties.clone(),
94        table_desc.secret_refs.clone(),
95    )
96    .context("failed to parse external table config")?;
97    let schema_table_name = SchemaTableName::from_properties(&table_desc.connect_properties);
98    let reader = table_type
99        .create_table_reader(
100            table_config,
101            table_schema,
102            table_pk_indices,
103            schema_table_name,
104            table_id.as_raw_id(),
105        )
106        .await?;
107    let stream = reader.get_parallel_cdc_splits(split_options);
108    let mut insert_batch = vec![];
109    let mut splits_num = 0;
110    let mut splits = vec![];
111    let txn = meta_store.conn.begin().await?;
112    pin_mut!(stream);
113    #[for_await]
114    for split in stream {
115        let split: CdcTableSnapshotSplit = split?;
116        splits_num += 1;
117        let left = split.left_bound_inclusive.value_serialize();
118        let right = split.right_bound_exclusive.value_serialize();
119        insert_batch.push(cdc_table_snapshot_split::ActiveModel {
120            table_id: Set(table_id.as_job_id()),
121            split_id: Set(split.split_id.to_owned()),
122            left: Set(left.clone()),
123            right: Set(right.clone()),
124            is_backfill_finished: Set(0),
125        });
126        splits.push(CdcTableSnapshotSplitRaw {
127            split_id: split.split_id,
128            left_bound_inclusive: left,
129            right_bound_exclusive: right,
130        });
131        if insert_batch.len() >= insert_batch_size as usize {
132            cdc_table_snapshot_split::Entity::insert_many(std::mem::take(&mut insert_batch))
133                .exec(&txn)
134                .await?;
135        }
136        if splits_num % sleep_split_interval == 0 {
137            tokio::time::sleep(Duration::from_millis(sleep_duration_millis)).await;
138        }
139    }
140    if !insert_batch.is_empty() {
141        cdc_table_snapshot_split::Entity::insert_many(std::mem::take(&mut insert_batch))
142            .exec(&txn)
143            .await?;
144    }
145    txn.commit().await?;
146    Ok(splits)
147}
148
149/// Returns some cdc scan node if the fragment is CDC scan and has parallelized backfill enabled.
150pub(crate) fn is_parallelized_backfill_enabled_cdc_scan_fragment(
151    fragment_type_mask: FragmentTypeMask,
152    nodes: &StreamNode,
153) -> Option<&StreamCdcScanNode> {
154    if !fragment_type_mask.contains(FragmentTypeFlag::StreamCdcScan) {
155        return None;
156    }
157    let mut stream_cdc_scan = None;
158    visit_stream_node_cont(nodes, |node| {
159        if let Some(NodeBody::StreamCdcScan(node)) = &node.node_body {
160            if is_parallelized_backfill_enabled(node) {
161                stream_cdc_scan = Some(&**node);
162            }
163            false
164        } else {
165            true
166        }
167    });
168    stream_cdc_scan
169}
170
171fn is_parallelized_backfill_enabled(node: &StreamCdcScanNode) -> bool {
172    if let Some(options) = &node.options
173        && CdcScanOptions::from_proto(options).is_parallelized_backfill()
174    {
175        return true;
176    }
177    false
178}
179
180pub(crate) fn parallel_cdc_table_backfill_fragment<'a>(
181    fragments: impl Iterator<Item = &'a Fragment>,
182) -> Option<(&'a Fragment, &'a StreamCdcScanNode)> {
183    let mut stream_scan_fragments = fragments.filter_map(|f| {
184        is_parallelized_backfill_enabled_cdc_scan_fragment(f.fragment_type_mask, &f.nodes)
185            .map(|cdc_scan| (f, cdc_scan))
186    });
187    let fragment = stream_scan_fragments.next()?;
188    assert_eq!(
189        stream_scan_fragments.count(),
190        0,
191        "Expect no remaining scan fragment",
192    );
193    Some(fragment)
194}
195
196pub(crate) fn assign_cdc_table_snapshot_splits(
197    actor_ids: HashSet<ActorId>,
198    splits: &[CdcTableSnapshotSplitRaw],
199    generation: u64,
200) -> MetaResult<HashMap<ActorId, PbCdcTableSnapshotSplits>> {
201    if actor_ids.is_empty() {
202        return Err(anyhow::anyhow!("Expect at least 1 actor, 0 was found.").into());
203    }
204    if splits.is_empty() {
205        return Err(
206            anyhow::anyhow!("Expect at least 1 CDC table snapshot splits, 0 was found.").into(),
207        );
208    }
209    let splits_per_actor = splits.len().div_ceil(actor_ids.len());
210    let mut assignments = HashMap::new();
211    for (actor_id, splits) in actor_ids.iter().copied().zip_eq_debug(
212        splits
213            .iter()
214            .map(build_cdc_table_snapshot_split)
215            .chunks(splits_per_actor)
216            .into_iter()
217            .map(|c| c.collect_vec())
218            .chain(iter::repeat(Vec::default()))
219            .take(actor_ids.len()),
220    ) {
221        assignments.insert(actor_id, PbCdcTableSnapshotSplits { splits, generation });
222    }
223    Ok(assignments)
224}
225
226#[derive(Debug)]
227pub enum CdcTableSnapshotSplits {
228    Backfilling(Vec<CdcTableSnapshotSplitRaw>),
229    Finished,
230}
231
232pub async fn reload_cdc_table_snapshot_splits(
233    txn: &impl ConnectionTrait,
234    database_id: Option<DatabaseId>,
235) -> MetaResult<HashMap<JobId, CdcTableSnapshotSplits>> {
236    let columns = [
237        cdc_table_snapshot_split::Column::TableId,
238        cdc_table_snapshot_split::Column::SplitId,
239        cdc_table_snapshot_split::Column::Left,
240        cdc_table_snapshot_split::Column::Right,
241        cdc_table_snapshot_split::Column::IsBackfillFinished,
242    ];
243    #[expect(clippy::type_complexity)]
244    let all_splits: Vec<(JobId, i64, Vec<u8>, Vec<u8>, i16)> =
245        if let Some(database_id) = database_id {
246            cdc_table_snapshot_split::Entity::find()
247                .join(JoinType::LeftJoin, Object.def())
248                .select_only()
249                .columns(columns)
250                .filter(object::Column::DatabaseId.eq(database_id))
251                .into_tuple()
252                .all(txn)
253                .await?
254        } else {
255            cdc_table_snapshot_split::Entity::find()
256                .select_only()
257                .columns(columns)
258                .into_tuple()
259                .all(txn)
260                .await?
261        };
262
263    let mut job_splits = HashMap::<_, Vec<_>>::new();
264    for (job_id, split_id, left, right, is_backfill_finished) in all_splits {
265        job_splits
266            .entry(job_id)
267            .or_default()
268            .push((split_id, left, right, is_backfill_finished));
269    }
270
271    job_splits
272        .into_iter()
273        .map(|(job_id, splits)| {
274            let splits = compose_job_splits(job_id, splits)?;
275            Ok((job_id, splits))
276        })
277        .try_collect()
278}
279
280pub fn compose_job_splits(
281    job_id: JobId,
282    splits: Vec<(i64, Vec<u8>, Vec<u8>, i16)>,
283) -> MetaResult<CdcTableSnapshotSplits> {
284    let split_completed_count = splits
285        .iter()
286        .filter(|(_, _, _, is_backfill_finished)| *is_backfill_finished == 1)
287        .count();
288    assert!(
289        split_completed_count <= 1,
290        "split_completed_count = {}",
291        split_completed_count
292    );
293    let is_backfill_finished = split_completed_count == 1;
294    let splits = if is_backfill_finished {
295        if splits.len() != 1 {
296            // CdcTableBackfillTracker::complete_job rewrites splits in a transaction.
297            // This error should only happen when the meta store reads uncommitted data.
298            tracing::error!(%job_id, ?splits, "unexpected split count");
299            bail!(
300                "unexpected split count: job_id={job_id}, split_total_count={}, split_completed_count={split_completed_count}",
301                splits.len()
302            );
303        }
304        CdcTableSnapshotSplits::Finished
305    } else {
306        let splits: Vec<_> = splits
307            .into_iter()
308            // The try_init_parallel_cdc_table_snapshot_splits ensures that split with a larger split_id will always have a larger left bound.
309            // Assigning consecutive splits to the same actor enables potential optimization in CDC backfill executor.
310            .sorted_by_key(|(split_id, _, _, _)| *split_id)
311            .map(
312                |(split_id, left_bound_inclusive, right_bound_exclusive, _)| {
313                    CdcTableSnapshotSplitRaw {
314                        split_id,
315                        left_bound_inclusive,
316                        right_bound_exclusive,
317                    }
318                },
319            )
320            .collect();
321        CdcTableSnapshotSplits::Backfilling(splits)
322    };
323    Ok(splits)
324}
325
326pub fn single_merged_split() -> CdcTableSnapshotSplitRaw {
327    CdcTableSnapshotSplitRaw {
328        // TODO(zw): remove magic number
329        split_id: 1,
330        left_bound_inclusive: OwnedRow::new(vec![None]).value_serialize(),
331        right_bound_exclusive: OwnedRow::new(vec![None]).value_serialize(),
332    }
333}