risingwave_storage/hummock/iterator/
merge_inner.rs

1// Copyright 2022 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::collections::binary_heap::PeekMut;
16use std::collections::{BinaryHeap, LinkedList};
17use std::ops::{Deref, DerefMut};
18
19use risingwave_hummock_sdk::key::FullKey;
20
21use crate::hummock::HummockResult;
22use crate::hummock::iterator::{
23    DirectionEnum, HummockIterator, HummockIteratorDirection, ValueMeta,
24};
25use crate::hummock::value::HummockValue;
26use crate::monitor::StoreLocalStatistic;
27
28pub struct Node<I: HummockIterator> {
29    iter: I,
30}
31
32impl<I: HummockIterator> Eq for Node<I> where Self: PartialEq {}
33impl<I: HummockIterator> PartialOrd for Node<I>
34where
35    Self: Ord,
36{
37    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
38        Some(self.cmp(other))
39    }
40}
41
42/// Implement `Ord` for unordered iter node. Only compare the key.
43impl<I: HummockIterator> Ord for Node<I> {
44    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
45        // Note: to implement min-heap by using max-heap internally, the comparing
46        // order should be reversed.
47
48        match I::Direction::direction() {
49            DirectionEnum::Forward => other.iter.key().cmp(&self.iter.key()),
50            DirectionEnum::Backward => self.iter.key().cmp(&other.iter.key()),
51        }
52    }
53}
54
55impl<I: HummockIterator> PartialEq for Node<I> {
56    fn eq(&self, other: &Self) -> bool {
57        self.iter.key() == other.iter.key()
58    }
59}
60
61/// Iterates on multiple iterators, a.k.a. `MergeIterator`.
62pub struct MergeIterator<I: HummockIterator> {
63    /// Invalid or non-initialized iterators.
64    unused_iters: LinkedList<Node<I>>,
65
66    /// The heap for merge sort.
67    heap: BinaryHeap<Node<I>>,
68
69    /// Statistics collected from iterators discarded after an error.
70    retired_stats: Option<Box<StoreLocalStatistic>>,
71}
72
73impl<I: HummockIterator> MergeIterator<I> {
74    fn collect_local_statistic_impl(&self, stats: &mut StoreLocalStatistic) {
75        if let Some(retired_stats) = &self.retired_stats {
76            stats.add(retired_stats);
77        }
78        for node in &self.heap {
79            node.iter.collect_local_statistic(stats);
80        }
81        for node in &self.unused_iters {
82            node.iter.collect_local_statistic(stats);
83        }
84    }
85}
86
87impl<I: HummockIterator> MergeIterator<I> {
88    pub fn new(iterators: impl IntoIterator<Item = I>) -> Self {
89        Self::create(iterators)
90    }
91
92    pub fn for_compactor(iterators: impl IntoIterator<Item = I>) -> Self {
93        Self::create(iterators)
94    }
95
96    fn create(iterators: impl IntoIterator<Item = I>) -> Self {
97        Self {
98            unused_iters: iterators.into_iter().map(|iter| Node { iter }).collect(),
99            heap: BinaryHeap::new(),
100            retired_stats: None,
101        }
102    }
103}
104
105impl<I: HummockIterator> MergeIterator<I>
106where
107    Node<I>: Ord,
108{
109    /// Moves all iterators from the `heap` to the linked list.
110    fn reset_heap(&mut self) {
111        self.unused_iters.extend(self.heap.drain());
112    }
113
114    /// After some iterators in `unused_iterators` are sought or rewound, calls this function
115    /// to construct a new heap using the valid ones.
116    fn build_heap(&mut self) {
117        assert!(self.heap.is_empty());
118
119        self.heap = self
120            .unused_iters
121            .extract_if(|i| i.iter.is_valid())
122            .collect();
123    }
124}
125
126/// This is a wrapper for the `PeekMut` of heap.
127///
128/// Several panics due to future cancellation are caused by calling `drop` on the `PeekMut` when
129/// futures holding the `PeekMut` are cancelled and dropped. Dropping a `PeekMut` will accidentally
130/// cause a comparison between the top node and the node below, and may call `key()` for top node
131/// iterators that are in some intermediate inconsistent states.
132///
133/// When a `PeekMut` is wrapped by this guard, when the guard is dropped, `PeekMut::pop` will be
134/// called on the `PeekMut`, and the popped node will be added to the linked list that collects the
135/// unused nodes. In this way, when the future holding the guard is dropped, the `PeekMut` will not
136/// be called `drop`, and there will not be unexpected `key()` called for heap comparison.
137///
138/// In normal usage, when we finished using the `PeekMut`, we should explicitly call `guard.used()`
139/// in every branch carefully. When we want to pop the `PeekMut`, we can simply call `guard.pop()`.
140struct PeekMutGuard<'a, T: Ord> {
141    peek: Option<PeekMut<'a, T>>,
142    unused: &'a mut LinkedList<T>,
143}
144
145impl<'a, T: Ord> PeekMutGuard<'a, T> {
146    /// Call `peek_mut` on the top of heap and return a guard over the `PeekMut` if the heap is not
147    /// empty.
148    fn peek_mut(heap: &'a mut BinaryHeap<T>, unused: &'a mut LinkedList<T>) -> Option<Self> {
149        heap.peek_mut().map(|peek| Self {
150            peek: Some(peek),
151            unused,
152        })
153    }
154
155    /// Call `pop` on the `PeekMut`.
156    fn pop(mut self) -> T {
157        PeekMut::pop(self.peek.take().expect("should not be None"))
158    }
159
160    /// Mark finish using the `PeekMut`. `drop` will be called on the `PeekMut` directly.
161    fn used(mut self) {
162        self.peek.take().expect("should not be None");
163    }
164}
165
166impl<T: Ord> Deref for PeekMutGuard<'_, T> {
167    type Target = T;
168
169    fn deref(&self) -> &Self::Target {
170        self.peek.as_ref().expect("should not be None")
171    }
172}
173
174impl<T: Ord> DerefMut for PeekMutGuard<'_, T> {
175    fn deref_mut(&mut self) -> &mut Self::Target {
176        self.peek.as_mut().expect("should not be None")
177    }
178}
179
180impl<T: Ord> Drop for PeekMutGuard<'_, T> {
181    /// When the guard is dropped, if `pop` or `used` is not called before it is dropped, we will
182    /// call `PeekMut::pop` on the `PeekMut` and recycle the node to the unused list.
183    fn drop(&mut self) {
184        if let Some(peek) = self.peek.take() {
185            tracing::debug!(
186                "PeekMut are dropped without used. May be caused by future cancellation"
187            );
188            let top = PeekMut::pop(peek);
189            self.unused.push_back(top);
190        }
191    }
192}
193
194impl<I: HummockIterator> HummockIterator for MergeIterator<I>
195where
196    Node<I>: Ord,
197{
198    type Direction = I::Direction;
199
200    async fn next(&mut self) -> HummockResult<()> {
201        let mut node =
202            PeekMutGuard::peek_mut(&mut self.heap, &mut self.unused_iters).expect("no inner iter");
203
204        // WARNING: within scope of BinaryHeap::PeekMut, we must carefully handle all places of
205        // return. Once the iterator enters an invalid state, we should remove it from heap
206        // before returning.
207
208        match node.iter.next().await {
209            Ok(_) => {}
210            Err(e) => {
211                // If the iterator returns error, we should clear the heap, so that this
212                // iterator becomes invalid. Collect their work before dropping them, but do not
213                // retain the failed/current heap iterators for a later seek or rewind.
214                let node = node.pop();
215                let retired_stats = self
216                    .retired_stats
217                    .get_or_insert_with(|| Box::new(StoreLocalStatistic::default()));
218                node.iter.collect_local_statistic(retired_stats);
219                for node in self.heap.drain() {
220                    node.iter.collect_local_statistic(retired_stats);
221                }
222                return Err(e);
223            }
224        }
225
226        if !node.iter.is_valid() {
227            // Put back to `unused_iters`
228            let node = node.pop();
229            self.unused_iters.push_back(node);
230        } else {
231            // This will update the heap top.
232            node.used();
233        }
234
235        Ok(())
236    }
237
238    fn key(&self) -> FullKey<&[u8]> {
239        self.heap.peek().expect("no inner iter").iter.key()
240    }
241
242    fn value(&self) -> HummockValue<&[u8]> {
243        self.heap.peek().expect("no inner iter").iter.value()
244    }
245
246    fn is_valid(&self) -> bool {
247        self.heap.peek().is_some_and(|n| n.iter.is_valid())
248    }
249
250    async fn rewind(&mut self) -> HummockResult<()> {
251        self.reset_heap();
252        futures::future::try_join_all(self.unused_iters.iter_mut().map(|x| x.iter.rewind()))
253            .await?;
254        self.build_heap();
255        Ok(())
256    }
257
258    async fn seek<'a>(&'a mut self, key: FullKey<&'a [u8]>) -> HummockResult<()> {
259        self.reset_heap();
260        futures::future::try_join_all(self.unused_iters.iter_mut().map(|x| x.iter.seek(key)))
261            .await?;
262        self.build_heap();
263        Ok(())
264    }
265
266    fn collect_local_statistic(&self, stats: &mut StoreLocalStatistic) {
267        self.collect_local_statistic_impl(stats);
268    }
269
270    fn value_meta(&self) -> ValueMeta {
271        self.heap.peek().expect("no inner iter").iter.value_meta()
272    }
273}
risingwave_storage/hummock/iterator/merge_inner.rs

risingwave_storage/hummock/iterator/
merge_inner.rs
