risingwave_common_estimate_size/collections/
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::ops::{Deref, DerefMut};
16
17use super::{EstimateSize, KvSize};
18
19pub mod vecdeque;
20pub use vecdeque::EstimatedVecDeque;
21pub mod hashmap;
22pub use hashmap::EstimatedHashMap;
23pub mod btreemap;
24pub use btreemap::EstimatedBTreeMap;
25pub mod vec;
26pub use vec::EstimatedVec;
27pub mod hashset;
28pub use hashset::EstimatedHashSet;
29
30mod private {
31    use super::*;
32
33    /// A trait that dispatches the size update method regarding the mutability of the reference
34    /// to the [`KvSize`].
35    pub trait GenericKvSize {
36        fn update_size(&mut self, from: usize, to: usize);
37    }
38
39    /// For mutable references, the size can be directly updated.
40    impl GenericKvSize for &'_ mut KvSize {
41        fn update_size(&mut self, from: usize, to: usize) {
42            self.add_size(to);
43            self.sub_size(from);
44        }
45    }
46
47    /// For immutable references, the size is updated atomically.
48    impl GenericKvSize for &'_ KvSize {
49        fn update_size(&mut self, from: usize, to: usize) {
50            self.update_size_atomic(from, to)
51        }
52    }
53}
54
55/// A guard holding a mutable reference to a value in a collection. When dropped, the size of the
56/// collection will be updated.
57pub struct MutGuard<'a, V, S = &'a mut KvSize>
58where
59    V: EstimateSize,
60    S: private::GenericKvSize,
61{
62    inner: &'a mut V,
63    // The size of the original value
64    original_val_size: usize,
65    // The total size of a collection
66    total_size: S,
67}
68
69/// Similar to [`MutGuard`], but the size is updated atomically. Useful for creating shared
70/// references to the entries in a collection.
71pub type AtomicMutGuard<'a, V> = MutGuard<'a, V, &'a KvSize>;
72
73impl<'a, V, S> MutGuard<'a, V, S>
74where
75    V: EstimateSize,
76    S: private::GenericKvSize,
77{
78    pub fn new(inner: &'a mut V, total_size: S) -> Self {
79        let original_val_size = inner.estimated_size();
80        Self {
81            inner,
82            original_val_size,
83            total_size,
84        }
85    }
86}
87
88impl<V, S> Drop for MutGuard<'_, V, S>
89where
90    V: EstimateSize,
91    S: private::GenericKvSize,
92{
93    fn drop(&mut self) {
94        self.total_size
95            .update_size(self.original_val_size, self.inner.estimated_size());
96    }
97}
98
99impl<V, S> Deref for MutGuard<'_, V, S>
100where
101    V: EstimateSize,
102    S: private::GenericKvSize,
103{
104    type Target = V;
105
106    fn deref(&self) -> &Self::Target {
107        self.inner
108    }
109}
110
111impl<V, S> DerefMut for MutGuard<'_, V, S>
112where
113    V: EstimateSize,
114    S: private::GenericKvSize,
115{
116    fn deref_mut(&mut self) -> &mut Self::Target {
117        self.inner
118    }
119}
risingwave_common_estimate_size/collections/mod.rs

risingwave_common_estimate_size/collections/
mod.rs
