iota_metric_checker/
query.rs

1// Copyright (c) Mysten Labs, Inc.
2// Modifications Copyright (c) 2024 IOTA Stiftung
3// SPDX-License-Identifier: Apache-2.0
4
5use anyhow::anyhow;
6use base64::{Engine, engine::general_purpose};
7use prometheus_http_query::Client;
8use reqwest::header::{AUTHORIZATION, HeaderValue};
9use tracing::{debug, info};
10
11use crate::unix_seconds_to_timestamp_string;
12
13pub async fn instant_query(
14    auth_header: &str,
15    client: Client,
16    query: &str,
17) -> Result<f64, anyhow::Error> {
18    debug!("Executing {query}");
19    let response = client
20        .query(query)
21        .header(
22            AUTHORIZATION,
23            HeaderValue::from_str(&format!(
24                "Basic {}",
25                general_purpose::STANDARD.encode(auth_header)
26            ))?,
27        )
28        .get()
29        .await?;
30
31    let result = response
32        .data()
33        .as_vector()
34        .unwrap_or_else(|| panic!("Expected result of type vector for {query}"));
35
36    if !result.is_empty() {
37        let first = result.first().unwrap();
38        debug!("Got value {}", first.sample().value());
39        Ok(first.sample().value())
40    } else {
41        Err(anyhow!(
42            "Did not get expected response from server for {query}"
43        ))
44    }
45}
46
47// This will return the median value of the queried metric over the given time
48// range.
49pub async fn range_query(
50    auth_header: &str,
51    client: Client,
52    query: &str,
53    start: i64,
54    end: i64,
55    step: f64,
56    percentile: u8,
57) -> Result<f64, anyhow::Error> {
58    debug!("Executing {query}");
59    let response = client
60        .query_range(query, start, end, step)
61        .header(
62            AUTHORIZATION,
63            HeaderValue::from_str(&format!(
64                "Basic {}",
65                general_purpose::STANDARD.encode(auth_header)
66            ))?,
67        )
68        .get()
69        .await?;
70
71    let result = response
72        .data()
73        .as_matrix()
74        .unwrap_or_else(|| panic!("Expected result of type matrix for {query}"));
75
76    if result.is_empty() {
77        return Err(anyhow!(
78            "Did not get expected response from server for {query}"
79        ));
80    }
81
82    let mut samples: Vec<f64> = result
83        .first()
84        .unwrap()
85        .samples()
86        .iter()
87        .filter_map(|sample| {
88            let v = sample.value();
89            if v.is_nan() { None } else { Some(v) }
90        })
91        .collect();
92    if samples.is_empty() {
93        return Err(anyhow!("Query returned zero data point! {query}"));
94    }
95    samples.sort_by(|a, b| a.partial_cmp(b).unwrap());
96
97    assert!(
98        (1..=100).contains(&percentile),
99        "Invalid percentile {percentile}"
100    );
101    let index = samples.len() * percentile as usize / 100;
102    let result = samples[index];
103    info!(
104        "{query}: got p{percentile} value {result}, over {} data points in time range {} - {}",
105        samples.len(),
106        unix_seconds_to_timestamp_string(start),
107        unix_seconds_to_timestamp_string(end)
108    );
109    Ok(result)
110}
iota_metric_checker/query.rs

iota_metric_checker/
query.rs
