risingwave_stream/executor/backfill/cdc/
state_v2.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 anyhow::anyhow;
16use risingwave_common::row;
17use risingwave_common::types::{JsonbVal, ScalarImpl};
18use risingwave_common::util::epoch::EpochPair;
19use risingwave_connector::source::cdc::external::CdcOffset;
20use risingwave_storage::StateStore;
21
22use crate::common::table::state_table::StateTable;
23use crate::executor::StreamExecutorResult;
24
25#[derive(Debug, Default)]
26pub struct CdcStateRecord {
27    pub is_finished: bool,
28    pub row_count: i64,
29    pub cdc_offset_low: Option<CdcOffset>,
30    pub cdc_offset_high: Option<CdcOffset>,
31}
32
33/// state schema: | `split_id` | `backfill_finished` | `row_count` | `cdc_offset_low` | `cdc_offset_high` |
34/// legacy state schema: | `split_id` | `backfill_finished` | `row_count` |
35pub struct ParallelizedCdcBackfillState<S: StateStore> {
36    state_table: StateTable<S>,
37    state_len: usize,
38    is_legacy_state: bool,
39}
40
41impl<S: StateStore> ParallelizedCdcBackfillState<S> {
42    pub fn new(state_table: StateTable<S>) -> Self {
43        let is_legacy_state = state_table.get_data_types().len() == 3;
44        let state_len = if is_legacy_state { 3 } else { 5 };
45        Self {
46            state_table,
47            state_len,
48            is_legacy_state,
49        }
50    }
51
52    pub async fn init_epoch(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
53        self.state_table.init_epoch(epoch).await
54    }
55
56    /// Restore the backfill state from storage
57    pub async fn restore_state(&mut self, split_id: i64) -> StreamExecutorResult<CdcStateRecord> {
58        let key = Some(split_id);
59        match self
60            .state_table
61            .get_row(row::once(key.map(ScalarImpl::from)))
62            .await?
63        {
64            Some(row) => {
65                tracing::info!("restored cdc backfill state: {:?}", row);
66                let state = row.into_inner().into_vec();
67                let is_finished = match state[1] {
68                    Some(ScalarImpl::Bool(val)) => val,
69                    _ => return Err(anyhow!("invalid backfill state: backfill_finished").into()),
70                };
71                let row_count = match state[2] {
72                    Some(ScalarImpl::Int64(val)) => val,
73                    _ => return Err(anyhow!("invalid backfill state: row_count").into()),
74                };
75                let (cdc_offset_low, cdc_offset_high) = if !self.is_legacy_state {
76                    let cdc_offset_low = match state[3] {
77                        Some(ScalarImpl::Jsonb(ref jsonb)) => {
78                            serde_json::from_value(jsonb.clone().take()).unwrap()
79                        }
80                        None => None,
81                        _ => return Err(anyhow!("invalid backfill state: cdc_offset_low").into()),
82                    };
83                    let cdc_offset_high = match state[4] {
84                        Some(ScalarImpl::Jsonb(ref jsonb)) => {
85                            serde_json::from_value(jsonb.clone().take()).unwrap()
86                        }
87                        None => None,
88                        _ => return Err(anyhow!("invalid backfill state: cdc_offset_high").into()),
89                    };
90                    (cdc_offset_low, cdc_offset_high)
91                } else {
92                    (None, None)
93                };
94
95                Ok(CdcStateRecord {
96                    is_finished,
97                    row_count,
98                    cdc_offset_low,
99                    cdc_offset_high,
100                })
101            }
102            None => Ok(CdcStateRecord::default()),
103        }
104    }
105
106    /// Modify the state of the corresponding split
107    pub async fn mutate_state(
108        &mut self,
109        split_id: i64,
110        is_finished: bool,
111        row_count: u64,
112        cdc_offset_low: Option<CdcOffset>,
113        cdc_offset_high: Option<CdcOffset>,
114    ) -> StreamExecutorResult<()> {
115        // schema: | `split_id` | `backfill_finished` | `row_count` | `cdc_offset_low` | `cdc_offset_high` |
116        let mut state = vec![None; self.state_len];
117        let split_id = Some(ScalarImpl::from(split_id));
118        state[0].clone_from(&split_id);
119        state[1] = Some(is_finished.into());
120        state[2] = Some((row_count as i64).into());
121        if !self.is_legacy_state {
122            state[3] = cdc_offset_low.map(|cdc_offset| {
123                let json = serde_json::to_value(cdc_offset).unwrap();
124                ScalarImpl::Jsonb(JsonbVal::from(json))
125            });
126            state[4] = cdc_offset_high.map(|cdc_offset| {
127                let json = serde_json::to_value(cdc_offset).unwrap();
128                ScalarImpl::Jsonb(JsonbVal::from(json))
129            });
130        }
131
132        match self.state_table.get_row(row::once(split_id)).await? {
133            Some(prev_row) => {
134                self.state_table.update(prev_row, state.as_slice());
135            }
136            None => {
137                self.state_table.insert(state.as_slice());
138            }
139        }
140        Ok(())
141    }
142
143    pub async fn init_state_if_absent(&mut self, split_id: i64) -> StreamExecutorResult<()> {
144        let key = Some(ScalarImpl::from(split_id));
145
146        if self
147            .state_table
148            .get_row(row::once(key.clone()))
149            .await?
150            .is_none()
151        {
152            let mut state = vec![None; self.state_len];
153            state[0] = key;
154            state[1] = Some(false.into());
155            state[2] = Some(0_i64.into());
156            self.state_table.insert(state.as_slice());
157        }
158
159        Ok(())
160    }
161
162    /// Persist the state to storage
163    pub async fn commit_state(&mut self, new_epoch: EpochPair) -> StreamExecutorResult<()> {
164        self.state_table
165            .commit_assert_no_update_vnode_bitmap(new_epoch)
166            .await
167    }
168
169    pub fn is_legacy_state(&self) -> bool {
170        self.is_legacy_state
171    }
172}
risingwave_stream/executor/backfill/cdc/state_v2.rs

risingwave_stream/executor/backfill/cdc/
state_v2.rs
