risingwave_connector/sink/
log_store.rs

1// Copyright 2023 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::cmp::Ordering;
16use std::collections::VecDeque;
17use std::fmt::Debug;
18use std::future::{Future, pending, poll_fn};
19use std::pin::pin;
20use std::sync::Arc;
21use std::task::Poll;
22use std::time::Instant;
23
24use await_tree::InstrumentAwait;
25use futures::future::BoxFuture;
26use futures::{TryFuture, TryFutureExt};
27use risingwave_common::array::StreamChunk;
28use risingwave_common::bail;
29use risingwave_common::bitmap::Bitmap;
30use risingwave_common::metrics::{LabelGuardedIntCounter, LabelGuardedIntGauge};
31use risingwave_common::util::epoch::{EpochPair, INVALID_EPOCH};
32use risingwave_common_estimate_size::EstimateSize;
33use risingwave_common_rate_limit::{RateLimit, RateLimiter};
34use risingwave_pb::stream_plan::PbSinkSchemaChange;
35use tokio::select;
36use tokio::sync::mpsc::UnboundedReceiver;
37
38pub type LogStoreResult<T> = Result<T, anyhow::Error>;
39pub type ChunkId = usize;
40
41#[derive(Debug, PartialEq, Eq, Copy, Clone)]
42pub enum TruncateOffset {
43    Chunk { epoch: u64, chunk_id: ChunkId },
44    Barrier { epoch: u64 },
45}
46
47impl Ord for TruncateOffset {
48    fn cmp(&self, other: &Self) -> Ordering {
49        let extract = |offset: &TruncateOffset| match offset {
50            TruncateOffset::Chunk { epoch, chunk_id } => (*epoch, *chunk_id),
51            TruncateOffset::Barrier { epoch } => (*epoch, usize::MAX),
52        };
53        extract(self).cmp(&extract(other))
54    }
55}
56
57impl PartialOrd for TruncateOffset {
58    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
59        Some(self.cmp(other))
60    }
61}
62
63impl TruncateOffset {
64    pub fn next_chunk_id(&self) -> ChunkId {
65        match self {
66            TruncateOffset::Chunk { chunk_id, .. } => chunk_id + 1,
67            TruncateOffset::Barrier { .. } => 0,
68        }
69    }
70
71    pub fn epoch(&self) -> u64 {
72        match self {
73            TruncateOffset::Chunk { epoch, .. } | TruncateOffset::Barrier { epoch } => *epoch,
74        }
75    }
76
77    pub fn check_next_offset(&self, next_offset: TruncateOffset) -> LogStoreResult<()> {
78        if *self >= next_offset {
79            bail!(
80                "next offset {:?} should be later than current offset {:?}",
81                next_offset,
82                self
83            )
84        } else {
85            Ok(())
86        }
87    }
88
89    pub fn check_next_item_epoch(&self, epoch: u64) -> LogStoreResult<()> {
90        match self {
91            TruncateOffset::Chunk {
92                epoch: offset_epoch,
93                ..
94            } => {
95                if epoch != *offset_epoch {
96                    bail!(
97                        "new item epoch {} does not match current chunk offset epoch {}",
98                        epoch,
99                        offset_epoch
100                    );
101                }
102            }
103            TruncateOffset::Barrier {
104                epoch: offset_epoch,
105            } => {
106                if epoch <= *offset_epoch {
107                    bail!(
108                        "new item epoch {} does not exceed barrier offset epoch {}",
109                        epoch,
110                        offset_epoch
111                    );
112                }
113            }
114        }
115        Ok(())
116    }
117}
118
119#[derive(Debug)]
120pub enum LogStoreReadItem {
121    StreamChunk {
122        chunk: StreamChunk,
123        chunk_id: ChunkId,
124    },
125    Barrier {
126        is_checkpoint: bool,
127        new_vnode_bitmap: Option<Arc<Bitmap>>,
128        is_stop: bool,
129        schema_change: Option<PbSinkSchemaChange>,
130    },
131}
132
133pub trait LogWriterPostFlushCurrentEpochFn<'a> = FnOnce() -> BoxFuture<'a, LogStoreResult<()>>;
134
135#[must_use]
136pub struct LogWriterPostFlushCurrentEpoch<'a>(
137    Box<dyn LogWriterPostFlushCurrentEpochFn<'a> + Send + 'a>,
138);
139
140impl<'a> LogWriterPostFlushCurrentEpoch<'a> {
141    pub fn new(f: impl LogWriterPostFlushCurrentEpochFn<'a> + Send + 'a) -> Self {
142        Self(Box::new(f))
143    }
144
145    pub async fn post_yield_barrier(self) -> LogStoreResult<()> {
146        self.0().await
147    }
148}
149
150pub struct FlushCurrentEpochOptions {
151    pub is_checkpoint: bool,
152    pub new_vnode_bitmap: Option<Arc<Bitmap>>,
153    pub is_stop: bool,
154    pub schema_change: Option<PbSinkSchemaChange>,
155    pub wait_log_store_flush: bool,
156}
157
158pub trait LogWriter: Send {
159    /// Initialize the log writer with an epoch
160    fn init(
161        &mut self,
162        epoch: EpochPair,
163        pause_read_on_bootstrap: bool,
164    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_;
165
166    /// Write a stream chunk to the log writer
167    fn write_chunk(
168        &mut self,
169        chunk: StreamChunk,
170    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_;
171
172    /// Mark current epoch as finished and sealed, and flush the unconsumed log data.
173    fn flush_current_epoch(
174        &mut self,
175        next_epoch: u64,
176        options: FlushCurrentEpochOptions,
177    ) -> impl Future<Output = LogStoreResult<LogWriterPostFlushCurrentEpoch<'_>>> + Send + '_;
178
179    fn pause(&mut self) -> LogStoreResult<()>;
180
181    fn resume(&mut self) -> LogStoreResult<()>;
182}
183
184pub trait LogReader: Send + Sized + 'static {
185    /// Initialize the log reader. Usually function as waiting for log writer to be initialized.
186    fn init(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_;
187
188    /// Consume log store from given `start_offset` or aligned start offset recorded previously.
189    fn start_from(
190        &mut self,
191        start_offset: Option<u64>,
192    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_;
193
194    /// Emit the next item.
195    ///
196    /// The implementation should ensure that the future is cancellation safe.
197    fn next_item(
198        &mut self,
199    ) -> impl Future<Output = LogStoreResult<(u64, LogStoreReadItem)>> + Send + '_;
200
201    /// Mark that all items emitted so far have been consumed and it is safe to truncate the log
202    /// from the current offset.
203    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()>;
204
205    /// Reset the log reader to after the latest truncate offset
206    ///
207    /// The return flag means whether the log store support rewind
208    fn rewind(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_;
209}
210
211pub trait LogStoreFactory: Send + 'static {
212    const ALLOW_REWIND: bool;
213    const REBUILD_SINK_ON_UPDATE_VNODE_BITMAP: bool;
214    type Reader: LogReader;
215    type Writer: LogWriter;
216
217    fn build(self) -> impl Future<Output = (Self::Reader, Self::Writer)> + Send;
218}
219
220pub struct TransformChunkLogReader<F: Fn(StreamChunk) -> StreamChunk, R: LogReader> {
221    f: F,
222    inner: R,
223}
224
225pub struct TruncateBarrierLogReader<R: LogReader> {
226    inner: R,
227    pending_barriers: VecDeque<u64>,
228}
229
230impl<R: LogReader> TruncateBarrierLogReader<R> {
231    pub fn new(inner: R) -> Self {
232        Self {
233            inner,
234            pending_barriers: VecDeque::new(),
235        }
236    }
237}
238
239impl<F: Fn(StreamChunk) -> StreamChunk + Send + 'static, R: LogReader> LogReader
240    for TransformChunkLogReader<F, R>
241{
242    fn init(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
243        self.inner.init()
244    }
245
246    async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
247        let (epoch, item) = self.inner.next_item().await?;
248        let item = match item {
249            LogStoreReadItem::StreamChunk { chunk, chunk_id } => LogStoreReadItem::StreamChunk {
250                chunk: (self.f)(chunk),
251                chunk_id,
252            },
253            other => other,
254        };
255        Ok((epoch, item))
256    }
257
258    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
259        self.inner.truncate(offset)
260    }
261
262    fn rewind(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
263        self.inner.rewind()
264    }
265
266    fn start_from(
267        &mut self,
268        start_offset: Option<u64>,
269    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
270        self.inner.start_from(start_offset)
271    }
272}
273
274impl<R: LogReader> LogReader for TruncateBarrierLogReader<R> {
275    async fn init(&mut self) -> LogStoreResult<()> {
276        self.pending_barriers.clear();
277        self.inner.init().await
278    }
279
280    async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
281        let (epoch, item) = self.inner.next_item().await?;
282        if matches!(item, LogStoreReadItem::Barrier { .. }) {
283            self.pending_barriers.push_back(epoch);
284        }
285        Ok((epoch, item))
286    }
287
288    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
289        while let Some(front_epoch) = self.pending_barriers.front().copied() {
290            match (TruncateOffset::Barrier { epoch: front_epoch }).cmp(&offset) {
291                Ordering::Less => {
292                    self.inner
293                        .truncate(TruncateOffset::Barrier { epoch: front_epoch })?;
294                    self.pending_barriers.pop_front();
295                }
296                Ordering::Equal => {
297                    self.inner.truncate(offset)?;
298                    self.pending_barriers.pop_front();
299                    return Ok(());
300                }
301                Ordering::Greater => {
302                    self.inner.truncate(offset)?;
303                    return Ok(());
304                }
305            }
306        }
307        self.inner.truncate(offset)?;
308        Ok(())
309    }
310
311    async fn rewind(&mut self) -> LogStoreResult<()> {
312        self.pending_barriers.clear();
313        self.inner.rewind().await
314    }
315
316    async fn start_from(&mut self, start_offset: Option<u64>) -> LogStoreResult<()> {
317        self.pending_barriers.clear();
318        self.inner.start_from(start_offset).await
319    }
320}
321
322pub struct BackpressureMonitoredLogReader<R: LogReader> {
323    inner: R,
324    /// Start time to wait for new future after poll ready
325    wait_new_future_start_time: Option<Instant>,
326    wait_new_future_duration_ns: LabelGuardedIntCounter,
327}
328
329impl<R: LogReader> BackpressureMonitoredLogReader<R> {
330    fn new(inner: R, wait_new_future_duration_ns: LabelGuardedIntCounter) -> Self {
331        Self {
332            inner,
333            wait_new_future_start_time: None,
334            wait_new_future_duration_ns,
335        }
336    }
337}
338
339impl<R: LogReader> LogReader for BackpressureMonitoredLogReader<R> {
340    fn init(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
341        self.wait_new_future_start_time = None;
342        self.inner.init()
343    }
344
345    fn next_item(
346        &mut self,
347    ) -> impl Future<Output = LogStoreResult<(u64, LogStoreReadItem)>> + Send + '_ {
348        if let Some(start_time) = self.wait_new_future_start_time.take() {
349            self.wait_new_future_duration_ns
350                .inc_by(start_time.elapsed().as_nanos() as _);
351        }
352        self.inner.next_item().inspect_ok(|_| {
353            // Set start time when return ready
354            self.wait_new_future_start_time = Some(Instant::now());
355        })
356    }
357
358    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
359        self.inner.truncate(offset)
360    }
361
362    fn rewind(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
363        self.inner.rewind().inspect_ok(|_| {
364            self.wait_new_future_start_time = None;
365        })
366    }
367
368    fn start_from(
369        &mut self,
370        start_offset: Option<u64>,
371    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
372        self.inner.start_from(start_offset)
373    }
374}
375
376pub struct MonitoredLogReader<R: LogReader> {
377    inner: R,
378    read_epoch: u64,
379    metrics: LogReaderMetrics,
380}
381
382pub struct LogReaderMetrics {
383    pub log_store_latest_read_epoch: LabelGuardedIntGauge,
384    pub log_store_read_rows: LabelGuardedIntCounter,
385    pub log_store_read_bytes: LabelGuardedIntCounter,
386    pub log_store_reader_wait_new_future_duration_ns: LabelGuardedIntCounter,
387}
388
389impl<R: LogReader> MonitoredLogReader<R> {
390    pub fn new(inner: R, metrics: LogReaderMetrics) -> Self {
391        Self {
392            inner,
393            read_epoch: INVALID_EPOCH,
394            metrics,
395        }
396    }
397}
398
399impl<R: LogReader> LogReader for MonitoredLogReader<R> {
400    async fn init(&mut self) -> LogStoreResult<()> {
401        self.inner.init().instrument_await("log_reader_init").await
402    }
403
404    async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
405        self.inner
406            .next_item()
407            .instrument_await("log_reader_next_item")
408            .await
409            .inspect(|(epoch, item)| {
410                if self.read_epoch != *epoch {
411                    self.read_epoch = *epoch;
412                    self.metrics.log_store_latest_read_epoch.set(*epoch as _);
413                }
414                if let LogStoreReadItem::StreamChunk { chunk, .. } = item {
415                    self.metrics
416                        .log_store_read_rows
417                        .inc_by(chunk.cardinality() as _);
418                    self.metrics
419                        .log_store_read_bytes
420                        .inc_by(chunk.estimated_size() as u64);
421                }
422            })
423    }
424
425    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
426        self.inner.truncate(offset)
427    }
428
429    fn rewind(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
430        self.inner.rewind().instrument_await("log_reader_rewind")
431    }
432
433    fn start_from(
434        &mut self,
435        start_offset: Option<u64>,
436    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
437        self.inner.start_from(start_offset)
438    }
439}
440
441#[derive(Copy, Clone, PartialOrd, PartialEq, Debug)]
442struct UpstreamChunkOffset(TruncateOffset);
443#[derive(Copy, Clone, PartialOrd, PartialEq)]
444struct DownstreamChunkOffset(TruncateOffset);
445
446struct RateLimitedLogReaderCore<R: LogReader> {
447    inner: R,
448    consuming_chunk: Option<(
449        UpstreamChunkOffset,
450        // Newer items at the front, push_front, pop_back
451        VecDeque<DownstreamChunkOffset>,
452        Vec<StreamChunk>, // split chunks
453    )>,
454    // Newer items at the front, push_front, pop_back
455    consumed_offset_queue: VecDeque<(UpstreamChunkOffset, VecDeque<DownstreamChunkOffset>)>,
456    next_chunk_id: usize,
457    rate_limiter: RateLimiter,
458}
459
460pub struct RateLimitedLogReader<R: LogReader> {
461    core: RateLimitedLogReaderCore<R>,
462    control_rx: UnboundedReceiver<RateLimit>,
463}
464
465impl<R: LogReader> RateLimitedLogReader<R> {
466    pub fn new(inner: R, control_rx: UnboundedReceiver<RateLimit>) -> Self {
467        Self {
468            core: RateLimitedLogReaderCore {
469                inner,
470                consuming_chunk: None,
471                consumed_offset_queue: VecDeque::new(),
472                next_chunk_id: 0,
473                rate_limiter: RateLimiter::new(RateLimit::Disabled),
474            },
475            control_rx,
476        }
477    }
478}
479
480impl<R: LogReader> RateLimitedLogReaderCore<R> {
481    fn peek_next_pending_chunk(&self) -> Option<&StreamChunk> {
482        self.consuming_chunk
483            .as_ref()
484            .and_then(|(_, _, chunk)| chunk.last())
485    }
486
487    fn consume_next_pending_chunk(&mut self) -> Option<(u64, StreamChunk, ChunkId)> {
488        let Some((upstream_offset, consumed_offsets, pending_chunk)) = &mut self.consuming_chunk
489        else {
490            return None;
491        };
492        let epoch = upstream_offset.0.epoch();
493
494        let item = pending_chunk.pop().map(|chunk| {
495            let chunk_id = self.next_chunk_id;
496            self.next_chunk_id += 1;
497            consumed_offsets.push_front(DownstreamChunkOffset(TruncateOffset::Chunk {
498                epoch,
499                chunk_id,
500            }));
501            (epoch, chunk, chunk_id)
502        });
503        if pending_chunk.is_empty() {
504            let (upstream_offset, consumed_offsets, _) =
505                self.consuming_chunk.take().expect("checked some");
506            self.consumed_offset_queue
507                .push_front((upstream_offset, consumed_offsets));
508        }
509        item
510    }
511
512    fn consume_single_upstream_item(
513        &mut self,
514        epoch: u64,
515        mut item: LogStoreReadItem,
516    ) -> (u64, LogStoreReadItem) {
517        assert!(self.consuming_chunk.is_none());
518        let (upstream_offset, downstream_offset) = match &mut item {
519            LogStoreReadItem::StreamChunk { chunk_id, .. } => {
520                let upstream_chunk_id = *chunk_id;
521                let downstream_chunk_id = self.next_chunk_id;
522                self.next_chunk_id += 1;
523                *chunk_id = downstream_chunk_id;
524                (
525                    UpstreamChunkOffset(TruncateOffset::Chunk {
526                        epoch,
527                        chunk_id: upstream_chunk_id,
528                    }),
529                    DownstreamChunkOffset(TruncateOffset::Chunk {
530                        epoch,
531                        chunk_id: downstream_chunk_id,
532                    }),
533                )
534            }
535            LogStoreReadItem::Barrier { .. } => (
536                UpstreamChunkOffset(TruncateOffset::Barrier { epoch }),
537                DownstreamChunkOffset(TruncateOffset::Barrier { epoch }),
538            ),
539        };
540        self.consumed_offset_queue
541            .push_front((upstream_offset, VecDeque::from_iter([downstream_offset])));
542        (epoch, item)
543    }
544
545    async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
546        match self.rate_limiter.rate_limit() {
547            RateLimit::Pause => pending().await,
548            RateLimit::Disabled => {
549                if let Some((epoch, chunk, chunk_id)) = self.consume_next_pending_chunk() {
550                    Ok((epoch, LogStoreReadItem::StreamChunk { chunk, chunk_id }))
551                } else {
552                    let (epoch, item) = self.inner.next_item().await?;
553                    Ok(self.consume_single_upstream_item(epoch, item))
554                }
555            }
556            RateLimit::Fixed(limit) => {
557                if self.peek_next_pending_chunk().is_none() {
558                    let (epoch, item) = self.inner.next_item().await?;
559                    match item {
560                        LogStoreReadItem::StreamChunk { chunk, chunk_id } => {
561                            let chunks = if chunk.rate_limit_permits() < limit.get() {
562                                vec![chunk]
563                            } else {
564                                let mut chunks = chunk.split(limit.get() as _);
565                                // reverse to make the first chunk to be popped first
566                                chunks.reverse();
567                                chunks
568                            };
569                            assert!(!chunks.is_empty());
570
571                            assert!(
572                                self.consuming_chunk
573                                    .replace((
574                                        UpstreamChunkOffset(TruncateOffset::Chunk {
575                                            epoch,
576                                            chunk_id
577                                        }),
578                                        VecDeque::new(),
579                                        chunks,
580                                    ))
581                                    .is_none()
582                            );
583                        }
584                        item @ LogStoreReadItem::Barrier { .. } => {
585                            return Ok(self.consume_single_upstream_item(epoch, item));
586                        }
587                    };
588                }
589                let chunk = self.peek_next_pending_chunk().expect("must Some");
590                self.rate_limiter.wait_chunk(chunk).await;
591                let (epoch, chunk, chunk_id) =
592                    self.consume_next_pending_chunk().expect("must Some");
593                Ok((epoch, LogStoreReadItem::StreamChunk { chunk, chunk_id }))
594            }
595        }
596    }
597}
598
599impl<R: LogReader> LogReader for RateLimitedLogReader<R> {
600    async fn init(&mut self) -> LogStoreResult<()> {
601        self.core.inner.init().await
602    }
603
604    async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
605        loop {
606            select! {
607                biased;
608                recv = pin!(self.control_rx.recv()) => {
609                    let new_rate_limit = match recv {
610                        Some(limit) => limit,
611                        None => bail!("rate limit control channel closed"),
612                    };
613                    let old_rate_limit = self.core.rate_limiter.update(new_rate_limit);
614                    let paused = matches!(new_rate_limit, RateLimit::Pause);
615                    tracing::info!("rate limit changed from {:?} to {:?}, paused = {paused}", old_rate_limit, new_rate_limit);
616                },
617                item = self.core.next_item() => {
618                    return item;
619                }
620            }
621        }
622    }
623
624    fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
625        let downstream_offset = DownstreamChunkOffset(offset);
626        let mut truncate_offset = None;
627        let mut stop = false;
628        'outer: while let Some((upstream_offset, downstream_offsets)) =
629            self.core.consumed_offset_queue.back_mut()
630        {
631            while let Some(prev_downstream_offset) = downstream_offsets.back() {
632                if *prev_downstream_offset <= downstream_offset {
633                    downstream_offsets.pop_back();
634                } else {
635                    stop = true;
636                    break 'outer;
637                }
638            }
639            truncate_offset = Some(*upstream_offset);
640            self.core.consumed_offset_queue.pop_back();
641        }
642        if !stop && let Some((_, downstream_offsets, _)) = &mut self.core.consuming_chunk {
643            while let Some(prev_downstream_offset) = downstream_offsets.back() {
644                if *prev_downstream_offset <= downstream_offset {
645                    downstream_offsets.pop_back();
646                } else {
647                    // stop = true;
648                    break;
649                }
650            }
651        }
652        tracing::trace!(
653            "rate limited log store reader truncate offset {:?}, downstream offset {:?}",
654            truncate_offset,
655            offset
656        );
657        if let Some(offset) = truncate_offset {
658            self.core.inner.truncate(offset.0)
659        } else {
660            Ok(())
661        }
662    }
663
664    fn rewind(&mut self) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
665        self.core.consuming_chunk = None;
666        self.core.consumed_offset_queue.clear();
667        self.core.next_chunk_id = 0;
668        self.core.inner.rewind()
669    }
670
671    fn start_from(
672        &mut self,
673        start_offset: Option<u64>,
674    ) -> impl Future<Output = LogStoreResult<()>> + Send + '_ {
675        self.core.inner.start_from(start_offset)
676    }
677}
678
679#[easy_ext::ext(LogReaderExt)]
680impl<T> T
681where
682    T: LogReader,
683{
684    pub fn transform_chunk<F: Fn(StreamChunk) -> StreamChunk + Sized>(
685        self,
686        f: F,
687    ) -> TransformChunkLogReader<F, Self> {
688        TransformChunkLogReader { f, inner: self }
689    }
690
691    pub fn monitored(self, metrics: LogReaderMetrics) -> impl LogReader {
692        // TODO: The `clone()` can be avoided if move backpressure inside `MonitoredLogReader`
693        let wait_new_future_duration = metrics.log_store_reader_wait_new_future_duration_ns.clone();
694        BackpressureMonitoredLogReader::new(
695            MonitoredLogReader::new(self, metrics),
696            wait_new_future_duration,
697        )
698    }
699
700    pub fn rate_limited(self, control_rx: UnboundedReceiver<RateLimit>) -> impl LogReader {
701        RateLimitedLogReader::new(self, control_rx)
702    }
703}
704
705pub struct MonitoredLogWriter<W: LogWriter> {
706    inner: W,
707    metrics: LogWriterMetrics,
708}
709
710pub struct LogWriterMetrics {
711    // Labels: [actor_id, sink_id, sink_name]
712    pub log_store_first_write_epoch: LabelGuardedIntGauge,
713    pub log_store_latest_write_epoch: LabelGuardedIntGauge,
714    pub log_store_write_rows: LabelGuardedIntCounter,
715}
716
717impl<W: LogWriter> LogWriter for MonitoredLogWriter<W> {
718    async fn init(
719        &mut self,
720        epoch: EpochPair,
721        pause_read_on_bootstrap: bool,
722    ) -> LogStoreResult<()> {
723        self.metrics
724            .log_store_first_write_epoch
725            .set(epoch.curr as _);
726        self.metrics
727            .log_store_latest_write_epoch
728            .set(epoch.curr as _);
729        self.inner.init(epoch, pause_read_on_bootstrap).await
730    }
731
732    async fn write_chunk(&mut self, chunk: StreamChunk) -> LogStoreResult<()> {
733        self.metrics
734            .log_store_write_rows
735            .inc_by(chunk.cardinality() as _);
736        self.inner.write_chunk(chunk).await
737    }
738
739    async fn flush_current_epoch(
740        &mut self,
741        next_epoch: u64,
742        options: FlushCurrentEpochOptions,
743    ) -> LogStoreResult<LogWriterPostFlushCurrentEpoch<'_>> {
744        let post_flush = self.inner.flush_current_epoch(next_epoch, options).await?;
745        self.metrics
746            .log_store_latest_write_epoch
747            .set(next_epoch as _);
748        Ok(post_flush)
749    }
750
751    fn pause(&mut self) -> LogStoreResult<()> {
752        self.inner.pause()
753    }
754
755    fn resume(&mut self) -> LogStoreResult<()> {
756        self.inner.resume()
757    }
758}
759
760#[easy_ext::ext(LogWriterExt)]
761impl<T> T
762where
763    T: LogWriter + Sized,
764{
765    pub fn monitored(self, metrics: LogWriterMetrics) -> MonitoredLogWriter<T> {
766        MonitoredLogWriter {
767            inner: self,
768            metrics,
769        }
770    }
771}
772
773enum DeliveryFutureManagerItem<F> {
774    Chunk {
775        chunk_id: ChunkId,
776        // earlier future at the front
777        futures: VecDeque<F>,
778    },
779    Barrier,
780}
781
782pub struct DeliveryFutureManager<F> {
783    future_count: usize,
784    max_future_count: usize,
785    // earlier items at the front
786    items: VecDeque<(u64, DeliveryFutureManagerItem<F>)>,
787}
788
789impl<F> DeliveryFutureManager<F> {
790    pub fn new(max_future_count: usize) -> Self {
791        Self {
792            future_count: 0,
793            max_future_count,
794            items: Default::default(),
795        }
796    }
797
798    pub fn add_barrier(&mut self, epoch: u64) {
799        if let Some((item_epoch, last_item)) = self.items.back() {
800            match last_item {
801                DeliveryFutureManagerItem::Chunk { .. } => {
802                    assert_eq!(*item_epoch, epoch)
803                }
804                DeliveryFutureManagerItem::Barrier => {
805                    assert!(
806                        epoch > *item_epoch,
807                        "new barrier epoch {} should be greater than prev barrier {}",
808                        epoch,
809                        item_epoch
810                    );
811                }
812            }
813        }
814        self.items
815            .push_back((epoch, DeliveryFutureManagerItem::Barrier));
816    }
817
818    pub fn start_write_chunk(
819        &mut self,
820        epoch: u64,
821        chunk_id: ChunkId,
822    ) -> DeliveryFutureManagerAddFuture<'_, F> {
823        if let Some((item_epoch, item)) = self.items.back() {
824            match item {
825                DeliveryFutureManagerItem::Chunk {
826                    chunk_id: item_chunk_id,
827                    ..
828                } => {
829                    assert_eq!(epoch, *item_epoch);
830                    assert!(
831                        chunk_id > *item_chunk_id,
832                        "new chunk id {} should be greater than prev chunk id {}",
833                        chunk_id,
834                        item_chunk_id
835                    );
836                }
837                DeliveryFutureManagerItem::Barrier => {
838                    assert!(
839                        epoch > *item_epoch,
840                        "new chunk epoch {} should be greater than prev barrier: {}",
841                        epoch,
842                        item_epoch
843                    );
844                }
845            }
846        }
847        self.items.push_back((
848            epoch,
849            DeliveryFutureManagerItem::Chunk {
850                chunk_id,
851                futures: VecDeque::new(),
852            },
853        ));
854        DeliveryFutureManagerAddFuture(self)
855    }
856}
857
858pub struct DeliveryFutureManagerAddFuture<'a, F>(&'a mut DeliveryFutureManager<F>);
859
860impl<F: TryFuture<Ok = ()> + Unpin + 'static> DeliveryFutureManagerAddFuture<'_, F> {
861    /// Add a new future to the latest started written chunk.
862    /// The returned bool value indicate whether we have awaited on any previous futures.
863    pub async fn add_future_may_await(&mut self, future: F) -> Result<bool, F::Error> {
864        let mut has_await = false;
865        while self.0.future_count >= self.0.max_future_count {
866            self.await_one_delivery().await?;
867            has_await = true;
868        }
869        match self.0.items.back_mut() {
870            Some((_, DeliveryFutureManagerItem::Chunk { futures, .. })) => {
871                futures.push_back(future);
872                self.0.future_count += 1;
873                Ok(has_await)
874            }
875            _ => unreachable!("should add future only after add a new chunk"),
876        }
877    }
878
879    pub async fn await_one_delivery(&mut self) -> Result<(), F::Error> {
880        for (_, item) in &mut self.0.items {
881            if let DeliveryFutureManagerItem::Chunk { futures, .. } = item
882                && let Some(mut delivery_future) = futures.pop_front()
883            {
884                self.0.future_count -= 1;
885                return poll_fn(|cx| delivery_future.try_poll_unpin(cx)).await;
886            } else {
887                continue;
888            }
889        }
890        Ok(())
891    }
892
893    pub fn future_count(&self) -> usize {
894        self.0.future_count
895    }
896
897    pub fn max_future_count(&self) -> usize {
898        self.0.max_future_count
899    }
900}
901
902impl<F: TryFuture<Ok = ()> + Unpin + 'static> DeliveryFutureManager<F> {
903    pub fn next_truncate_offset(
904        &mut self,
905    ) -> impl Future<Output = Result<TruncateOffset, F::Error>> + '_ {
906        poll_fn(move |cx| {
907            let mut latest_offset: Option<TruncateOffset> = None;
908            'outer: while let Some((epoch, item)) = self.items.front_mut() {
909                match item {
910                    DeliveryFutureManagerItem::Chunk { chunk_id, futures } => {
911                        while let Some(future) = futures.front_mut() {
912                            match future.try_poll_unpin(cx) {
913                                Poll::Ready(result) => match result {
914                                    Ok(()) => {
915                                        self.future_count -= 1;
916                                        futures.pop_front();
917                                    }
918                                    Err(e) => {
919                                        return Poll::Ready(Err(e));
920                                    }
921                                },
922                                Poll::Pending => {
923                                    break 'outer;
924                                }
925                            }
926                        }
927
928                        // when we reach here, there must not be any pending or error future.
929                        // Which means all futures of this stream chunk have been finished
930                        assert!(futures.is_empty());
931                        latest_offset = Some(TruncateOffset::Chunk {
932                            epoch: *epoch,
933                            chunk_id: *chunk_id,
934                        });
935                        self.items.pop_front().expect("items not empty");
936                    }
937                    DeliveryFutureManagerItem::Barrier => {
938                        latest_offset = Some(TruncateOffset::Barrier { epoch: *epoch });
939                        self.items.pop_front().expect("items not empty");
940                        // Barrier will be yielded anyway
941                        break 'outer;
942                    }
943                }
944            }
945            if let Some(offset) = latest_offset {
946                Poll::Ready(Ok(offset))
947            } else {
948                Poll::Pending
949            }
950        })
951    }
952}
953
954#[cfg(test)]
955mod tests {
956    use std::collections::VecDeque;
957    use std::future::{Future, poll_fn};
958    use std::pin::pin;
959    use std::task::Poll;
960
961    use futures::{FutureExt, TryFuture};
962    use risingwave_common::array::StreamChunk;
963    use risingwave_common::util::epoch::test_epoch;
964    use tokio::sync::oneshot;
965    use tokio::sync::oneshot::Receiver;
966
967    use super::{LogReader, LogStoreReadItem, LogStoreResult, TruncateBarrierLogReader};
968    use crate::sink::log_store::{DeliveryFutureManager, TruncateOffset};
969
970    struct MockLogReader {
971        items: VecDeque<(u64, LogStoreReadItem)>,
972        truncate_calls: Vec<TruncateOffset>,
973    }
974
975    impl MockLogReader {
976        fn new(items: impl IntoIterator<Item = (u64, LogStoreReadItem)>) -> Self {
977            Self {
978                items: items.into_iter().collect(),
979                truncate_calls: Vec::new(),
980            }
981        }
982    }
983
984    impl LogReader for MockLogReader {
985        async fn init(&mut self) -> LogStoreResult<()> {
986            Ok(())
987        }
988
989        async fn start_from(&mut self, _start_offset: Option<u64>) -> LogStoreResult<()> {
990            Ok(())
991        }
992
993        async fn next_item(&mut self) -> LogStoreResult<(u64, LogStoreReadItem)> {
994            self.items
995                .pop_front()
996                .ok_or_else(|| anyhow::anyhow!("no item"))
997        }
998
999        fn truncate(&mut self, offset: TruncateOffset) -> LogStoreResult<()> {
1000            self.truncate_calls.push(offset);
1001            Ok(())
1002        }
1003
1004        async fn rewind(&mut self) -> LogStoreResult<()> {
1005            Ok(())
1006        }
1007    }
1008
1009    #[test]
1010    fn test_truncate_offset_cmp() {
1011        assert!(
1012            TruncateOffset::Barrier { epoch: 232 }
1013                < TruncateOffset::Chunk {
1014                    epoch: 233,
1015                    chunk_id: 1
1016                }
1017        );
1018        assert_eq!(
1019            TruncateOffset::Chunk {
1020                epoch: 1,
1021                chunk_id: 1
1022            },
1023            TruncateOffset::Chunk {
1024                epoch: 1,
1025                chunk_id: 1
1026            }
1027        );
1028        assert!(
1029            TruncateOffset::Chunk {
1030                epoch: 1,
1031                chunk_id: 1
1032            } < TruncateOffset::Chunk {
1033                epoch: 1,
1034                chunk_id: 2
1035            }
1036        );
1037        assert!(
1038            TruncateOffset::Barrier { epoch: 1 }
1039                > TruncateOffset::Chunk {
1040                    epoch: 1,
1041                    chunk_id: 2
1042                }
1043        );
1044        assert!(
1045            TruncateOffset::Chunk {
1046                epoch: 1,
1047                chunk_id: 2
1048            } < TruncateOffset::Barrier { epoch: 1 }
1049        );
1050        assert!(
1051            TruncateOffset::Chunk {
1052                epoch: 2,
1053                chunk_id: 2
1054            } > TruncateOffset::Barrier { epoch: 1 }
1055        );
1056        assert!(TruncateOffset::Barrier { epoch: 2 } > TruncateOffset::Barrier { epoch: 1 });
1057    }
1058
1059    type TestFuture = impl TryFuture<Ok = (), Error = anyhow::Error> + Unpin + 'static;
1060
1061    #[define_opaque(TestFuture)]
1062    fn to_test_future(rx: Receiver<LogStoreResult<()>>) -> TestFuture {
1063        async move { rx.await.unwrap() }.boxed()
1064    }
1065
1066    #[tokio::test]
1067    async fn test_empty() {
1068        let mut manager = DeliveryFutureManager::<TestFuture>::new(2);
1069        let mut future = pin!(manager.next_truncate_offset());
1070        assert!(
1071            poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx)))
1072                .await
1073                .is_pending()
1074        );
1075    }
1076
1077    #[tokio::test]
1078    async fn test_future_delivery_manager_basic() {
1079        let mut manager = DeliveryFutureManager::new(2);
1080        let epoch1 = 233;
1081        let chunk_id1 = 1;
1082        let (tx1_1, rx1_1) = oneshot::channel();
1083        let mut write_chunk = manager.start_write_chunk(epoch1, chunk_id1);
1084        assert!(
1085            !write_chunk
1086                .add_future_may_await(to_test_future(rx1_1))
1087                .await
1088                .unwrap()
1089        );
1090        assert_eq!(manager.future_count, 1);
1091        {
1092            let mut next_truncate_offset = pin!(manager.next_truncate_offset());
1093            assert!(
1094                poll_fn(|cx| Poll::Ready(next_truncate_offset.as_mut().poll(cx)))
1095                    .await
1096                    .is_pending()
1097            );
1098            tx1_1.send(Ok(())).unwrap();
1099            assert_eq!(
1100                next_truncate_offset.await.unwrap(),
1101                TruncateOffset::Chunk {
1102                    epoch: epoch1,
1103                    chunk_id: chunk_id1
1104                }
1105            );
1106        }
1107        assert_eq!(manager.future_count, 0);
1108        manager.add_barrier(epoch1);
1109        assert_eq!(
1110            manager.next_truncate_offset().await.unwrap(),
1111            TruncateOffset::Barrier { epoch: epoch1 }
1112        );
1113    }
1114
1115    #[tokio::test]
1116    async fn test_future_delivery_manager_compress_chunk() {
1117        let mut manager = DeliveryFutureManager::new(10);
1118        let epoch1 = test_epoch(233);
1119        let chunk_id1 = 1;
1120        let chunk_id2 = chunk_id1 + 1;
1121        let chunk_id3 = chunk_id2 + 1;
1122        let (tx1_1, rx1_1) = oneshot::channel();
1123        let (tx1_2, rx1_2) = oneshot::channel();
1124        let (tx1_3, rx1_3) = oneshot::channel();
1125        let epoch2 = test_epoch(234);
1126        let (tx2_1, rx2_1) = oneshot::channel();
1127        assert!(
1128            !manager
1129                .start_write_chunk(epoch1, chunk_id1)
1130                .add_future_may_await(to_test_future(rx1_1))
1131                .await
1132                .unwrap()
1133        );
1134        assert!(
1135            !manager
1136                .start_write_chunk(epoch1, chunk_id2)
1137                .add_future_may_await(to_test_future(rx1_2))
1138                .await
1139                .unwrap()
1140        );
1141        assert!(
1142            !manager
1143                .start_write_chunk(epoch1, chunk_id3)
1144                .add_future_may_await(to_test_future(rx1_3))
1145                .await
1146                .unwrap()
1147        );
1148        manager.add_barrier(epoch1);
1149        assert!(
1150            !manager
1151                .start_write_chunk(epoch2, chunk_id1)
1152                .add_future_may_await(to_test_future(rx2_1))
1153                .await
1154                .unwrap()
1155        );
1156        assert_eq!(manager.future_count, 4);
1157        {
1158            let mut next_truncate_offset = pin!(manager.next_truncate_offset());
1159            assert!(
1160                poll_fn(|cx| Poll::Ready(next_truncate_offset.as_mut().poll(cx)))
1161                    .await
1162                    .is_pending()
1163            );
1164            tx1_2.send(Ok(())).unwrap();
1165            assert!(
1166                poll_fn(|cx| Poll::Ready(next_truncate_offset.as_mut().poll(cx)))
1167                    .await
1168                    .is_pending()
1169            );
1170            tx1_1.send(Ok(())).unwrap();
1171            // The offset of chunk1 and chunk2 are compressed
1172            assert_eq!(
1173                next_truncate_offset.await.unwrap(),
1174                TruncateOffset::Chunk {
1175                    epoch: epoch1,
1176                    chunk_id: chunk_id2
1177                }
1178            );
1179        }
1180        assert_eq!(manager.future_count, 2);
1181        {
1182            let mut next_truncate_offset = pin!(manager.next_truncate_offset());
1183            assert!(
1184                poll_fn(|cx| Poll::Ready(next_truncate_offset.as_mut().poll(cx)))
1185                    .await
1186                    .is_pending()
1187            );
1188            tx1_3.send(Ok(())).unwrap();
1189            tx2_1.send(Ok(())).unwrap();
1190            // Emit barrier though later chunk has finished.
1191            assert_eq!(
1192                next_truncate_offset.await.unwrap(),
1193                TruncateOffset::Barrier { epoch: epoch1 }
1194            );
1195        }
1196        assert_eq!(manager.future_count, 1);
1197        assert_eq!(
1198            manager.next_truncate_offset().await.unwrap(),
1199            TruncateOffset::Chunk {
1200                epoch: epoch2,
1201                chunk_id: chunk_id1
1202            }
1203        );
1204    }
1205
1206    #[tokio::test]
1207    async fn test_future_delivery_manager_await_future() {
1208        let mut manager = DeliveryFutureManager::new(2);
1209        let epoch = 233;
1210        let chunk_id1 = 1;
1211        let chunk_id2 = chunk_id1 + 1;
1212        let (tx1_1, rx1_1) = oneshot::channel();
1213        let (tx1_2, rx1_2) = oneshot::channel();
1214        let (tx2_1, rx2_1) = oneshot::channel();
1215        let (tx2_2, rx2_2) = oneshot::channel();
1216
1217        {
1218            let mut write_chunk = manager.start_write_chunk(epoch, chunk_id1);
1219            assert!(
1220                !write_chunk
1221                    .add_future_may_await(to_test_future(rx1_1))
1222                    .await
1223                    .unwrap()
1224            );
1225            assert!(
1226                !write_chunk
1227                    .add_future_may_await(to_test_future(rx1_2))
1228                    .await
1229                    .unwrap()
1230            );
1231            assert_eq!(manager.future_count, 2);
1232        }
1233
1234        {
1235            let mut write_chunk = manager.start_write_chunk(epoch, chunk_id2);
1236            {
1237                let mut future1 = pin!(write_chunk.add_future_may_await(to_test_future(rx2_1)));
1238                assert!(
1239                    poll_fn(|cx| Poll::Ready(future1.as_mut().poll(cx)))
1240                        .await
1241                        .is_pending()
1242                );
1243                tx1_1.send(Ok(())).unwrap();
1244                assert!(future1.await.unwrap());
1245            }
1246            assert_eq!(2, write_chunk.future_count());
1247            {
1248                let mut future2 = pin!(write_chunk.add_future_may_await(to_test_future(rx2_2)));
1249                assert!(
1250                    poll_fn(|cx| Poll::Ready(future2.as_mut().poll(cx)))
1251                        .await
1252                        .is_pending()
1253                );
1254                tx1_2.send(Ok(())).unwrap();
1255                assert!(future2.await.unwrap());
1256            }
1257            assert_eq!(2, write_chunk.future_count());
1258            {
1259                let mut future3 = pin!(write_chunk.await_one_delivery());
1260                assert!(
1261                    poll_fn(|cx| Poll::Ready(future3.as_mut().poll(cx)))
1262                        .await
1263                        .is_pending()
1264                );
1265                tx2_1.send(Ok(())).unwrap();
1266                future3.await.unwrap();
1267            }
1268            assert_eq!(1, write_chunk.future_count());
1269        }
1270
1271        assert_eq!(
1272            manager.next_truncate_offset().await.unwrap(),
1273            TruncateOffset::Chunk {
1274                epoch,
1275                chunk_id: chunk_id1
1276            }
1277        );
1278
1279        assert_eq!(1, manager.future_count);
1280
1281        {
1282            let mut future = pin!(manager.next_truncate_offset());
1283            assert!(
1284                poll_fn(|cx| Poll::Ready(future.as_mut().poll(cx)))
1285                    .await
1286                    .is_pending()
1287            );
1288            tx2_2.send(Ok(())).unwrap();
1289            assert_eq!(
1290                future.await.unwrap(),
1291                TruncateOffset::Chunk {
1292                    epoch,
1293                    chunk_id: chunk_id2
1294                }
1295            );
1296        }
1297
1298        assert_eq!(0, manager.future_count);
1299    }
1300
1301    #[tokio::test]
1302    async fn test_truncate_barrier_reader_truncates_previous_barriers_first() {
1303        let inner = MockLogReader::new([
1304            (
1305                1,
1306                LogStoreReadItem::Barrier {
1307                    is_checkpoint: false,
1308                    new_vnode_bitmap: None,
1309                    is_stop: false,
1310                    schema_change: None,
1311                },
1312            ),
1313            (
1314                2,
1315                LogStoreReadItem::Barrier {
1316                    is_checkpoint: false,
1317                    new_vnode_bitmap: None,
1318                    is_stop: false,
1319                    schema_change: None,
1320                },
1321            ),
1322        ]);
1323        let mut reader = TruncateBarrierLogReader::new(inner);
1324
1325        reader.init().await.unwrap();
1326        reader.start_from(None).await.unwrap();
1327        reader.next_item().await.unwrap();
1328        reader.next_item().await.unwrap();
1329        reader
1330            .truncate(TruncateOffset::Barrier { epoch: 2 })
1331            .unwrap();
1332
1333        assert_eq!(
1334            reader.inner.truncate_calls,
1335            vec![
1336                TruncateOffset::Barrier { epoch: 1 },
1337                TruncateOffset::Barrier { epoch: 2 },
1338            ]
1339        );
1340    }
1341
1342    #[tokio::test]
1343    async fn test_truncate_barrier_reader_keeps_later_barrier_pending() {
1344        let inner = MockLogReader::new([
1345            (
1346                1,
1347                LogStoreReadItem::Barrier {
1348                    is_checkpoint: false,
1349                    new_vnode_bitmap: None,
1350                    is_stop: false,
1351                    schema_change: None,
1352                },
1353            ),
1354            (
1355                2,
1356                LogStoreReadItem::StreamChunk {
1357                    chunk: StreamChunk::default(),
1358                    chunk_id: 0,
1359                },
1360            ),
1361            (
1362                2,
1363                LogStoreReadItem::Barrier {
1364                    is_checkpoint: false,
1365                    new_vnode_bitmap: None,
1366                    is_stop: false,
1367                    schema_change: None,
1368                },
1369            ),
1370        ]);
1371        let mut reader = TruncateBarrierLogReader::new(inner);
1372
1373        reader.init().await.unwrap();
1374        reader.start_from(None).await.unwrap();
1375        reader.next_item().await.unwrap();
1376        reader.next_item().await.unwrap();
1377        reader.next_item().await.unwrap();
1378        reader
1379            .truncate(TruncateOffset::Chunk {
1380                epoch: 2,
1381                chunk_id: 0,
1382            })
1383            .unwrap();
1384
1385        assert_eq!(
1386            reader.inner.truncate_calls,
1387            vec![
1388                TruncateOffset::Barrier { epoch: 1 },
1389                TruncateOffset::Chunk {
1390                    epoch: 2,
1391                    chunk_id: 0,
1392                },
1393            ]
1394        );
1395
1396        reader
1397            .truncate(TruncateOffset::Barrier { epoch: 2 })
1398            .unwrap();
1399        assert_eq!(
1400            reader.inner.truncate_calls,
1401            vec![
1402                TruncateOffset::Barrier { epoch: 1 },
1403                TruncateOffset::Chunk {
1404                    epoch: 2,
1405                    chunk_id: 0,
1406                },
1407                TruncateOffset::Barrier { epoch: 2 },
1408            ]
1409        );
1410    }
1411}
risingwave_connector/sink/log_store.rs

risingwave_connector/sink/
log_store.rs
