risingwave_stream/executor/backfill/cdc/
state.rs

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

risingwave_stream/executor/backfill/cdc/
state.rs
