mod types;
use std::collections::BTreeMap;
use chrono::Utc;
use reqwest::header::{
CONTENT_ENCODING, CONTENT_TYPE, HeaderMap, HeaderName, HeaderValue,
USER_AGENT,
};
use crate::types::{Request, Sample, Timeseries};
#[derive(Clone)]
pub struct Mertic {
http: reqwest::Client,
url: String,
headers: HeaderMap,
}
impl Mertic {
pub fn new(url: String) -> Self {
let cli = reqwest::Client::new();
let mut headers = HeaderMap::new();
headers.insert(
CONTENT_TYPE,
HeaderValue::from_static(
"application/x-protobuf;proto=io.prometheus.write.v2.Request",
),
);
headers.insert(CONTENT_ENCODING, HeaderValue::from_static("snappy"));
headers.insert(
HeaderName::from_static("x-prometheus-remote-write-version"),
HeaderValue::from_static("2.0.0"),
);
headers.insert(USER_AGENT, HeaderValue::from_static("promwrite/0.0.1"));
Self {
http: cli,
url: url,
headers: headers,
}
}
pub fn name(&self, name: String, unit: String, help: String) -> Label {
let mut labels: BTreeMap<String, String> = BTreeMap::new();
let mut n = name;
if !unit.is_empty() {
n = n + "_" + &unit;
}
Label {
metric: self.clone(),
name: n,
uint: unit,
help: help,
lables: labels,
samples: vec![],
}
}
}
#[derive(Clone)]
pub struct Label {
metric: Mertic,
pub name: String,
pub uint: String,
pub help: String,
pub lables: BTreeMap<String, String>,
pub samples: Vec<Sample>,
}
impl Label {
pub fn label(&mut self, name: String, value: String) -> &mut Self {
self.lables.insert(name, value);
self
}
pub fn value(&mut self, v: f64) -> &mut Self {
self.samples.push(Sample {
value: v,
timestamp: Utc::now().timestamp_millis(),
});
self
}
pub async fn send(&mut self) {
let mut req = Request {
symbols: vec![],
timeseries: vec![],
};
self.lables
.insert("__name__".to_string(), self.name.to_string());
let mut labelref: Vec<u32> = vec![];
for (pos, key) in self.lables.keys().enumerate() {
req.symbols.push(key.to_string());
req.symbols
.push(self.lables.get(&key.to_string()).unwrap().to_string());
labelref.push(pos as u32 * 2);
labelref.push(pos as u32 * 2 + 1);
}
req.timeseries.push(Timeseries {
labels_refs: labelref,
samples: self.samples.clone(),
});
if let Ok(r) = snap::raw::Encoder::new()
.compress_vec(&prost::Message::encode_to_vec(&req))
{
self.metric
.http
.post(&self.metric.url)
.headers(self.metric.headers.clone())
.body(r)
.send()
.await
.unwrap();
}
self.lables.clear();
self.samples.clear();
}
}
#[cfg(test)]
mod tests {
use crate::Mertic;
#[tokio::test]
async fn test_request() {
let m = Mertic::new("http://127.0.0.1:9090/api/v1/write".to_string());
let mut up =
m.name("rust".to_string(), "up".to_string(), "".to_string());
up.label("hostname".to_string(), "host1".to_string())
.value(20.0)
.send()
.await;
up.value(10.0).send().await;
let mut test =
m.name("test".to_string(), "up".to_string(), "".to_string());
test.label("hostname".to_string(), "host1".to_string())
.label("ip".to_string(), "192.168.1.1".to_string())
.value(30.0)
.send()
.await;
}
}