hesione 0.1.1

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt::Write,
    sync::atomic::Ordering,
};

use atomic_traits::{Atomic, NumOps};
use num_traits::One;
use parking_lot::{lock_api::RwLockUpgradableReadGuard, RwLock};

use crate::{Label, Metric, NameError, Value};

/// Gauge metric
///
/// From the Prometheus docs:
///
/// > A gauge is a metric that represents a single numerical value that can
/// > arbitrarily go up and down.
/// >
/// > Gauges are typically used for measured values like temperatures or current
/// > memory usage, but also "counts" that can go up and down, like the number
/// > of concurrent requests.
///
/// # Examples
///
/// ```
/// use std::sync::atomic::{Ordering, AtomicIsize};
/// use hesione::{labels, Gauge, Metric};
///
/// // Create an example gauge
/// let gauge = Gauge::<AtomicIsize>::new("example").unwrap();
///
/// // Mess with the value a bit
/// gauge.inc(Ordering::SeqCst, None);
/// gauge.dec(Ordering::SeqCst, None);
/// gauge.dec(Ordering::SeqCst, None);
///
/// // When we convert the metric into Prometheus' format, it'll look like this:
/// let expected = concat!(
///     "# TYPE example gauge\n",
///     "example -1\n",
///     "\n",
/// );
///
/// let mut actual = String::with_capacity(64);
/// gauge.to_prometheus_lines(&mut actual, Ordering::SeqCst);
///
/// assert_eq!(actual, expected);
/// ```
pub struct Gauge<T> {
    // Name of the metric
    name: String,

    // Longer description of what metric measures
    help: Option<String>,

    // Little explanation of this data structure:
    //
    // * (None, T) is `metric_name 1337`
    // * (Some(BTreeSet<Label>), T) would be something like
    //   `metric_name{foo="bar", baz="qux"} 1337`
    values: RwLock<BTreeMap<Option<BTreeSet<Label>>, T>>,
}

impl<T> Gauge<T> {
    /// Create a new [`Gauge`][Gauge]-type metric
    pub fn new<S: ToString>(name: S) -> Result<Self, NameError> {
        let name = name.to_string();

        if crate::is_valid_name(&name) {
            Ok(Self {
                name,
                help: Default::default(),
                values: Default::default(),
            })
        } else {
            Err(NameError)
        }
    }

    /// Set the help text for this gauge
    pub fn help<S: ToString>(mut self, help: S) -> Self {
        self.help = Some(help.to_string());

        self
    }
}

impl<T> Gauge<T>
where
    T: Default + Atomic,
{
    /// Gets the current value of the gauge
    pub fn get(
        &self,
        order: Ordering,
        labels: Option<BTreeSet<Label>>,
    ) -> <T as Atomic>::Type {
        let values = self.values.upgradable_read();

        match values.get(&labels) {
            Some(x) => x.load(order),
            None => {
                let mut values = RwLockUpgradableReadGuard::upgrade(values);
                let x = values.entry(labels).or_default();

                x.load(order)
            }
        }
    }
}

impl<T> Gauge<T>
where
    T: Default + NumOps + Atomic,
    <T as Atomic>::Type: One,
{
    /// Increments the gauge by 1
    pub fn inc(&self, order: Ordering, labels: Option<BTreeSet<Label>>) {
        let values = self.values.upgradable_read();

        match values.get(&labels) {
            Some(x) => {
                x.fetch_add(<T as Atomic>::Type::one(), order);
            }
            None => {
                let mut values = RwLockUpgradableReadGuard::upgrade(values);
                let x = values.entry(labels).or_default();

                x.fetch_add(<T as Atomic>::Type::one(), order);
            }
        }
    }

    /// Decrements the gauge by 1
    pub fn dec(&self, order: Ordering, labels: Option<BTreeSet<Label>>) {
        let values = self.values.upgradable_read();

        match values.get(&labels) {
            Some(x) => {
                x.fetch_sub(<T as Atomic>::Type::one(), order);
            }
            None => {
                let mut values = RwLockUpgradableReadGuard::upgrade(values);
                let x = values.entry(labels).or_default();

                x.fetch_sub(<T as Atomic>::Type::one(), order);
            }
        }
    }

    /// Sets the gauge to an arbitrary value
    pub fn set(
        &self,
        order: Ordering,
        labels: Option<BTreeSet<Label>>,
        set: <T as Atomic>::Type,
    ) {
        let values = self.values.upgradable_read();

        match values.get(&labels) {
            Some(x) => {
                x.store(set, order);
            }
            None => {
                let mut values = RwLockUpgradableReadGuard::upgrade(values);
                let x = values.entry(labels).or_default();

                x.store(set, order);
            }
        }
    }
}

impl<T, W> Metric<W> for Gauge<T>
where
    T: Value,
    W: Write,
{
    fn to_prometheus_lines(&self, buf: &mut W, order: Ordering) {
        if let Some(help) = &self.help {
            writeln!(buf, "# HELP {} {}", self.name, help).unwrap();
        }

        writeln!(buf, "# TYPE {} gauge", self.name).unwrap();

        for (k, v) in self.values.read().iter() {
            let text = v.to_prometheus_string(order);

            if let Some(labels) = k {
                write!(buf, "{}", self.name).unwrap();

                if !labels.is_empty() {
                    write!(buf, "{{").unwrap();
                }

                let len = labels.len();
                for (i, label) in labels.iter().enumerate() {
                    let maybe_separator = if i + 1 == len {
                        // This is the end
                        ""
                    } else {
                        // We need another separator
                        ", "
                    };
                    write!(
                        buf,
                        "{}=\"{}\"{}",
                        label.name, label.value, maybe_separator
                    )
                    .unwrap();
                }

                if !labels.is_empty() {
                    write!(buf, "}}").unwrap();
                }

                writeln!(buf, " {}", text).unwrap();
            } else {
                writeln!(buf, "{} {}", self.name, text).unwrap();
            }
        }

        // Just for readability in the output
        writeln!(buf).unwrap();
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn labelled_increment() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::{labels, Gauge};

        let gauge = Gauge::<AtomicU8>::new("tests").unwrap();

        let labels = labels! {
            "foo": "bar".to_string(),
        }
        .unwrap();

        gauge.inc(Ordering::SeqCst, Some(labels.clone()));
        gauge.dec(Ordering::SeqCst, Some(labels.clone()));
        gauge.inc(Ordering::SeqCst, Some(labels.clone()));

        // internal stuff to get the string value
        let x = gauge.get(Ordering::SeqCst, Some(labels));
        assert_eq!(x, 1);
    }

    #[test]
    fn no_values_to_lines() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::{Gauge, Metric};

        let gauge = Gauge::<AtomicU8>::new("tests")
            .unwrap()
            .help("The number of times this test has been run");

        let expected = concat!(
            "# HELP tests The number of times this test has been run\n",
            "# TYPE tests gauge\n",
            "\n",
        );

        let mut buf = String::with_capacity(128);
        gauge.to_prometheus_lines(&mut buf, Ordering::SeqCst);
        assert_eq!(expected, buf);
    }

    #[test]
    fn one_value_no_labels_to_lines() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::{Gauge, Metric};

        let gauge = Gauge::<AtomicU8>::new("tests")
            .unwrap()
            .help("The number of times this test has been run");

        gauge.inc(Ordering::SeqCst, None);
        gauge.dec(Ordering::SeqCst, None);
        gauge.inc(Ordering::SeqCst, None);

        let expected = concat!(
            "# HELP tests The number of times this test has been run\n",
            "# TYPE tests gauge\n",
            "tests 1\n",
            "\n",
        );

        let mut buf = String::with_capacity(128);
        gauge.to_prometheus_lines(&mut buf, Ordering::SeqCst);
        assert_eq!(expected, buf);
    }

    #[test]
    fn one_value_one_label_to_lines() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::{labels, Gauge, Metric};

        let gauge = Gauge::<AtomicU8>::new("tests")
            .unwrap()
            .help("The number of times this test has been run");

        let labels = labels! {
            "foo": "bar".to_string(),
        }
        .unwrap();

        gauge.inc(Ordering::SeqCst, Some(labels.clone()));
        gauge.dec(Ordering::SeqCst, Some(labels.clone()));
        gauge.inc(Ordering::SeqCst, Some(labels));

        let expected = concat!(
            "# HELP tests The number of times this test has been run\n",
            "# TYPE tests gauge\n",
            "tests{foo=\"bar\"} 1\n",
            "\n",
        );

        let mut buf = String::with_capacity(128);
        gauge.to_prometheus_lines(&mut buf, Ordering::SeqCst);
        assert_eq!(expected, buf);
    }

    #[test]
    fn one_value_two_labels_to_lines() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::{labels, Gauge, Metric};

        let gauge = Gauge::<AtomicU8>::new("tests")
            .unwrap()
            .help("The number of times this test has been run");

        let labels = labels! {
            "foo": "bar".to_string(),
            "baz": "qux".to_string(),
        }
        .unwrap();

        gauge.inc(Ordering::SeqCst, Some(labels.clone()));
        gauge.dec(Ordering::SeqCst, Some(labels.clone()));
        gauge.inc(Ordering::SeqCst, Some(labels));

        let expected = concat!(
            "# HELP tests The number of times this test has been run\n",
            "# TYPE tests gauge\n",
            // This order is deterministic because it's a BTreeSet and not a
            // HashSet
            "tests{baz=\"qux\", foo=\"bar\"} 1\n",
            "\n",
        );

        let mut buf = String::with_capacity(128);
        gauge.to_prometheus_lines(&mut buf, Ordering::SeqCst);
        assert_eq!(expected, buf);
    }

    #[test]
    fn set() {
        use std::sync::atomic::{AtomicU8, Ordering};

        use crate::Gauge;

        let gauge = Gauge::<AtomicU8>::new("tests")
            .unwrap()
            .help("The number of times this test has been run");

        let expected = 10;
        gauge.set(Ordering::SeqCst, None, expected);

        assert_eq!(expected, gauge.get(Ordering::SeqCst, None));
    }
}