sirun 0.1.9

A benchmark test runner
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

use crate::metric_value::*;
use anyhow::*;
use async_std::{
    net::UdpSocket,
    sync::{Arc, Barrier, RwLock},
};
use std::{collections::HashMap, env};

pub(crate) async fn statsd_listener(
    barrier: Arc<Barrier>,
    statsd_buf: Arc<RwLock<String>>,
) -> Result<()> {
    // If the env var is set, we'll use it, otherwise use 0 to grab an available port.
    let port: u16 = env::var("SIRUN_STATSD_PORT").map_or(0, |p| p.parse().unwrap_or(0));
    let socket = UdpSocket::bind(format!("127.0.0.1:{}", port)).await;
    let socket = match socket {
        Ok(s) => s,
        Err(error) => panic!("Cannot bind to 127.0.0.1:{}: {}", port, error),
    };
    let port = socket.local_addr()?.port();
    env::set_var("SIRUN_STATSD_PORT", format!("{}", port));
    barrier.wait().await; // indicates to main task that socket is listening

    loop {
        let mut buf = vec![0u8; 4096];
        let (recv, _peer) = socket.recv_from(&mut buf).await?;

        let datum = String::from_utf8(buf[..recv].into()).unwrap_or_else(|_| String::new());
        statsd_buf.write().await.push_str(&datum);
    }
}

pub(crate) async fn get_statsd_metrics(
    udp_data: Arc<RwLock<String>>,
) -> Result<HashMap<String, MetricValue>> {
    let mut metrics = HashMap::new();
    let udp_string = udp_data.read().await.clone();
    let lines = udp_string.trim().lines();
    udp_data.write().await.clear();
    for line in lines {
        let metric: Vec<&str> = match line.split('|').next() {
            None => continue,
            Some(metric) => metric.split(':').collect(),
        };
        if metric.len() < 2 {
            continue;
        }
        metrics.insert(metric[0].into(), metric[1].parse::<f64>()?.into());
    }
    Ok(metrics)
}