risingwave_common_estimate_size/collections/
hashmap.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::collections::HashMap;
16use std::ops::Deref;
17
18use super::{AtomicMutGuard, MutGuard};
19use crate::{EstimateSize, KvSize};
20
21pub struct EstimatedHashMap<K, V> {
22    inner: HashMap<K, V>,
23    heap_size: KvSize,
24}
25
26impl<K, V> EstimateSize for EstimatedHashMap<K, V> {
27    fn estimated_heap_size(&self) -> usize {
28        self.heap_size.size()
29    }
30}
31
32impl<K, V> Deref for EstimatedHashMap<K, V> {
33    type Target = HashMap<K, V>;
34
35    fn deref(&self) -> &Self::Target {
36        &self.inner
37    }
38}
39
40impl<K, V> Default for EstimatedHashMap<K, V> {
41    fn default() -> Self {
42        Self {
43            inner: HashMap::default(),
44            heap_size: KvSize::default(),
45        }
46    }
47}
48
49impl<K, V> EstimatedHashMap<K, V>
50where
51    K: EstimateSize + Eq + std::hash::Hash,
52    V: EstimateSize,
53{
54    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
55        self.heap_size.add_val(&key);
56        self.heap_size.add_val(&value);
57        self.inner.insert(key, value)
58    }
59
60    pub fn get_mut(&mut self, key: &K) -> Option<MutGuard<'_, V>> {
61        self.inner
62            .get_mut(key)
63            .map(|v| MutGuard::new(v, &mut self.heap_size))
64    }
65
66    pub fn values_mut(&mut self) -> impl Iterator<Item = AtomicMutGuard<'_, V>> + '_ {
67        let heap_size = &self.heap_size;
68        self.inner
69            .values_mut()
70            .map(move |v| AtomicMutGuard::new(v, heap_size))
71    }
72
73    pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ {
74        self.heap_size.set(0);
75        self.inner.drain()
76    }
77}
risingwave_common_estimate_size/collections/hashmap.rs

risingwave_common_estimate_size/collections/
hashmap.rs
