openvm_stark_sdk/bench/
mod.rs

Help
1use std::{collections::BTreeMap, ffi::OsStr};
2
3use metrics_tracing_context::{MetricsLayer, TracingContextLayer};
4use metrics_util::{
5    debugging::{DebugValue, DebuggingRecorder, Snapshot},
6    layers::Layer,
7    CompositeKey, MetricKind,
8};
9use serde_json::json;
10use tracing_forest::ForestLayer;
11use tracing_subscriber::{layer::SubscriberExt, EnvFilter, Registry};
12
13/// Run a function with metric collection enabled. The metrics will be written to a file specified
14/// by an environment variable which name is `output_path_envar`.
15pub fn run_with_metric_collection<R>(
16    output_path_envar: impl AsRef<OsStr>,
17    f: impl FnOnce() -> R,
18) -> R {
19    let file = std::env::var(output_path_envar).map(|path| std::fs::File::create(path).unwrap());
20    // Set up tracing:
21    let env_filter =
22        EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info,p3_=warn"));
23    // Plonky3 logging is more verbose, so we set default to debug.
24    let subscriber = Registry::default()
25        .with(env_filter)
26        .with(ForestLayer::default())
27        .with(MetricsLayer::new());
28    // Prepare tracing.
29    tracing::subscriber::set_global_default(subscriber).unwrap();
30
31    // Prepare metrics.
32    let recorder = DebuggingRecorder::new();
33    let snapshotter = recorder.snapshotter();
34    let recorder = TracingContextLayer::all().layer(recorder);
35    // Install the registry as the global recorder
36    metrics::set_global_recorder(recorder).unwrap();
37    let res = f();
38
39    if let Ok(file) = file {
40        serde_json::to_writer_pretty(&file, &serialize_metric_snapshot(snapshotter.snapshot()))
41            .unwrap();
42    }
43    res
44}
45
46/// Serialize a gauge/counter metric into a JSON object. The object has the following structure:
47/// {
48///    "metric": <Metric Name>,
49///    "labels": [
50///       (<key1>, <value1>),
51///       (<key2>, <value2>),
52///     ],
53///    "value": <float value if gauge | integer value if counter>
54/// }
55///
56fn serialize_metric(ckey: CompositeKey, value: DebugValue) -> serde_json::Value {
57    let (_kind, key) = ckey.into_parts();
58    let (key_name, labels) = key.into_parts();
59    let value = match value {
60        DebugValue::Gauge(v) => v.into_inner().to_string(),
61        DebugValue::Counter(v) => v.to_string(),
62        DebugValue::Histogram(_) => todo!("Histograms not supported yet."),
63    };
64    let labels = labels
65        .into_iter()
66        .map(|label| {
67            let (k, v) = label.into_parts();
68            (k.as_ref().to_owned(), v.as_ref().to_owned())
69        })
70        .collect::<Vec<_>>();
71
72    json!({
73        "metric": key_name.as_str(),
74        "labels": labels,
75        "value": value,
76    })
77}
78
79/// Serialize a metric snapshot into a JSON object. The object has the following structure:
80/// {
81///   "gauge": [
82///     {
83///         "metric": <Metric Name>,
84///         "labels": [
85///             (<key1>, <value1>),
86///             (<key2>, <value2>),
87///         ],
88///         "value": <float value>
89///     },
90///     ...
91///   ],
92///   ...
93/// }
94///
95pub fn serialize_metric_snapshot(snapshot: Snapshot) -> serde_json::Value {
96    let mut ret = BTreeMap::<_, Vec<serde_json::Value>>::new();
97    for (ckey, _, _, value) in snapshot.into_vec() {
98        match ckey.kind() {
99            MetricKind::Gauge => {
100                ret.entry("gauge")
101                    .or_default()
102                    .push(serialize_metric(ckey, value));
103            }
104            MetricKind::Counter => {
105                ret.entry("counter")
106                    .or_default()
107                    .push(serialize_metric(ckey, value));
108            }
109            MetricKind::Histogram => todo!(),
110        }
111    }
112    json!(ret)
113}
openvm_stark_sdk/bench/mod.rs

openvm_stark_sdk/bench/
mod.rs
