risingwave_stream/executor/watermark/
mod.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 std::cmp::Reverse;
16use std::collections::{BTreeMap, BinaryHeap, HashSet, VecDeque};
17use std::hash::Hash;
18
19use super::Watermark;
20
21#[derive(Default, Debug)]
22pub(super) struct StagedWatermarks {
23    in_heap: bool,
24    staged: VecDeque<Watermark>,
25}
26
27pub(super) struct BufferedWatermarks<Id> {
28    /// We store the smallest watermark of each upstream, because the next watermark to emit is
29    /// among them.
30    pub first_buffered_watermarks: BinaryHeap<Reverse<(Watermark, Id)>>,
31    /// We buffer other watermarks of each upstream. The next-to-smallest one will become the
32    /// smallest when the smallest is emitted and be moved into heap.
33    pub other_buffered_watermarks: BTreeMap<Id, StagedWatermarks>,
34}
35
36impl<Id: Ord + Hash + std::fmt::Debug> BufferedWatermarks<Id> {
37    pub fn with_ids(buffer_ids: impl IntoIterator<Item = Id>) -> Self {
38        let other_buffered_watermarks: BTreeMap<_, _> = buffer_ids
39            .into_iter()
40            .map(|id| (id, Default::default()))
41            .collect();
42        let first_buffered_watermarks = BinaryHeap::with_capacity(other_buffered_watermarks.len());
43
44        BufferedWatermarks {
45            first_buffered_watermarks,
46            other_buffered_watermarks,
47        }
48    }
49
50    pub fn add_buffers(&mut self, buffer_ids: impl IntoIterator<Item = Id>) {
51        buffer_ids.into_iter().for_each(|id| {
52            self.other_buffered_watermarks
53                .try_insert(id, Default::default())
54                .unwrap();
55        });
56    }
57
58    pub fn clear(&mut self) {
59        self.first_buffered_watermarks.clear();
60        self.other_buffered_watermarks
61            .values_mut()
62            .for_each(|staged_watermarks| {
63                std::mem::take(staged_watermarks);
64            });
65    }
66
67    /// Handle a new watermark message. Optionally returns the watermark message to emit and the
68    /// buffer id.
69    pub fn handle_watermark(&mut self, buffer_id: Id, watermark: Watermark) -> Option<Watermark> {
70        // Note: The staged watermark buffer should be created before handling the watermark.
71        let staged = self.other_buffered_watermarks.get_mut(&buffer_id).unwrap();
72
73        if staged.in_heap {
74            staged.staged.push_back(watermark);
75            None
76        } else {
77            staged.in_heap = true;
78            self.first_buffered_watermarks
79                .push(Reverse((watermark, buffer_id)));
80            self.check_watermark_heap()
81        }
82    }
83
84    /// Check the watermark heap and decide whether to emit a watermark message.
85    pub fn check_watermark_heap(&mut self) -> Option<Watermark> {
86        let len = self.other_buffered_watermarks.len();
87        let mut watermark_to_emit = None;
88        while !self.first_buffered_watermarks.is_empty()
89            && (self.first_buffered_watermarks.len() == len
90                || watermark_to_emit.as_ref().is_some_and(|watermark| {
91                    watermark == &self.first_buffered_watermarks.peek().unwrap().0.0
92                }))
93        {
94            let Reverse((watermark, id)) = self.first_buffered_watermarks.pop().unwrap();
95            watermark_to_emit = Some(watermark);
96            let staged = self.other_buffered_watermarks.get_mut(&id).unwrap();
97            if let Some(first) = staged.staged.pop_front() {
98                self.first_buffered_watermarks.push(Reverse((first, id)));
99            } else {
100                staged.in_heap = false;
101            }
102        }
103        watermark_to_emit
104    }
105
106    /// Remove buffers and return watermark to emit.
107    pub fn remove_buffer(&mut self, buffer_ids_to_remove: HashSet<Id>) -> Option<Watermark> {
108        self.first_buffered_watermarks
109            .retain(|Reverse((_, id))| !buffer_ids_to_remove.contains(id));
110        self.other_buffered_watermarks
111            .retain(|id, _| !buffer_ids_to_remove.contains(id));
112        // Call `check_watermark_heap` in case the only buffers(s) that does not have watermark in
113        // heap is removed
114        self.check_watermark_heap()
115    }
116}
risingwave_stream/executor/watermark/mod.rs

risingwave_stream/executor/watermark/
mod.rs
