Skip to main content

memscope_rs/render_engine/export/
task_graph.rs

1//! Task graphs and accounting projected from explicit, immutable inputs.
2
3use 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
16/// Export the current registry graph. Use the snapshot API for a replayable report.
17pub 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
22/// Export only facts in the supplied events.
23///
24/// Events contain task IDs but not names, status or parent relationships. Those fields
25/// remain unknown unless the caller supplies a captured graph to the snapshot API.
26pub 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
33/// Export a fixed event cut and optional task metadata captured with that report.
34/// All byte counters come from events; registry counters are never used as a fallback.
35pub 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            // Compatibility flag: any gap means consumers must not treat the row as exact.
95            // Lifecycle ambiguity can affect live bytes in either direction.
96            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}