Repository: risingwavelabs/risingwave
File: src/stream/src/common/log_store_impl/kv_log_store/buffer.rs (Rust)

Code before:
    pub(crate) async fn await_barrier_truncation(
        &self,
        curr_epoch: u64,
    ) -> ReaderTruncationOffsetType {
        let mut ret = None;
        loop {
            {
                let mut inner = self.buffer.inner();
                while let Some((epoch, _)) = inner.truncation_list.front()
                    && *epoch <= curr_epoch
                {
                    ret = inner.truncation_list.pop_front();
                }
            }
            if let Some((epoch, None)) = ret
                && epoch == curr_epoch
            {
                return (epoch, None);
            }

            let notified = self.truncate_notify.notified();
            notified
                .instrument_await("Wait For Barrier Truncation")
                .await;
        }
    }

    pub(crate) fn flush_all_unflushed(
        &mut self,
        mut flush_fn: impl FnMut(&StreamChunk, u64, SeqId, SeqId) -> LogStoreResult<()>,
    ) -> LogStoreResult<()> {
        let mut inner_guard = self.buffer.inner();
        let inner = inner_guard.deref_mut();
        for (epoch, item) in inner
            .unconsumed_queue
            .iter_mut()
            .chain(inner.consumed_queue.iter_mut())
        {
            if let LogStoreBufferItem::StreamChunk {
                chunk,
                start_seq_id,
                end_seq_id,
                flushed,
                ..
            } = item
            {
                if *flushed {
                    // Since we iterate from new data to old data, when we meet a flushed data, the
                    // rest should have been flushed.
                    break;
                }

Review comment (bug):
The `await_barrier_truncation` method has a race condition. After checking the condition at line 339-342 and before setting up the notifier at line 345, a truncation notification could be added (lines 467 or 474). This would cause the notification to be missed, and the method would wait indefinitely.

**Recommendation**: Call `self.truncate_notify.notified()` before checking the condition to ensure notifications aren't missed:

```rust
pub(crate) async fn await_barrier_truncation(
    &self,
    curr_epoch: u64,
) -> ReaderTruncationOffsetType {
    loop {
        let notified = self.truncate_notify.notified();
        let mut ret = None;
        {
            let mut inner = self.buffer.inner();
            while let Some((epoch, _)) = inner.truncation_list.front()
                && *epoch <= curr_epoch
            {
                ret = inner.truncation_list.pop_front();
            }
        }
        if let Some((epoch, None)) = ret
            && epoch == curr_epoch
        {
            return (epoch, None);
        }

        notified
            .instrument_await("Wait For Barrier Truncation")
            .await;
    }
}
```

Code after:
    pub(crate) async fn wait_for_barrier_truncation(
        &self,
        curr_epoch: u64,
    ) -> LogStoreResult<ReaderTruncationOffsetType> {
        loop {
            let notified = self.truncate_notify.notified();

            {
                let mut inner = self.buffer.inner();
                while let Some((epoch, seq_id)) = inner.truncation_list.pop_front() {
                    if epoch > curr_epoch {
                        // TODO: should panic, after we confirm the correctness
                        return Err(anyhow::anyhow!(
                            "truncation epoch {} should not be larger than current epoch {}",
                            epoch,
                            curr_epoch
                        ));
                    }
                    if epoch == curr_epoch && seq_id.is_none() {
                        return Ok((epoch, seq_id));
                    }
                }
            }

            notified
                .instrument_await("Wait For Barrier Truncation")
                .await;
        }
    }

    pub(crate) fn flush_all_unflushed(
        &mut self,
        mut flush_fn: impl FnMut(&StreamChunk, u64, SeqId, SeqId) -> LogStoreResult<()>,
    ) -> LogStoreResult<()> {
        let mut inner_guard = self.buffer.inner();
        let inner = inner_guard.deref_mut();
        for (epoch, item) in inner
            .unconsumed_queue
            .iter_mut()
            .chain(inner.consumed_queue.iter_mut())
        {
            if let LogStoreBufferItem::StreamChunk {
                chunk,
                start_seq_id,
                end_seq_id,
                flushed,
                ..
            } = item
            {
                if *flushed {
                    // Since we iterate from new data to old data, when we meet a flushed data, the