metrics.rs

//! Utilities to work with metrics up for tests.
use std::sync::atomic::Ordering;

use eyre::{Result, eyre};
use metrics::{
    Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString,
    Unit,
};
use metrics_util::registry::{AtomicStorage, Registry, Storage};
use tokio::time;

/// A memory default implementation of the metrics::Recorder.
pub struct MemoryRecorder {
    registry: Registry<Key, AtomicStorage>,
}

// I've not implemented a `reset` method here. This is because I'd like to use
// `set_default_local_recorder` as long as it is possible. I can use `{}` scope
// blocks to handle the resets and get a new registry.
//
// Note: reset can be implemented by `retain_counters` with a callback that
// always returns false.
impl MemoryRecorder {
    /// Creates a new instance of `MemoryRecorder`.
    pub fn new() -> MemoryRecorder {
        MemoryRecorder {
            registry: Registry::atomic(),
        }
    }

    /// Tries to fetch a key
    pub fn maybe_get_counter(&self, key: &Key) -> Option<u64> {
        self.registry
            .get_counter(key)
            .map(|v| v.load(Ordering::Relaxed))
    }

    /// Fetches a key
    pub fn get_counter(&self, key: &Key) -> Result<u64> {
        Ok(self
            .maybe_get_counter(key)
            .ok_or(eyre!("key does not exist"))?)
    }

    /// Tries to fetch a key
    pub fn maybe_get_gauge(&self, key: &Key) -> Option<f64> {
        self.registry
            .get_gauge(key)
            .map(|v| f64::from_bits(v.load(Ordering::Relaxed)))
    }

    pub fn get_gauge(&self, key: &Key) -> Result<f64> {
        Ok(self
            .maybe_get_gauge(key)
            .ok_or(eyre!("key does not exist"))?)
    }

    /// Returns a histogram for the given key.
    pub fn get_histogram(
        &self,
        _key: &Key,
    ) -> Option<<AtomicStorage as Storage<Key>>::Histogram> {
        todo!()
    }

    /// Asserts that a gauge has the expected value by the deadline.
    pub async fn assert_gauge(&self, key: &Key, value: f64) -> Result<()> {
        on_deadline(|| {
            if self.get_gauge(key)? == value {
                return Ok(true);
            }

            Ok(false)
        })
        .await?;

        Ok(())
    }
}

impl Recorder for MemoryRecorder {
    fn describe_counter(
        &self,
        _key: KeyName,
        _unit: Option<Unit>,
        _description: SharedString,
    ) {
    }

    fn describe_gauge(
        &self,
        _key: KeyName,
        _unit: Option<Unit>,
        _description: SharedString,
    ) {
    }

    fn describe_histogram(
        &self,
        _key: KeyName,
        _unit: Option<Unit>,
        _description: SharedString,
    ) {
    }

    fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
        self.registry
            .get_or_create_counter(key, |c| Counter::from_arc(c.clone()))
    }

    fn register_gauge(&self, key: &Key, _metadata: &Metadata<'_>) -> Gauge {
        self.registry
            .get_or_create_gauge(key, |g| Gauge::from_arc(g.clone()))
    }

    fn register_histogram(
        &self,
        key: &Key,
        _metadata: &Metadata<'_>,
    ) -> Histogram {
        self.registry
            .get_or_create_histogram(key, |h| Histogram::from_arc(h.clone()))
    }
}

async fn on_deadline<F>(check: F) -> Result<()>
where
    F: Fn() -> Result<bool>,
{
    let deadline = time::Duration::from_secs(1);
    time::timeout(deadline, async move {
        let mut tick = time::interval(time::Duration::from_millis(10));
        loop {
            tick.tick().await;

            if check()? {
                break;
            }
        }

        Ok::<_, eyre::Report>(())
    })
    .await??;

    tracing::info!("finished");

    Ok(())
}