OXYGENHUB

Coverage / crates/flow/src/metrics.rs 92.3%

96 of 104 lines run, 8 of 14 functions · run #187 on main · run missed

1
//! The runner's own metrics (oxygenflow#35, oxygenhub#212): a Prometheus
2
//! text page on an address of the runner's choosing, off unless
3
//! `metrics = "127.0.0.1:9464"` is set under `[runner]`. With
4
//! `metrics_token`, a scrape must send it as a bearer token.
5
6
use std::io::{BufRead, BufReader, Write};
7
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
8
use std::time::Instant;
9
10
/// Jobs this runner claimed.
11
pub static CLAIMS: AtomicU64 = AtomicU64::new(0);
12
/// Jobs declined before running (the preflight said no).
13
pub static DECLINED: AtomicU64 = AtomicU64::new(0);
14
/// Finished jobs by result.
15
pub static PASSED: AtomicU64 = AtomicU64::new(0);
16
pub static FAILED: AtomicU64 = AtomicU64::new(0);
17
pub static ERRORS: AtomicU64 = AtomicU64::new(0);
18
pub static CANCELLED: AtomicU64 = AtomicU64::new(0);
19
/// Milliseconds finished jobs ran, together.
20
pub static JOB_MS: AtomicU64 = AtomicU64::new(0);
21
/// Jobs running now.
22
pub static RUNNING: AtomicU64 = AtomicU64::new(0);
23
24
/// Count a job's end: by exit code, its error, and how long it ran.
251×
pub fn finished(exit_code: Option<i32>, error: Option<&str>, duration_ms: u64, cancelled: bool) {
261×
    let counter = if cancelled {
270×
        &CANCELLED
281×
    } else if error.is_some() {
290×
        &ERRORS
301×
    } else if exit_code == Some(0) {
311×
        &PASSED
32
    } else {
330×
        &FAILED
34
    };
351×
    counter.fetch_add(1, Relaxed);
361×
    JOB_MS.fetch_add(duration_ms, Relaxed);
371×
}
38
39
/// The page.
401×
pub fn render(name: &str, started: Instant) -> String {
411×
    let n = name.replace('\\', "\\\\").replace('"', "\\\"");
421×
    let mut out = String::new();
437×
    let mut family = |metric: &str, help: &str, kind: &str, samples: &[(String, u64)]| {
447×
        out.push_str(&format!("# HELP {metric} {help}\n# TYPE {metric} {kind}\n"));
4510×
        for (labels, v) in samples {
4610×
            let sep = if labels.is_empty() { "" } else { "," };
4710×
            out.push_str(&format!("{metric}{{runner=\"{n}\"{sep}{labels}}} {v}\n"));
48
        }
497×
    };
501×
    family("flow_runner_build_info", "This runner's version.", "gauge", &[(format!("version=\"{}\"", env!("CARGO_PKG_VERSION")), 1)]);
511×
    family("flow_runner_uptime_seconds", "Seconds since the runner started.", "gauge", &[(String::new(), started.elapsed().as_secs())]);
521×
    family("flow_runner_claims_total", "Jobs claimed from the hub.", "counter", &[(String::new(), CLAIMS.load(Relaxed))]);
531×
    family("flow_runner_declined_total", "Claimed jobs declined before running.", "counter", &[(String::new(), DECLINED.load(Relaxed))]);
541×
    family(
551×
        "flow_runner_jobs_finished_total",
561×
        "Finished jobs by result.",
571×
        "counter",
581×
        &[
591×
            ("result=\"passed\"".into(), PASSED.load(Relaxed)),
601×
            ("result=\"failed\"".into(), FAILED.load(Relaxed)),
611×
            ("result=\"error\"".into(), ERRORS.load(Relaxed)),
621×
            ("result=\"cancelled\"".into(), CANCELLED.load(Relaxed)),
631×
        ],
641×
    );
651×
    family(
661×
        "flow_runner_job_seconds_total",
671×
        "Seconds finished jobs ran, together.",
681×
        "counter",
691×
        &[(String::new(), JOB_MS.load(Relaxed) / 1000)],
701×
    );
711×
    family("flow_runner_jobs_running", "Jobs running now.", "gauge", &[(String::new(), RUNNING.load(Relaxed))]);
721×
    out
731×
}
74
75
/// Serve the page at `addr` on a thread of its own; a failure to listen is
76
/// said and the runner goes on without it.
771×
pub fn serve(addr: &str, token: Option<String>, name: String) {
781×
    let listener = match std::net::TcpListener::bind(addr) {
791×
        Ok(l) => l,
800×
        Err(e) => {
810×
            eprintln!("flow runner {name}: not serving metrics on {addr}: {e}");
820×
            return;
83
        }
84
    };
851×
    eprintln!("flow runner {name}: metrics on http://{addr}/metrics");
861×
    let started = Instant::now();
871×
    std::thread::spawn(move || {
882×
        for stream in listener.incoming().flatten() {
892×
            let _ = answer(stream, token.as_deref(), &name, started);
902×
        }
911×
    });
921×
}
93
942×
fn answer(mut stream: std::net::TcpStream, token: Option<&str>, name: &str, started: Instant) -> std::io::Result<()> {
952×
    stream.set_read_timeout(Some(std::time::Duration::from_secs(5)))?;
962×
    let mut reader = BufReader::new(stream.try_clone()?);
972×
    let mut line = String::new();
982×
    reader.read_line(&mut line)?;
992×
    let path = line.split_whitespace().nth(1).unwrap_or("").to_string();
1002×
    let mut authorized = token.is_none();
101
    loop {
1025×
        let mut h = String::new();
1035×
        if reader.read_line(&mut h)? == 0 || h.trim().is_empty() {
1042×
            break;
1053×
        }
1063×
        if let (Some(t), Some((k, v))) = (token, h.split_once(':')) {
1073×
            if k.trim().eq_ignore_ascii_case("authorization") && v.trim() == format!("Bearer {t}") {
1081×
                authorized = true;
1092×
            }
1100×
        }
111
    }
1122×
    let (status, body) = match (path.as_str(), authorized) {
1132×
        ("/metrics", true) => ("200 OK", render(name, started)),
1141×
        ("/metrics", false) => ("401 Unauthorized", "a bearer token is needed\n".to_string()),
1150×
        _ => ("404 Not Found", "not found\n".to_string()),
116
    };
1172×
    write!(
1182×
        stream,
119
        "HTTP/1.1 {status}\r\nContent-Type: text/plain; version=0.0.4\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
1202×
        body.len()
121
    )
1222×
}
123
124
#[cfg(test)]
125
mod tests {
126
    use super::*;
127
    use std::io::Read;
128
129
    #[test]
1301×
    fn a_scrape_needs_the_token_and_reads_the_counters() {
1311×
        let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
1321×
        let addr = l.local_addr().unwrap().to_string();
1331×
        drop(l);
1341×
        serve(&addr, Some("t0k".into()), "earth-1".into());
1351×
        CLAIMS.fetch_add(1, Relaxed);
1361×
        finished(Some(0), None, 2500, false);
1372×
        let get = |auth: &str| {
1382×
            let mut s = std::net::TcpStream::connect(&addr).unwrap();
1392×
            write!(s, "GET /metrics HTTP/1.1\r\nHost: x\r\n{auth}\r\n").unwrap();
1402×
            let mut body = String::new();
1412×
            s.read_to_string(&mut body).unwrap();
1422×
            body
1432×
        };
1441×
        assert!(get("").starts_with("HTTP/1.1 401"));
1451×
        let page = get("Authorization: Bearer t0k\r\n");
1461×
        assert!(page.starts_with("HTTP/1.1 200"), "{page}");
1471×
        assert!(page.contains("flow_runner_claims_total{runner=\"earth-1\"} ") && page.contains("result=\"passed\"}"), "{page}");
1481×
    }
149
}