risingwave_connector/sink/
coordinate.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 std::cmp::Ordering;
16use std::future::pending;
17use std::num::NonZeroU64;
18use std::time::Instant;
19
20use anyhow::anyhow;
21use async_trait::async_trait;
22use risingwave_common::bail;
23use risingwave_common::bitmap::Bitmap;
24use risingwave_pb::connector_service::SinkMetadata;
25use tracing::{info, warn};
26
27use super::{
28    LogSinker, SinkCoordinationRpcClientEnum, SinkLogReader, SinkWriterMetrics, SinkWriterParam,
29};
30use crate::sink::decouple_checkpoint_log_sink::should_force_commit_on_checkpoint_barrier;
31use crate::sink::writer::SinkWriter;
32use crate::sink::{LogStoreReadItem, Result, SinkError, SinkParam, TruncateOffset};
33
34pub struct CoordinatedLogSinker<W: SinkWriter<CommitMetadata = Option<SinkMetadata>>> {
35    writer: W,
36    sink_coordinate_client: SinkCoordinationRpcClientEnum,
37    param: SinkParam,
38    vnode_bitmap: Bitmap,
39    commit_checkpoint_interval: NonZeroU64,
40    sink_writer_metrics: SinkWriterMetrics,
41}
42
43impl<W: SinkWriter<CommitMetadata = Option<SinkMetadata>>> CoordinatedLogSinker<W> {
44    pub async fn new(
45        writer_param: &SinkWriterParam,
46        param: SinkParam,
47        writer: W,
48        commit_checkpoint_interval: NonZeroU64,
49    ) -> Result<Self> {
50        Ok(Self {
51            writer,
52            sink_coordinate_client: writer_param
53                .meta_client
54                .as_ref()
55                .ok_or_else(|| anyhow!("should have meta client"))?
56                .clone()
57                .sink_coordinate_client()
58                .await,
59            param,
60            vnode_bitmap: writer_param.vnode_bitmap.clone(),
61            commit_checkpoint_interval,
62            sink_writer_metrics: SinkWriterMetrics::new(writer_param),
63        })
64    }
65}
66
67#[async_trait]
68impl<W: SinkWriter<CommitMetadata = Option<SinkMetadata>>> LogSinker for CoordinatedLogSinker<W> {
69    async fn consume_log_and_sink(self, mut log_reader: impl SinkLogReader) -> Result<!> {
70        let (mut coordinator_stream_handle, log_store_rewind_start_epoch) = self
71            .sink_coordinate_client
72            .new_stream_handle(&self.param, self.vnode_bitmap)
73            .await?;
74        let mut sink_writer = self.writer;
75        log_reader.start_from(log_store_rewind_start_epoch).await?;
76        let mut first_item = log_reader.next_item().await?;
77        if let (Some(log_store_rewind_start_epoch), (first_epoch, _)) =
78            (log_store_rewind_start_epoch, &first_item)
79        {
80            if log_store_rewind_start_epoch >= *first_epoch {
81                bail!(
82                    "log_store_rewind_start_epoch {} not later than first_epoch {}",
83                    log_store_rewind_start_epoch,
84                    first_epoch
85                );
86            }
87        } else {
88            let &(initial_epoch, _) = &first_item;
89            let aligned_initial_epoch = coordinator_stream_handle
90                .align_initial_epoch(initial_epoch)
91                .await?;
92            if initial_epoch != aligned_initial_epoch {
93                warn!(
94                    initial_epoch,
95                    aligned_initial_epoch,
96                    sink_id = %self.param.sink_id,
97                    "initial epoch not matched aligned initial epoch"
98                );
99                let mut peeked_first = Some(first_item);
100                first_item = loop {
101                    let (epoch, item) = if let Some(peeked_first) = peeked_first.take() {
102                        peeked_first
103                    } else {
104                        log_reader.next_item().await?
105                    };
106                    match epoch.cmp(&aligned_initial_epoch) {
107                        Ordering::Less => {
108                            continue;
109                        }
110                        Ordering::Equal => {
111                            break (epoch, item);
112                        }
113                        Ordering::Greater => {
114                            return Err(anyhow!(
115                                "initial epoch {} greater than aligned initial epoch {}",
116                                initial_epoch,
117                                aligned_initial_epoch
118                            )
119                            .into());
120                        }
121                    }
122                };
123            }
124        }
125
126        let mut first_item = Some(first_item);
127
128        #[derive(Debug)]
129        enum LogConsumerState {
130            /// Mark that the log consumer is not initialized yet
131            Uninitialized,
132
133            /// Mark that a new epoch has begun.
134            EpochBegun { curr_epoch: u64 },
135
136            /// Mark that the consumer has just received a barrier
137            BarrierReceived { prev_epoch: u64 },
138        }
139
140        let mut state = LogConsumerState::Uninitialized;
141
142        let mut current_checkpoint: u64 = 0;
143        let commit_checkpoint_interval = self.commit_checkpoint_interval;
144        let sink_writer_metrics = self.sink_writer_metrics;
145
146        loop {
147            let (epoch, item) = if let Some(item) = first_item.take() {
148                item
149            } else {
150                log_reader.next_item().await?
151            };
152
153            // begin_epoch when not previously began
154            state = match state {
155                LogConsumerState::Uninitialized => {
156                    sink_writer.begin_epoch(epoch).await?;
157                    LogConsumerState::EpochBegun { curr_epoch: epoch }
158                }
159                LogConsumerState::EpochBegun { curr_epoch } => {
160                    assert!(
161                        epoch >= curr_epoch,
162                        "new epoch {} should not be below the current epoch {}",
163                        epoch,
164                        curr_epoch
165                    );
166                    LogConsumerState::EpochBegun { curr_epoch: epoch }
167                }
168                LogConsumerState::BarrierReceived { prev_epoch, .. } => {
169                    assert!(
170                        epoch > prev_epoch,
171                        "new epoch {} should be greater than prev epoch {}",
172                        epoch,
173                        prev_epoch
174                    );
175
176                    sink_writer.begin_epoch(epoch).await?;
177                    LogConsumerState::EpochBegun { curr_epoch: epoch }
178                }
179            };
180            match item {
181                LogStoreReadItem::StreamChunk { chunk, .. } => {
182                    if let Err(e) = sink_writer.write_batch(chunk).await {
183                        sink_writer.abort().await?;
184                        return Err(e);
185                    }
186                }
187                LogStoreReadItem::Barrier {
188                    is_checkpoint,
189                    new_vnode_bitmap,
190                    is_stop,
191                    schema_change,
192                } => {
193                    let prev_epoch = match state {
194                        LogConsumerState::EpochBegun { curr_epoch } => curr_epoch,
195                        _ => unreachable!("epoch must have begun before handling barrier"),
196                    };
197                    if is_checkpoint {
198                        current_checkpoint += 1;
199                        if current_checkpoint >= commit_checkpoint_interval.get()
200                            || should_force_commit_on_checkpoint_barrier(
201                                new_vnode_bitmap.is_some(),
202                                is_stop,
203                                schema_change.is_some(),
204                            )
205                        {
206                            let start_time = Instant::now();
207                            let metadata = sink_writer.barrier(true).await?;
208                            let metadata = metadata.ok_or_else(|| {
209                                SinkError::Coordinator(anyhow!(
210                                    "should get metadata on checkpoint barrier"
211                                ))
212                            })?;
213                            if schema_change.is_some() {
214                                tracing::info!(
215                                    sink_id = %self.param.sink_id,
216                                    ?schema_change,
217                                    "schema change received for coordinated log sinker"
218                                );
219                                assert!(
220                                    is_stop,
221                                    "schema change should stop current sink for sink {}",
222                                    self.param.sink_id
223                                );
224                            }
225                            coordinator_stream_handle
226                                .commit(epoch, metadata, schema_change)
227                                .await?;
228                            sink_writer_metrics
229                                .sink_commit_duration
230                                .observe(start_time.elapsed().as_secs_f64());
231
232                            current_checkpoint = 0;
233                            if let Some(new_vnode_bitmap) = new_vnode_bitmap {
234                                let epoch = coordinator_stream_handle
235                                    .update_vnode_bitmap(&new_vnode_bitmap)
236                                    .await?;
237                                if epoch != prev_epoch {
238                                    bail!(
239                                        "newly start epoch {} after update vnode bitmap not matched with prev_epoch {}",
240                                        epoch,
241                                        prev_epoch
242                                    );
243                                }
244                            }
245                            if is_stop {
246                                coordinator_stream_handle.stop().await?;
247                                info!(
248                                    sink_id = %self.param.sink_id,
249                                    "coordinated log sinker stops"
250                                );
251                                log_reader.truncate(TruncateOffset::Barrier { epoch })?;
252                                return pending().await;
253                            }
254                            log_reader.truncate(TruncateOffset::Barrier { epoch })?;
255                        } else {
256                            let metadata = sink_writer.barrier(false).await?;
257                            if let Some(metadata) = metadata {
258                                warn!(?metadata, "get metadata on non-checkpoint barrier");
259                            }
260                        }
261                    } else {
262                        let metadata = sink_writer.barrier(false).await?;
263                        if let Some(metadata) = metadata {
264                            warn!(?metadata, "get metadata on non-checkpoint barrier");
265                        }
266                    }
267                    state = LogConsumerState::BarrierReceived { prev_epoch }
268                }
269            }
270        }
271    }
272}
risingwave_connector/sink/coordinate.rs

risingwave_connector/sink/
coordinate.rs
