risingwave_connector/sink/
batching_log_sink.rs

1// Copyright 2026 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 async_trait::async_trait;
16use risingwave_common::array::StreamChunk;
17
18use crate::sink::log_store::{LogStoreReadItem, TruncateOffset};
19use crate::sink::{LogSinker, Result, SinkLogReader};
20
21/// A sink writer that buffers rows across chunks and commits them in batches, driven by
22/// [`BatchingLogSinker`].
23#[async_trait]
24pub trait BatchingSinkWriter: Send + 'static {
25    async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()>;
26
27    /// Commits buffered data if a batch is ready. Returning `true` means everything received so
28    /// far is visible downstream, allowing the log store to truncate up to this point.
29    async fn try_commit(&mut self) -> Result<bool>;
30
31    /// Called at a barrier. Returns whether the barrier may be truncated, i.e. everything received
32    /// so far is committed or was never buffered. Batching across barriers by returning `false`
33    /// while data is pending is only safe for sinks guaranteed to run decoupled: on the in-memory
34    /// log store, an untruncated checkpoint barrier blocks the checkpoint from completing. Sinks
35    /// that may run non-decoupled must flush here and return `true`.
36    async fn commit_on_barrier(&mut self) -> Result<bool>;
37}
38
39/// Log sinker for sinks that batch rows across chunks: an offset is truncated only once a commit
40/// has made its rows visible downstream, preserving at-least-once delivery on restart.
41pub struct BatchingLogSinker<W> {
42    writer: W,
43}
44
45impl<W> BatchingLogSinker<W> {
46    pub fn new(writer: W) -> Self {
47        BatchingLogSinker { writer }
48    }
49}
50
51#[async_trait]
52impl<W: BatchingSinkWriter> LogSinker for BatchingLogSinker<W> {
53    async fn consume_log_and_sink(self, mut log_reader: impl SinkLogReader) -> Result<!> {
54        log_reader.start_from(None).await?;
55        let mut sink_writer = self.writer;
56        enum LogConsumerState {
57            Uninitialized,
58            EpochBegun { curr_epoch: u64 },
59            BarrierReceived { prev_epoch: u64 },
60        }
61
62        let mut state = LogConsumerState::Uninitialized;
63        loop {
64            let (epoch, item): (u64, LogStoreReadItem) = log_reader.next_item().await?;
65            state = match state {
66                LogConsumerState::Uninitialized => {
67                    LogConsumerState::EpochBegun { curr_epoch: epoch }
68                }
69                LogConsumerState::EpochBegun { curr_epoch } => {
70                    assert!(
71                        epoch >= curr_epoch,
72                        "new epoch {} should not be below the current epoch {}",
73                        epoch,
74                        curr_epoch
75                    );
76                    LogConsumerState::EpochBegun { curr_epoch: epoch }
77                }
78                LogConsumerState::BarrierReceived { prev_epoch } => {
79                    assert!(
80                        epoch > prev_epoch,
81                        "new epoch {} should be greater than prev epoch {}",
82                        epoch,
83                        prev_epoch
84                    );
85                    LogConsumerState::EpochBegun { curr_epoch: epoch }
86                }
87            };
88            match item {
89                LogStoreReadItem::StreamChunk { chunk, chunk_id } => {
90                    sink_writer.write_batch(chunk).await?;
91                    if sink_writer.try_commit().await? {
92                        // A chunk truncation also covers all preceding barriers.
93                        log_reader.truncate(TruncateOffset::Chunk { epoch, chunk_id })?;
94                    }
95                }
96                LogStoreReadItem::Barrier { .. } => {
97                    let prev_epoch = match state {
98                        LogConsumerState::EpochBegun { curr_epoch } => curr_epoch,
99                        _ => unreachable!("epoch must have begun before handling barrier"),
100                    };
101
102                    // Truncating in idle periods keeps barriers from accumulating in the log store.
103                    if sink_writer.commit_on_barrier().await? {
104                        log_reader.truncate(TruncateOffset::Barrier { epoch: prev_epoch })?;
105                    }
106
107                    state = LogConsumerState::BarrierReceived { prev_epoch }
108                }
109            }
110        }
111    }
112}
risingwave_connector/sink/batching_log_sink.rs

risingwave_connector/sink/
batching_log_sink.rs
