Skip to content

Commit

Permalink
refine sink await-tree
Browse files Browse the repository at this point in the history 
Signed-off-by: Bugen Zhao <[email protected]>
  • Loading branch information
@BugenZhao
BugenZhao committed Jun 13, 2024
1 parent 2c6f7c3 commit 40eb907
Show file tree
Hide file tree
Showing 8 changed files with 49 additions and 40 deletions.
3 changes: 2 additions & 1 deletion src/connector/src/sink/decouple_checkpoint_log_sink.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::convert::Infallible;
use std::num::NonZeroU64;
use std::time::Instant;

Expand Down Expand Up @@ -48,7 +49,7 @@ impl<W> DecoupleCheckpointLogSinkerOf<W> {

#[async_trait]
impl<W: SinkWriter<CommitMetadata = ()>> LogSinker for DecoupleCheckpointLogSinkerOf<W> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<()> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<Infallible> {
let mut sink_writer = self.writer;
let sink_metrics = self.sink_metrics;
#[derive(Debug)]
Expand Down
31 changes: 18 additions & 13 deletions src/connector/src/sink/log_store.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ use std::sync::Arc;
use std::task::Poll;
use std::time::Instant;

use await_tree::InstrumentAwait;
use futures::{TryFuture, TryFutureExt};
use risingwave_common::array::StreamChunk;
use risingwave_common::bail;
Expand Down Expand Up @@ -282,21 +283,25 @@ impl<R: LogReader> MonitoredLogReader<R> {

impl<R: LogReader> LogReader for MonitoredLogReader<R> {
async fn init(&mut self) -> LogStoreResult<()> {
self.inner.init().await
self.inner.init().instrument_await("log_reader_init").await
}

async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
self.inner.next_item().await.inspect(|(epoch, item)| {
if self.read_epoch != *epoch {
self.read_epoch = *epoch;
self.metrics.log_store_latest_read_epoch.set(*epoch as _);
}
if let LogStoreReadItem::StreamChunk { chunk, .. } = item {
self.metrics
.log_store_read_rows
.inc_by(chunk.cardinality() as _);
}
})
self.inner
.next_item()
.instrument_await("log_reader_next_item")
.await
.inspect(|(epoch, item)| {
if self.read_epoch != *epoch {
self.read_epoch = *epoch;
self.metrics.log_store_latest_read_epoch.set(*epoch as _);
}
if let LogStoreReadItem::StreamChunk { chunk, .. } = item {
self.metrics
.log_store_read_rows
.inc_by(chunk.cardinality() as _);
}
})
}

fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
Expand All @@ -306,7 +311,7 @@ impl<R: LogReader> LogReader for MonitoredLogReader<R> {
fn rewind(
&mut self,
) -> impl Future<Output = LogStoreResult<(bool, Option<Bitmap>)>> + Send + '_ {
self.inner.rewind()
self.inner.rewind().instrument_await("log_reader_rewind")
}
}

Expand Down
3 changes: 2 additions & 1 deletion src/connector/src/sink/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ pub mod utils;
pub mod writer;

use std::collections::BTreeMap;
use std::convert::Infallible;
use std::future::Future;

use ::clickhouse::error::Error as ClickHouseError;
Expand Down Expand Up @@ -386,7 +387,7 @@ impl<R: LogReader> SinkLogReader for R {

#[async_trait]
pub trait LogSinker: 'static {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<()>;
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<Infallible>;
}

#[async_trait]
Expand Down
26 changes: 9 additions & 17 deletions src/connector/src/sink/remote.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.

use std::collections::{BTreeMap, VecDeque};
use std::convert::Infallible;
use std::marker::PhantomData;
use std::ops::Deref;
use std::pin::pin;
Expand Down Expand Up @@ -301,7 +302,7 @@ impl RemoteLogSinker {

#[async_trait]
impl LogSinker for RemoteLogSinker {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<()> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<Infallible> {
let mut request_tx = self.request_sender;
let mut response_err_stream_rx = self.response_stream;
let sink_metrics = self.sink_metrics;
Expand All @@ -313,7 +314,7 @@ impl LogSinker for RemoteLogSinker {
let result = response_err_stream_rx
.stream
.try_next()
.instrument_await("Wait Response Stream")
.instrument_await("log_sinker_wait_next_response")
.await;
match result {
Ok(Some(response)) => {
Expand Down Expand Up @@ -368,20 +369,12 @@ impl LogSinker for RemoteLogSinker {
let mut sent_offset_queue: VecDeque<(TruncateOffset, Option<Instant>)> =
VecDeque::new();

let mut curr_epoch = 0;

loop {
let either_result: futures::future::Either<
Option<SinkWriterStreamResponse>,
LogStoreResult<(u64, LogStoreReadItem)>,
> = drop_either_future(
select(
pin!(response_rx.recv()),
pin!(log_reader
.next_item()
.instrument_await(format!("Wait Next Item: {}", curr_epoch))),
)
.await,
select(pin!(response_rx.recv()), pin!(log_reader.next_item())).await,
);
match either_result {
futures::future::Either::Left(opt) => {
Expand Down Expand Up @@ -435,7 +428,6 @@ impl LogSinker for RemoteLogSinker {
}
futures::future::Either::Right(result) => {
let (epoch, item): (u64, LogStoreReadItem) = result?;
curr_epoch = epoch;

match item {
LogStoreReadItem::StreamChunk { chunk, chunk_id } => {
Expand All @@ -456,8 +448,7 @@ impl LogSinker for RemoteLogSinker {
chunk,
})
.instrument_await(format!(
"Send Chunk Request: {} {}",
curr_epoch, chunk_id
"log_sinker_send_chunk (chunk {chunk_id})"
))
.await?;
prev_offset = Some(offset);
Expand All @@ -473,15 +464,16 @@ impl LogSinker for RemoteLogSinker {
let start_time = Instant::now();
request_tx
.barrier(epoch, true)
.instrument_await(format!("Commit: {}", curr_epoch))
.instrument_await(format!(
"log_sinker_commit_checkpoint (epoch {epoch})"
))
.await?;
Some(start_time)
} else {
request_tx
.barrier(epoch, false)
.instrument_await(format!(
"Send Barrier Request: {}",
curr_epoch
"log_sinker_send_barrier (epoch {epoch})"
))
.await?;
None
Expand Down
3 changes: 2 additions & 1 deletion src/connector/src/sink/trivial.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::convert::Infallible;
use std::marker::PhantomData;

use async_trait::async_trait;
Expand Down Expand Up @@ -75,7 +76,7 @@ impl<T: TrivialSinkName> Sink for TrivialSink<T> {

#[async_trait]
impl<T: TrivialSinkName> LogSinker for TrivialSink<T> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<()> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<Infallible> {
loop {
let (epoch, item) = log_reader.next_item().await?;
match item {
Expand Down
8 changes: 6 additions & 2 deletions src/connector/src/sink/writer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::convert::Infallible;
use std::future::{Future, Ready};
use std::pin::pin;
use std::sync::Arc;
Expand Down Expand Up @@ -126,7 +127,7 @@ impl<W> LogSinkerOf<W> {

#[async_trait]
impl<W: SinkWriter<CommitMetadata = ()>> LogSinker for LogSinkerOf<W> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<()> {
async fn consume_log_and_sink(self, log_reader: &mut impl SinkLogReader) -> Result<Infallible> {
let mut sink_writer = self.writer;
let sink_metrics = self.sink_metrics;
#[derive(Debug)]
Expand Down Expand Up @@ -242,7 +243,10 @@ impl<W: AsyncTruncateSinkWriter> AsyncTruncateLogSinkerOf<W> {

#[async_trait]
impl<W: AsyncTruncateSinkWriter> LogSinker for AsyncTruncateLogSinkerOf<W> {
async fn consume_log_and_sink(mut self, log_reader: &mut impl SinkLogReader) -> Result<()> {
async fn consume_log_and_sink(
mut self,
log_reader: &mut impl SinkLogReader,
) -> Result<Infallible> {
loop {
let select_result = drop_either_future(
select(
Expand Down
5 changes: 5 additions & 0 deletions src/stream/src/common/log_store_impl/in_mem.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
use std::sync::Arc;

use anyhow::{anyhow, Context};
use await_tree::InstrumentAwait;
use risingwave_common::array::StreamChunk;
use risingwave_common::buffer::Bitmap;
use risingwave_common::util::epoch::{EpochExt, EpochPair, INVALID_EPOCH};
Expand Down Expand Up @@ -260,6 +261,7 @@ impl LogWriter for BoundedInMemLogStoreWriter {
async fn write_chunk(&mut self, chunk: StreamChunk) -> LogStoreResult<()> {
self.item_tx
.send(InMemLogStoreItem::StreamChunk(chunk))
.instrument_await("in_mem_send_item_chunk")
.await
.map_err(|_| anyhow!("unable to send stream chunk"))?;
Ok(())
Expand All @@ -275,6 +277,7 @@ impl LogWriter for BoundedInMemLogStoreWriter {
next_epoch,
is_checkpoint,
})
.instrument_await("in_mem_send_item_barrier")
.await
.map_err(|_| anyhow!("unable to send barrier"))?;

Expand All @@ -287,6 +290,7 @@ impl LogWriter for BoundedInMemLogStoreWriter {
let truncated_epoch = self
.truncated_epoch_rx
.recv()
.instrument_await("in_mem_recv_truncated_epoch")
.await
.ok_or_else(|| anyhow!("cannot get truncated epoch"))?;
assert_eq!(truncated_epoch, prev_epoch);
Expand All @@ -298,6 +302,7 @@ impl LogWriter for BoundedInMemLogStoreWriter {
async fn update_vnode_bitmap(&mut self, new_vnodes: Arc<Bitmap>) -> LogStoreResult<()> {
self.item_tx
.send(InMemLogStoreItem::UpdateVnodeBitmap(new_vnodes))
.instrument_await("in_mem_send_item_vnode_bitmap")
.await
.map_err(|_| anyhow!("unable to send vnode bitmap"))
}
Expand Down
10 changes: 5 additions & 5 deletions src/stream/src/executor/sink.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::convert::Infallible;
use std::mem;

use anyhow::anyhow;
Expand Down Expand Up @@ -217,10 +218,9 @@ impl<F: LogStoreFactory> SinkExecutor<F> {
self.sink_writer_param,
self.actor_context,
)
.instrument_await(format!(
"Consume Log: sink_id: {} actor_id: {}, executor_id: {}",
sink_id, actor_id, executor_id,
));
.instrument_await(format!("consume_log (sink_id {sink_id})"))
.map_ok(|f| match f {}); // unify return type to `Message`

// TODO: may try to remove the boxed
select(consume_log_stream.into_stream(), write_log_stream).boxed()
})
Expand Down Expand Up @@ -403,7 +403,7 @@ impl<F: LogStoreFactory> SinkExecutor<F> {
sink_param: SinkParam,
mut sink_writer_param: SinkWriterParam,
actor_context: ActorContextRef,
) -> StreamExecutorResult<Message> {
) -> StreamExecutorResult<Infallible> {
let metrics = sink_writer_param.sink_metrics.clone();

let visible_columns = columns
Expand Down

0 comments on commit 40eb907

Please sign in to comment.