memscope_rs/export/
data_localizer.rs

1//! data localizer - reduce global state access overhead
2//!
3//! This module implements data localization functionality,
4//! fetching all export data at once to avoid repeated access to global state,
5//! thus significantly improving export performance.
6
7use crate::analysis::unsafe_ffi_tracker::{
8    get_global_unsafe_ffi_tracker, EnhancedAllocationInfo, UnsafeFFIStats,
9};
10use crate::core::scope_tracker::get_global_scope_tracker;
11use crate::core::tracker::get_global_tracker;
12use crate::core::types::MemoryStats;
13use crate::core::types::ScopeInfo;
14use crate::core::types::{AllocationInfo, TrackingError, TrackingResult};
15use std::time::{Duration, Instant};
16
17/// data localizer - fetch all export data at once to avoid repeated access to global state
18pub struct DataLocalizer {
19    /// cached basic allocation data
20    cached_allocations: Option<Vec<AllocationInfo>>,
21    /// cached FFI enhanced data
22    cached_ffi_data: Option<Vec<EnhancedAllocationInfo>>,
23    /// cached stats
24    cached_stats: Option<MemoryStats>,
25    /// cached FFI stats
26    cached_ffi_stats: Option<UnsafeFFIStats>,
27    /// cached scope info
28    cached_scope_info: Option<Vec<ScopeInfo>>,
29    /// last update time
30    last_update: Instant,
31    /// cache ttl
32    cache_ttl: Duration,
33}
34
35/// localized export data, containing all necessary information
36#[derive(Debug, Clone)]
37pub struct LocalizedExportData {
38    /// basic memory allocation info
39    pub allocations: Vec<AllocationInfo>,
40    /// FFI enhanced allocation info
41    pub enhanced_allocations: Vec<EnhancedAllocationInfo>,
42    /// stats
43    pub stats: MemoryStats,
44    /// FFI stats
45    pub ffi_stats: UnsafeFFIStats,
46    /// scope info
47    pub scope_info: Vec<ScopeInfo>,
48    /// timestamp
49    pub timestamp: Instant,
50}
51
52/// data gathering stats
53#[derive(Debug, Clone)]
54pub struct DataGatheringStats {
55    /// total time ms
56    pub total_time_ms: u64,
57    /// basic data time ms
58    pub basic_data_time_ms: u64,
59    /// FFI data time ms
60    pub ffi_data_time_ms: u64,
61    /// scope data time ms
62    pub scope_data_time_ms: u64,
63    /// allocation count
64    pub allocation_count: usize,
65    /// ffi allocation count
66    pub ffi_allocation_count: usize,
67    /// scope count
68    pub scope_count: usize,
69}
70
71impl DataLocalizer {
72    /// create new data localizer
73    pub fn new() -> Self {
74        Self {
75            cached_allocations: None,
76            cached_ffi_data: None,
77            cached_stats: None,
78            cached_ffi_stats: None,
79            cached_scope_info: None,
80            last_update: Instant::now(),
81            cache_ttl: Duration::from_millis(100), // 100ms cache, avoid too frequent data fetching
82        }
83    }
84
85    /// create data localizer with custom cache ttl
86    pub fn with_cache_ttl(cache_ttl: Duration) -> Self {
87        Self {
88            cached_allocations: None,
89            cached_ffi_data: None,
90            cached_stats: None,
91            cached_ffi_stats: None,
92            cached_scope_info: None,
93            last_update: Instant::now(),
94            cache_ttl,
95        }
96    }
97
98    /// gather all export data at once to avoid repeated access to global state
99    pub fn gather_all_export_data(
100        &mut self,
101    ) -> TrackingResult<(LocalizedExportData, DataGatheringStats)> {
102        let total_start = Instant::now();
103
104        println!("šŸ”„ start data localization to reduce global state access...");
105
106        // check if cache is still valid
107        if self.is_cache_valid() {
108            println!("āœ… using cached data, skipping repeated fetching");
109            return self.get_cached_data();
110        }
111
112        // step 1: get basic memory tracking data
113        let basic_start = Instant::now();
114        let tracker = get_global_tracker();
115        let allocations = tracker.get_active_allocations().map_err(|e| {
116            TrackingError::ExportError(format!("get active allocations failed: {}", e))
117        })?;
118        let stats = tracker
119            .get_stats()
120            .map_err(|e| TrackingError::ExportError(format!("get stats failed: {}", e)))?;
121        let basic_time = basic_start.elapsed();
122
123        // step 2: get ffi related data
124        let ffi_start = Instant::now();
125        let ffi_tracker = get_global_unsafe_ffi_tracker();
126        let enhanced_allocations = ffi_tracker.get_enhanced_allocations().unwrap_or_else(|e| {
127            eprintln!(
128                "sāš ļø get enhanced allocations failed: {}, using empty data",
129                e
130            );
131            Vec::new()
132        });
133        let ffi_stats = ffi_tracker.get_stats();
134        let ffi_time = ffi_start.elapsed();
135
136        // step 3: get scope data
137        let scope_start = Instant::now();
138        let scope_tracker = get_global_scope_tracker();
139        let scope_info = scope_tracker.get_all_scopes();
140        let scope_time = scope_start.elapsed();
141
142        let total_time = total_start.elapsed();
143
144        // update cache
145        self.cached_allocations = Some(allocations.clone());
146        self.cached_ffi_data = Some(enhanced_allocations.clone());
147        self.cached_stats = Some(stats.clone());
148        self.cached_ffi_stats = Some(ffi_stats.clone());
149        self.cached_scope_info = Some(scope_info.clone());
150        self.last_update = Instant::now();
151
152        let localized_data = LocalizedExportData {
153            allocations: allocations.clone(),
154            enhanced_allocations: enhanced_allocations.clone(),
155            stats,
156            ffi_stats,
157            scope_info: scope_info.clone(),
158            timestamp: total_start,
159        };
160
161        let gathering_stats = DataGatheringStats {
162            total_time_ms: total_time.as_millis() as u64,
163            basic_data_time_ms: basic_time.as_millis() as u64,
164            ffi_data_time_ms: ffi_time.as_millis() as u64,
165            scope_data_time_ms: scope_time.as_millis() as u64,
166            allocation_count: allocations.len(),
167            ffi_allocation_count: enhanced_allocations.len(),
168            scope_count: scope_info.len(),
169        };
170
171        // print performance stats
172        println!("āœ… data localization completed:");
173        println!("   total time: {:?}", total_time);
174        println!(
175            "   basic data: {:?} ({} allocations)",
176            basic_time, gathering_stats.allocation_count
177        );
178        println!(
179            "   ffi data: {:?} ({} enhanced allocations)",
180            ffi_time, gathering_stats.ffi_allocation_count
181        );
182        println!(
183            "   scope data: {:?} ({} scopes)",
184            scope_time, gathering_stats.scope_count
185        );
186        println!(
187            "   data localization avoided {} global state accesses",
188            self.estimate_avoided_global_accesses(&gathering_stats)
189        );
190
191        Ok((localized_data, gathering_stats))
192    }
193
194    /// refresh cache and gather all export data
195    pub fn refresh_cache(&mut self) -> TrackingResult<(LocalizedExportData, DataGatheringStats)> {
196        self.invalidate_cache();
197        self.gather_all_export_data()
198    }
199
200    /// check if cache is still valid
201    fn is_cache_valid(&self) -> bool {
202        self.cached_allocations.is_some()
203            && self.cached_ffi_data.is_some()
204            && self.cached_stats.is_some()
205            && self.cached_ffi_stats.is_some()
206            && self.cached_scope_info.is_some()
207            && self.last_update.elapsed() < self.cache_ttl
208    }
209
210    /// get cached data
211    fn get_cached_data(&self) -> TrackingResult<(LocalizedExportData, DataGatheringStats)> {
212        let localized_data = LocalizedExportData {
213            allocations: self.cached_allocations.as_ref().unwrap().clone(),
214            enhanced_allocations: self.cached_ffi_data.as_ref().unwrap().clone(),
215            stats: self.cached_stats.as_ref().unwrap().clone(),
216            ffi_stats: self.cached_ffi_stats.as_ref().unwrap().clone(),
217            scope_info: self.cached_scope_info.as_ref().unwrap().clone(),
218            timestamp: self.last_update,
219        };
220
221        let gathering_stats = DataGatheringStats {
222            total_time_ms: 0, // cache hit, no time
223            basic_data_time_ms: 0,
224            ffi_data_time_ms: 0,
225            scope_data_time_ms: 0,
226            allocation_count: localized_data.allocations.len(),
227            ffi_allocation_count: localized_data.enhanced_allocations.len(),
228            scope_count: localized_data.scope_info.len(),
229        };
230
231        Ok((localized_data, gathering_stats))
232    }
233
234    /// invalidate cache
235    pub fn invalidate_cache(&mut self) {
236        self.cached_allocations = None;
237        self.cached_ffi_data = None;
238        self.cached_stats = None;
239        self.cached_ffi_stats = None;
240        self.cached_scope_info = None;
241    }
242
243    /// estimate avoided global accesses
244    fn estimate_avoided_global_accesses(&self, stats: &DataGatheringStats) -> usize {
245        // In the traditional export process, each allocation may need multiple accesses to global state
246        // Here we estimate how many accesses we avoided through data localization
247        let basic_accesses = stats.allocation_count * 2; // Each allocation needs to access tracker 2 times
248        let ffi_accesses = stats.ffi_allocation_count * 3; // FFI allocations need more accesses
249        let scope_accesses = stats.scope_count * 1; // scope access
250
251        basic_accesses + ffi_accesses + scope_accesses
252    }
253
254    /// get cache stats
255    pub fn get_cache_stats(&self) -> CacheStats {
256        CacheStats {
257            is_cached: self.is_cache_valid(),
258            cache_age_ms: self.last_update.elapsed().as_millis() as u64,
259            cache_ttl_ms: self.cache_ttl.as_millis() as u64,
260            cached_allocation_count: self
261                .cached_allocations
262                .as_ref()
263                .map(|v| v.len())
264                .unwrap_or(0),
265            cached_ffi_count: self.cached_ffi_data.as_ref().map(|v| v.len()).unwrap_or(0),
266            cached_scope_count: self
267                .cached_scope_info
268                .as_ref()
269                .map(|v| v.len())
270                .unwrap_or(0),
271        }
272    }
273}
274
275/// cache stats
276#[derive(Debug, Clone)]
277pub struct CacheStats {
278    /// whether there is valid cache
279    pub is_cached: bool,
280    /// cache age (milliseconds)
281    pub cache_age_ms: u64,
282    /// cache ttl (milliseconds)
283    pub cache_ttl_ms: u64,
284    /// cached allocation count
285    pub cached_allocation_count: usize,
286    /// cached ffi count
287    pub cached_ffi_count: usize,
288    /// cached scope count
289    pub cached_scope_count: usize,
290}
291
292impl Default for DataLocalizer {
293    fn default() -> Self {
294        Self::new()
295    }
296}
297
298impl LocalizedExportData {
299    /// get data age
300    pub fn age(&self) -> Duration {
301        self.timestamp.elapsed()
302    }
303
304    /// check if data is still fresh
305    pub fn is_fresh(&self, max_age: Duration) -> bool {
306        self.age() < max_age
307    }
308
309    /// get total allocation count (basic + ffi)
310    pub fn total_allocation_count(&self) -> usize {
311        self.allocations.len() + self.enhanced_allocations.len()
312    }
313
314    /// get data summary
315    pub fn get_summary(&self) -> String {
316        format!(
317            "LocalizedExportData {{ allocations: {}, ffi_allocations: {}, scopes: {}, age: {:?} }}",
318            self.allocations.len(),
319            self.enhanced_allocations.len(),
320            self.scope_info.len(),
321            self.age()
322        )
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    #[test]
331    fn test_data_localizer_creation() {
332        let localizer = DataLocalizer::new();
333        assert!(!localizer.is_cache_valid());
334
335        let cache_stats = localizer.get_cache_stats();
336        assert!(!cache_stats.is_cached);
337        assert_eq!(cache_stats.cached_allocation_count, 0);
338    }
339
340    #[test]
341    fn test_cache_ttl() {
342        let short_ttl = Duration::from_millis(1);
343        let mut localizer = DataLocalizer::with_cache_ttl(short_ttl);
344
345        // simulate cached data
346        localizer.cached_allocations = Some(vec![]);
347        localizer.cached_ffi_data = Some(vec![]);
348        localizer.cached_stats = Some(MemoryStats::default());
349        localizer.cached_ffi_stats = Some(UnsafeFFIStats::default());
350        localizer.cached_scope_info = Some(vec![]);
351        localizer.last_update = Instant::now();
352
353        assert!(localizer.is_cache_valid());
354
355        // wait for cache to expire
356        std::thread::sleep(Duration::from_millis(2));
357        assert!(!localizer.is_cache_valid());
358    }
359
360    #[test]
361    fn test_localized_export_data() {
362        let data = LocalizedExportData {
363            allocations: vec![],
364            enhanced_allocations: vec![],
365            stats: MemoryStats::default(),
366            ffi_stats: UnsafeFFIStats::default(),
367            scope_info: vec![],
368            timestamp: Instant::now(),
369        };
370
371        assert_eq!(data.total_allocation_count(), 0);
372        assert!(data.is_fresh(Duration::from_secs(1)));
373
374        let summary = data.get_summary();
375        assert!(summary.contains("allocations: 0"));
376    }
377}