risingwave_common/memory/
monitored_heap.rs

1// Copyright 2024 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::BinaryHeap;
16use std::mem::size_of;
17
18use risingwave_common_estimate_size::EstimateSize;
19
20use crate::memory::{MemoryContext, MonitoredGlobalAlloc};
21
22pub struct MemMonitoredHeap<T> {
23    inner: BinaryHeap<T>,
24    mem_ctx: MemoryContext,
25}
26
27impl<T: Ord + EstimateSize> MemMonitoredHeap<T> {
28    pub fn new_with(mem_ctx: MemoryContext) -> Self {
29        Self {
30            inner: BinaryHeap::new(),
31            mem_ctx,
32        }
33    }
34
35    pub fn with_capacity(capacity: usize, mem_ctx: MemoryContext) -> Self {
36        let inner = BinaryHeap::with_capacity(capacity);
37        // The allocation already succeeded; a budget overrun must not discard its charge.
38        mem_ctx.add_unchecked((capacity * size_of::<T>()) as i64);
39        Self { inner, mem_ctx }
40    }
41
42    pub fn push(&mut self, item: T) {
43        let prev_cap = self.inner.capacity();
44        let item_heap = item.estimated_heap_size();
45        self.inner.push(item);
46        let new_cap = self.inner.capacity();
47        self.mem_ctx
48            .add_unchecked(((new_cap - prev_cap) * size_of::<T>() + item_heap) as i64);
49    }
50
51    pub fn pop(&mut self) -> Option<T> {
52        let prev_cap = self.inner.capacity();
53        let item = self.inner.pop();
54        let item_heap = item.as_ref().map(|i| i.estimated_heap_size()).unwrap_or(0);
55        let new_cap = self.inner.capacity();
56        self.mem_ctx
57            .add_unchecked(-(((prev_cap - new_cap) * size_of::<T>() + item_heap) as i64));
58
59        item
60    }
61
62    pub fn is_empty(&self) -> bool {
63        self.inner.is_empty()
64    }
65
66    pub fn len(&self) -> usize {
67        self.inner.len()
68    }
69
70    pub fn peek(&self) -> Option<&T> {
71        self.inner.peek()
72    }
73
74    /// Moves the elements into a sorted vector and releases the old heap's backing-storage charge.
75    ///
76    /// # Warning
77    ///
78    /// `deallocate` only subtracts the backing-storage size. The elements may already have been
79    /// dropped, so it cannot determine the size of their separately allocated memory. Those charges
80    /// need separate cleanup, either explicitly or through a private memory context. A future
81    /// redesign should make this handling automatic.
82    ///
83    /// In this example, the heap uses a new private `mem_ctx`. Values are dropped as they are
84    /// consumed, but their payload charges remain until the context is dropped on normal return,
85    /// errors, or cancellation.
86    ///
87    /// ```rust
88    /// # #![feature(allocator_api)]
89    /// # use prometheus::core::Atomic;
90    /// # use risingwave_common::memory::{MemMonitoredHeap, MemoryContext};
91    /// # use risingwave_common::metrics::TrAdderAtomic;
92    /// # let parent = MemoryContext::none();
93    /// let mem_ctx = MemoryContext::new(Some(parent), TrAdderAtomic::new(0));
94    /// let mut heap = MemMonitoredHeap::new_with(mem_ctx.clone());
95    /// heap.push(String::from("value"));
96    /// for value in heap.into_sorted_vec() {
97    ///     drop(value);
98    /// }
99    /// ```
100    pub fn into_sorted_vec(self) -> Vec<T, MonitoredGlobalAlloc> {
101        let old_cap = self.inner.capacity();
102        let alloc = MonitoredGlobalAlloc::with_memory_context(self.mem_ctx.clone());
103        let vec = self.inner.into_iter_sorted();
104
105        let mut ret = Vec::with_capacity_in(vec.len(), alloc);
106        ret.extend(vec);
107
108        self.mem_ctx
109            .add_unchecked(-((old_cap * size_of::<T>()) as i64));
110        ret
111    }
112
113    pub fn mem_context(&self) -> &MemoryContext {
114        &self.mem_ctx
115    }
116}
117
118impl<T> Extend<T> for MemMonitoredHeap<T>
119where
120    T: Ord + EstimateSize,
121{
122    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
123        let old_cap = self.inner.capacity();
124        let mut items_heap_size = 0usize;
125        let items = iter.into_iter();
126        self.inner.reserve_exact(items.size_hint().0);
127        for item in items {
128            items_heap_size += item.estimated_heap_size();
129            self.inner.push(item);
130        }
131
132        let new_cap = self.inner.capacity();
133
134        let diff = (new_cap - old_cap) * size_of::<T>() + items_heap_size;
135        self.mem_ctx.add_unchecked(diff as i64);
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::MemMonitoredHeap;
142    use crate::memory::MemoryContext;
143    use crate::metrics::LabelGuardedIntGauge;
144
145    #[test]
146    fn test_heap() {
147        let gauge = LabelGuardedIntGauge::test_int_gauge::<4>();
148        let mem_ctx = MemoryContext::root(gauge.clone(), u64::MAX);
149
150        let mut heap = MemMonitoredHeap::<u8>::new_with(mem_ctx);
151        assert_eq!(0, gauge.get());
152
153        heap.push(9u8);
154        heap.push(1u8);
155        assert_eq!(heap.inner.capacity() as i64, gauge.get());
156
157        heap.pop().unwrap();
158        assert_eq!(heap.inner.capacity() as i64, gauge.get());
159
160        assert!(!heap.is_empty());
161    }
162
163    #[test]
164    fn test_heap_drop() {
165        let gauge = LabelGuardedIntGauge::test_int_gauge::<4>();
166        let mem_ctx = MemoryContext::root(gauge.clone(), u64::MAX);
167
168        let vec = {
169            let mut heap = MemMonitoredHeap::<u8>::new_with(mem_ctx);
170            assert_eq!(0, gauge.get());
171
172            heap.push(9u8);
173            heap.push(1u8);
174            assert_eq!(heap.inner.capacity() as i64, gauge.get());
175
176            heap.into_sorted_vec()
177        };
178
179        assert_eq!(2, gauge.get());
180
181        drop(vec);
182
183        assert_eq!(0, gauge.get());
184    }
185}
risingwave_common/memory/monitored_heap.rs

risingwave_common/memory/
monitored_heap.rs
