risingwave_stream/executor/wrapper/
epoch_provide.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use futures::{TryStreamExt, pin_mut};
16use futures_async_stream::try_stream;
17use risingwave_common::util::epoch;
18
19use crate::executor::error::StreamExecutorError;
20use crate::executor::{Message, MessageStream};
21
22/// Streams wrapped by `epoch_provide` is able to retrieve the current epoch pair from the functions
23/// from [`epoch::task_local`].
24#[try_stream(ok = Message, error = StreamExecutorError)]
25pub async fn epoch_provide(input: impl MessageStream) {
26    pin_mut!(input);
27
28    let mut epoch = None;
29
30    while let Some(message) = if let Some(epoch) = epoch {
31        epoch::task_local::scope(epoch, input.try_next()).await?
32    } else {
33        input.try_next().await?
34    } {
35        // The inner executor has yielded a new barrier message. In next polls, we will provide the
36        // updated epoch pair.
37        if let Message::Barrier(b) = &message {
38            epoch = Some(b.epoch);
39        }
40        yield message;
41    }
42}
risingwave_stream/executor/wrapper/epoch_provide.rs

risingwave_stream/executor/wrapper/
epoch_provide.rs
