epimetheus 0.9.0

/*! An easy-to-use prometheus-compatible metrics library

# Writing metrics

A "metric" is a named value of type `f64`.  There is a single global set of
metrics; you can update it from any thread.

```
use epimetheus::metric;

metric!(foobar).set(12.3);
metric!(foobar).add(0.7);
```

If you increment a metric which has never been set, it is considered to
start from zero.

```
# use epimetheus::metric;
metric!(barqux).add(6.5);
// now barqux = 6.5
```

## Labels

The base part of the name is fixed statically at at compile-time.  However,
a metric's name may also include "labels", which are dynamic.

```
# use epimetheus::metric;
let user_id = 7;
metric!(login_attempts{user=user_id}).add(1.0);
```

The label values can be anything which implements `Display`.

```
# use epimetheus::metric;
# let user_id = 0;
# let passwd = 0;
# let try_log_in = |_, _| 0;
// enum LoginResult { Success, BadUsername, BadPassword }
// impl Display for LoginResult { ... }

let result = try_log_in(user_id, passwd);
metric!(login_attempts{user=user_id, result=result}).add(1.0);
```

Labels can be useful, but they come at a performance cost (see README).

# Seeing your metrics

## ...via a function call

You can call `query()` to see the current value of the metrics:

```
# use epimetheus::metric;
# metric!(foobar).set(12.3);
# metric!(foobar).add(0.7);
# metric!(barqux).set(6.5);
# metric!(login_attempts{result="Success",user=7}).set(1.);
# metric!(login_attempts{user=7}).set(1.);
let mut metrics = epimetheus::query();
assert_eq!(metrics.next(), Some(("barqux".to_string(), 6.5)));
assert_eq!(metrics.next(), Some(("epimetheus_total_flushes".to_string(), 1.)));
assert_eq!(metrics.next(), Some(("epimetheus_total_updates".to_string(), 5.)));
assert_eq!(metrics.next(), Some(("foobar".to_string(), 13.)));
assert_eq!(metrics.next(), Some(("login_attempts{result=\"Success\",user=\"7\"}".to_string(), 1.)));
assert_eq!(metrics.next(), Some(("login_attempts{user=\"7\"}".to_string(), 1.)));
```

Note the "epimetheus_*" lines: these are metrics exposed by epimetheus itself.

## ...via HTTP

Set the `RUST_METRICS_PORT` environment variable when starting your program.
An HTTP server will be spawned the first time a metric is written.
Clients are sent the current state of all metrics in the Prometheus exposistion format.

```console
$ RUST_METRICS_PORT=9898 cargo run &
$ curl localhost:9898
barqux 6.5
epimetheus_total_flushes 2
epimetheus_total_updates 5
foobar 13
login_attempts{result="Success",user="7"} 1
login_attempts{user="7"} 1
```

## ...via systemd-report

Set the `RUST_METRICS_PATH` environment variable when starting your program.
A varlink server will be spawned the first time a metric is written.
The server speaks the `io.systemd.Metrics` varlink protocol.

```console
$ RUST_METRICS_PATH=/run/systemd/report/com.example.my_app cargo run &
$ systemd-report metrics
FAMILY                                          OBJECT FIELDS                         VALUE
com.example.my_app.barqux                       -      -                              6.5
com.example.my_app.epimetheus_total_flushes     -      -                              12
com.example.my_app.epimetheus_total_updates     -      -                              286571
com.example.my_app.foobar                       -      -                              1.0725e+03
com.example.my_app.login_attempts               -      {"user":"7"}                   715
com.example.my_app.login_attempts               -      {"result":"failed","user":"1"} 355000
```

Note: if you want systemd-report to pick up the metrics automatically,
you need to choose a path in /run/systemd/report/ (if you're root) or
/run/user/$UID/systemd/report/ (otherwise).

Tip: If you're running in via systemd, you can put something like this in the
service file, and the metrics will be namespaced by the name of the unit:
```systemd
Environment=RUST_METRICS_PATH=/run/systemd/report/com.example.%N
```

*/

use std::collections::{BTreeMap, BTreeSet};
use std::io::{self, prelude::*};
use std::net::{Ipv4Addr, TcpListener, TcpStream};
use std::os::unix::fs::PermissionsExt;
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::{Path, PathBuf};
use std::sync::mpsc::{Receiver, Sender, channel};
use std::sync::{LazyLock, Mutex};
use std::time::Duration;
use std::{fmt, thread};
use tracing::*;

struct State {
    chan: Sender<(Metric, Action)>,
    tracker: Mutex<Tracker>,
}

static STATE: LazyLock<State> = LazyLock::new(|| {
    let (tx, rx) = channel();
    let state = State {
        chan: tx,
        tracker: Mutex::new(Tracker {
            metrics: BTreeMap::default(),
            chan: rx,
        }),
    };
    // We want to ensure that the channel is regularly drained, even when
    // there are no new connections coming in.  (Otherwise, we'd have a
    // memory leak and - worse - the metric update latency would suffer.)
    // Therefore we spawn a thread which regularly drains the channel.
    thread::Builder::new()
        .name("epimetheus-drainer".into())
        .spawn(move || {
            loop {
                // We sleep first to avoid recursing
                thread::sleep(Duration::from_secs(20));
                STATE.tracker.lock().unwrap().update();
            }
        })
        .expect("Failed to spawn drainer thread");
    if let Some(port) = http_port() {
        match try_spawn_http_server_on(port) {
            Ok(()) => (),
            Err(e) => warn!("{e:#}"),
        }
    }
    if let Some(path) = systemd_path() {
        match try_spawn_systemd_server_at(&path) {
            Ok(()) => (),
            Err(e) => warn!("{e:#}"),
        }
    }
    state
});

/// A named metric;  it has a associated global mutable `f64` value.
///
/// You can create these by hand, but you might find it more convenient to
/// use the `metric!()` macro.
pub struct Metric {
    pub name: &'static str, // aka. "family"
    pub labels: Labels,     // aka. "fields"
}
type Labels = Vec<(&'static str, Box<dyn fmt::Display + Send>)>;
enum Action {
    Inc(f64),
    Set(f64),
    Min(f64),
    Max(f64),
}

impl Metric {
    /// Set the metric to the specified value.
    #[inline]
    pub fn set(self, x: f64) {
        send_chan((self, Action::Set(x)));
    }

    /// Increment the metric by the specified amount.
    #[inline]
    pub fn add(self, x: f64) {
        send_chan((self, Action::Inc(x)));
    }

    /// Set the metric to the specified value if it is smaller than the
    /// current value.
    #[inline]
    pub fn min(self, x: f64) {
        send_chan((self, Action::Min(x)));
    }

    /// Set the metric to the specified value if it is larger than the
    /// current value.
    #[inline]
    pub fn max(self, x: f64) {
        send_chan((self, Action::Max(x)));
    }
}

impl From<&'static str> for Metric {
    fn from(name: &'static str) -> Self {
        Metric {
            name,
            labels: vec![],
        }
    }
}

impl Metric {
    fn prepare(self) -> PreparedMetric {
        let labels = self
            .labels
            .into_iter()
            .map(|(k, v)| (k, v.to_string()))
            .collect::<BTreeMap<_, _>>();
        PreparedMetric {
            name: self.name,
            labels,
        }
    }
}

#[inline]
fn send_chan(x: (Metric, Action)) {
    STATE.chan.send(x).unwrap();
}

/// Refer to a metric.
#[macro_export]
macro_rules! metric {
    ($name:ident) => {
        $crate::Metric {
            name: stringify!($name),
            labels: Vec::new(),
        }
    };
    ($name:ident{$($key:ident = $val:expr),*}) => {
        $crate::Metric {
            name: stringify!($name),
            labels: vec![$((stringify!($key), Box::new($val))),*],
        }
    };
}

/// A partially pre-rendered version of [`Metric`]
#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
struct PreparedMetric {
    name: &'static str,
    labels: BTreeMap<&'static str, String>,
}
impl PreparedMetric {
    fn prometheus<'a>(&'a self) -> PrometheusMetric<'a> {
        PrometheusMetric {
            name: self.name,
            labels: self.labels.iter(),
        }
    }

    fn json<'a>(&'a self, namespace: &'a str) -> JsonMetric<'a> {
        JsonMetric {
            name: self.name,
            labels: self.labels.iter(),
            namespace,
        }
    }
}

struct PrometheusMetric<'a> {
    name: &'a str,
    labels: std::collections::btree_map::Iter<'a, &'static str, String>,
}
impl<'a> fmt::Display for PrometheusMetric<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.write_str(self.name)?;
        let mut labels = self.labels.clone();
        if let Some((k, v)) = labels.next() {
            write!(f, "{{{k}=\"{v}\"")?;
            for (k, v) in labels {
                write!(f, ",{k}=\"{v}\"")?;
            }
            f.write_str("}")?;
        }
        Ok(())
    }
}

struct JsonMetric<'a> {
    name: &'a str,
    labels: std::collections::btree_map::Iter<'a, &'static str, String>,
    namespace: &'a str,
}
impl<'a> fmt::Display for JsonMetric<'a> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // We use String's Debug impl to escape label values, just in case
        // they contain quotes or newlines.  Apparently Debug's escaping is
        // not _exactly_ the same as JSON escaping, in the case of non_ASCII
        // characters.  Proper JSON-conpliant escaping is a TODO.
        write!(f, r#""name":"{}.{}""#, self.namespace, self.name)?;
        let mut labels = self.labels.clone();
        if let Some(head) = labels.next() {
            f.write_str(r#","fields":"#)?;
            f.write_str("{\"")?;
            f.write_str(head.0)?;
            f.write_str("\":")?;
            fmt::Debug::fmt(head.1, f)?;
            for (k, v) in labels {
                f.write_str(",\"")?;
                f.write_str(k)?;
                f.write_str("\":")?;
                fmt::Debug::fmt(v, f)?;
            }
            f.write_str("}")?;
        }
        Ok(())
    }
}

/// Tracks the current state of the metrics
struct Tracker {
    /// Stores the current value of all the metrics in a map.  (We use a
    /// BTreeMap so the metrics are nicely sorted when we print them.)
    metrics: BTreeMap<PreparedMetric, f64>,
    chan: Receiver<(Metric, Action)>,
}

impl Tracker {
    fn update(&mut self) {
        let mut n = 0.;
        for (metric, action) in self.chan.try_iter() {
            let metric = metric.prepare();
            let entry = self.metrics.entry(metric).or_insert(0.0);
            match action {
                Action::Inc(x) => *entry += x,
                Action::Set(x) => *entry = x,
                Action::Min(x) => *entry = entry.min(x),
                Action::Max(x) => *entry = entry.max(x),
            }
            n += 1.;
        }
        let total_updates = Metric::from("epimetheus_total_updates").prepare();
        let total_flushes = Metric::from("epimetheus_total_flushes").prepare();
        *self.metrics.entry(total_updates).or_insert(0.) += n;
        *self.metrics.entry(total_flushes).or_insert(0.) += 1.;
    }
}

/// Get the current state of the metrics.
///
/// ```
/// use epimetheus::metric;
/// metric!(a_metric).set(42.0);
/// assert_eq!(
///     epimetheus::query().next(),
///     Some(("a_metric".to_string(), 42.0))
/// );
/// ```
pub fn query() -> impl Iterator<Item = (String, f64)> {
    get_metrics()
        .into_iter()
        .map(|(metric, val)| (metric.prometheus().to_string(), val))
}

fn get_metrics() -> BTreeMap<PreparedMetric, f64> {
    let mut tracker = STATE.tracker.lock().unwrap();
    tracker.update();
    tracker.metrics.clone()
}

fn http_port() -> Option<u16> {
    match std::env::var("RUST_METRICS_PORT") {
        Ok(x) => match x.parse::<u16>() {
            Ok(port) => return Some(port),
            Err(_) => warn!("RUST_METRICS_PORT present but not a valid port number"),
        },
        Err(std::env::VarError::NotPresent) => (),
        Err(std::env::VarError::NotUnicode(_)) => {
            warn!("RUST_METRICS_PORT present but not a valid port number")
        }
    }
    None
}

fn systemd_path() -> Option<PathBuf> {
    match std::env::var("RUST_METRICS_PATH") {
        Ok(x) => return Some(PathBuf::from(x)),
        Err(std::env::VarError::NotPresent) => (),
        Err(std::env::VarError::NotUnicode(_)) => {
            warn!("RUST_METRICS_PATH present but not a valid path")
        }
    }
    None
}

/// Bind a socket and listen for incoming connections from HTTP clients.
/// When we get a connection, we:
///
/// 1. drain any updates from the channel and apply them to the global
///    metrics map; then
/// 2. render the map to prometheus exposition format and send it to
///    the client.
fn try_spawn_http_server_on(port: u16) -> std::io::Result<()> {
    let sock = TcpListener::bind((Ipv4Addr::LOCALHOST, port))?;
    info!("Listening on port {port}");
    std::thread::Builder::new()
        .name("epimetheus-http".into())
        .spawn(move || {
            for conn in sock.incoming() {
                if let Err(e) = conn.and_then(|conn| handle_http_client(conn, get_metrics())) {
                    warn!("{}", e);
                }
            }
        })?;
    Ok(())
}

/// This is the world's simplest HTTP implementation.  It completely
/// ignores the request, unconditionally sending the same response.
/// This response comes with no headers or anything - just a body.
fn handle_http_client(
    conn: TcpStream,
    metrics: BTreeMap<PreparedMetric, f64>,
) -> Result<(), std::io::Error> {
    // We don't care about the request, but some HTTP clients get upset if you
    // don't at least read it.  Unfortunate.
    //
    // HTTP clients keep connection open: to know when to stop reading, you have
    // to look for \r\n\r\n in the stream. We do this below.
    let mut conn = std::io::BufReader::with_capacity(128, conn);
    let mut progress = 0;
    for b in std::io::Read::by_ref(&mut conn).bytes() {
        match b {
            Ok(b) => match progress {
                0 if b == b'\r' => progress = 1,
                1 if b == b'\n' => progress = 2,
                2 if b == b'\r' => progress = 3,
                3 if b == b'\n' => break,
                _ => progress = 0,
            },
            Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => (),
            Err(e) => return Err(e),
        }
    }
    let mut conn = conn.into_inner();
    writeln!(conn, "HTTP/1.1 200 OK\r\n")?;
    for (metric, val) in metrics {
        writeln!(conn, "{} {val}", metric.prometheus())?;
    }
    Ok(())
}

fn try_spawn_systemd_server_at(path: &Path) -> std::io::Result<()> {
    if !path.starts_with("/run/systemd/report/") {
        warn!(
            "{}: Not in system-report's socket directory",
            path.display()
        );
    }
    let namespace = path
        .file_name()
        .and_then(|x| x.to_str())
        .ok_or_else(|| io::Error::other(format!("{}: No filename", path.display())))?
        .to_owned();
    match std::fs::remove_file(path) {
        Err(e) if e.kind() == io::ErrorKind::NotFound => (), // expected
        Ok(()) => warn!("{}: already exists.  Removing...", path.display()),
        Err(e) => Err(e)?,
    }
    let sock = UnixListener::bind(path)?;
    std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o666))?;
    info!("Listening at {}", path.display());
    std::thread::Builder::new()
        .name("epimetheus-systemd".into())
        .spawn(move || {
            for conn in sock.incoming() {
                if let Err(e) =
                    conn.and_then(|conn| handle_systemd_client(conn, &namespace, get_metrics()))
                {
                    warn!("{}", e);
                }
            }
        })?;
    Ok(())
}

/// This is the world's simplest varlink implementation.  It
/// only responds to the two methods required by systemd-report
/// (io.systemd.Metrics.{Describe,List}).
fn handle_systemd_client(
    mut conn: UnixStream,
    namespace: &str,
    metrics: BTreeMap<PreparedMetric, f64>,
) -> Result<(), std::io::Error> {
    let mut method = vec![];
    std::io::BufReader::new(&conn).read_until(0, &mut method)?;
    method.pop(); // Drop the null
    let method = String::from_utf8(method).map_err(|e| io::Error::other(e.to_string()))?;
    debug!(method, "Got connection");
    // TODO: Be more lenient with the JSON.  Use gjson?
    let Some(method) = method
        .strip_prefix(r#"{"method":"io.systemd.Metrics."#)
        .and_then(|x| x.strip_suffix(r#"","more":true}"#))
    else {
        return Err(io::Error::other(format!("{method}: Unexpected interface")));
    };
    match method {
        "Describe" => {
            let names = metrics.keys().map(|x| x.name).collect::<BTreeSet<_>>();
            for (i, name) in names.iter().enumerate() {
                let continues = i + 1 != names.len();
                let desc = "(Epimetheus doesn't allow users to describe their metrics)";
                write!(
                    conn,
                    r#"{{"parameters":{{"name":"{namespace}.{}","description":"{desc}","type":"gauge"}},"continues":{continues}}}"#,
                    name,
                )?;
                conn.write_all(&[0])?;
            }
        }
        "List" => {
            for (i, (name, val)) in metrics.iter().enumerate() {
                let continues = i + 1 != metrics.len();
                write!(
                    conn,
                    r#"{{"parameters":{{{},"value":{val}}},"continues":{continues}}}"#,
                    name.json(namespace),
                )?;
                conn.write_all(&[0])?;
            }
        }
        _ => return Err(io::Error::other(format!("{method}: Unexpected method"))),
    }
    let mut buf = vec![];
    std::io::BufReader::new(&conn).read_until(0, &mut buf)?;
    debug!("Client said: {buf:?}");
    Ok(())
}

#[cfg(test)]
mod tests {
    use crate::*;

    #[test]
    fn test_non_http() -> Result<(), Box<dyn std::error::Error>> {
        metric!(foo).set(1.0);
        metric!(bar).add(1.0);
        metric!(bar).add(2.0);
        assert_eq!(query().find(|(k, _)| k == "foo").map(|x| x.1), Some(1.0));
        assert_eq!(query().find(|(k, _)| k == "bar").map(|x| x.1), Some(3.0));
        assert_eq!(query().find(|(k, _)| k == "qux"), None);
        metric!(bar).max(1.5);
        assert_eq!(query().find(|(k, _)| k == "bar").map(|x| x.1), Some(3.0));
        metric!(bar).min(1.5);
        assert_eq!(query().find(|(k, _)| k == "bar").map(|x| x.1), Some(1.5));
        Ok(())
    }

    #[test]
    fn labels() -> Result<(), Box<dyn std::error::Error>> {
        metric!(labels).set(1.0);
        metric!(labels{user=1, admin=false}).set(1.0);
        metric!(labels{user=2, name="Pete"}).set(1.0);
        assert_eq!(
            query()
                .filter(|x| x.0.starts_with("labels"))
                .collect::<Vec<_>>(),
            [
                r#"labels"#,
                r#"labels{admin="false",user="1"}"#,
                r#"labels{name="Pete",user="2"}"#,
            ]
            .into_iter()
            .map(|x| (x.to_string(), 1.0))
            .collect::<Vec<_>>(),
        );
        Ok(())
    }
}