casper-node 0.6.3

use prometheus::{self, IntGauge, Registry};

use crate::reactor::{EventQueueHandle, QueueKind};
use itertools::Itertools;
use std::collections::HashMap;
use tracing::debug;

/// Metrics for event queue sizes.
#[derive(Debug)]
pub(super) struct EventQueueMetrics {
    /// Per queue kind gauges that measure number of event in the queue.
    event_queue_gauges: HashMap<QueueKind, IntGauge>,
    /// Total events count.
    event_total: IntGauge,
    /// Instance of registry to unregister from when being dropped.
    registry: Registry,
}

impl EventQueueMetrics {
    /// Initializes event queue sizes metrics.
    pub(super) fn new<REv: 'static>(
        registry: Registry,
        event_queue_handle: EventQueueHandle<REv>,
    ) -> Result<Self, prometheus::Error> {
        let mut event_queue_gauges: HashMap<QueueKind, IntGauge> = HashMap::new();
        for queue_kind in event_queue_handle.event_queues_counts().keys() {
            let key = format!("scheduler_queue_{}_count", queue_kind.metrics_name());
            let queue_event_counter = IntGauge::new(key, "Event in the queue.".to_string())?;
            registry.register(Box::new(queue_event_counter.clone()))?;
            let result = event_queue_gauges.insert(*queue_kind, queue_event_counter);
            assert!(result.is_none(), "Map keys should not be overwritten.");
        }

        let event_total = IntGauge::new(
            "scheduler_queue_total_count",
            "total count of events in queues.",
        )?;
        registry.register(Box::new(event_total.clone()))?;

        Ok(EventQueueMetrics {
            event_queue_gauges,
            event_total,
            registry,
        })
    }

    /// Updates the event queues size metrics.
    /// NOTE: Count may be off by one b/c of the way locking works when elements are popped.
    /// It's fine for its purposes.
    pub(super) fn record_event_queue_counts<REv: 'static>(
        &self,
        event_queue_handle: &EventQueueHandle<REv>,
    ) {
        let event_queue_count = event_queue_handle.event_queues_counts();

        let total = event_queue_count.values().sum::<usize>() as i64;
        self.event_total.set(total);

        let event_counts: String = event_queue_count
            .iter()
            .sorted_by_key(|k| k.0)
            .map(|(queue, event_count)| {
                let _ = self
                    .event_queue_gauges
                    .get(queue)
                    .map(|gauge| gauge.set(*event_count as i64))
                    .expect("queue exists.");
                format!("{}={}", queue, event_count)
            })
            .join(",");

        debug!(%total, %event_counts, "Collected new set of event queue sizes metrics.")
    }
}

impl Drop for EventQueueMetrics {
    fn drop(&mut self) {
        self.registry
            .unregister(Box::new(self.event_total.clone()))
            .expect("did not expect de-registering of scheduler_queue_total_count to fail.");
        self.event_queue_gauges
            .iter()
            .for_each(|(key, queue_gauge)| {
                self.registry
                    .unregister(Box::new(queue_gauge.clone()))
                    .unwrap_or_else(|_| {
                        panic!("did not expect de-registering of {:?} to fail.", key)
                    })
            });
    }
}