risingwave_connector/sink/
boxed.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 std::future::Future;
16use std::ops::DerefMut;
17
18use async_trait::async_trait;
19use futures::FutureExt;
20use futures::future::BoxFuture;
21use risingwave_pb::connector_service::SinkMetadata;
22use risingwave_pb::stream_plan::PbSinkSchemaChange;
23
24use crate::sink::log_store::{LogStoreReadItem, LogStoreResult, TruncateOffset};
25use crate::sink::{
26    LogSinker, SinglePhaseCommitCoordinator, SinkLogReader, TwoPhaseCommitCoordinator,
27};
28
29pub type BoxSinglePhaseCoordinator = Box<dyn SinglePhaseCommitCoordinator + Send + 'static>;
30pub type BoxTwoPhaseCoordinator = Box<dyn TwoPhaseCommitCoordinator + Send + 'static>;
31
32pub type BoxLogSinker = Box<
33    dyn for<'a> FnOnce(&'a mut dyn DynLogReader) -> BoxFuture<'a, crate::sink::Result<!>>
34        + Send
35        + 'static,
36>;
37
38#[async_trait]
39pub trait DynLogReader: Send {
40    async fn dyn_start_from(&mut self, start_offset: Option<u64>) -> LogStoreResult<()>;
41    async fn dyn_next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)>;
42
43    fn dyn_truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()>;
44}
45
46#[async_trait]
47impl<R: SinkLogReader> DynLogReader for R {
48    async fn dyn_start_from(&mut self, start_offset: Option<u64>) -> LogStoreResult<()> {
49        R::start_from(self, start_offset).await
50    }
51
52    async fn dyn_next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
53        R::next_item(self).await
54    }
55
56    fn dyn_truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
57        R::truncate(self, offset)
58    }
59}
60
61impl SinkLogReader for &mut dyn DynLogReader {
62    fn start_from(
63        &mut self,
64        start_offset: Option<u64>,
65    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
66        (*self).dyn_start_from(start_offset)
67    }
68
69    fn next_item(
70        &mut self,
71    ) -> impl Future<Output = LogStoreResult<(u64, LogStoreReadItem)>> + Send + '_ {
72        (*self).dyn_next_item()
73    }
74
75    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
76        (*self).dyn_truncate(offset)
77    }
78}
79
80pub fn boxed_log_sinker(log_sinker: impl LogSinker) -> BoxLogSinker {
81    fn make_future<'a>(
82        log_sinker: impl LogSinker,
83        log_reader: &'a mut dyn DynLogReader,
84    ) -> BoxFuture<'a, crate::sink::Result<!>> {
85        log_sinker.consume_log_and_sink(log_reader).boxed()
86    }
87
88    // Note: it's magical that the following expression can be cast to the expected return type
89    // without any explicit conversion, such as `<expr> as _` or `<expr>.into()`.
90    // TODO: may investigate the reason. The currently successful compilation seems volatile to future compatibility.
91    Box::new(move |log_reader: &mut dyn DynLogReader| make_future(log_sinker, log_reader))
92}
93
94#[async_trait]
95impl LogSinker for BoxLogSinker {
96    async fn consume_log_and_sink(
97        self,
98        mut log_reader: impl SinkLogReader,
99    ) -> crate::sink::Result<!> {
100        (self)(&mut log_reader).await
101    }
102}
103
104#[async_trait]
105impl SinglePhaseCommitCoordinator for BoxSinglePhaseCoordinator {
106    async fn init(&mut self) -> crate::sink::Result<()> {
107        self.deref_mut().init().await
108    }
109
110    async fn commit(
111        &mut self,
112        epoch: u64,
113        metadata: Vec<SinkMetadata>,
114        schema_change: Option<PbSinkSchemaChange>,
115    ) -> crate::sink::Result<()> {
116        self.deref_mut()
117            .commit(epoch, metadata, schema_change)
118            .await
119    }
120}
121
122#[async_trait]
123impl TwoPhaseCommitCoordinator for BoxTwoPhaseCoordinator {
124    async fn init(&mut self) -> crate::sink::Result<()> {
125        self.deref_mut().init().await
126    }
127
128    async fn pre_commit(
129        &mut self,
130        epoch: u64,
131        metadata: Vec<SinkMetadata>,
132        schema_change: Option<PbSinkSchemaChange>,
133    ) -> crate::sink::Result<Vec<u8>> {
134        self.deref_mut()
135            .pre_commit(epoch, metadata, schema_change)
136            .await
137    }
138
139    async fn commit(
140        &mut self,
141        epoch: u64,
142        commit_metadata: Vec<u8>,
143        schema_change: Option<PbSinkSchemaChange>,
144    ) -> crate::sink::Result<()> {
145        self.deref_mut()
146            .commit(epoch, commit_metadata, schema_change)
147            .await
148    }
149
150    async fn abort(&mut self, epoch: u64, commit_metadata: Vec<u8>) {
151        self.deref_mut().abort(epoch, commit_metadata).await;
152    }
153}
risingwave_connector/sink/boxed.rs

risingwave_connector/sink/
boxed.rs
