1use 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 { 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 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 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(()), |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 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 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 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 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 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 sink.begin_epoch(2023).await.unwrap();
941 assert_eq!(sink.epoch, Some(2023));
942
943 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}