risingwave_common_estimate_size/
lib.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
15#![feature(allocator_api)]
16#![feature(btree_cursors)]
17
18pub mod collections;
19
20use std::cmp::Reverse;
21use std::marker::PhantomData;
22use std::sync::atomic::{AtomicUsize, Ordering};
23
24use bytes::Bytes;
25use fixedbitset::FixedBitSet;
26pub use risingwave_common_proc_macro::EstimateSize;
27
28/// The trait for estimating the actual memory usage of a struct.
29///
30/// Used for cache eviction now.
31pub trait EstimateSize {
32    /// The estimated heap size of the current struct in bytes.
33    fn estimated_heap_size(&self) -> usize;
34
35    /// The estimated total size of the current struct in bytes, including the `estimated_heap_size`
36    /// and the size of `Self`.
37    fn estimated_size(&self) -> usize
38    where
39        Self: Sized,
40    {
41        self.estimated_heap_size() + std::mem::size_of::<Self>()
42    }
43}
44
45impl EstimateSize for FixedBitSet {
46    fn estimated_heap_size(&self) -> usize {
47        std::mem::size_of_val(self.as_slice())
48    }
49}
50
51impl EstimateSize for String {
52    fn estimated_heap_size(&self) -> usize {
53        self.capacity()
54    }
55}
56
57impl<T: EstimateSize> EstimateSize for Option<T> {
58    fn estimated_heap_size(&self) -> usize {
59        if let Some(inner) = self {
60            inner.estimated_heap_size()
61        } else {
62            0
63        }
64    }
65}
66
67/// SAFETY: `Bytes` can store a pointer in some cases, that may cause the size
68/// of a `Bytes` be calculated more than one and when memory stats is larger than the real value.
69impl EstimateSize for Bytes {
70    fn estimated_heap_size(&self) -> usize {
71        self.len()
72    }
73}
74
75impl EstimateSize for Box<str> {
76    fn estimated_heap_size(&self) -> usize {
77        self.len()
78    }
79}
80
81impl EstimateSize for serde_json::Value {
82    fn estimated_heap_size(&self) -> usize {
83        // FIXME: implement correct size
84        // https://github.com/risingwavelabs/risingwave/issues/9377
85        match self {
86            Self::Null => 0,
87            Self::Bool(_) => 0,
88            Self::Number(_) => 0,
89            Self::String(s) => s.estimated_heap_size(),
90            Self::Array(v) => std::mem::size_of::<Self>() * v.capacity(),
91            Self::Object(map) => std::mem::size_of::<Self>() * map.len(),
92        }
93    }
94}
95
96impl EstimateSize for jsonbb::Value {
97    fn estimated_heap_size(&self) -> usize {
98        self.capacity()
99    }
100}
101
102impl EstimateSize for jsonbb::Builder {
103    fn estimated_heap_size(&self) -> usize {
104        self.capacity()
105    }
106}
107
108impl<T1: EstimateSize, T2: EstimateSize> EstimateSize for (T1, T2) {
109    fn estimated_heap_size(&self) -> usize {
110        self.0.estimated_heap_size() + self.1.estimated_heap_size()
111    }
112}
113
114macro_rules! primitive_estimate_size_impl {
115    ($($t:ty)*) => ($(
116        impl ZeroHeapSize for $t {}
117    )*)
118}
119
120primitive_estimate_size_impl! { () usize u8 u16 u32 u64 u128 isize i8 i16 i32 i64 i128 f32 f64 bool }
121
122pub trait ZeroHeapSize {}
123
124impl<T: ZeroHeapSize> EstimateSize for T {
125    fn estimated_heap_size(&self) -> usize {
126        0
127    }
128}
129
130impl<T: ZeroHeapSize> EstimateSize for Vec<T> {
131    fn estimated_heap_size(&self) -> usize {
132        std::mem::size_of::<T>() * self.capacity()
133    }
134}
135
136impl<T: ZeroHeapSize> EstimateSize for Box<[T]> {
137    fn estimated_heap_size(&self) -> usize {
138        std::mem::size_of::<T>() * self.len()
139    }
140}
141
142impl<T: EstimateSize> EstimateSize for Reverse<T> {
143    fn estimated_heap_size(&self) -> usize {
144        self.0.estimated_heap_size()
145    }
146}
147
148impl<T: ZeroHeapSize, const LEN: usize> EstimateSize for [T; LEN] {
149    fn estimated_heap_size(&self) -> usize {
150        0
151    }
152}
153
154impl ZeroHeapSize for rust_decimal::Decimal {}
155
156impl ZeroHeapSize for ethnum::I256 {}
157
158impl<T> ZeroHeapSize for PhantomData<T> {}
159
160/// The size of the collection.
161///
162/// We use an atomic value here to enable operating the size without a mutable reference.
163/// See [`collections::AtomicMutGuard`] for more details.
164///
165/// In the most cases, we have the mutable reference of this struct, so we can directly
166/// operate the underlying value.
167#[derive(Default)]
168pub struct KvSize(AtomicUsize);
169
170/// Clone the [`KvSize`] will duplicate the underlying value.
171impl Clone for KvSize {
172    fn clone(&self) -> Self {
173        Self(self.size().into())
174    }
175}
176
177impl KvSize {
178    pub fn new() -> Self {
179        Self(0.into())
180    }
181
182    pub fn with_size(size: usize) -> Self {
183        Self(size.into())
184    }
185
186    pub fn add<K: EstimateSize, V: EstimateSize>(&mut self, key: &K, val: &V) {
187        self.add_size(key.estimated_size());
188        self.add_size(val.estimated_size());
189    }
190
191    pub fn sub<K: EstimateSize, V: EstimateSize>(&mut self, key: &K, val: &V) {
192        self.sub_size(key.estimated_size());
193        self.sub_size(val.estimated_size());
194    }
195
196    /// Add the size of `val` and return it.
197    pub fn add_val<V: EstimateSize>(&mut self, val: &V) -> usize {
198        let size = val.estimated_size();
199        self.add_size(size);
200        size
201    }
202
203    pub fn sub_val<V: EstimateSize>(&mut self, val: &V) {
204        self.sub_size(val.estimated_size());
205    }
206
207    pub fn add_size(&mut self, size: usize) {
208        let this = self.0.get_mut(); // get the underlying value since we have a mutable reference
209        *this = this.saturating_add(size);
210    }
211
212    pub fn sub_size(&mut self, size: usize) {
213        let this = self.0.get_mut(); // get the underlying value since we have a mutable reference
214        *this = this.saturating_sub(size);
215    }
216
217    /// Update the size of the collection by `to - from` atomically, i.e., without a mutable reference.
218    pub fn update_size_atomic(&self, from: usize, to: usize) {
219        let _ = (self.0).fetch_update(Ordering::Relaxed, Ordering::Relaxed, |this| {
220            Some(this.saturating_add(to).saturating_sub(from))
221        });
222    }
223
224    pub fn set(&mut self, size: usize) {
225        self.0 = size.into();
226    }
227
228    pub fn size(&self) -> usize {
229        self.0.load(Ordering::Relaxed)
230    }
231}
risingwave_common_estimate_size/lib.rs

risingwave_common_estimate_size/
lib.rs
