use anyhow::{Result, anyhow};
use reqwest::Client;
use serde::Deserialize;
use serde::de::DeserializeOwned;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::Duration;
type QueryCache = Arc<Mutex<HashMap<String, (i64, i64, Duration, Vec<Series>)>>>;
type QueryWaiter = tokio::sync::oneshot::Sender<Result<Vec<Series>, String>>;
type InflightQueries = Arc<Mutex<HashMap<String, Vec<QueryWaiter>>>>;
#[derive(Debug, Clone)]
pub(crate) struct PromClient {
pub(crate) base: String,
client: reqwest::Client,
cache: QueryCache,
inflight: InflightQueries,
}
impl PromClient {
pub(crate) fn new(base: String) -> Self {
let http = Client::builder()
.timeout(Duration::from_secs(10))
.connect_timeout(Duration::from_secs(5))
.build()
.unwrap_or_else(|e| {
eprintln!(
"Warning: Failed to configure HTTP client with timeouts: {}",
e
);
eprintln!(" Falling back to default client (requests may hang).");
Client::new()
});
Self {
base,
client: http,
cache: Arc::new(Mutex::new(HashMap::new())),
inflight: Arc::new(Mutex::new(HashMap::new())),
}
}
pub(crate) fn build_query_range_url(
&self,
expr: &str,
start: i64,
end: i64,
step: Duration,
) -> String {
let step_s = step.as_secs().max(1);
let step_param = format!("{}s", step_s);
format!(
"{}/api/v1/query_range?query={}&start={}&end={}&step={}",
self.base.trim_end_matches('/'),
urlencoding::encode(expr),
start,
end,
step_param
)
}
pub(crate) fn build_query_url(&self, expr: &str, time: i64) -> String {
format!(
"{}/api/v1/query?query={}&time={}",
self.base.trim_end_matches('/'),
urlencoding::encode(expr),
time
)
}
pub(crate) async fn query_range(
&self,
expr: &str,
start: i64,
end: i64,
step: Duration,
) -> Result<Vec<Series>> {
{
let cache = self.cache.lock().unwrap();
if let Some((c_start, c_end, c_step, data)) = cache.get(expr) {
if *c_start == start && *c_end == end && *c_step == step {
return Ok(data.clone());
}
}
}
let inflight_key = format!("{}|{}|{}|{}", expr, start, end, step.as_secs());
let rx = {
let mut inflight = self.inflight.lock().unwrap();
if let Some(waiters) = inflight.get_mut(&inflight_key) {
let (tx, rx) = tokio::sync::oneshot::channel();
waiters.push(tx);
Some(rx)
} else {
inflight.insert(inflight_key.clone(), Vec::new());
None
}
};
if let Some(rx) = rx {
return match rx.await {
Ok(Ok(res)) => Ok(res),
Ok(Err(s)) => Err(anyhow!(s)),
Err(_) => Err(anyhow!("inflight request cancelled")),
};
}
let url = self.build_query_range_url(expr, start, end, step);
let max_retries = 3;
let mut last_err = anyhow!("unknown error");
let mut final_res = Err(anyhow!("unknown error"));
for attempt in 0..=max_retries {
if attempt > 0 {
tokio::time::sleep(Duration::from_millis(100 * (1 << attempt))).await;
}
match self.perform_request(&url).await {
Ok(series) => {
{
let mut cache = self.cache.lock().unwrap();
cache.insert(expr.to_string(), (start, end, step, series.clone()));
}
final_res = Ok(series);
break;
}
Err(e) => last_err = e,
}
}
if final_res.is_err() {
final_res = Err(last_err);
}
{
let mut inflight = self.inflight.lock().unwrap();
if let Some(waiters) = inflight.remove(&inflight_key) {
for tx in waiters {
let _ = tx.send(match &final_res {
Ok(v) => Ok(v.clone()),
Err(e) => Err(e.to_string()),
});
}
}
}
final_res
}
async fn perform_request(&self, url: &str) -> Result<Vec<Series>> {
let text = self.get_text(url).await?;
let body: PromResponse<QueryRangeData> = serde_json::from_str(&text)
.map_err(|e| anyhow!("parsing json: {} (body: {})", e, text))?;
if body.status != "success" {
return Err(anyhow!(
"prometheus error status: {} — body: {}",
body.status,
text
));
}
Ok(body.data.result)
}
pub(crate) async fn label_values(&self, label: &str) -> Result<Vec<String>> {
let url = format!(
"{}/api/v1/label/{}/values",
self.base.trim_end_matches('/'),
urlencoding::encode(label)
);
let body: PromResponse<Vec<String>> = self.get_json(&url).await?;
ensure_success(&body.status)?;
Ok(body.data)
}
pub(crate) async fn series_label_values(
&self,
metric: &str,
label: &str,
start: i64,
end: i64,
) -> Result<Vec<String>> {
let url = format!(
"{}/api/v1/series?match[]={}&start={}&end={}",
self.base.trim_end_matches('/'),
urlencoding::encode(metric),
start,
end
);
let body: PromResponse<Vec<HashMap<String, String>>> = self.get_json(&url).await?;
ensure_success(&body.status)?;
Ok(body
.data
.into_iter()
.filter_map(|series| series.get(label).cloned())
.collect())
}
pub(crate) async fn query_instant_result_strings(
&self,
expr: &str,
time: i64,
) -> Result<Vec<String>> {
let url = self.build_query_url(expr, time);
let body: PromResponse<QueryInstantData> = self.get_json(&url).await?;
ensure_success(&body.status)?;
Ok(body.data.result_strings())
}
pub(crate) async fn query_instant_series(&self, expr: &str, time: i64) -> Result<Vec<Series>> {
let url = self.build_query_url(expr, time);
let body: PromResponse<QueryInstantData> = self.get_json(&url).await?;
ensure_success(&body.status)?;
Ok(body.data.into_series(time))
}
async fn get_json<T: DeserializeOwned>(&self, url: &str) -> Result<T> {
let text = self.get_text(url).await?;
serde_json::from_str(&text).map_err(|e| anyhow!("parsing json: {} (body: {})", e, text))
}
async fn get_text(&self, url: &str) -> Result<String> {
let resp = self
.client
.get(url)
.send()
.await
.map_err(|e| anyhow!("request failed: {}", e))?;
let status = resp.status();
let text = resp
.text()
.await
.map_err(|e| anyhow!("reading text: {}", e))?;
if !status.is_success() {
return Err(anyhow!("prometheus {}: {}", status, text));
}
Ok(text)
}
}
fn ensure_success(status: &str) -> Result<()> {
if status == "success" {
Ok(())
} else {
Err(anyhow!("prometheus error status: {}", status))
}
}
#[derive(Debug, Deserialize, Clone)]
struct PromResponse<T> {
status: String,
data: T,
}
#[derive(Debug, Deserialize, Clone)]
pub(crate) struct QueryRangeData {
#[serde(rename = "resultType")]
#[allow(dead_code)]
pub(crate) result_type: String,
pub(crate) result: Vec<Series>,
}
#[derive(Debug, Deserialize, Clone)]
pub(crate) struct Series {
pub(crate) metric: std::collections::HashMap<String, String>,
pub(crate) values: Vec<(f64, String)>, }
#[derive(Debug, Deserialize, Clone)]
struct QueryInstantData {
#[serde(rename = "resultType")]
result_type: String,
result: serde_json::Value,
}
impl QueryInstantData {
fn result_strings(self) -> Vec<String> {
match self.result_type.as_str() {
"vector" => vector_result_strings(&self.result),
"scalar" | "string" => scalar_result_string(&self.result).into_iter().collect(),
_ => Vec::new(),
}
}
fn into_series(self, time: i64) -> Vec<Series> {
match self.result_type.as_str() {
"vector" => vector_result_series(&self.result, time),
"scalar" => scalar_result_series(&self.result, time)
.into_iter()
.collect(),
_ => Vec::new(),
}
}
}
fn vector_result_series(result: &serde_json::Value, time: i64) -> Vec<Series> {
result
.as_array()
.into_iter()
.flatten()
.filter_map(|sample| {
let metric = sample.get("metric")?.as_object()?;
let value = sample
.get("value")
.and_then(|value| value.as_array())
.and_then(|value| value.get(1))
.and_then(|value| value.as_str())?;
Some(Series {
metric: metric
.iter()
.filter_map(|(label, value)| {
value
.as_str()
.map(|value| (label.clone(), value.to_string()))
})
.collect(),
values: vec![(time as f64, value.to_string())],
})
})
.collect()
}
fn scalar_result_series(result: &serde_json::Value, time: i64) -> Option<Series> {
let value = result
.as_array()
.and_then(|value| value.get(1))
.and_then(|value| value.as_str())?;
Some(Series {
metric: HashMap::new(),
values: vec![(time as f64, value.to_string())],
})
}
fn vector_result_strings(result: &serde_json::Value) -> Vec<String> {
result
.as_array()
.into_iter()
.flatten()
.filter_map(|sample| {
let metric = sample.get("metric")?.as_object()?;
let value = sample
.get("value")
.and_then(|value| value.as_array())
.and_then(|value| value.get(1))
.and_then(|value| value.as_str())
.unwrap_or_default();
let mut labels: Vec<_> = metric
.iter()
.filter_map(|(label, value)| value.as_str().map(|value| (label, value)))
.collect();
labels.sort_by(|a, b| a.0.cmp(b.0));
let labels = labels
.into_iter()
.map(|(label, value)| format!("{}=\"{}\"", label, value))
.collect::<Vec<_>>()
.join(", ");
if labels.is_empty() {
Some(value.to_string())
} else {
Some(format!("{{{}}} {}", labels, value))
}
})
.collect()
}
fn scalar_result_string(result: &serde_json::Value) -> Option<String> {
result
.as_array()
.and_then(|value| value.get(1))
.and_then(|value| value.as_str())
.map(ToString::to_string)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_build_query_range_url() {
let client = PromClient::new("http://localhost:9090".to_string());
let expr = "up{job=\"node\"}";
let start = 1600000000;
let end = 1600003600;
let step = Duration::from_secs(60);
let url = client.build_query_range_url(expr, start, end, step);
assert_eq!(
url,
"http://localhost:9090/api/v1/query_range?query=up%7Bjob%3D%22node%22%7D&start=1600000000&end=1600003600&step=60s"
);
}
#[test]
fn test_build_query_url_preserves_path_prefix() {
let client = PromClient::new("http://localhost:9090/prometheus/".to_string());
let url = client.build_query_url("up{job=\"node\"}", 1600003600);
assert_eq!(
url,
"http://localhost:9090/prometheus/api/v1/query?query=up%7Bjob%3D%22node%22%7D&time=1600003600"
);
}
#[test]
fn test_deserialize_query_range_response() {
let json = r#"
{
"status": "success",
"data": {
"resultType": "matrix",
"result": [
{
"metric": {
"__name__": "up",
"job": "prometheus"
},
"values": [
[1435781451.781, "1"],
[1435781466.781, "1"]
]
}
]
}
}
"#;
let resp: PromResponse<QueryRangeData> = serde_json::from_str(json).unwrap();
assert_eq!(resp.status, "success");
assert_eq!(resp.data.result_type, "matrix");
assert_eq!(resp.data.result.len(), 1);
assert_eq!(resp.data.result[0].metric.get("job").unwrap(), "prometheus");
assert_eq!(resp.data.result[0].values.len(), 2);
}
#[test]
fn test_query_instant_vector_result_strings() {
let json = r#"
{
"resultType": "vector",
"result": [
{
"metric": { "instance": "node-1", "job": "node" },
"value": [1435781451.781, "1"]
}
]
}
"#;
let data: QueryInstantData = serde_json::from_str(json).unwrap();
assert_eq!(
data.result_strings(),
vec![r#"{instance="node-1", job="node"} 1"#]
);
}
#[test]
fn test_query_instant_vector_converts_to_series() {
let json = r#"
{
"resultType": "vector",
"result": [
{
"metric": { "instance": "node-1", "job": "node" },
"value": [1435781451.781, "1"]
},
{
"metric": { "instance": "node-2", "job": "node" },
"value": [1435781451.781, "2.5"]
}
]
}
"#;
let data: QueryInstantData = serde_json::from_str(json).unwrap();
let series = data.into_series(1_435_781_451);
assert_eq!(series.len(), 2);
assert_eq!(series[0].metric.get("instance").unwrap(), "node-1");
assert_eq!(series[0].values, vec![(1_435_781_451.0, "1".to_string())]);
assert_eq!(series[1].metric.get("instance").unwrap(), "node-2");
assert_eq!(series[1].values, vec![(1_435_781_451.0, "2.5".to_string())]);
}
#[test]
fn test_query_instant_scalar_converts_to_unlabeled_series() {
let json = r#"
{
"resultType": "scalar",
"result": [1435781451.781, "42"]
}
"#;
let data: QueryInstantData = serde_json::from_str(json).unwrap();
let series = data.into_series(1_435_781_451);
assert_eq!(series.len(), 1);
assert!(series[0].metric.is_empty());
assert_eq!(series[0].values, vec![(1_435_781_451.0, "42".to_string())]);
}
}