khora_telemetry/metrics/
registry.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Registry for managing metrics.
16
17use crate::storage::{backend::MetricsBackend, memory_backend::InMemoryBackend};
18use khora_core::telemetry::metrics::{Metric, MetricId, MetricType, MetricsError, MetricsResult};
19use std::sync::Arc;
20
21/// Central registry for metrics in the KhoraEngine
22///
23/// This registry provides a high-level API for metrics management and
24/// serves as the main entry point for the metrics system. It handles
25/// metric registration, updates, and queries while providing type safety
26/// and convenient helper methods.
27#[derive(Debug)]
28pub struct MetricsRegistry {
29    backend: Arc<dyn MetricsBackend>,
30}
31
32impl MetricsRegistry {
33    /// Create a new metrics registry with the default in-memory backend
34    pub fn new() -> Self {
35        Self {
36            backend: Arc::new(InMemoryBackend::new()),
37        }
38    }
39
40    /// Create a new metrics registry with a custom backend
41    pub fn with_backend(backend: Arc<dyn MetricsBackend>) -> Self {
42        Self { backend }
43    }
44
45    /// Register a new counter metric
46    pub fn register_counter(
47        &self,
48        namespace: impl Into<String>,
49        name: impl Into<String>,
50        description: impl Into<String>,
51    ) -> MetricsResult<CounterHandle> {
52        let id = MetricId::new(namespace, name);
53        let metric = Metric::new_counter(id.clone(), description, 0);
54        self.backend.put_metric(metric)?;
55        Ok(CounterHandle::new(id, self.backend.clone()))
56    }
57
58    /// Register a new counter metric with labels
59    pub fn register_counter_with_labels(
60        &self,
61        namespace: impl Into<String>,
62        name: impl Into<String>,
63        description: impl Into<String>,
64        labels: Vec<(String, String)>,
65    ) -> MetricsResult<CounterHandle> {
66        let mut id = MetricId::new(namespace, name);
67        for (key, value) in labels {
68            id = id.with_label(key, value);
69        }
70        let metric = Metric::new_counter(id.clone(), description, 0);
71        self.backend.put_metric(metric)?;
72        Ok(CounterHandle::new(id, self.backend.clone()))
73    }
74
75    /// Register a new gauge metric
76    pub fn register_gauge(
77        &self,
78        namespace: impl Into<String>,
79        name: impl Into<String>,
80        description: impl Into<String>,
81        unit: impl Into<String>,
82    ) -> MetricsResult<GaugeHandle> {
83        let id = MetricId::new(namespace, name);
84        let metric = Metric::new_gauge(id.clone(), description, unit, 0.0);
85        self.backend.put_metric(metric)?;
86        Ok(GaugeHandle::new(id, self.backend.clone()))
87    }
88
89    /// Register a new gauge metric with labels
90    pub fn register_gauge_with_labels(
91        &self,
92        namespace: impl Into<String>,
93        name: impl Into<String>,
94        description: impl Into<String>,
95        unit: impl Into<String>,
96        labels: Vec<(String, String)>,
97    ) -> MetricsResult<GaugeHandle> {
98        let mut id = MetricId::new(namespace, name);
99        for (key, value) in labels {
100            id = id.with_label(key, value);
101        }
102        let metric = Metric::new_gauge(id.clone(), description, unit, 0.0);
103        self.backend.put_metric(metric)?;
104        Ok(GaugeHandle::new(id, self.backend.clone()))
105    }
106
107    /// Register a new histogram metric
108    pub fn register_histogram(
109        &self,
110        namespace: impl Into<String>,
111        name: impl Into<String>,
112        description: impl Into<String>,
113        unit: impl Into<String>,
114        buckets: Vec<f64>,
115    ) -> MetricsResult<HistogramHandle> {
116        let id = MetricId::new(namespace, name);
117        let metric = Metric::new_histogram(id.clone(), description, unit, buckets);
118        self.backend.put_metric(metric)?;
119        Ok(HistogramHandle::new(id, self.backend.clone()))
120    }
121
122    /// Get a metric by ID
123    pub fn get_metric(&self, id: &MetricId) -> MetricsResult<Metric> {
124        self.backend.get_metric(id)
125    }
126
127    /// Check if a metric exists
128    pub fn contains_metric(&self, id: &MetricId) -> bool {
129        self.backend.contains_metric(id)
130    }
131
132    /// Get all metrics in a namespace
133    pub fn get_namespace_metrics(&self, namespace: &str) -> Vec<Metric> {
134        // Try to cast to InMemoryBackend for more efficient operation
135        if let Some(memory_backend) = self
136            .backend
137            .as_ref()
138            .as_any()
139            .downcast_ref::<InMemoryBackend>()
140        {
141            memory_backend.get_metrics_by_namespace(namespace)
142        } else {
143            // Fallback for other backends
144            self.backend
145                .list_all_metrics()
146                .into_iter()
147                .filter(|m| m.metadata.id.namespace == namespace)
148                .collect()
149        }
150    }
151
152    /// Get all counters
153    pub fn get_all_counters(&self) -> Vec<Metric> {
154        if let Some(memory_backend) = self
155            .backend
156            .as_ref()
157            .as_any()
158            .downcast_ref::<InMemoryBackend>()
159        {
160            memory_backend.get_metrics_by_type(MetricType::Counter)
161        } else {
162            self.backend
163                .list_all_metrics()
164                .into_iter()
165                .filter(|m| m.metadata.metric_type == MetricType::Counter)
166                .collect()
167        }
168    }
169
170    /// Get all gauges
171    pub fn get_all_gauges(&self) -> Vec<Metric> {
172        if let Some(memory_backend) = self
173            .backend
174            .as_ref()
175            .as_any()
176            .downcast_ref::<InMemoryBackend>()
177        {
178            memory_backend.get_metrics_by_type(MetricType::Gauge)
179        } else {
180            self.backend
181                .list_all_metrics()
182                .into_iter()
183                .filter(|m| m.metadata.metric_type == MetricType::Gauge)
184                .collect()
185        }
186    }
187
188    /// Get the total number of metrics
189    pub fn metric_count(&self) -> usize {
190        self.backend.metric_count()
191    }
192
193    /// Clear all metrics
194    pub fn clear_all(&self) -> MetricsResult<()> {
195        self.backend.clear_all()
196    }
197
198    /// Get direct access to the backend (for advanced operations)
199    pub fn backend(&self) -> &Arc<dyn MetricsBackend> {
200        &self.backend
201    }
202}
203
204impl Default for MetricsRegistry {
205    fn default() -> Self {
206        Self::new()
207    }
208}
209
210/// Handle for efficient counter operations
211#[derive(Debug, Clone)]
212pub struct CounterHandle {
213    id: MetricId,
214    backend: Arc<dyn MetricsBackend>,
215}
216
217impl CounterHandle {
218    fn new(id: MetricId, backend: Arc<dyn MetricsBackend>) -> Self {
219        Self { id, backend }
220    }
221
222    /// Increment the counter by 1
223    pub fn increment(&self) -> MetricsResult<u64> {
224        self.backend.increment_counter(&self.id, 1)
225    }
226
227    /// Increment the counter by a specific amount
228    pub fn increment_by(&self, amount: u64) -> MetricsResult<u64> {
229        self.backend.increment_counter(&self.id, amount)
230    }
231
232    /// Get the current counter value
233    pub fn get(&self) -> MetricsResult<u64> {
234        let metric = self.backend.get_metric(&self.id)?;
235        metric
236            .value
237            .as_counter()
238            .ok_or_else(|| MetricsError::TypeMismatch {
239                expected: MetricType::Counter,
240                found: metric.value.metric_type(),
241            })
242    }
243
244    /// Get the metric ID
245    pub fn id(&self) -> &MetricId {
246        &self.id
247    }
248}
249
250/// Handle for efficient gauge operations
251#[derive(Debug, Clone)]
252pub struct GaugeHandle {
253    id: MetricId,
254    backend: Arc<dyn MetricsBackend>,
255}
256
257impl GaugeHandle {
258    fn new(id: MetricId, backend: Arc<dyn MetricsBackend>) -> Self {
259        Self { id, backend }
260    }
261
262    /// Set the gauge to a specific value
263    pub fn set(&self, value: f64) -> MetricsResult<()> {
264        self.backend.set_gauge(&self.id, value)
265    }
266
267    /// Increment the gauge by a specific amount
268    pub fn add(&self, delta: f64) -> MetricsResult<f64> {
269        let current = self.get()?;
270        let new_value = current + delta;
271        self.set(new_value)?;
272        Ok(new_value)
273    }
274
275    /// Decrement the gauge by a specific amount
276    pub fn sub(&self, delta: f64) -> MetricsResult<f64> {
277        self.add(-delta)
278    }
279
280    /// Get the current gauge value
281    pub fn get(&self) -> MetricsResult<f64> {
282        let metric = self.backend.get_metric(&self.id)?;
283        metric
284            .value
285            .as_gauge()
286            .ok_or_else(|| MetricsError::TypeMismatch {
287                expected: MetricType::Gauge,
288                found: metric.value.metric_type(),
289            })
290    }
291
292    /// Get the metric ID
293    pub fn id(&self) -> &MetricId {
294        &self.id
295    }
296}
297
298/// Handle for efficient histogram operations
299#[derive(Debug, Clone)]
300pub struct HistogramHandle {
301    id: MetricId,
302    backend: Arc<dyn MetricsBackend>,
303}
304
305impl HistogramHandle {
306    fn new(id: MetricId, backend: Arc<dyn MetricsBackend>) -> Self {
307        Self { id, backend }
308    }
309
310    /// Record a sample in the histogram
311    pub fn observe(&self, value: f64) -> MetricsResult<()> {
312        self.backend.record_histogram_sample(&self.id, value)
313    }
314
315    /// Get the metric ID
316    pub fn id(&self) -> &MetricId {
317        &self.id
318    }
319
320    /// Get the full histogram metric (for analysis)
321    pub fn get_metric(&self) -> MetricsResult<Metric> {
322        self.backend.get_metric(&self.id)
323    }
324}
325
326#[cfg(test)]
327mod tests {
328    use super::*;
329
330    #[test]
331    fn test_registry_creation() {
332        let registry = MetricsRegistry::new();
333        assert_eq!(registry.metric_count(), 0);
334    }
335
336    #[test]
337    fn test_counter_registration_and_operations() {
338        let registry = MetricsRegistry::new();
339
340        let counter = registry
341            .register_counter("engine", "frame_count", "Total number of frames rendered")
342            .unwrap();
343
344        // Test increment operations
345        assert_eq!(counter.increment().unwrap(), 1);
346        assert_eq!(counter.increment_by(5).unwrap(), 6);
347        assert_eq!(counter.get().unwrap(), 6);
348
349        // Verify it's stored in the registry
350        assert!(registry.contains_metric(counter.id()));
351        assert_eq!(registry.metric_count(), 1);
352    }
353
354    #[test]
355    fn test_gauge_registration_and_operations() {
356        let registry = MetricsRegistry::new();
357
358        let gauge = registry
359            .register_gauge("memory", "heap_usage", "Current heap usage", "MB")
360            .unwrap();
361
362        // Test gauge operations
363        gauge.set(100.5).unwrap();
364        assert_eq!(gauge.get().unwrap(), 100.5);
365
366        assert_eq!(gauge.add(50.0).unwrap(), 150.5);
367        assert_eq!(gauge.sub(25.0).unwrap(), 125.5);
368
369        // Verify it's stored in the registry
370        assert!(registry.contains_metric(gauge.id()));
371    }
372
373    #[test]
374    fn test_histogram_registration_and_operations() {
375        let registry = MetricsRegistry::new();
376
377        let histogram = registry
378            .register_histogram(
379                "renderer",
380                "frame_time",
381                "Frame rendering time distribution",
382                "ms",
383                vec![1.0, 5.0, 10.0, 50.0, 100.0],
384            )
385            .unwrap();
386
387        // Test histogram operations
388        histogram.observe(2.5).unwrap();
389        histogram.observe(15.0).unwrap();
390        histogram.observe(75.0).unwrap();
391
392        // Verify it's stored
393        assert!(registry.contains_metric(histogram.id()));
394
395        let metric = histogram.get_metric().unwrap();
396        if let khora_core::telemetry::metrics::MetricValue::Histogram { samples, .. } = metric.value
397        {
398            assert_eq!(samples.len(), 3);
399            assert!(samples.contains(&2.5));
400            assert!(samples.contains(&15.0));
401            assert!(samples.contains(&75.0));
402        } else {
403            panic!("Expected histogram metric");
404        }
405    }
406
407    #[test]
408    fn test_metrics_with_labels() {
409        let registry = MetricsRegistry::new();
410
411        let counter = registry
412            .register_counter_with_labels(
413                "renderer",
414                "triangles_rendered",
415                "Number of triangles rendered",
416                vec![
417                    ("quality".to_string(), "high".to_string()),
418                    ("pass".to_string(), "main".to_string()),
419                ],
420            )
421            .unwrap();
422
423        counter.increment_by(1000).unwrap();
424
425        let id_str = counter.id().to_string_formatted();
426        assert!(id_str.contains("quality=high"));
427        assert!(id_str.contains("pass=main"));
428    }
429
430    #[test]
431    fn test_namespace_filtering() {
432        let registry = MetricsRegistry::new();
433
434        registry
435            .register_counter("engine", "frames", "Frame count")
436            .unwrap();
437        registry
438            .register_counter("engine", "updates", "Update count")
439            .unwrap();
440        registry
441            .register_counter("renderer", "draws", "Draw calls")
442            .unwrap();
443        registry
444            .register_gauge("memory", "heap", "Heap usage", "MB")
445            .unwrap();
446
447        let engine_metrics = registry.get_namespace_metrics("engine");
448        assert_eq!(engine_metrics.len(), 2);
449
450        let renderer_metrics = registry.get_namespace_metrics("renderer");
451        assert_eq!(renderer_metrics.len(), 1);
452
453        let memory_metrics = registry.get_namespace_metrics("memory");
454        assert_eq!(memory_metrics.len(), 1);
455    }
456
457    #[test]
458    fn test_type_filtering() {
459        let registry = MetricsRegistry::new();
460
461        registry
462            .register_counter("test", "c1", "Counter 1")
463            .unwrap();
464        registry
465            .register_counter("test", "c2", "Counter 2")
466            .unwrap();
467        registry
468            .register_gauge("test", "g1", "Gauge 1", "unit")
469            .unwrap();
470
471        let counters = registry.get_all_counters();
472        assert_eq!(counters.len(), 2);
473
474        let gauges = registry.get_all_gauges();
475        assert_eq!(gauges.len(), 1);
476    }
477
478    #[test]
479    fn test_clear_all() {
480        let registry = MetricsRegistry::new();
481
482        registry
483            .register_counter("test", "counter", "Test counter")
484            .unwrap();
485        registry
486            .register_gauge("test", "gauge", "Test gauge", "unit")
487            .unwrap();
488
489        assert_eq!(registry.metric_count(), 2);
490
491        registry.clear_all().unwrap();
492        assert_eq!(registry.metric_count(), 0);
493    }
494}
khora_telemetry/metrics/registry.rs

khora_telemetry/metrics/
registry.rs
