memscope_rs/capture/backends/
task_profile.rs

1//! Task-level memory profiling for the unified tracker
2//!
3//! This module provides task-aware memory tracking capabilities,
4//! allowing users to track memory usage patterns at the granularity
5//! of individual tasks or workloads.
6
7use serde::{Deserialize, Serialize};
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::{Arc, Mutex};
11use std::time::Duration;
12
13/// Task type classification for categorizing different workloads
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
15pub enum TaskType {
16    /// CPU-intensive tasks (e.g., matrix multiplication, data processing)
17    CpuIntensive,
18    /// IO-intensive tasks (e.g., file operations, database queries)
19    IoIntensive,
20    /// Network-intensive tasks (e.g., HTTP requests, RPC calls)
21    NetworkIntensive,
22    /// Memory-intensive tasks (e.g., large data structures, caching)
23    MemoryIntensive,
24    /// GPU compute tasks (e.g., CUDA, OpenCL operations)
25    GpuCompute,
26    /// Mixed workload with balanced resource usage
27    #[default]
28    Mixed,
29    /// Streaming data processing tasks
30    Streaming,
31    /// Background maintenance tasks
32    Background,
33}
34
35impl TaskType {
36    /// Get human-readable description of task type
37    pub fn description(&self) -> &'static str {
38        match self {
39            Self::CpuIntensive => "CPU-intensive workload",
40            Self::IoIntensive => "IO-intensive workload",
41            Self::NetworkIntensive => "Network-intensive workload",
42            Self::MemoryIntensive => "Memory-intensive workload",
43            Self::GpuCompute => "GPU compute workload",
44            Self::Mixed => "Mixed workload",
45            Self::Streaming => "Streaming workload",
46            Self::Background => "Background task",
47        }
48    }
49
50    /// Get resource priority for this task type
51    pub fn resource_priority(&self) -> (f64, f64, f64, f64) {
52        match self {
53            Self::CpuIntensive => (1.0, 0.3, 0.2, 0.1),
54            Self::IoIntensive => (0.3, 1.0, 0.2, 0.1),
55            Self::NetworkIntensive => (0.3, 0.2, 1.0, 0.1),
56            Self::MemoryIntensive => (0.2, 0.3, 0.2, 1.0),
57            Self::GpuCompute => (0.5, 0.2, 0.1, 0.8),
58            Self::Mixed => (0.5, 0.5, 0.5, 0.5),
59            Self::Streaming => (0.4, 0.4, 0.6, 0.4),
60            Self::Background => (0.2, 0.2, 0.2, 0.2),
61        }
62    }
63}
64
65/// Memory usage profile for a single task
66#[derive(Debug, Clone, Serialize, Deserialize)]
67pub struct TaskMemoryProfile {
68    /// Unique task identifier (auto-generated, never reused)
69    pub task_id: u64,
70    /// Tokio task ID (if running in tokio context)
71    pub tokio_task_id: Option<u64>,
72    /// Task name for identification
73    pub task_name: String,
74    /// Task type classification
75    pub task_type: TaskType,
76    /// Task creation timestamp (milliseconds since Unix epoch)
77    pub created_at_ms: u64,
78    /// Task completion timestamp (milliseconds since Unix epoch, if completed)
79    pub completed_at_ms: Option<u64>,
80    /// Total bytes allocated by this task
81    pub total_bytes: u64,
82    /// Current memory usage (allocated - deallocated)
83    pub current_memory: u64,
84    /// Peak memory usage observed
85    pub peak_memory: u64,
86    /// Number of allocation operations
87    pub total_allocations: u64,
88    /// Number of deallocation operations
89    pub total_deallocations: u64,
90    /// Task duration in nanoseconds
91    pub duration_ns: u64,
92    /// Memory allocation rate (bytes/second)
93    pub allocation_rate: f64,
94    /// Memory efficiency score (0.0 to 1.0)
95    pub efficiency_score: f64,
96    /// Average allocation size
97    pub average_allocation_size: f64,
98    /// Numeric thread id where the task was created or last observed.
99    ///
100    /// Recorded when the information is genuinely available so reports never have
101    /// to guess: using the thread that *builds the snapshot* would label every task
102    /// with the exporter's thread and silently merge them.
103    #[serde(default)]
104    pub thread_id_u64: Option<u64>,
105}
106
107impl TaskMemoryProfile {
108    /// Create new task profile
109    pub fn new(task_id: u64, task_name: String, task_type: TaskType) -> Self {
110        Self {
111            task_id,
112            tokio_task_id: None,
113            task_name,
114            task_type,
115            created_at_ms: Self::now_ms(),
116            completed_at_ms: None,
117            total_bytes: 0,
118            current_memory: 0,
119            peak_memory: 0,
120            total_allocations: 0,
121            total_deallocations: 0,
122            duration_ns: 0,
123            allocation_rate: 0.0,
124            efficiency_score: 1.0,
125            average_allocation_size: 0.0,
126            thread_id_u64: None,
127        }
128    }
129
130    /// Record the numeric thread id (callers that genuinely know it).
131    pub fn set_thread_id_u64(&mut self, thread_id: Option<u64>) {
132        self.thread_id_u64 = thread_id;
133    }
134
135    /// Create new task profile with tokio task ID
136    pub fn with_tokio_id(
137        task_id: u64,
138        tokio_task_id: u64,
139        task_name: String,
140        task_type: TaskType,
141    ) -> Self {
142        let mut profile = Self::new(task_id, task_name, task_type);
143        profile.tokio_task_id = Some(tokio_task_id);
144        profile
145    }
146
147    /// Get current timestamp in milliseconds
148    fn now_ms() -> u64 {
149        std::time::SystemTime::now()
150            .duration_since(std::time::UNIX_EPOCH)
151            .unwrap_or_default()
152            .as_millis() as u64
153    }
154
155    /// Mark task as completed
156    pub fn mark_completed(&mut self) {
157        self.completed_at_ms.get_or_insert_with(Self::now_ms);
158        self.update_metrics();
159    }
160
161    /// Record allocation event
162    pub fn record_allocation(&mut self, size: u64) {
163        self.total_bytes += size;
164        self.current_memory += size;
165        self.total_allocations += 1;
166
167        if self.current_memory > self.peak_memory {
168            self.peak_memory = self.current_memory;
169        }
170
171        self.update_metrics();
172    }
173
174    /// Record an in-place reallocation of an already counted block.
175    ///
176    /// The old size leaves the live set, the new size enters it, and the block is
177    /// not counted again as a separate allocation.
178    pub fn record_reallocation(&mut self, old_size: u64, new_size: u64) {
179        self.total_bytes += new_size;
180        self.current_memory = self
181            .current_memory
182            .saturating_sub(old_size)
183            .saturating_add(new_size);
184        if self.current_memory > self.peak_memory {
185            self.peak_memory = self.current_memory;
186        }
187        self.update_metrics();
188    }
189
190    /// Record deallocation event
191    pub fn record_deallocation(&mut self, size: u64) {
192        self.current_memory = self.current_memory.saturating_sub(size);
193        self.total_deallocations += 1;
194        self.update_metrics();
195    }
196
197    /// Check if task is completed
198    pub fn is_completed(&self) -> bool {
199        self.completed_at_ms.is_some()
200    }
201
202    /// Get task lifetime duration
203    pub fn lifetime(&self) -> Duration {
204        let end_ms = self.completed_at_ms.unwrap_or_else(Self::now_ms);
205        let start_ms = self.created_at_ms;
206        Duration::from_millis(end_ms.saturating_sub(start_ms))
207    }
208
209    /// Calculate memory efficiency (deallocated / allocated)
210    pub fn memory_efficiency(&self) -> f64 {
211        if self.total_bytes == 0 {
212            1.0
213        } else {
214            let deallocated = self.total_bytes - self.current_memory;
215            deallocated as f64 / self.total_bytes as f64
216        }
217    }
218
219    /// Check if task has potential memory leak
220    pub fn has_potential_leak(&self) -> bool {
221        self.is_completed() && self.current_memory > 0
222    }
223
224    /// Update derived metrics
225    fn update_metrics(&mut self) {
226        self.refresh_at(Self::now_ms().saturating_mul(1_000_000));
227    }
228
229    pub(crate) fn refresh_at(&mut self, captured_at_ns: u64) {
230        let end = self
231            .completed_at_ms
232            .map(|ms| ms.saturating_mul(1_000_000))
233            .unwrap_or(captured_at_ns);
234        self.duration_ns = end.saturating_sub(self.created_at_ms.saturating_mul(1_000_000));
235        let lifetime_secs = self.duration_ns as f64 / 1_000_000_000.0;
236
237        self.allocation_rate = if lifetime_secs > 0.0 {
238            self.total_bytes as f64 / lifetime_secs
239        } else {
240            0.0
241        };
242
243        self.efficiency_score = self.memory_efficiency();
244
245        self.average_allocation_size = if self.total_allocations > 0 {
246            self.total_bytes as f64 / self.total_allocations as f64
247        } else {
248            0.0
249        };
250    }
251
252    /// Get memory usage summary
253    pub fn summary(&self) -> String {
254        format!(
255            "Task '{}' (ID: {}, Type: {:?}): {} allocations, {:.2} MB total, {:.2} MB peak, {:.1}% efficiency",
256            self.task_name,
257            self.task_id,
258            self.task_type,
259            self.total_allocations,
260            self.total_bytes as f64 / 1_048_576.0,
261            self.peak_memory as f64 / 1_048_576.0,
262            self.efficiency_score * 100.0
263        )
264    }
265}
266
267/// Manager for tracking multiple task profiles
268#[derive(Debug, Clone)]
269pub struct TaskProfileManager {
270    profiles: Arc<Mutex<HashMap<u64, TaskMemoryProfile>>>,
271    next_task_id: Arc<AtomicU64>,
272}
273
274impl TaskProfileManager {
275    /// Create new task profile manager
276    pub fn new() -> Self {
277        Self {
278            profiles: Arc::new(Mutex::new(HashMap::new())),
279            next_task_id: Arc::new(AtomicU64::new(1)),
280        }
281    }
282
283    /// Create a new task profile
284    pub fn create_task(&self, task_name: String, task_type: TaskType) -> u64 {
285        let task_id = self.next_task_id.fetch_add(1, Ordering::Relaxed);
286
287        let profile = TaskMemoryProfile::new(task_id, task_name, task_type);
288
289        if let Ok(mut profiles) = self.profiles.lock() {
290            profiles.insert(task_id, profile);
291        }
292
293        task_id
294    }
295
296    /// Record allocation for a task
297    pub fn record_allocation(&self, task_id: u64, size: u64) {
298        if let Ok(mut profiles) = self.profiles.lock() {
299            if let Some(profile) = profiles.get_mut(&task_id) {
300                profile.record_allocation(size);
301            }
302        }
303    }
304
305    /// Record deallocation for a task
306    pub fn record_deallocation(&self, task_id: u64, size: u64) {
307        if let Ok(mut profiles) = self.profiles.lock() {
308            if let Some(profile) = profiles.get_mut(&task_id) {
309                profile.record_deallocation(size);
310            }
311        }
312    }
313
314    /// Mark task as completed
315    pub fn complete_task(&self, task_id: u64) {
316        if let Ok(mut profiles) = self.profiles.lock() {
317            if let Some(profile) = profiles.get_mut(&task_id) {
318                profile.mark_completed();
319            }
320        }
321    }
322
323    /// Get profile for a specific task
324    pub fn get_profile(&self, task_id: u64) -> Option<TaskMemoryProfile> {
325        if let Ok(profiles) = self.profiles.lock() {
326            profiles.get(&task_id).cloned()
327        } else {
328            None
329        }
330    }
331
332    /// Get all task profiles
333    pub fn get_all_profiles(&self) -> Vec<TaskMemoryProfile> {
334        if let Ok(profiles) = self.profiles.lock() {
335            profiles.values().cloned().collect()
336        } else {
337            Vec::new()
338        }
339    }
340
341    /// Get profiles by task type
342    pub fn get_profiles_by_type(&self, task_type: TaskType) -> Vec<TaskMemoryProfile> {
343        if let Ok(profiles) = self.profiles.lock() {
344            profiles
345                .values()
346                .filter(|p| p.task_type == task_type)
347                .cloned()
348                .collect()
349        } else {
350            Vec::new()
351        }
352    }
353
354    /// Get aggregated statistics
355    pub fn get_aggregated_stats(&self) -> AggregatedTaskStats {
356        let profiles = self.get_all_profiles();
357
358        let total_tasks = profiles.len();
359        let completed_tasks = profiles.iter().filter(|p| p.is_completed()).count();
360
361        let total_memory_allocated: u64 = profiles.iter().map(|p| p.total_bytes).sum();
362        let current_memory_usage: u64 = profiles.iter().map(|p| p.current_memory).sum();
363        let peak_memory_usage: u64 = profiles.iter().map(|p| p.peak_memory).max().unwrap_or(0);
364
365        let total_duration: Duration = profiles.iter().map(|p| p.lifetime()).sum::<Duration>();
366        let average_lifetime = if total_tasks > 0 {
367            let total_secs = total_duration.as_secs_f64();
368            let avg_secs = total_secs / total_tasks as f64;
369            Duration::from_secs_f64(avg_secs)
370        } else {
371            Duration::ZERO
372        };
373
374        let overall_efficiency = if total_memory_allocated > 0 {
375            let total_deallocated = total_memory_allocated - current_memory_usage;
376            total_deallocated as f64 / total_memory_allocated as f64
377        } else {
378            1.0
379        };
380
381        let potential_leaks = profiles.iter().filter(|p| p.has_potential_leak()).count();
382
383        AggregatedTaskStats {
384            total_tasks,
385            completed_tasks,
386            total_memory_allocated,
387            current_memory_usage,
388            peak_memory_usage,
389            average_lifetime,
390            overall_efficiency,
391            potential_leaks,
392        }
393    }
394
395    /// Clear all profiles
396    pub fn clear(&self) {
397        if let Ok(mut profiles) = self.profiles.lock() {
398            profiles.clear();
399        }
400    }
401
402    /// Get number of active tasks
403    pub fn active_task_count(&self) -> usize {
404        if let Ok(profiles) = self.profiles.lock() {
405            profiles.iter().filter(|(_, p)| !p.is_completed()).count()
406        } else {
407            0
408        }
409    }
410}
411
412impl Default for TaskProfileManager {
413    fn default() -> Self {
414        Self::new()
415    }
416}
417
418/// Aggregated statistics across multiple tasks
419#[derive(Debug, Clone, Serialize, Deserialize)]
420pub struct AggregatedTaskStats {
421    /// Total number of tasks tracked
422    pub total_tasks: usize,
423    /// Number of completed tasks
424    pub completed_tasks: usize,
425    /// Total memory allocated across all tasks
426    pub total_memory_allocated: u64,
427    /// Current memory usage across all active tasks
428    pub current_memory_usage: u64,
429    /// Peak memory usage observed
430    pub peak_memory_usage: u64,
431    /// Average task lifetime
432    pub average_lifetime: Duration,
433    /// Memory efficiency across all tasks
434    pub overall_efficiency: f64,
435    /// Tasks with potential memory leaks
436    pub potential_leaks: usize,
437}
438
439impl AggregatedTaskStats {
440    /// Create empty statistics
441    pub fn new() -> Self {
442        Self {
443            total_tasks: 0,
444            completed_tasks: 0,
445            total_memory_allocated: 0,
446            current_memory_usage: 0,
447            peak_memory_usage: 0,
448            average_lifetime: Duration::ZERO,
449            overall_efficiency: 1.0,
450            potential_leaks: 0,
451        }
452    }
453
454    /// Get memory usage summary
455    pub fn memory_summary(&self) -> String {
456        format!(
457            "Tasks: {} ({}% complete), Memory: {:.2}MB allocated, {:.2}MB current, {:.1}% efficiency, {} potential leaks",
458            self.total_tasks,
459            self.completed_tasks
460                .checked_mul(100)
461                .and_then(|v| v.checked_div(self.total_tasks))
462                .unwrap_or(0),
463            self.total_memory_allocated as f64 / 1_048_576.0,
464            self.current_memory_usage as f64 / 1_048_576.0,
465            self.overall_efficiency * 100.0,
466            self.potential_leaks
467        )
468    }
469}
470
471impl Default for AggregatedTaskStats {
472    fn default() -> Self {
473        Self::new()
474    }
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480
481    #[test]
482    fn test_task_type_description() {
483        assert_eq!(
484            TaskType::CpuIntensive.description(),
485            "CPU-intensive workload"
486        );
487        assert_eq!(TaskType::IoIntensive.description(), "IO-intensive workload");
488        assert_eq!(
489            TaskType::NetworkIntensive.description(),
490            "Network-intensive workload"
491        );
492    }
493
494    #[test]
495    fn test_task_type_resource_priority() {
496        let (cpu, io, net, mem) = TaskType::CpuIntensive.resource_priority();
497        assert!(cpu > io && cpu > net && cpu > mem);
498
499        let (cpu, io, net, mem) = TaskType::IoIntensive.resource_priority();
500        assert!(io > cpu && io > net && io > mem);
501    }
502
503    #[test]
504    fn test_task_memory_profile_basic() {
505        let profile = TaskMemoryProfile::new(1, "test_task".to_string(), TaskType::CpuIntensive);
506
507        assert_eq!(profile.task_id, 1);
508        assert_eq!(profile.task_name, "test_task");
509        assert_eq!(profile.task_type, TaskType::CpuIntensive);
510        assert!(!profile.is_completed());
511        assert_eq!(profile.current_memory, 0);
512        assert_eq!(profile.total_bytes, 0);
513    }
514
515    #[test]
516    fn test_record_allocation() {
517        let mut profile = TaskMemoryProfile::new(1, "test".to_string(), TaskType::Mixed);
518
519        profile.record_allocation(1024);
520        assert_eq!(profile.current_memory, 1024);
521        assert_eq!(profile.total_bytes, 1024);
522        assert_eq!(profile.peak_memory, 1024);
523        assert_eq!(profile.total_allocations, 1);
524
525        profile.record_allocation(2048);
526        assert_eq!(profile.current_memory, 3072);
527        assert_eq!(profile.total_bytes, 3072);
528        assert_eq!(profile.peak_memory, 3072);
529        assert_eq!(profile.total_allocations, 2);
530    }
531
532    #[test]
533    fn test_record_deallocation() {
534        let mut profile = TaskMemoryProfile::new(1, "test".to_string(), TaskType::Mixed);
535
536        profile.record_allocation(3072);
537        profile.record_deallocation(1024);
538        assert_eq!(profile.current_memory, 2048);
539        assert_eq!(profile.total_bytes, 3072);
540        assert_eq!(profile.peak_memory, 3072);
541        assert_eq!(profile.total_deallocations, 1);
542    }
543
544    #[test]
545    fn test_memory_efficiency() {
546        let mut profile = TaskMemoryProfile::new(1, "test".to_string(), TaskType::Mixed);
547
548        profile.record_allocation(1000);
549        profile.record_deallocation(1000);
550        assert_eq!(profile.memory_efficiency(), 1.0);
551
552        profile.record_allocation(1000);
553        assert_eq!(profile.memory_efficiency(), 0.5);
554    }
555
556    #[test]
557    fn test_potential_leak_detection() {
558        let mut profile = TaskMemoryProfile::new(1, "test".to_string(), TaskType::Mixed);
559
560        profile.record_allocation(1000);
561        assert!(!profile.has_potential_leak());
562
563        profile.record_deallocation(1000);
564        profile.mark_completed();
565        assert!(!profile.has_potential_leak());
566
567        let mut profile2 = TaskMemoryProfile::new(2, "test2".to_string(), TaskType::Mixed);
568        profile2.record_allocation(1000);
569        profile2.mark_completed();
570        assert!(profile2.has_potential_leak());
571    }
572
573    #[test]
574    fn test_task_profile_manager() {
575        let manager = TaskProfileManager::new();
576
577        let task_id = manager.create_task("test_task".to_string(), TaskType::CpuIntensive);
578        assert!(task_id > 0);
579
580        manager.record_allocation(task_id, 1024);
581        manager.record_allocation(task_id, 2048);
582
583        let profile = manager.get_profile(task_id);
584        assert!(profile.is_some());
585        assert_eq!(profile.unwrap().total_bytes, 3072);
586    }
587
588    #[test]
589    fn test_aggregated_stats() {
590        let manager = TaskProfileManager::new();
591
592        let task1 = manager.create_task("task1".to_string(), TaskType::Mixed);
593        manager.record_allocation(task1, 1000);
594        manager.record_deallocation(task1, 500);
595        manager.complete_task(task1);
596
597        let task2 = manager.create_task("task2".to_string(), TaskType::Mixed);
598        manager.record_allocation(task2, 2000);
599
600        let stats = manager.get_aggregated_stats();
601        assert_eq!(stats.total_tasks, 2);
602        assert_eq!(stats.completed_tasks, 1);
603        assert_eq!(stats.total_memory_allocated, 3000);
604        assert_eq!(stats.current_memory_usage, 2500);
605    }
606
607    #[test]
608    fn test_active_task_count() {
609        let manager = TaskProfileManager::new();
610
611        let task1 = manager.create_task("task1".to_string(), TaskType::Mixed);
612        let task2 = manager.create_task("task2".to_string(), TaskType::Mixed);
613
614        assert_eq!(manager.active_task_count(), 2);
615
616        manager.complete_task(task1);
617        assert_eq!(manager.active_task_count(), 1);
618
619        manager.complete_task(task2);
620        assert_eq!(manager.active_task_count(), 0);
621    }
622
623    #[test]
624    fn test_profiles_by_type() {
625        let manager = TaskProfileManager::new();
626
627        let _ = manager.create_task("cpu_task".to_string(), TaskType::CpuIntensive);
628        let _ = manager.create_task("io_task".to_string(), TaskType::IoIntensive);
629        let _ = manager.create_task("cpu_task2".to_string(), TaskType::CpuIntensive);
630
631        let cpu_profiles = manager.get_profiles_by_type(TaskType::CpuIntensive);
632        assert_eq!(cpu_profiles.len(), 2);
633
634        let io_profiles = manager.get_profiles_by_type(TaskType::IoIntensive);
635        assert_eq!(io_profiles.len(), 1);
636    }
637}
memscope_rs/capture/backends/task_profile.rs

memscope_rs/capture/backends/
task_profile.rs
