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

pub fn finished(exit_code: Option<i32>, error: Option<&str>, duration_ms: u64, cancelled: bool) {

    let counter = if cancelled {

    } else if error.is_some() {

    } else if exit_code == Some(0) {

        &PASSED

    counter.fetch_add(1, Relaxed);

    JOB_MS.fetch_add(duration_ms, Relaxed);

}

pub fn render(name: &str, started: Instant) -> String {

    let n = name.replace('\\', "\\\\").replace('"', "\\\"");

    let mut out = String::new();

    let mut family = |metric: &str, help: &str, kind: &str, samples: &[(String, u64)]| {

        out.push_str(&format!("# HELP {metric} {help}\n# TYPE {metric} {kind}\n"));

        for (labels, v) in samples {

            let sep = if labels.is_empty() { "" } else { "," };

            out.push_str(&format!("{metric}{{runner=\"{n}\"{sep}{labels}}} {v}\n"));

    };

    family("flow_runner_build_info", "This runner's version.", "gauge", &[(format!("version=\"{}\"", env!("CARGO_PKG_VERSION")), 1)]);

    family("flow_runner_uptime_seconds", "Seconds since the runner started.", "gauge", &[(String::new(), started.elapsed().as_secs())]);

    family("flow_runner_claims_total", "Jobs claimed from the hub.", "counter", &[(String::new(), CLAIMS.load(Relaxed))]);

    family("flow_runner_declined_total", "Claimed jobs declined before running.", "counter", &[(String::new(), DECLINED.load(Relaxed))]);

    family(

        "flow_runner_jobs_finished_total",

        "Finished jobs by result.",

        "counter",

        &[

            ("result=\"passed\"".into(), PASSED.load(Relaxed)),

            ("result=\"failed\"".into(), FAILED.load(Relaxed)),

            ("result=\"error\"".into(), ERRORS.load(Relaxed)),

            ("result=\"cancelled\"".into(), CANCELLED.load(Relaxed)),

        ],

    );

    family(

        "flow_runner_job_seconds_total",

        "Seconds finished jobs ran, together.",

        "counter",

        &[(String::new(), JOB_MS.load(Relaxed) / 1000)],

    );

    family("flow_runner_jobs_running", "Jobs running now.", "gauge", &[(String::new(), RUNNING.load(Relaxed))]);

    out

}

pub fn serve(addr: &str, token: Option<String>, name: String) {

    let listener = match std::net::TcpListener::bind(addr) {

        Ok(l) => l,

    eprintln!("flow runner {name}: metrics on http://{addr}/metrics");

    let started = Instant::now();

    std::thread::spawn(move || {

        for stream in listener.incoming().flatten() {

            let _ = answer(stream, token.as_deref(), &name, started);

        }

    });

}

fn answer(mut stream: std::net::TcpStream, token: Option<&str>, name: &str, started: Instant) -> std::io::Result<()> {

    stream.set_read_timeout(Some(std::time::Duration::from_secs(5)))?;

    let mut reader = BufReader::new(stream.try_clone()?);

    let mut line = String::new();

    reader.read_line(&mut line)?;

    let path = line.split_whitespace().nth(1).unwrap_or("").to_string();

    let mut authorized = token.is_none();

        let mut h = String::new();

        if reader.read_line(&mut h)? == 0 || h.trim().is_empty() {

            break;

        }

        if let (Some(t), Some((k, v))) = (token, h.split_once(':')) {

            if k.trim().eq_ignore_ascii_case("authorization") && v.trim() == format!("Bearer {t}") {

                authorized = true;

            }

    let (status, body) = match (path.as_str(), authorized) {

        ("/metrics", true) => ("200 OK", render(name, started)),

        ("/metrics", false) => ("401 Unauthorized", "a bearer token is needed\n".to_string()),

    write!(

        stream,

        body.len()

}

    fn a_scrape_needs_the_token_and_reads_the_counters() {

        let l = std::net::TcpListener::bind("127.0.0.1:0").unwrap();

        let addr = l.local_addr().unwrap().to_string();

        drop(l);

        serve(&addr, Some("t0k".into()), "earth-1".into());

        CLAIMS.fetch_add(1, Relaxed);

        finished(Some(0), None, 2500, false);

        let get = |auth: &str| {

            let mut s = std::net::TcpStream::connect(&addr).unwrap();

            write!(s, "GET /metrics HTTP/1.1\r\nHost: x\r\n{auth}\r\n").unwrap();

            let mut body = String::new();

            s.read_to_string(&mut body).unwrap();

            body

        };

        assert!(get("").starts_with("HTTP/1.1 401"));

        let page = get("Authorization: Bearer t0k\r\n");

        assert!(page.starts_with("HTTP/1.1 200"), "{page}");

        assert!(page.contains("flow_runner_claims_total{runner=\"earth-1\"} ") && page.contains("result=\"passed\"}"), "{page}");

    }