iota_util_mem/
lib.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5// Copyright 2020 Parity Technologies
6//
7// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
8// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
9// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
10// option. This file may not be copied, modified, or distributed
11// except according to those terms.
12
13//! Crate for parity memory management related utilities.
14//! It includes global allocator choice, heap measurement and
15//! memory erasure.
16
17#![allow(clippy::all)]
18#![cfg_attr(not(feature = "std"), no_std)]
19
20#[cfg(not(feature = "std"))]
21extern crate alloc;
22
23// default allocator used
24mod memory_stats_noop;
25
26pub mod allocators;
27
28use memory_stats_noop as memory_stats;
29
30#[cfg(any(
31    any(target_os = "macos", target_os = "ios"),
32    feature = "estimate-heapsize"
33))]
34pub mod sizeof;
35
36/// This is a copy of patched crate `malloc_size_of` as a module.
37/// We need to have it as an inner module to be able to define our own traits
38/// implementation, if at some point the trait become standard enough we could
39/// use the right way of doing it by implementing it in our type traits crates.
40/// At this time moving this trait to the primitive types level would impact too
41/// much of the dependencies to be easily manageable.
42#[macro_use]
43mod malloc_size;
44
45pub mod external_impls;
46
47pub use allocators::MallocSizeOfExt;
48pub use iota_util_mem_derive::*;
49pub use malloc_size::{MallocShallowSizeOf, MallocSizeOf, MallocSizeOfOps};
50
51/// Heap size of structure.
52///
53/// Structure can be anything that implements MallocSizeOf.
54pub fn malloc_size<T: MallocSizeOf + ?Sized>(t: &T) -> usize {
55    MallocSizeOf::size_of(t, &mut allocators::new_malloc_size_ops())
56}
57
58/// An error related to the memory stats gathering.
59#[derive(Clone, Debug)]
60pub struct MemoryStatsError(memory_stats::Error);
61
62#[cfg(feature = "std")]
63impl std::fmt::Display for MemoryStatsError {
64    fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
65        self.0.fmt(fmt)
66    }
67}
68
69#[cfg(feature = "std")]
70impl std::error::Error for MemoryStatsError {}
71
72/// Snapshot of collected memory metrics.
73#[non_exhaustive]
74#[derive(Debug, Clone)]
75pub struct MemoryAllocationSnapshot {
76    /// Total resident memory, in bytes.
77    pub resident: u64,
78    /// Total allocated memory, in bytes.
79    pub allocated: u64,
80}
81
82/// Accessor to the allocator internals.
83#[derive(Clone)]
84pub struct MemoryAllocationTracker(self::memory_stats::MemoryAllocationTracker);
85
86impl MemoryAllocationTracker {
87    /// Create an instance of an allocation tracker.
88    pub fn new() -> Result<Self, MemoryStatsError> {
89        self::memory_stats::MemoryAllocationTracker::new()
90            .map(MemoryAllocationTracker)
91            .map_err(MemoryStatsError)
92    }
93
94    /// Create an allocation snapshot.
95    pub fn snapshot(&self) -> Result<MemoryAllocationSnapshot, MemoryStatsError> {
96        self.0.snapshot().map_err(MemoryStatsError)
97    }
98}
99
100#[cfg(feature = "std")]
101#[cfg(test)]
102mod test {
103    use std::sync::Arc;
104
105    use super::{MallocSizeOf, MallocSizeOfExt, malloc_size};
106
107    #[test]
108    fn test_arc() {
109        let val = Arc::new("test".to_string());
110        let s = val.malloc_size_of();
111        assert!(s > 0);
112    }
113
114    #[test]
115    fn test_dyn() {
116        trait Augmented: MallocSizeOf {}
117        impl Augmented for Vec<u8> {}
118        let val: Arc<dyn Augmented> = Arc::new(vec![0u8; 1024]);
119        assert!(malloc_size(&*val) > 1000);
120    }
121}