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

risingwave_common_estimate_size/
lib.rs
