openvm_prof/
lib.rs

1use std::{collections::HashMap, fs::File, path::Path};
2
3use aggregate::{
4    EXECUTE_TIME_LABEL, PROOF_TIME_LABEL, PROVE_EXCL_TRACE_TIME_LABEL, TRACE_GEN_TIME_LABEL,
5};
6use eyre::Result;
7use memmap2::Mmap;
8
9use crate::types::{Labels, Metric, MetricDb, MetricsFile};
10
11pub mod aggregate;
12pub mod summary;
13pub mod types;
14
15impl MetricDb {
16    pub fn new(metrics_file: impl AsRef<Path>) -> Result<Self> {
17        let file = File::open(metrics_file)?;
18        let mmap = unsafe { Mmap::map(&file)? };
19        let metrics: MetricsFile = serde_json::from_slice(&mmap)?;
20
21        let mut db = MetricDb::default();
22
23        // Process counters
24        for entry in metrics.counter {
25            if entry.value == 0.0 {
26                continue;
27            }
28            let labels = Labels::from(entry.labels);
29            db.add_to_flat_dict(labels, entry.metric, entry.value);
30        }
31
32        // Process gauges
33        for entry in metrics.gauge {
34            let labels = Labels::from(entry.labels);
35            db.add_to_flat_dict(labels, entry.metric, entry.value);
36        }
37
38        db.apply_aggregations();
39        db.separate_by_label_types();
40
41        Ok(db)
42    }
43
44    // Currently hardcoding aggregations
45    pub fn apply_aggregations(&mut self) {
46        for metrics in self.flat_dict.values_mut() {
47            let get = |key: &str| metrics.iter().find(|m| m.name == key).map(|m| m.value);
48            let execute_time = get(EXECUTE_TIME_LABEL);
49            let trace_gen_time = get(TRACE_GEN_TIME_LABEL);
50            let prove_excl_trace_time = get(PROVE_EXCL_TRACE_TIME_LABEL);
51            if let (Some(execute_time), Some(trace_gen_time), Some(prove_excl_trace_time)) =
52                (execute_time, trace_gen_time, prove_excl_trace_time)
53            {
54                let total_time = execute_time + trace_gen_time + prove_excl_trace_time;
55                metrics.push(Metric::new(PROOF_TIME_LABEL.to_string(), total_time));
56            }
57        }
58    }
59
60    pub fn add_to_flat_dict(&mut self, labels: Labels, metric: String, value: f64) {
61        self.flat_dict
62            .entry(labels)
63            .or_default()
64            .push(Metric::new(metric, value));
65    }
66
67    // Custom sorting function that ensures 'group' comes first.
68    // Other keys are sorted alphabetically.
69    pub fn custom_sort_label_keys(label_keys: &mut [String]) {
70        // Prioritize 'group' by giving it the lowest possible sort value
71        label_keys.sort_by_key(|key| {
72            if key == "group" {
73                (0, key.clone()) // Lowest priority for 'group'
74            } else {
75                (1, key.clone()) // Normal priority for other keys
76            }
77        });
78    }
79
80    pub fn separate_by_label_types(&mut self) {
81        self.dict_by_label_types.clear();
82
83        for (labels, metrics) in &self.flat_dict {
84            // Get sorted label keys
85            let mut label_keys: Vec<String> = labels.0.iter().map(|(key, _)| key.clone()).collect();
86            Self::custom_sort_label_keys(&mut label_keys);
87
88            // Create label_values based on sorted keys
89            let label_dict: HashMap<String, String> = labels.0.iter().cloned().collect();
90
91            let label_values: Vec<String> = label_keys
92                .iter()
93                .map(|key| label_dict.get(key).unwrap().clone())
94                .collect();
95
96            // Add to dict_by_label_types
97            self.dict_by_label_types
98                .entry(label_keys)
99                .or_default()
100                .entry(label_values)
101                .or_default()
102                .extend(metrics.clone());
103        }
104    }
105
106    pub fn generate_markdown_tables(&self) -> String {
107        let mut markdown_output = String::new();
108        // Get sorted keys to iterate in consistent order
109        let mut sorted_keys: Vec<_> = self.dict_by_label_types.keys().cloned().collect();
110        sorted_keys.sort();
111
112        for label_keys in sorted_keys {
113            if label_keys.contains(&"cycle_tracker_span".to_string()) {
114                // Skip cycle_tracker_span as it is too long for markdown and visualized in
115                // flamegraphs
116                continue;
117            }
118            let metrics_dict = &self.dict_by_label_types[&label_keys];
119            let mut metric_names: Vec<String> = metrics_dict
120                .values()
121                .flat_map(|metrics| metrics.iter().map(|m| m.name.clone()))
122                .collect::<std::collections::HashSet<_>>()
123                .into_iter()
124                .collect();
125            metric_names.sort_by(|a, b| b.cmp(a));
126
127            // Create table header
128            let header = format!(
129                "| {} | {} |",
130                label_keys.join(" | "),
131                metric_names.join(" | ")
132            );
133
134            let separator = "| ".to_string()
135                + &vec!["---"; label_keys.len() + metric_names.len()].join(" | ")
136                + " |";
137
138            markdown_output.push_str(&header);
139            markdown_output.push('\n');
140            markdown_output.push_str(&separator);
141            markdown_output.push('\n');
142
143            // Fill table rows
144            for (label_values, metrics) in metrics_dict {
145                let mut row = String::new();
146                row.push_str("| ");
147                row.push_str(&label_values.join(" | "));
148                row.push_str(" | ");
149
150                // Add metric values
151                for metric_name in &metric_names {
152                    let metric_value = metrics
153                        .iter()
154                        .find(|m| &m.name == metric_name)
155                        .map(|m| Self::format_number(m.value))
156                        .unwrap_or_default();
157
158                    row.push_str(&format!("{} | ", metric_value));
159                }
160
161                markdown_output.push_str(&row);
162                markdown_output.push('\n');
163            }
164
165            markdown_output.push('\n');
166        }
167
168        markdown_output
169    }
170}
openvm_prof/lib.rs

openvm_prof/
lib.rs
