kaspa_utils/
mem_size.rs

1//! Defines a [`MemSizeEstimator`] trait and a companying [`MemMode`] which are used to
2//! estimate sizes of run-time objects in memory, including deep heap allocations. See
3//! struct-level docs for more details.
4
5use std::{collections::HashSet, sync::Arc};
6
7use parking_lot::RwLock;
8
9/// The memory mode of the tracked object
10#[derive(Debug, Clone, Copy)]
11pub enum MemMode {
12    Bytes,
13    Units,
14    Undefined,
15}
16
17/// The contract for estimating deep memory size owned by this object. Implementors
18/// are expected to support only a single function - bytes or units. Objects which are
19/// containers of items with pre-compilation known static size should implement the `_units`
20/// estimation and return the number of logical items (i.e. number of items in
21/// the container). Objects with more complex and varying runtime sizes should implement the `_bytes` estimation.
22///
23/// By panicking on the remaining unimplemented function we ensure that tests will catch any inconsistency over the
24/// used units between the object implementing the contract and the code using its size for various purposes (e.g. cache
25/// size tracking).
26/// Exceptions to the above are objects which delegate the estimation to an underlying inner field (such as Arc or RwLock),
27/// which should implement both methods and rely on the inner object to be implemented correctly
28pub trait MemSizeEstimator {
29    /// Estimates the size of this object depending on the passed mem mode
30    fn estimate_size(&self, mem_mode: MemMode) -> usize {
31        match mem_mode {
32            MemMode::Bytes => self.estimate_mem_bytes(),
33            MemMode::Units => self.estimate_mem_units(),
34            MemMode::Undefined => unimplemented!(),
35        }
36    }
37
38    /// Estimates the (deep) size of this object in bytes (including heap owned inner data)
39    fn estimate_mem_bytes(&self) -> usize {
40        unimplemented!()
41    }
42
43    /// Estimates the number of units this object holds in memory where the unit byte size is usually
44    /// a constant known to the caller as well (and hence we avoid computing it over and over)
45    fn estimate_mem_units(&self) -> usize {
46        unimplemented!()
47    }
48}
49
50impl MemSizeEstimator for u64 {}
51impl MemSizeEstimator for u32 {}
52impl MemSizeEstimator for u16 {}
53impl MemSizeEstimator for u8 {}
54impl MemSizeEstimator for i64 {}
55impl MemSizeEstimator for i32 {}
56impl MemSizeEstimator for i16 {}
57impl MemSizeEstimator for i8 {}
58
59impl<T> MemSizeEstimator for Vec<T> {
60    fn estimate_mem_units(&self) -> usize {
61        self.len()
62    }
63}
64
65impl<T, S> MemSizeEstimator for HashSet<T, S> {
66    fn estimate_mem_units(&self) -> usize {
67        self.len()
68    }
69}
70
71impl<T: MemSizeEstimator> MemSizeEstimator for Arc<T> {
72    fn estimate_mem_bytes(&self) -> usize {
73        self.as_ref().estimate_mem_bytes() + size_of::<Self>()
74    }
75
76    fn estimate_mem_units(&self) -> usize {
77        self.as_ref().estimate_mem_units()
78    }
79}
80
81impl<T: MemSizeEstimator> MemSizeEstimator for RwLock<T> {
82    fn estimate_mem_bytes(&self) -> usize {
83        self.read().estimate_mem_bytes() + size_of::<Self>()
84    }
85
86    fn estimate_mem_units(&self) -> usize {
87        self.read().estimate_mem_units()
88    }
89}
kaspa_utils/mem_size.rs

kaspa_utils/
mem_size.rs
