memscope_rs/
lib.rs

1//! Memory tracking and visualization tools for Rust applications.
2//!
3//! This crate provides tools for tracking memory allocations and visualizing
4//! memory usage in Rust applications. It includes a custom global allocator
5//! that tracks all heap allocations and deallocations, and provides utilities
6//! for exporting memory usage data in various formats.
7
8#![warn(missing_docs)]
9
10pub mod allocator;
11pub mod export;
12pub mod export_enhanced;
13pub mod tracker;
14pub mod types;
15
16// Re-export main types for easier use
17pub use allocator::TrackingAllocator;
18pub use tracker::{get_global_tracker, MemoryTracker};
19pub use types::{AllocationInfo, TrackingError, TrackingResult};
20
21// Set up the global allocator when the tracking-allocator feature is enabled
22#[cfg(feature = "tracking-allocator")]
23#[global_allocator]
24/// Global tracking allocator instance used when the tracking-allocator feature is enabled.
25pub static GLOBAL: TrackingAllocator = TrackingAllocator::new();
26
27/// Trait for types that can be tracked by the memory tracker.
28pub trait Trackable {
29    /// Get the pointer to the heap allocation for this value.
30    fn get_heap_ptr(&self) -> Option<usize>;
31
32    /// Get the type name for this value.
33    fn get_type_name(&self) -> &'static str;
34}
35
36// Implement Trackable for common heap-allocated types
37impl<T> Trackable for Vec<T> {
38    fn get_heap_ptr(&self) -> Option<usize> {
39        if self.capacity() > 0 {
40            Some(self.as_ptr() as usize)
41        } else {
42            None
43        }
44    }
45
46    fn get_type_name(&self) -> &'static str {
47        std::any::type_name::<Vec<T>>()
48    }
49}
50
51impl Trackable for String {
52    fn get_heap_ptr(&self) -> Option<usize> {
53        if self.capacity() > 0 {
54            Some(self.as_ptr() as usize)
55        } else {
56            None
57        }
58    }
59
60    fn get_type_name(&self) -> &'static str {
61        "String"
62    }
63}
64
65impl<T> Trackable for Box<T> {
66    fn get_heap_ptr(&self) -> Option<usize> {
67        Some(self.as_ref() as *const T as usize)
68    }
69
70    fn get_type_name(&self) -> &'static str {
71        std::any::type_name::<Box<T>>()
72    }
73}
74
75impl<T> Trackable for std::rc::Rc<T> {
76    fn get_heap_ptr(&self) -> Option<usize> {
77        // For Rc, the allocation tracking is complex because Rc uses a control block
78        // We'll track the Rc itself rather than the inner data to avoid pointer issues
79        Some(std::rc::Rc::as_ptr(self) as usize)
80    }
81
82    fn get_type_name(&self) -> &'static str {
83        std::any::type_name::<std::rc::Rc<T>>()
84    }
85}
86
87impl<T> Trackable for std::sync::Arc<T> {
88    fn get_heap_ptr(&self) -> Option<usize> {
89        // For Arc, the allocation tracking is complex because Arc uses a control block
90        // We'll track the Arc itself rather than the inner data to avoid pointer issues
91        Some(std::sync::Arc::as_ptr(self) as usize)
92    }
93
94    fn get_type_name(&self) -> &'static str {
95        std::any::type_name::<std::sync::Arc<T>>()
96    }
97}
98
99/// Macro to track a variable's memory allocation.
100///
101/// This macro associates a variable name with its heap allocation,
102/// allowing the memory tracker to provide meaningful names in reports.
103///
104/// # Example
105/// ```rust
106/// use memscope_rs::track_var;
107///
108/// let my_vec = vec![1, 2, 3, 4, 5];
109/// track_var!(my_vec);
110/// ```
111#[macro_export]
112macro_rules! track_var {
113    ($var:ident) => {
114        $crate::_track_var_impl(&$var, stringify!($var))
115    };
116}
117
118/// Internal implementation function for the track_var! macro.
119/// This function should not be called directly.
120#[doc(hidden)]
121pub fn _track_var_impl<T: Trackable>(var: &T, var_name: &str) -> TrackingResult<()> {
122    if let Some(ptr) = var.get_heap_ptr() {
123        let tracker = get_global_tracker();
124        tracker.associate_var(ptr, var_name.to_string(), var.get_type_name().to_string())
125    } else {
126        // Variable doesn't have a heap allocation (e.g., empty Vec)
127        Ok(())
128    }
129}
130
131/// Initialize the memory tracking system.
132///
133/// This function sets up the tracing subscriber and prepares the global tracker.
134/// Call this early in your application, typically in main().
135///
136/// # Example
137/// ```rust
138/// memscope_rs::init();
139/// // Your application code here
140/// ```
141pub fn init() {
142    use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
143
144    tracing_subscriber::registry()
145        .with(
146            tracing_subscriber::EnvFilter::try_from_default_env()
147                .unwrap_or_else(|_| "memscope_rs=info".into()),
148        )
149        .with(tracing_subscriber::fmt::layer())
150        .init();
151
152    tracing::info!("memscope-rs initialized");
153}