memscope_rs/core/
optimized_tracker.rs

1//! Optimized memory tracker using Arc sharing to reduce clone overhead
2//!
3//! This module provides an optimized version of the memory tracker that uses
4//! Arc-based sharing to minimize expensive clone operations while maintaining
5//! full API compatibility.
6
7use crate::core::clone_monitor;
8use crate::core::optimized_types::OptimizedAllocationInfo;
9use crate::core::shared_types::{SharedAllocationCollection, SharedAllocationInfo};
10use crate::core::types::{MemoryStats, TrackingResult};
11use std::collections::HashMap;
12use std::sync::{Arc, Mutex};
13
14/// Optimized memory tracker that uses Arc sharing
15pub struct OptimizedMemoryTracker {
16    /// Active allocations using Arc sharing
17    active_allocations: Mutex<HashMap<usize, Arc<OptimizedAllocationInfo>>>,
18    /// Memory statistics
19    stats: Mutex<MemoryStats>,
20    /// Allocation history using Arc sharing
21    allocation_history: Mutex<Vec<Arc<OptimizedAllocationInfo>>>,
22}
23
24impl OptimizedMemoryTracker {
25    /// Create a new optimized memory tracker
26    pub fn new() -> Self {
27        Self {
28            active_allocations: Mutex::new(HashMap::new()),
29            stats: Mutex::new(MemoryStats::default()),
30            allocation_history: Mutex::new(Vec::new()),
31        }
32    }
33
34    /// Track an allocation with Arc sharing
35    pub fn track_allocation_shared(&self, ptr: usize, size: usize) -> TrackingResult<()> {
36        let allocation = OptimizedAllocationInfo::new(ptr, size);
37        let shared_allocation = Arc::new(allocation);
38
39        // Record that we're using Arc sharing instead of cloning
40        clone_monitor::record_avoided_clone("AllocationInfo", size);
41
42        match (self.active_allocations.try_lock(), self.stats.try_lock()) {
43            (Ok(mut active), Ok(mut stats)) => {
44                // Insert the Arc - no cloning needed
45                active.insert(ptr, Arc::clone(&shared_allocation));
46
47                // Update statistics
48                stats.total_allocations = stats.total_allocations.saturating_add(1);
49                stats.active_allocations = stats.active_allocations.saturating_add(1);
50                stats.active_memory = stats.active_memory.saturating_add(size);
51
52                // Add to history - sharing the same Arc
53                if let Ok(mut history) = self.allocation_history.try_lock() {
54                    history.push(shared_allocation);
55                }
56
57                Ok(())
58            }
59            _ => Err(crate::core::types::TrackingError::ThreadSafetyError(
60                "Lock contention".to_string(),
61            )),
62        }
63    }
64
65    /// Get active allocations as shared collection
66    pub fn get_active_allocations_shared(&self) -> TrackingResult<SharedAllocationCollection> {
67        match self.active_allocations.lock() {
68            Ok(active) => {
69                let shared_infos: Vec<SharedAllocationInfo> = active
70                    .values()
71                    .map(|arc_info| SharedAllocationInfo::new((**arc_info).clone()))
72                    .collect();
73
74                Ok(SharedAllocationCollection::new(shared_infos))
75            }
76            Err(_) => Err(crate::core::types::TrackingError::ThreadSafetyError(
77                "Lock contention".to_string(),
78            )),
79        }
80    }
81
82    /// Get allocation history as shared collection
83    pub fn get_allocation_history_shared(&self) -> TrackingResult<SharedAllocationCollection> {
84        match self.allocation_history.lock() {
85            Ok(history) => {
86                let shared_infos: Vec<SharedAllocationInfo> = history
87                    .iter()
88                    .map(|arc_info| SharedAllocationInfo::new((**arc_info).clone()))
89                    .collect();
90
91                Ok(SharedAllocationCollection::new(shared_infos))
92            }
93            Err(_) => Err(crate::core::types::TrackingError::ThreadSafetyError(
94                "Lock contention".to_string(),
95            )),
96        }
97    }
98
99    /// Associate variable with existing allocation (Arc-optimized)
100    pub fn associate_var_shared(
101        &self,
102        ptr: usize,
103        var_name: String,
104        type_name: String,
105    ) -> TrackingResult<()> {
106        match self.active_allocations.lock() {
107            Ok(mut active) => {
108                if let Some(arc_allocation) = active.get_mut(&ptr) {
109                    // Create a new allocation with updated info
110                    let mut updated_allocation = (**arc_allocation).clone();
111                    updated_allocation.var_name =
112                        Some(crate::core::string_pool::intern_string(&var_name));
113                    updated_allocation.type_name =
114                        Some(crate::core::string_pool::intern_string(&type_name));
115
116                    // Replace with new Arc
117                    *arc_allocation = Arc::new(updated_allocation);
118
119                    Ok(())
120                } else {
121                    Err(crate::core::types::TrackingError::InvalidPointer(
122                        "Allocation not found".to_string(),
123                    ))
124                }
125            }
126            Err(_) => Err(crate::core::types::TrackingError::ThreadSafetyError(
127                "Lock contention".to_string(),
128            )),
129        }
130    }
131
132    /// Get memory statistics
133    pub fn get_stats(&self) -> TrackingResult<MemoryStats> {
134        match self.stats.lock() {
135            Ok(stats) => Ok(stats.clone()),
136            Err(_) => Err(crate::core::types::TrackingError::ThreadSafetyError(
137                "Lock contention".to_string(),
138            )),
139        }
140    }
141
142    /// Get clone optimization statistics
143    pub fn get_clone_stats(&self) -> clone_monitor::CloneMonitorStats {
144        clone_monitor::get_clone_stats()
145    }
146}
147
148impl Default for OptimizedMemoryTracker {
149    fn default() -> Self {
150        Self::new()
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157
158    #[test]
159    fn test_new_optimized_tracker() {
160        // Test creating a new optimized tracker
161        let tracker = OptimizedMemoryTracker::new();
162        let stats = tracker.get_stats().expect("Failed to get stats");
163
164        assert_eq!(stats.active_allocations, 0);
165        assert_eq!(stats.active_memory, 0);
166        assert_eq!(stats.total_allocations, 0);
167    }
168
169    #[test]
170    fn test_track_allocation_shared() {
171        // Test tracking an allocation with Arc sharing
172        let tracker = OptimizedMemoryTracker::new();
173        let ptr = 0x1000;
174        let size = 1024;
175
176        tracker
177            .track_allocation_shared(ptr, size)
178            .expect("Failed to track allocation");
179
180        let stats = tracker.get_stats().expect("Failed to get stats");
181        assert_eq!(stats.active_allocations, 1);
182        assert_eq!(stats.active_memory, size);
183        assert_eq!(stats.total_allocations, 1);
184    }
185
186    #[test]
187    fn test_track_multiple_allocations() {
188        // Test tracking multiple allocations
189        let tracker = OptimizedMemoryTracker::new();
190
191        for i in 0..5 {
192            let ptr = 0x1000 + (i * 0x100);
193            let size = 256 * (i + 1);
194
195            tracker
196                .track_allocation_shared(ptr, size)
197                .expect("Failed to track allocation");
198        }
199
200        let stats = tracker.get_stats().expect("Failed to get stats");
201        assert_eq!(stats.active_allocations, 5);
202        assert_eq!(stats.total_allocations, 5);
203
204        // Total memory: 256 + 512 + 768 + 1024 + 1280 = 3840
205        assert_eq!(stats.active_memory, 3840);
206    }
207
208    #[test]
209    fn test_get_active_allocations_shared() {
210        // Test getting active allocations as shared collection
211        let tracker = OptimizedMemoryTracker::new();
212
213        // Track some allocations
214        tracker
215            .track_allocation_shared(0x2000, 512)
216            .expect("Failed to track allocation");
217        tracker
218            .track_allocation_shared(0x3000, 1024)
219            .expect("Failed to track allocation");
220
221        let shared_collection = tracker
222            .get_active_allocations_shared()
223            .expect("Failed to get active allocations");
224
225        assert_eq!(shared_collection.len(), 2);
226
227        // Check that allocations are present
228        let ptrs: Vec<usize> = shared_collection.iter().map(|info| info.ptr()).collect();
229        assert!(ptrs.contains(&0x2000));
230        assert!(ptrs.contains(&0x3000));
231    }
232
233    #[test]
234    fn test_get_allocation_history_shared() {
235        // Test getting allocation history as shared collection
236        let tracker = OptimizedMemoryTracker::new();
237
238        // Track some allocations
239        for i in 0..3 {
240            tracker
241                .track_allocation_shared(0x4000 + i * 0x100, 256)
242                .expect("Failed to track allocation");
243        }
244
245        let history = tracker
246            .get_allocation_history_shared()
247            .expect("Failed to get allocation history");
248
249        assert_eq!(history.len(), 3);
250    }
251
252    #[test]
253    fn test_associate_var_shared() {
254        // Test associating variable with existing allocation
255        let tracker = OptimizedMemoryTracker::new();
256        let ptr = 0x5000;
257
258        // First track the allocation
259        tracker
260            .track_allocation_shared(ptr, 1024)
261            .expect("Failed to track allocation");
262
263        // Then associate a variable with it
264        tracker
265            .associate_var_shared(ptr, "test_variable".to_string(), "TestType".to_string())
266            .expect("Failed to associate variable");
267
268        // Verify the association
269        let allocations = tracker
270            .get_active_allocations_shared()
271            .expect("Failed to get allocations");
272
273        let found = allocations.iter().find(|info| info.ptr() == ptr);
274
275        assert!(found.is_some());
276        let allocation = found.unwrap();
277        assert!(allocation.var_name_str().is_some());
278        assert!(allocation.type_name_str().is_some());
279    }
280
281    #[test]
282    fn test_associate_var_non_existing_allocation() {
283        // Test associating variable with non-existing allocation
284        let tracker = OptimizedMemoryTracker::new();
285
286        let result = tracker.associate_var_shared(
287            0x9999,
288            "non_existing".to_string(),
289            "TestType".to_string(),
290        );
291
292        assert!(result.is_err());
293        match result {
294            Err(crate::core::types::TrackingError::InvalidPointer(_)) => {}
295            _ => panic!("Expected InvalidPointer error"),
296        }
297    }
298
299    #[test]
300    fn test_default_implementation() {
301        // Test the Default trait implementation
302        let tracker = OptimizedMemoryTracker::default();
303        let stats = tracker.get_stats().expect("Failed to get stats");
304
305        assert_eq!(stats.active_allocations, 0);
306        assert_eq!(stats.active_memory, 0);
307    }
308
309    #[test]
310    fn test_clone_stats_retrieval() {
311        // Test getting clone optimization statistics
312        let tracker = OptimizedMemoryTracker::new();
313
314        // Track some allocations which should trigger clone monitoring
315        tracker
316            .track_allocation_shared(0x6000, 2048)
317            .expect("Failed to track allocation");
318
319        let clone_stats = tracker.get_clone_stats();
320
321        // Should have recorded avoided clones
322        assert!(clone_stats.avoided_clones > 0);
323    }
324
325    #[test]
326    fn test_arc_sharing_efficiency() {
327        // Test that Arc sharing is actually being used efficiently
328        let tracker = OptimizedMemoryTracker::new();
329
330        // Track an allocation
331        tracker
332            .track_allocation_shared(0x7000, 1024)
333            .expect("Failed to track allocation");
334
335        // Get the allocation multiple times
336        let collection1 = tracker
337            .get_active_allocations_shared()
338            .expect("Failed to get allocations");
339        let collection2 = tracker
340            .get_active_allocations_shared()
341            .expect("Failed to get allocations");
342
343        // Both collections should contain data
344        assert_eq!(collection1.len(), 1);
345        assert_eq!(collection2.len(), 1);
346    }
347
348    #[test]
349    fn test_large_allocation_tracking() {
350        // Test tracking very large allocations
351        let tracker = OptimizedMemoryTracker::new();
352        let large_size = 1024 * 1024 * 10; // 10 MB
353
354        tracker
355            .track_allocation_shared(0x8000, large_size)
356            .expect("Failed to track large allocation");
357
358        let stats = tracker.get_stats().expect("Failed to get stats");
359        assert_eq!(stats.active_memory, large_size);
360    }
361
362    #[test]
363    fn test_allocation_history_persistence() {
364        // Test that allocation history persists across operations
365        let tracker = OptimizedMemoryTracker::new();
366
367        // Track allocations
368        for i in 0..5 {
369            tracker
370                .track_allocation_shared(0x9000 + i * 0x100, 512)
371                .expect("Failed to track allocation");
372        }
373
374        // Get history multiple times
375        let history1 = tracker
376            .get_allocation_history_shared()
377            .expect("Failed to get history");
378        let history2 = tracker
379            .get_allocation_history_shared()
380            .expect("Failed to get history");
381
382        assert_eq!(history1.len(), history2.len());
383        assert_eq!(history1.len(), 5);
384    }
385}
memscope_rs/core/optimized_tracker.rs

memscope_rs/core/
optimized_tracker.rs
