openvm_prof/
lib.rs

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

openvm_prof/
lib.rs
