Skip to main content

risingwave_connector/sink/
decouple_checkpoint_log_sink.rs

1// Copyright 2024 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::num::NonZeroU64;
16use std::time::Instant;
17
18use async_trait::async_trait;
19
20use crate::sink::log_store::{LogStoreReadItem, TruncateOffset};
21use crate::sink::writer::SinkWriter;
22use crate::sink::{LogSinker, Result, SinkLogReader, SinkWriterMetrics};
23
24pub const DEFAULT_COMMIT_CHECKPOINT_INTERVAL_WITH_SINK_DECOUPLE: u64 = 10;
25pub const DEFAULT_COMMIT_CHECKPOINT_INTERVAL_WITHOUT_SINK_DECOUPLE: u64 = 1;
26
27pub const ICEBERG_DEFAULT_COMMIT_CHECKPOINT_INTERVAL: u64 = 60;
28pub const COMMIT_CHECKPOINT_INTERVAL: &str = "commit_checkpoint_interval";
29
30pub fn default_commit_checkpoint_interval() -> u64 {
31    DEFAULT_COMMIT_CHECKPOINT_INTERVAL_WITH_SINK_DECOUPLE
32}
33
34pub fn iceberg_default_commit_checkpoint_interval() -> u64 {
35    ICEBERG_DEFAULT_COMMIT_CHECKPOINT_INTERVAL
36}
37
38/// The `LogSinker` implementation used for commit-decoupled sinks (such as `Iceberg`, `DeltaLake` and `StarRocks`).
39/// The concurrent/frequent commit capability of these sinks is poor, so by leveraging the decoupled log reader,
40/// we delay the checkpoint barrier to make commits less frequent.
41pub struct DecoupleCheckpointLogSinkerOf<W> {
42    writer: W,
43    sink_writer_metrics: SinkWriterMetrics,
44    commit_checkpoint_interval: NonZeroU64,
45}
46
47impl<W> DecoupleCheckpointLogSinkerOf<W> {
48    /// Create a log sinker with a commit checkpoint interval. The sinker should be used with a
49    /// decouple log reader `KvLogStoreReader`.
50    pub fn new(
51        writer: W,
52        sink_writer_metrics: SinkWriterMetrics,
53        commit_checkpoint_interval: NonZeroU64,
54    ) -> Self {
55        DecoupleCheckpointLogSinkerOf {
56            writer,
57            sink_writer_metrics,
58            commit_checkpoint_interval,
59        }
60    }
61}
62
63pub(crate) fn should_force_commit_on_checkpoint_barrier(
64    vnode_bitmap_updated: bool,
65    is_stop: bool,
66    has_schema_change: bool,
67) -> bool {
68    vnode_bitmap_updated || is_stop || has_schema_change
69}
70
71#[async_trait]
72impl<W: SinkWriter<CommitMetadata = ()>> LogSinker for DecoupleCheckpointLogSinkerOf<W> {
73    async fn consume_log_and_sink(self, mut log_reader: impl SinkLogReader) -> Result<!> {
74        let mut sink_writer = self.writer;
75        log_reader.start_from(None).await?;
76        #[derive(Debug)]
77        enum LogConsumerState {
78            /// Mark that the log consumer is not initialized yet
79            Uninitialized,
80
81            /// Mark that a new epoch has begun.
82            EpochBegun { curr_epoch: u64 },
83
84            /// Mark that the consumer has just received a barrier
85            BarrierReceived { prev_epoch: u64 },
86        }
87
88        let mut state = LogConsumerState::Uninitialized;
89
90        let mut current_checkpoint: u64 = 0;
91        let commit_checkpoint_interval = self.commit_checkpoint_interval;
92        let sink_writer_metrics = self.sink_writer_metrics;
93
94        loop {
95            let (epoch, item): (u64, LogStoreReadItem) = log_reader.next_item().await?;
96            // begin_epoch when not previously began
97            state = match state {
98                LogConsumerState::Uninitialized => {
99                    sink_writer.begin_epoch(epoch).await?;
100                    LogConsumerState::EpochBegun { curr_epoch: epoch }
101                }
102                LogConsumerState::EpochBegun { curr_epoch } => {
103                    assert!(
104                        epoch >= curr_epoch,
105                        "new epoch {} should not be below the current epoch {}",
106                        epoch,
107                        curr_epoch
108                    );
109                    LogConsumerState::EpochBegun { curr_epoch: epoch }
110                }
111                LogConsumerState::BarrierReceived { prev_epoch, .. } => {
112                    assert!(
113                        epoch > prev_epoch,
114                        "new epoch {} should be greater than prev epoch {}",
115                        epoch,
116                        prev_epoch
117                    );
118
119                    sink_writer.begin_epoch(epoch).await?;
120                    LogConsumerState::EpochBegun { curr_epoch: epoch }
121                }
122            };
123            match item {
124                LogStoreReadItem::StreamChunk { chunk, .. } => {
125                    if let Err(e) = sink_writer.write_batch(chunk).await {
126                        sink_writer.abort().await?;
127                        return Err(e);
128                    }
129                }
130                LogStoreReadItem::Barrier {
131                    is_checkpoint,
132                    new_vnode_bitmap,
133                    is_stop,
134                    schema_change,
135                } => {
136                    let prev_epoch = match state {
137                        LogConsumerState::EpochBegun { curr_epoch } => curr_epoch,
138                        _ => unreachable!("epoch must have begun before handling barrier"),
139                    };
140                    if is_checkpoint {
141                        current_checkpoint += 1;
142                        if current_checkpoint >= commit_checkpoint_interval.get()
143                            || should_force_commit_on_checkpoint_barrier(
144                                new_vnode_bitmap.is_some(),
145                                is_stop,
146                                schema_change.is_some(),
147                            )
148                        {
149                            let start_time = Instant::now();
150                            sink_writer.barrier(true).await?;
151                            sink_writer_metrics
152                                .sink_commit_duration
153                                .observe(start_time.elapsed().as_secs_f64());
154                            log_reader.truncate(TruncateOffset::Barrier { epoch })?;
155
156                            current_checkpoint = 0;
157                        } else {
158                            sink_writer.barrier(false).await?;
159                        }
160                    } else {
161                        assert!(new_vnode_bitmap.is_none());
162                        sink_writer.barrier(false).await?;
163                    }
164                    state = LogConsumerState::BarrierReceived { prev_epoch }
165                }
166            }
167        }
168    }
169}
170
171#[cfg(test)]
172mod tests {
173    use super::should_force_commit_on_checkpoint_barrier;
174
175    #[test]
176    fn test_should_force_commit_on_checkpoint_barrier() {
177        assert!(!should_force_commit_on_checkpoint_barrier(
178            false, false, false
179        ));
180        assert!(should_force_commit_on_checkpoint_barrier(
181            true, false, false
182        ));
183        assert!(should_force_commit_on_checkpoint_barrier(
184            false, true, false
185        ));
186        assert!(should_force_commit_on_checkpoint_barrier(
187            false, false, true
188        ));
189    }
190}