memscope_rs/render_engine/export/
task_graph.rs1use crate::core::{MemScopeError, MemScopeResult};
4use crate::event_store::MemoryEvent;
5use crate::task_registry::{self, TaskGraph, TaskIdRegistry, TaskLifecycleQuality};
6use serde_json::{json, Value};
7use std::collections::BTreeMap;
8use std::path::Path;
9
10fn lifecycle_json(quality: TaskLifecycleQuality) -> Value {
11 let mut value = json!(quality);
12 value["complete"] = json!(quality.is_complete());
13 value
14}
15
16pub fn export_task_graph_json<P: AsRef<Path>>(base_path: P) -> MemScopeResult<()> {
18 let graph = task_registry::global_registry().export_graph();
19 write_graph(base_path.as_ref(), &json!(graph))
20}
21
22pub fn export_task_graph_json_with_accounting<P: AsRef<Path>>(
27 base_path: P,
28 events: &[MemoryEvent],
29) -> MemScopeResult<()> {
30 export_task_graph_snapshot(base_path, events, None)
31}
32
33pub fn export_task_graph_snapshot<P: AsRef<Path>>(
36 base_path: P,
37 events: &[MemoryEvent],
38 graph: Option<&TaskGraph>,
39) -> MemScopeResult<()> {
40 let accounting = TaskIdRegistry::task_accounting_from_events(events);
41 let (charge, task_charges) = TaskIdRegistry::charge_quality_from_events(events);
42 let mut nodes = BTreeMap::new();
43 let mut metadata_complete = graph.is_some();
44 if let Some(graph) = graph {
45 for node in &graph.nodes {
46 let mut node = json!(node);
47 node["metadata_source"] = json!("task_graph_snapshot");
48 nodes.insert(
49 node["id"].as_u64().expect("serialized task ID is u64"),
50 node,
51 );
52 }
53 }
54 for task in events.iter().filter_map(|event| event.task_id) {
55 nodes.entry(task).or_insert_with(|| {
56 metadata_complete = false;
57 json!({
58 "id": task, "name": null, "status": null, "metadata_source": "unavailable",
59 })
60 });
61 }
62 for (id, node) in &mut nodes {
63 let stats = accounting.stats.get(id).copied().unwrap_or_default();
64 node["memory_usage"] = json!(stats.live_bytes);
65 node["allocation_count"] = json!(stats.allocation_count);
66 node["stats_source"] = json!("event_snapshot");
67 }
68 let tasks: Vec<_> = nodes
69 .keys()
70 .map(|id| {
71 let stats = accounting.stats.get(id).copied().unwrap_or_default();
72 let charge = task_charges.get(id).copied().unwrap_or_default();
73 let lifecycle = accounting
74 .task_lifecycle_quality
75 .get(id)
76 .copied()
77 .unwrap_or_default();
78 let mut row = json!(stats);
79 let fields = row
80 .as_object_mut()
81 .expect("serialized task stats are an object");
82 fields.extend(
83 json!(charge)
84 .as_object()
85 .expect("serialized quality is an object")
86 .clone(),
87 );
88 fields.insert("task_id".to_string(), json!(id));
89 fields.insert("stats_source".to_string(), json!("event_snapshot"));
90 fields.insert(
91 "complete".to_string(),
92 json!(charge.is_complete() && lifecycle.is_complete()),
93 );
94 fields.insert(
97 "lower_bound".to_string(),
98 json!(!charge.is_complete() || !lifecycle.is_complete()),
99 );
100 fields.insert(
101 "live_bytes_exact".to_string(),
102 json!(lifecycle.is_complete()),
103 );
104 fields.insert("lifecycle_quality".to_string(), lifecycle_json(lifecycle));
105 row
106 })
107 .collect();
108 let complete = charge.is_complete() && accounting.lifecycle_quality.is_complete();
109 let mut charge_json = json!(charge);
110 charge_json["complete"] = json!(charge.is_complete());
111 let mut edges = graph.map_or_else(Vec::new, |graph| graph.edges.clone());
112 edges.sort_by_key(|edge| (edge.from, edge.to));
113 write_graph(
114 base_path.as_ref(),
115 &json!({
116 "nodes": nodes.into_values().collect::<Vec<_>>(),
117 "edges": edges,
118 "metadata_complete": metadata_complete,
119 "accounting": {
120 "stats_source": "event_snapshot", "complete": complete, "lower_bound": !complete,
121 "charge_quality": charge_json,
122 "lifecycle_quality": lifecycle_json(accounting.lifecycle_quality), "tasks": tasks,
123 },
124 }),
125 )
126}
127
128fn write_graph(base_path: &Path, value: &Value) -> MemScopeResult<()> {
129 let content = serde_json::to_vec_pretty(value).map_err(|error| {
130 MemScopeError::error("export", "export_task_graph_json", error.to_string())
131 })?;
132 std::fs::write(base_path.join("task_graph.json"), content).map_err(|error| {
133 MemScopeError::error("export", "export_task_graph_json", error.to_string())
134 })
135}