risingwave_connector/sink/
remote.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::collections::VecDeque;
16use std::marker::PhantomData;
17use std::num::NonZero;
18use std::ops::Deref;
19use std::pin::pin;
20use std::time::Instant;
21
22use anyhow::{Context, anyhow};
23use async_trait::async_trait;
24use await_tree::{InstrumentAwait, span};
25use futures::TryStreamExt;
26use futures::future::select;
27use jni::JavaVM;
28use phf::phf_set;
29use prost::Message;
30use risingwave_common::array::StreamChunk;
31use risingwave_common::bail;
32use risingwave_common::catalog::{ColumnDesc, ColumnId};
33use risingwave_common::session_config::sink_decouple::SinkDecouple;
34use risingwave_common::types::DataType;
35use risingwave_jni_core::jvm_runtime::{JVM, execute_with_jni_env};
36use risingwave_jni_core::{
37    JniReceiverType, JniSenderType, JniSinkWriterStreamRequest, call_static_method, gen_class_name,
38};
39use risingwave_pb::connector_service::sink_coordinator_stream_request::StartCoordinator;
40use risingwave_pb::connector_service::sink_writer_stream_request::{
41    Request as SinkRequest, StartSink,
42};
43use risingwave_pb::connector_service::{
44    PbSinkParam, SinkCoordinatorStreamRequest, SinkCoordinatorStreamResponse, SinkMetadata,
45    SinkWriterStreamRequest, SinkWriterStreamResponse, TableSchema, ValidateSinkRequest,
46    ValidateSinkResponse, sink_coordinator_stream_request, sink_coordinator_stream_response,
47    sink_writer_stream_response,
48};
49use risingwave_rpc_client::error::RpcError;
50use risingwave_rpc_client::{
51    BidiStreamReceiver, BidiStreamSender, DEFAULT_BUFFER_SIZE, SinkCoordinatorStreamHandle,
52    SinkWriterStreamHandle,
53};
54use rw_futures_util::drop_either_future;
55use sea_orm::DatabaseConnection;
56use thiserror_ext::AsReport;
57use tokio::sync::mpsc;
58use tokio::sync::mpsc::{Receiver, unbounded_channel};
59use tokio::task::spawn_blocking;
60use tokio_stream::wrappers::ReceiverStream;
61use tracing::warn;
62
63use super::SinkCommittedEpochSubscriber;
64use super::elasticsearch_opensearch::elasticsearch_converter::{
65    StreamChunkConverter, is_remote_es_sink,
66};
67use super::elasticsearch_opensearch::elasticsearch_opensearch_config::ES_OPTION_DELIMITER;
68use crate::enforce_secret::EnforceSecret;
69use crate::error::ConnectorResult;
70use crate::sink::coordinate::CoordinatedLogSinker;
71use crate::sink::log_store::{LogStoreReadItem, LogStoreResult, TruncateOffset};
72use crate::sink::writer::SinkWriter;
73use crate::sink::{
74    DummySinkCommitCoordinator, LogSinker, Result, Sink, SinkCommitCoordinator, SinkError,
75    SinkLogReader, SinkParam, SinkWriterMetrics, SinkWriterParam,
76};
77
78macro_rules! def_remote_sink {
79    () => {
80        def_remote_sink! {
81            //todo!, delete java impl
82            // { ElasticSearchJava, ElasticSearchJavaSink, "elasticsearch_v1" }
83            // { OpensearchJava, OpenSearchJavaSink, "opensearch_v1"}
84            { Cassandra, CassandraSink, "cassandra", [ "cassandra.url" ] }
85            { Jdbc, JdbcSink, "jdbc", [ "jdbc.url" ] }
86        }
87    };
88    ({ $variant_name:ident, $sink_type_name:ident, $sink_name:expr, [ $($enforce_secret_prop:expr),* ] }) => {
89        #[derive(Debug)]
90        pub struct $variant_name;
91        impl RemoteSinkTrait for $variant_name {
92            const SINK_NAME: &'static str = $sink_name;
93        }
94        impl EnforceSecret for $variant_name {
95            const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf_set! {
96                $($enforce_secret_prop),*
97            };
98        }
99        pub type $sink_type_name = RemoteSink<$variant_name>;
100    };
101    ({ $variant_name:ident, $sink_type_name:ident, $sink_name:expr, [ $($enforce_secret_prop:expr),* ], |$desc:ident| $body:expr }) => {
102        #[derive(Debug)]
103        pub struct $variant_name;
104        impl RemoteSinkTrait for $variant_name {
105            const SINK_NAME: &'static str = $sink_name;
106            fn default_sink_decouple($desc: &SinkDesc) -> bool {
107                $body
108            }
109        }
110        impl EnforceSecret for $variant_name {
111            const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = phf_set! {
112                $($enforce_secret_prop),*
113            };
114        }
115        pub type $sink_type_name = RemoteSink<$variant_name>;
116    };
117    ({ $($first:tt)+ } $({$($rest:tt)+})*) => {
118        def_remote_sink! {
119            {$($first)+}
120        }
121        def_remote_sink! {
122            $({$($rest)+})*
123        }
124    };
125    ($($invalid:tt)*) => {
126        compile_error! {concat! {"invalid `", stringify!{$($invalid)*}, "`"}}
127    }
128}
129
130def_remote_sink!();
131
132pub trait RemoteSinkTrait: EnforceSecret + Send + Sync + 'static {
133    const SINK_NAME: &'static str;
134    fn default_sink_decouple() -> bool {
135        true
136    }
137}
138
139#[derive(Debug)]
140pub struct RemoteSink<R: RemoteSinkTrait> {
141    param: SinkParam,
142    _phantom: PhantomData<R>,
143}
144
145impl<R: RemoteSinkTrait> EnforceSecret for RemoteSink<R> {
146    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = R::ENFORCE_SECRET_PROPERTIES;
147}
148
149impl<R: RemoteSinkTrait> TryFrom<SinkParam> for RemoteSink<R> {
150    type Error = SinkError;
151
152    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
153        Ok(Self {
154            param,
155            _phantom: PhantomData,
156        })
157    }
158}
159
160impl<R: RemoteSinkTrait> Sink for RemoteSink<R> {
161    type Coordinator = DummySinkCommitCoordinator;
162    type LogSinker = RemoteLogSinker;
163
164    const SINK_NAME: &'static str = R::SINK_NAME;
165
166    fn is_sink_decouple(user_specified: &SinkDecouple) -> Result<bool> {
167        match user_specified {
168            SinkDecouple::Default => Ok(R::default_sink_decouple()),
169            SinkDecouple::Enable => Ok(true),
170            SinkDecouple::Disable => Ok(false),
171        }
172    }
173
174    async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
175        RemoteLogSinker::new(self.param.clone(), writer_param, Self::SINK_NAME).await
176    }
177
178    async fn validate(&self) -> Result<()> {
179        validate_remote_sink(&self.param, Self::SINK_NAME).await?;
180        Ok(())
181    }
182}
183
184async fn validate_remote_sink(param: &SinkParam, sink_name: &str) -> ConnectorResult<()> {
185    // if sink_name == OpenSearchJavaSink::SINK_NAME {
186    //     risingwave_common::license::Feature::OpenSearchSink
187    //         .check_available()
188    //         .map_err(|e| anyhow::anyhow!(e))?;
189    // }
190    if is_remote_es_sink(sink_name)
191        && param.downstream_pk.len() > 1
192        && !param.properties.contains_key(ES_OPTION_DELIMITER)
193    {
194        bail!("Es sink only supports single pk or pk with delimiter option");
195    }
196    // FIXME: support struct and array in stream sink
197    param.columns.iter().try_for_each(|col| {
198        match &col.data_type {
199            DataType::Int16
200                    | DataType::Int32
201                    | DataType::Int64
202                    | DataType::Float32
203                    | DataType::Float64
204                    | DataType::Boolean
205                    | DataType::Decimal
206                    | DataType::Timestamp
207                    | DataType::Timestamptz
208                    | DataType::Varchar
209                    | DataType::Date
210                    | DataType::Time
211                    | DataType::Interval
212                    | DataType::Jsonb
213                    | DataType::Bytea => Ok(()),
214            DataType::List(list) => {
215                if is_remote_es_sink(sink_name) || matches!(list.as_ref(), DataType::Int16 | DataType::Int32 | DataType::Int64 | DataType::Float32 | DataType::Float64 | DataType::Varchar){
216                    Ok(())
217                } else{
218                    Err(SinkError::Remote(anyhow!(
219                        "Remote sink only supports list<int16, int32, int64, float, double, varchar>, got {:?}: {:?}",
220                        col.name,
221                        col.data_type,
222                    )))
223                }
224            },
225            DataType::Struct(_) => {
226                if is_remote_es_sink(sink_name){
227                    Ok(())
228                }else{
229                    Err(SinkError::Remote(anyhow!(
230                        "Only Es sink supports struct, got {:?}: {:?}",
231                        col.name,
232                        col.data_type,
233                    )))
234                }
235            },
236            DataType::Serial | DataType::Int256 | DataType::Map(_) => Err(SinkError::Remote(anyhow!(
237                            "remote sink supports Int16, Int32, Int64, Float32, Float64, Boolean, Decimal, Time, Date, Interval, Jsonb, Timestamp, Timestamptz, Bytea, List and Varchar, (Es sink support Struct) got {:?}: {:?}",
238                            col.name,
239                            col.data_type,
240                        )))}})?;
241
242    let jvm = JVM.get_or_init()?;
243    let sink_param = param.to_proto();
244
245    spawn_blocking(move || -> anyhow::Result<()> {
246        execute_with_jni_env(jvm, |env| {
247            let validate_sink_request = ValidateSinkRequest {
248                sink_param: Some(sink_param),
249            };
250            let validate_sink_request_bytes =
251                env.byte_array_from_slice(&Message::encode_to_vec(&validate_sink_request))?;
252
253            let validate_sink_response_bytes = call_static_method!(
254                env,
255                {com.risingwave.connector.JniSinkValidationHandler},
256                {byte[] validate(byte[] validateSourceRequestBytes)},
257                &validate_sink_request_bytes
258            )?;
259
260            let validate_sink_response: ValidateSinkResponse = Message::decode(
261                risingwave_jni_core::to_guarded_slice(&validate_sink_response_bytes, env)?.deref(),
262            )?;
263
264            validate_sink_response.error.map_or_else(
265                || Ok(()), // If there is no error message, return Ok here.
266                |err| bail!("sink cannot pass validation: {}", err.error_message),
267            )
268        })
269    })
270    .await
271    .context("JoinHandle returns error")??;
272
273    Ok(())
274}
275
276pub struct RemoteLogSinker {
277    request_sender: BidiStreamSender<JniSinkWriterStreamRequest>,
278    response_stream: BidiStreamReceiver<SinkWriterStreamResponse>,
279    stream_chunk_converter: StreamChunkConverter,
280    sink_writer_metrics: SinkWriterMetrics,
281}
282
283impl RemoteLogSinker {
284    async fn new(
285        sink_param: SinkParam,
286        writer_param: SinkWriterParam,
287        sink_name: &str,
288    ) -> Result<Self> {
289        let sink_proto = sink_param.to_proto();
290        let payload_schema = if is_remote_es_sink(sink_name) {
291            let columns = vec![
292                ColumnDesc::unnamed(ColumnId::from(0), DataType::Varchar).to_protobuf(),
293                ColumnDesc::unnamed(ColumnId::from(1), DataType::Varchar).to_protobuf(),
294                ColumnDesc::unnamed(ColumnId::from(2), DataType::Jsonb).to_protobuf(),
295                ColumnDesc::unnamed(ColumnId::from(2), DataType::Varchar).to_protobuf(),
296            ];
297            Some(TableSchema {
298                columns,
299                pk_indices: vec![],
300            })
301        } else {
302            sink_proto.table_schema.clone()
303        };
304
305        let SinkWriterStreamHandle {
306            request_sender,
307            response_stream,
308        } = EmbeddedConnectorClient::new()?
309            .start_sink_writer_stream(payload_schema, sink_proto)
310            .await?;
311
312        Ok(RemoteLogSinker {
313            request_sender,
314            response_stream,
315            sink_writer_metrics: SinkWriterMetrics::new(&writer_param),
316            stream_chunk_converter: StreamChunkConverter::new(
317                sink_name,
318                sink_param.schema(),
319                &sink_param.downstream_pk,
320                &sink_param.properties,
321                sink_param.sink_type.is_append_only(),
322            )?,
323        })
324    }
325}
326
327#[async_trait]
328impl LogSinker for RemoteLogSinker {
329    async fn consume_log_and_sink(self, mut log_reader: impl SinkLogReader) -> Result<!> {
330        log_reader.start_from(None).await?;
331        let mut request_tx = self.request_sender;
332        let mut response_err_stream_rx = self.response_stream;
333        let sink_writer_metrics = self.sink_writer_metrics;
334
335        let (response_tx, mut response_rx) = unbounded_channel();
336
337        let poll_response_stream = async move {
338            loop {
339                let result = response_err_stream_rx
340                    .stream
341                    .try_next()
342                    .instrument_await("log_sinker_wait_next_response")
343                    .await;
344                match result {
345                    Ok(Some(response)) => {
346                        response_tx.send(response).map_err(|err| {
347                            SinkError::Remote(anyhow!("unable to send response: {:?}", err.0))
348                        })?;
349                    }
350                    Ok(None) => return Err(SinkError::Remote(anyhow!("end of response stream"))),
351                    Err(e) => return Err(SinkError::Remote(anyhow!(e))),
352                }
353            }
354        };
355
356        let poll_consume_log_and_sink = async move {
357            fn truncate_matched_offset(
358                queue: &mut VecDeque<(TruncateOffset, Option<Instant>)>,
359                persisted_offset: TruncateOffset,
360                log_reader: &mut impl SinkLogReader,
361                sink_writer_metrics: &SinkWriterMetrics,
362            ) -> Result<()> {
363                while let Some((sent_offset, _)) = queue.front()
364                    && sent_offset < &persisted_offset
365                {
366                    queue.pop_front();
367                }
368
369                let (sent_offset, start_time) = queue.pop_front().ok_or_else(|| {
370                    anyhow!("get unsent offset {:?} in response", persisted_offset)
371                })?;
372                if sent_offset != persisted_offset {
373                    bail!(
374                        "new response offset {:?} does not match the buffer offset {:?}",
375                        persisted_offset,
376                        sent_offset
377                    );
378                }
379
380                if let (TruncateOffset::Barrier { .. }, Some(start_time)) =
381                    (persisted_offset, start_time)
382                {
383                    sink_writer_metrics
384                        .sink_commit_duration
385                        .observe(start_time.elapsed().as_millis() as f64);
386                }
387
388                log_reader.truncate(persisted_offset)?;
389                Ok(())
390            }
391
392            let mut prev_offset: Option<TruncateOffset> = None;
393            // Push from back and pop from front
394            let mut sent_offset_queue: VecDeque<(TruncateOffset, Option<Instant>)> =
395                VecDeque::new();
396
397            loop {
398                let either_result: futures::future::Either<
399                    Option<SinkWriterStreamResponse>,
400                    LogStoreResult<(u64, LogStoreReadItem)>,
401                > = drop_either_future(
402                    select(pin!(response_rx.recv()), pin!(log_reader.next_item())).await,
403                );
404                match either_result {
405                    futures::future::Either::Left(opt) => {
406                        let response = opt.context("end of response stream")?;
407                        match response {
408                            SinkWriterStreamResponse {
409                                response:
410                                    Some(sink_writer_stream_response::Response::Batch(
411                                        sink_writer_stream_response::BatchWrittenResponse {
412                                            epoch,
413                                            batch_id,
414                                        },
415                                    )),
416                            } => {
417                                truncate_matched_offset(
418                                    &mut sent_offset_queue,
419                                    TruncateOffset::Chunk {
420                                        epoch,
421                                        chunk_id: batch_id as _,
422                                    },
423                                    &mut log_reader,
424                                    &sink_writer_metrics,
425                                )?;
426                            }
427                            SinkWriterStreamResponse {
428                                response:
429                                    Some(sink_writer_stream_response::Response::Commit(
430                                        sink_writer_stream_response::CommitResponse {
431                                            epoch,
432                                            metadata,
433                                        },
434                                    )),
435                            } => {
436                                if let Some(metadata) = metadata {
437                                    warn!("get unexpected non-empty metadata: {:?}", metadata);
438                                }
439                                truncate_matched_offset(
440                                    &mut sent_offset_queue,
441                                    TruncateOffset::Barrier { epoch },
442                                    &mut log_reader,
443                                    &sink_writer_metrics,
444                                )?;
445                            }
446                            response => {
447                                return Err(SinkError::Remote(anyhow!(
448                                    "get unexpected response: {:?}",
449                                    response
450                                )));
451                            }
452                        }
453                    }
454                    futures::future::Either::Right(result) => {
455                        let (epoch, item): (u64, LogStoreReadItem) = result?;
456
457                        match item {
458                            LogStoreReadItem::StreamChunk { chunk, chunk_id } => {
459                                let offset = TruncateOffset::Chunk { epoch, chunk_id };
460                                if let Some(prev_offset) = &prev_offset {
461                                    prev_offset.check_next_offset(offset)?;
462                                }
463                                let cardinality = chunk.cardinality();
464                                sink_writer_metrics
465                                    .connector_sink_rows_received
466                                    .inc_by(cardinality as _);
467
468                                let chunk = self.stream_chunk_converter.convert_chunk(chunk)?;
469                                request_tx
470                                    .send_request(JniSinkWriterStreamRequest::Chunk {
471                                        epoch,
472                                        batch_id: chunk_id as u64,
473                                        chunk,
474                                    })
475                                    .instrument_await(span!(
476                                        "log_sinker_send_chunk (chunk {chunk_id})"
477                                    ))
478                                    .await?;
479                                prev_offset = Some(offset);
480                                sent_offset_queue
481                                    .push_back((TruncateOffset::Chunk { epoch, chunk_id }, None));
482                            }
483                            LogStoreReadItem::Barrier { is_checkpoint, .. } => {
484                                let offset = TruncateOffset::Barrier { epoch };
485                                if let Some(prev_offset) = &prev_offset {
486                                    prev_offset.check_next_offset(offset)?;
487                                }
488                                let start_time = if is_checkpoint {
489                                    let start_time = Instant::now();
490                                    request_tx
491                                        .barrier(epoch, true)
492                                        .instrument_await(span!(
493                                            "log_sinker_commit_checkpoint (epoch {epoch})"
494                                        ))
495                                        .await?;
496                                    Some(start_time)
497                                } else {
498                                    request_tx
499                                        .barrier(epoch, false)
500                                        .instrument_await(span!(
501                                            "log_sinker_send_barrier (epoch {epoch})"
502                                        ))
503                                        .await?;
504                                    None
505                                };
506                                prev_offset = Some(offset);
507                                sent_offset_queue
508                                    .push_back((TruncateOffset::Barrier { epoch }, start_time));
509                            }
510                        }
511                    }
512                }
513            }
514        };
515
516        select(pin!(poll_response_stream), pin!(poll_consume_log_and_sink))
517            .await
518            .factor_first()
519            .0
520    }
521}
522
523#[derive(Debug)]
524pub struct CoordinatedRemoteSink<R: RemoteSinkTrait> {
525    param: SinkParam,
526    _phantom: PhantomData<R>,
527}
528
529impl<R: RemoteSinkTrait> EnforceSecret for CoordinatedRemoteSink<R> {
530    const ENFORCE_SECRET_PROPERTIES: phf::Set<&'static str> = R::ENFORCE_SECRET_PROPERTIES;
531}
532
533impl<R: RemoteSinkTrait> TryFrom<SinkParam> for CoordinatedRemoteSink<R> {
534    type Error = SinkError;
535
536    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
537        Ok(Self {
538            param,
539            _phantom: PhantomData,
540        })
541    }
542}
543
544impl<R: RemoteSinkTrait> Sink for CoordinatedRemoteSink<R> {
545    type Coordinator = RemoteCoordinator;
546    type LogSinker = CoordinatedLogSinker<CoordinatedRemoteSinkWriter>;
547
548    const SINK_NAME: &'static str = R::SINK_NAME;
549
550    async fn validate(&self) -> Result<()> {
551        validate_remote_sink(&self.param, Self::SINK_NAME).await?;
552        Ok(())
553    }
554
555    async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
556        let metrics = SinkWriterMetrics::new(&writer_param);
557        CoordinatedLogSinker::new(
558            &writer_param,
559            self.param.clone(),
560            CoordinatedRemoteSinkWriter::new(self.param.clone(), metrics.clone()).await?,
561            NonZero::new(1).unwrap(),
562        )
563        .await
564    }
565
566    fn is_coordinated_sink(&self) -> bool {
567        true
568    }
569
570    async fn new_coordinator(&self, _db: DatabaseConnection) -> Result<Self::Coordinator> {
571        RemoteCoordinator::new::<R>(self.param.clone()).await
572    }
573}
574
575pub struct CoordinatedRemoteSinkWriter {
576    epoch: Option<u64>,
577    batch_id: u64,
578    stream_handle: SinkWriterStreamHandle<JniSinkWriterStreamRequest>,
579    metrics: SinkWriterMetrics,
580}
581
582impl CoordinatedRemoteSinkWriter {
583    pub async fn new(param: SinkParam, metrics: SinkWriterMetrics) -> Result<Self> {
584        let sink_proto = param.to_proto();
585        let stream_handle = EmbeddedConnectorClient::new()?
586            .start_sink_writer_stream(sink_proto.table_schema.clone(), sink_proto)
587            .await?;
588
589        Ok(Self {
590            epoch: None,
591            batch_id: 0,
592            stream_handle,
593            metrics,
594        })
595    }
596
597    #[cfg(test)]
598    fn for_test(
599        response_receiver: Receiver<ConnectorResult<SinkWriterStreamResponse>>,
600        request_sender: tokio::sync::mpsc::Sender<JniSinkWriterStreamRequest>,
601    ) -> CoordinatedRemoteSinkWriter {
602        use futures::StreamExt;
603
604        let stream_handle = SinkWriterStreamHandle::for_test(
605            request_sender,
606            ReceiverStream::new(response_receiver)
607                .map_err(RpcError::from)
608                .boxed(),
609        );
610
611        CoordinatedRemoteSinkWriter {
612            epoch: None,
613            batch_id: 0,
614            stream_handle,
615            metrics: SinkWriterMetrics::for_test(),
616        }
617    }
618}
619
620#[async_trait]
621impl SinkWriter for CoordinatedRemoteSinkWriter {
622    type CommitMetadata = Option<SinkMetadata>;
623
624    async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
625        let cardinality = chunk.cardinality();
626        self.metrics
627            .connector_sink_rows_received
628            .inc_by(cardinality as _);
629
630        let epoch = self.epoch.ok_or_else(|| {
631            SinkError::Remote(anyhow!("epoch has not been initialize, call `begin_epoch`"))
632        })?;
633        let batch_id = self.batch_id;
634        self.stream_handle
635            .request_sender
636            .send_request(JniSinkWriterStreamRequest::Chunk {
637                chunk,
638                epoch,
639                batch_id,
640            })
641            .await?;
642        self.batch_id += 1;
643        Ok(())
644    }
645
646    async fn begin_epoch(&mut self, epoch: u64) -> Result<()> {
647        self.epoch = Some(epoch);
648        Ok(())
649    }
650
651    async fn barrier(&mut self, is_checkpoint: bool) -> Result<Option<SinkMetadata>> {
652        let epoch = self.epoch.ok_or_else(|| {
653            SinkError::Remote(anyhow!("epoch has not been initialize, call `begin_epoch`"))
654        })?;
655        if is_checkpoint {
656            // TODO: add metrics to measure commit time
657            let rsp = self.stream_handle.commit(epoch).await?;
658            rsp.metadata
659                .ok_or_else(|| {
660                    SinkError::Remote(anyhow!(
661                        "get none metadata in commit response for coordinated sink writer"
662                    ))
663                })
664                .map(Some)
665        } else {
666            self.stream_handle.barrier(epoch).await?;
667            Ok(None)
668        }
669    }
670}
671
672pub struct RemoteCoordinator {
673    stream_handle: SinkCoordinatorStreamHandle,
674}
675
676impl RemoteCoordinator {
677    pub async fn new<R: RemoteSinkTrait>(param: SinkParam) -> Result<Self> {
678        let stream_handle = EmbeddedConnectorClient::new()?
679            .start_sink_coordinator_stream(param.clone())
680            .await?;
681
682        tracing::trace!("{:?} RemoteCoordinator started", R::SINK_NAME,);
683
684        Ok(RemoteCoordinator { stream_handle })
685    }
686}
687
688#[async_trait]
689impl SinkCommitCoordinator for RemoteCoordinator {
690    async fn init(&mut self, _subscriber: SinkCommittedEpochSubscriber) -> Result<Option<u64>> {
691        Ok(None)
692    }
693
694    async fn commit(&mut self, epoch: u64, metadata: Vec<SinkMetadata>) -> Result<()> {
695        Ok(self.stream_handle.commit(epoch, metadata).await?)
696    }
697}
698
699struct EmbeddedConnectorClient {
700    jvm: &'static JavaVM,
701}
702
703impl EmbeddedConnectorClient {
704    fn new() -> Result<Self> {
705        let jvm = JVM
706            .get_or_init()
707            .context("failed to create EmbeddedConnectorClient")?;
708        Ok(EmbeddedConnectorClient { jvm })
709    }
710
711    async fn start_sink_writer_stream(
712        &self,
713        payload_schema: Option<TableSchema>,
714        sink_proto: PbSinkParam,
715    ) -> Result<SinkWriterStreamHandle<JniSinkWriterStreamRequest>> {
716        let (handle, first_rsp) = SinkWriterStreamHandle::initialize(
717            SinkWriterStreamRequest {
718                request: Some(SinkRequest::Start(StartSink {
719                    sink_param: Some(sink_proto),
720                    payload_schema,
721                })),
722            },
723            |rx| async move {
724                let rx = self.start_jvm_worker_thread(
725                    gen_class_name!(com.risingwave.connector.JniSinkWriterHandler),
726                    "runJniSinkWriterThread",
727                    rx,
728                );
729                Ok(ReceiverStream::new(rx).map_err(RpcError::from))
730            },
731        )
732        .await?;
733
734        match first_rsp {
735            SinkWriterStreamResponse {
736                response: Some(sink_writer_stream_response::Response::Start(_)),
737            } => Ok(handle),
738            msg => Err(SinkError::Internal(anyhow!(
739                "should get start response but get {:?}",
740                msg
741            ))),
742        }
743    }
744
745    async fn start_sink_coordinator_stream(
746        &self,
747        param: SinkParam,
748    ) -> Result<SinkCoordinatorStreamHandle> {
749        let (handle, first_rsp) = SinkCoordinatorStreamHandle::initialize(
750            SinkCoordinatorStreamRequest {
751                request: Some(sink_coordinator_stream_request::Request::Start(
752                    StartCoordinator {
753                        param: Some(param.to_proto()),
754                    },
755                )),
756            },
757            |rx| async move {
758                let rx = self.start_jvm_worker_thread(
759                    gen_class_name!(com.risingwave.connector.JniSinkCoordinatorHandler),
760                    "runJniSinkCoordinatorThread",
761                    rx,
762                );
763                Ok(ReceiverStream::new(rx).map_err(RpcError::from))
764            },
765        )
766        .await?;
767
768        match first_rsp {
769            SinkCoordinatorStreamResponse {
770                response: Some(sink_coordinator_stream_response::Response::Start(_)),
771            } => Ok(handle),
772            msg => Err(SinkError::Internal(anyhow!(
773                "should get start response but get {:?}",
774                msg
775            ))),
776        }
777    }
778
779    fn start_jvm_worker_thread<REQ: Send + 'static, RSP: Send + 'static>(
780        &self,
781        class_name: &'static str,
782        method_name: &'static str,
783        mut request_rx: JniReceiverType<REQ>,
784    ) -> Receiver<std::result::Result<RSP, anyhow::Error>> {
785        let (mut response_tx, response_rx): (JniSenderType<RSP>, _) =
786            mpsc::channel(DEFAULT_BUFFER_SIZE);
787
788        let jvm = self.jvm;
789        std::thread::spawn(move || {
790            let result = execute_with_jni_env(jvm, |env| {
791                let result = call_static_method!(
792                    env,
793                    class_name,
794                    method_name,
795                    {{void}, {long requestRx, long responseTx}},
796                    &mut request_rx as *mut JniReceiverType<REQ>,
797                    &mut response_tx as *mut JniSenderType<RSP>
798                );
799
800                match result {
801                    Ok(_) => {
802                        tracing::debug!("end of jni call {}::{}", class_name, method_name);
803                    }
804                    Err(e) => {
805                        tracing::error!(error = %e.as_report(), "jni call error");
806                    }
807                };
808
809                Ok(())
810            });
811
812            if let Err(e) = result {
813                let _ = response_tx.blocking_send(Err(e));
814            }
815        });
816        response_rx
817    }
818}
819
820#[cfg(test)]
821mod test {
822    use std::time::Duration;
823
824    use risingwave_common::array::StreamChunk;
825    use risingwave_common::test_prelude::StreamChunkTestExt;
826    use risingwave_jni_core::JniSinkWriterStreamRequest;
827    use risingwave_pb::connector_service::sink_writer_stream_request::{Barrier, Request};
828    use risingwave_pb::connector_service::sink_writer_stream_response::{CommitResponse, Response};
829    use risingwave_pb::connector_service::{SinkWriterStreamRequest, SinkWriterStreamResponse};
830    use tokio::sync::mpsc;
831
832    use crate::sink::SinkWriter;
833    use crate::sink::remote::CoordinatedRemoteSinkWriter;
834
835    #[tokio::test]
836    async fn test_epoch_check() {
837        let (request_sender, mut request_recv) = mpsc::channel(16);
838        let (_, resp_recv) = mpsc::channel(16);
839
840        let mut sink = CoordinatedRemoteSinkWriter::for_test(resp_recv, request_sender);
841        let chunk = StreamChunk::from_pretty(
842            " i T
843            + 1 Ripper
844        ",
845        );
846
847        // test epoch check
848        assert!(
849            tokio::time::timeout(Duration::from_secs(10), sink.barrier(true))
850                .await
851                .expect("test failed: should not commit without epoch")
852                .is_err(),
853            "test failed: no epoch check for commit()"
854        );
855        assert!(
856            request_recv.try_recv().is_err(),
857            "test failed: unchecked epoch before request"
858        );
859
860        assert!(
861            tokio::time::timeout(Duration::from_secs(1), sink.write_batch(chunk))
862                .await
863                .expect("test failed: should not write without epoch")
864                .is_err(),
865            "test failed: no epoch check for write_batch()"
866        );
867        assert!(
868            request_recv.try_recv().is_err(),
869            "test failed: unchecked epoch before request"
870        );
871    }
872
873    #[tokio::test]
874    async fn test_remote_sink() {
875        let (request_sender, mut request_receiver) = mpsc::channel(16);
876        let (response_sender, response_receiver) = mpsc::channel(16);
877        let mut sink = CoordinatedRemoteSinkWriter::for_test(response_receiver, request_sender);
878
879        let chunk_a = StreamChunk::from_pretty(
880            " i T
881            + 1 Alice
882            + 2 Bob
883            + 3 Clare
884        ",
885        );
886        let chunk_b = StreamChunk::from_pretty(
887            " i T
888            + 4 David
889            + 5 Eve
890            + 6 Frank
891        ",
892        );
893
894        // test write batch
895        sink.begin_epoch(2022).await.unwrap();
896        assert_eq!(sink.epoch, Some(2022));
897
898        sink.write_batch(chunk_a.clone()).await.unwrap();
899        assert_eq!(sink.epoch, Some(2022));
900        assert_eq!(sink.batch_id, 1);
901        match request_receiver.recv().await.unwrap() {
902            JniSinkWriterStreamRequest::Chunk {
903                epoch,
904                batch_id,
905                chunk,
906            } => {
907                assert_eq!(epoch, 2022);
908                assert_eq!(batch_id, 0);
909                assert_eq!(chunk, chunk_a);
910            }
911            _ => panic!("test failed: failed to construct write request"),
912        }
913
914        // test commit
915        response_sender
916            .send(Ok(SinkWriterStreamResponse {
917                response: Some(Response::Commit(CommitResponse {
918                    epoch: 2022,
919                    metadata: None,
920                })),
921            }))
922            .await
923            .expect("test failed: failed to sync epoch");
924        sink.barrier(false).await.unwrap();
925        let commit_request = request_receiver.recv().await.unwrap();
926        match commit_request {
927            JniSinkWriterStreamRequest::PbRequest(SinkWriterStreamRequest {
928                request:
929                    Some(Request::Barrier(Barrier {
930                        epoch,
931                        is_checkpoint: false,
932                    })),
933            }) => {
934                assert_eq!(epoch, 2022);
935            }
936            _ => panic!("test failed: failed to construct sync request "),
937        };
938
939        // begin another epoch
940        sink.begin_epoch(2023).await.unwrap();
941        assert_eq!(sink.epoch, Some(2023));
942
943        // test another write
944        sink.write_batch(chunk_b.clone()).await.unwrap();
945        assert_eq!(sink.epoch, Some(2023));
946        assert_eq!(sink.batch_id, 2);
947        match request_receiver.recv().await.unwrap() {
948            JniSinkWriterStreamRequest::Chunk {
949                epoch,
950                batch_id,
951                chunk,
952            } => {
953                assert_eq!(epoch, 2023);
954                assert_eq!(batch_id, 1);
955                assert_eq!(chunk, chunk_b);
956            }
957            _ => panic!("test failed: failed to construct write request"),
958        }
959    }
960}