metrics_dashboard/
recorder.rs

1use metrics::{Key, Metadata, Recorder};
2use parking_lot::RwLock;
3use serde::Serialize;
4use std::{collections::HashMap, sync::Arc};
5
6use crate::DashboardOptions;
7
8use self::{counter::SimpleCounter, gauge::SimpleGauge, histogram::SimpleHistogram};
9
10mod counter;
11mod gauge;
12mod histogram;
13
14#[derive(Debug, Serialize, Clone)]
15pub enum MetricType {
16    Counter,
17    Gauge,
18    Histogram,
19}
20
21#[derive(Debug, Serialize, Clone)]
22pub struct MetricMeta {
23    pub key: String,
24    typ: MetricType,
25    pub desc: Option<String>,
26    pub unit: Option<String>,
27}
28
29#[derive(Debug, Serialize, Clone)]
30pub struct MetricValue {
31    pub key: String,
32    #[serde(rename = "value", skip_serializing_if = "Option::is_none")]
33    pub value_u64: Option<u64>,
34    #[serde(rename = "value", skip_serializing_if = "Option::is_none")]
35    pub value_f64: Option<f64>,
36}
37
38#[derive(Default)]
39struct DashboardStorage {
40    counters: HashMap<String, SimpleCounter>,
41    gauges: HashMap<String, SimpleGauge>,
42    histograms: HashMap<String, SimpleHistogram>,
43}
44
45impl DashboardStorage {
46    fn get_counter(&mut self, key: &str) -> SimpleCounter {
47        let entry = self.counters.entry(key.to_string()).or_default();
48        entry.clone()
49    }
50
51    fn get_gauge(&mut self, key: &str) -> SimpleGauge {
52        let entry = self.gauges.entry(key.to_string()).or_default();
53        entry.clone()
54    }
55
56    fn get_histogram(&mut self, key: &str) -> SimpleHistogram {
57        let entry = self.histograms.entry(key.to_string()).or_default();
58        entry.clone()
59    }
60}
61
62#[derive(Clone)]
63pub struct DashboardRecorder {
64    pub options: DashboardOptions,
65    storage: Arc<RwLock<DashboardStorage>>,
66    metrics: Arc<RwLock<HashMap<String, MetricMeta>>>,
67}
68
69/// The `DashboardRecorder` struct represents a recorder for metrics dashboard.
70/// It provides methods for adding bound keys, retrieving metrics, and retrieving metric values.
71impl DashboardRecorder {
72    /// Creates a new instance of `DashboardRecorder`.
73    ///
74    /// # Returns
75    ///
76    /// A new instance of `DashboardRecorder`.
77    pub fn new(opts: DashboardOptions) -> Self {
78        Self {
79            options: opts,
80            storage: Default::default(),
81            metrics: Arc::new(RwLock::new(HashMap::new())),
82        }
83    }
84
85    /// Retrieves the metrics as a vector of `MetricMeta`.
86    ///
87    /// # Returns
88    ///
89    /// A vector of `MetricMeta`.
90    pub fn metrics(&self) -> Vec<MetricMeta> {
91        let mut res = vec![];
92        let metrics = &*self.metrics.read();
93        for (_key, meta) in metrics.iter() {
94            res.push(meta.clone());
95        }
96        res.sort_by_cached_key(|m: &MetricMeta| m.key.clone());
97        res
98    }
99
100    /// Retrieves the metric values for the specified keys.
101    ///
102    /// # Arguments
103    ///
104    /// * `keys` - The keys to retrieve metric values for.
105    ///
106    /// # Returns
107    ///
108    /// A vector of `MetricValue`.
109    pub fn metrics_value(&self, keys: Vec<&str>) -> Vec<MetricValue> {
110        let mut storage = self.storage.write();
111        let metrics = self.metrics.read();
112        let mut data = vec![];
113        for key in keys {
114            if let Some(meta) = metrics.get(key) {
115                match meta.typ {
116                    MetricType::Counter => {
117                        let counter = storage.get_counter(key);
118                        data.push(MetricValue {
119                            key: key.to_string(),
120                            value_u64: Some(counter.value()),
121                            value_f64: None,
122                        });
123                    }
124                    MetricType::Gauge => {
125                        let gauge = storage.get_gauge(key);
126                        data.push(MetricValue {
127                            key: key.to_string(),
128                            value_u64: None,
129                            value_f64: Some((gauge.value() * 100.0).round() / 100.0),
130                        });
131                    }
132                    MetricType::Histogram => {
133                        // let _histogram = self
134                        //     .registry
135                        //     .get_or_create_histogram(&metrics::Key::from(key.to_string()), |a| {
136                        //         a.clone()
137                        //     });
138                        // TODO
139                    }
140                };
141            }
142        }
143        data
144    }
145}
146
147impl Recorder for DashboardRecorder {
148    fn describe_counter(
149        &self,
150        key: metrics::KeyName,
151        unit: Option<metrics::Unit>,
152        description: metrics::SharedString,
153    ) {
154        let mut metrics = self.metrics.write();
155        if let Some(metric) = metrics.get_mut(key.as_str()) {
156            metric.desc = Some(description.to_string());
157        } else {
158            metrics.insert(
159                key.as_str().to_string(),
160                MetricMeta {
161                    key: key.as_str().to_string(),
162                    typ: MetricType::Counter,
163                    desc: Some(description.to_string()),
164                    unit: unit.map(|u| u.as_canonical_label().to_string()),
165                },
166            );
167        }
168    }
169
170    fn describe_gauge(
171        &self,
172        key: metrics::KeyName,
173        unit: Option<metrics::Unit>,
174        description: metrics::SharedString,
175    ) {
176        let mut metrics = self.metrics.write();
177        if let Some(metric) = metrics.get_mut(key.as_str()) {
178            metric.desc = Some(description.to_string())
179        } else {
180            metrics.insert(
181                key.as_str().to_string(),
182                MetricMeta {
183                    key: key.as_str().to_string(),
184                    typ: MetricType::Gauge,
185                    desc: Some(description.to_string()),
186                    unit: unit.map(|u| u.as_canonical_label().to_string()),
187                },
188            );
189        }
190    }
191
192    fn describe_histogram(
193        &self,
194        key: metrics::KeyName,
195        unit: Option<metrics::Unit>,
196        description: metrics::SharedString,
197    ) {
198        let mut metrics = self.metrics.write();
199        if let Some(metric) = metrics.get_mut(key.as_str()) {
200            metric.desc = Some(description.to_string())
201        } else {
202            metrics.insert(
203                key.as_str().to_string(),
204                MetricMeta {
205                    key: key.as_str().to_string(),
206                    typ: MetricType::Histogram,
207                    desc: Some(description.to_string()),
208                    unit: unit.map(|u| u.as_canonical_label().to_string()),
209                },
210            );
211        }
212    }
213
214    fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> metrics::Counter {
215        let mut metrics = self.metrics.write();
216        if !metrics.contains_key(key.name()) {
217            metrics.insert(
218                key.name().to_string(),
219                MetricMeta {
220                    key: key.name().to_string(),
221                    typ: MetricType::Counter,
222                    desc: None,
223                    unit: None,
224                },
225            );
226        }
227        drop(metrics);
228
229        metrics::Counter::from_arc(self.storage.write().get_counter(key.name()).into())
230    }
231
232    fn register_gauge(&self, key: &Key, _metadata: &Metadata<'_>) -> metrics::Gauge {
233        let mut metrics = self.metrics.write();
234        if !metrics.contains_key(key.name()) {
235            metrics.insert(
236                key.name().to_string(),
237                MetricMeta {
238                    key: key.name().to_string(),
239                    typ: MetricType::Gauge,
240                    desc: None,
241                    unit: None,
242                },
243            );
244        }
245        drop(metrics);
246
247        metrics::Gauge::from_arc(self.storage.write().get_gauge(key.name()).into())
248    }
249
250    fn register_histogram(&self, key: &Key, _metadata: &Metadata<'_>) -> metrics::Histogram {
251        let mut metrics = self.metrics.write();
252        if !metrics.contains_key(key.name()) {
253            metrics.insert(
254                key.name().to_string(),
255                MetricMeta {
256                    key: key.name().to_string(),
257                    typ: MetricType::Histogram,
258                    desc: None,
259                    unit: None,
260                },
261            );
262        }
263        drop(metrics);
264
265        metrics::Histogram::from_arc(self.storage.write().get_histogram(key.name()).into())
266    }
267}