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

risingwave_common_estimate_size/
lib.rs
