gitdataai/apps/metrics/src/scrape.rs

136 lines
4.0 KiB
Rust

use awc::Client;
use std::collections::HashMap;
use crate::target::ScrapeTarget;
#[derive(Clone)]
pub struct HttpClient {
client: Client,
}
impl HttpClient {
pub fn new(timeout_secs: u64) -> Self {
let client = Client::builder()
.timeout(std::time::Duration::from_secs(timeout_secs))
.finish();
Self { client }
}
pub async fn scrape(&self, target: &ScrapeTarget) -> ScrapeResult {
let start = std::time::Instant::now();
let url = target.url();
let mut resp = match self.client.get(url).send().await {
Ok(resp) => resp,
Err(e) => {
let msg = e.to_string();
if msg.contains("timeout") || msg.contains("TimedOut") || msg.contains("timed out")
{
return ScrapeResult::Timeout;
}
return ScrapeResult::ConnectionError(msg);
}
};
if !resp.status().is_success() {
return ScrapeResult::HttpError(resp.status().as_u16());
}
let body = match resp.body().await {
Ok(bytes) => String::from_utf8_lossy(&bytes).into_owned(),
Err(e) => return ScrapeResult::ConnectionError(e.to_string()),
};
let scrape_ms = start.elapsed().as_millis() as f64;
ScrapeResult::Success(body, scrape_ms)
}
}
pub enum ScrapeResult {
Success(String, f64),
Timeout,
ConnectionError(String),
HttpError(u16),
}
#[derive(Clone, Debug)]
pub struct PromMetric {
pub name: String,
pub value: f64,
pub labels: HashMap<String, String>,
}
pub fn parse_prometheus(body: &str) -> Vec<PromMetric> {
let mut metrics = Vec::new();
for line in body.lines() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
let (name_and_labels, value_str) = match line.find(' ') {
Some(pos) => (&line[..pos], &line[pos + 1..]),
None => continue,
};
let value: f64 = match value_str
.split_whitespace()
.next()
.and_then(|v| v.parse().ok())
{
Some(v) => v,
None => continue,
};
let (metric_name, labels) = if let Some(brace) = name_and_labels.find('{') {
let name = &name_and_labels[..brace];
let label_str = &name_and_labels[brace + 1..name_and_labels.len() - 1];
let labels = parse_labels(label_str);
(name.to_string(), labels)
} else {
(name_and_labels.to_string(), HashMap::new())
};
metrics.push(PromMetric {
name: metric_name,
value,
labels,
});
}
metrics
}
pub fn parse_labels(s: &str) -> HashMap<String, String> {
let mut labels = HashMap::new();
let mut remaining = s;
while !remaining.is_empty() {
if let Some(eq) = remaining.find('=') {
let key = remaining[..eq].trim().to_string();
remaining = &remaining[eq + 1..];
let (value, rest) = if remaining.starts_with('"') {
let end = remaining[1..]
.find('"')
.map(|p| p + 1)
.unwrap_or(remaining.len());
(&remaining[1..end], &remaining[end + 1..])
} else if remaining.starts_with('\'') {
let end = remaining[1..]
.find('\'')
.map(|p| p + 1)
.unwrap_or(remaining.len());
(&remaining[1..end], &remaining[end + 1..])
} else {
let end = remaining
.find(|c: char| !c.is_alphanumeric() && c != '_' && c != '-')
.unwrap_or(remaining.len());
(&remaining[..end], &remaining[end..])
};
labels.insert(key, value.to_string());
remaining = rest.trim_start_matches(',').trim_start();
} else {
break;
}
}
labels
}