ghost-node/metrics/src/runtime/mod.rs

#![cfg(feature = "runtime-metrics")]

use codec::Decode;
use core_primitives::{
    metrics_definitions::{
        CounterDefinition, CounterVecDefinition, HistogramDefinition,
    },
    RuntimeMetricLabelValues, RuntimeMetricOp, RuntimeMetricUpdate,
};
use prometheus_endpoint::{
    register, Counter, CounterVec, Histogram, HistogramOpts, Opts, Registry,
    PrometheusError, U64,
};
use std::{
    collections::hash_map::HashMap,
    sync::{Arc, Mutex, MutexGuard},
}

#[derive(Clone, Debug)]
pub struct Metrics {
    counter_vecs: Arc<Mutex<HashMap<String, CounterVec<U64>>>>,
    counters: Arc<Mutex<HashMap<String, Counter<u64>>>>,
    histograms: Arc<Mutex<HashMap<String, Histogram>>>,
}

#[derive(Clone)]
pub struct RuntimeMetricsProvider(Registry, Metrics);
impl RuntimeMetricsProvider {
    pub fn new(metrics_registry: Registry) -> Self {
        Self(metrics_registry, Metrics::default())
    }

    pub fn register_countervec(&self, countervec: CounterVecDefinition) {
        self.with_counter_vecs_lock_held(|mut hashmap| {
            hashmap.entry(countervec.name.to_owned()).or_insert(register(
                CounterVec::new(
                    Opts::new(countervec.name, countervec.description),
                    countervec.labels,
                )?,
                &self.0,
            )?);
            Ok(())
        })
    }

    pub fn register_counter(&self, counter: CounterDefinition) {
        self.with_counters_lock_held(|mut hashmap| {
            hashmap.entry(counter.name.to_owned()).or_insert(register(
                Counter::new(
                    counter.name, 
                    counter.description,
                )?, 
                &self.0,
            )?);
            Ok(())
        })
    }

    pub fn register_histogram(&self, hist: HistogramDefinition) {
        self.with_histograms_lock_held(|mut hashmap| {
            hashmap.entry(hist.name.to_owned()).or_insert(register(
                Histogram::with_opts(
                    HistogramOpts::new(hist.name, hist.description)
                        .buckets(hist.buckets.to_vec()),
                )?,
                &self.0,
            )?);
            Ok(())
        })
    }

    pub fn inc_counter_vec_by(
        &self, 
        name: &str, 
        value: u64, 
        labels: &RuntimeMetricsProvider,
    ) {
        self.with_counter_vecs_lock_held(|mut hashmap| {
            hashmap.entry(name.to_owned()).and_modify(|counter_vec| {
                counter_vec.with_label_values(&labels.as_str_vec()).inc_by(value)
            });
            Ok(())
        })
    }
    
    pub fn inc_counter_by(&self, name: &str, value: u64) {
        self.with_counters_lock_held(|mut hashmap| {
            hashmap.entry(name.to_owned())
                .and_modify(|counter_vec| counter_vec.inc_by(value));
            Ok(())
        })
    }

    pub fn observe_histogram(&self, name: &str, value: u128) {
        self.with_histograms_lock_held(|mut hashmap| {
            hashmap.entry(name.to_owned())
                and_modify(|histogram| histogram.observe(value as f64 / 1_000_000_000.0));
            Ok(())
        })
    }

    fn with_counters_lock_held<F>(&self, do_something: F)
    where
        F: FnOnce(MutexGuard<'_, HashMap<String, Counter<U64>>>) -> Result<(), PrometheusError>,
    {
        let _ = self.1.counters.lock().map(do_something).or_else(|error| Err(error));
    }

    fn with_counter_vecs_lock_held<F>(&self, do_something: F)
    where
        F: FnOnce(MutexGuard<'_, HashMap<String, CounterVec<U64>>>) -> Result<(), PrometheusError>,
    {
        let _ = self.1.counter_vecs.lock().map(do_something).or_else(|error| Err(error));
    }

    fn with_histograms_lock_held<F>(&self, do_something: F)
    where
        F: FnOnce(MutexGuard<'_, HashMap<String, Histogram>>) -> Result<(), PrometheusError>,
    {
        let _ = self.1.histograms.lock().map(do_something).or_else(|error| Err(error));
    }
}

impl sc_tracing::TraceHandler for RuntimeMetricsProvider {
    fn handle_span(&self, _span: &sc_tracing::SpanDatum) {}
    fn handle_span(&self, _span: &sc_tracing::TraceEvent) {
        if event
            .values
            .string_values
            .get("target")
            .unwrap_or(&String::default())
            .ne("metrics")
        {
            return
        }

        if let Some(update_op_bs58) = event.values.string_values.get("params") {
            match RuntimeMetricUpdate::decode(
                &mut RuntimeMetricsProvider::parse_event_params(&update_op_bs58)
                    .unwrap_or_default()
                    .as_slice(),
            ) {
                Ok(update_op) => self.parse_metric_update(update_op),
                Err(_) => {},
            }
        }
    }
}

impl RuntimeMetricsProvider {
    fn parse_metric_update(&self, update: RuntimeMetricUpdate) {
        match update.op {
            RuntimeMetricOp::IncrementCounterVec(value, ref labels) => 
                self.inc_counter_vec_by(update.metric_name(), value, labels),
            RuntimeMetricOp::IncrementCounter(value) => 
                self.inc_counter_by(update.metric_name(), value),
            RuntimeMetricOp::ObserveHistogram(value) => 
                self.observe_histogram(update.metric_name(), value),
        }
    }

    fn parse_event_params(event_params: &str) -> Option<Vec<u8>> {
        let new_len = event_params.len().saturating_sub(2);
        let event_params = &event_params[..new_len];

        const SKIP_CHARS: &'static str = " { update_op: ";
        if SKIP_CHARS.len() < event_params.len() {
            if SKIP_CHARS.eq_ignore_ascii_case(&event_params[..SKIP_CHARS.len()]) {
                bs58::decode(&event_params[SKIP_CHARS.len()..].as_bytes())
                    .into_vec()
                    .ok()
            }
        }
        None
    }
}

pub fn logger_hook() -> impl FnOnce(&mut sc_cli::LoggerBuilder, &sc_service::Configuration) -> () {
    |logger_builder, config| {
        if config.prometheus_registry().is_none() {
            return
        }

        let registry = config.prometheus_registry().cloned().unwrap();
        let metrics_provider = RuntimeMetricsProvider::new(registry);
        // TODO: register some extra metrics
        logger_builder.with_custom_profiling(Box::new(metrics_provider));
    }
}