mod aggregation;
mod binary;
pub(crate) mod function;
mod selector;
use std::collections::{BTreeMap, HashMap};
use std::sync::Arc;
use regex::Regex;
use crate::promql::ast::{
Duration, Expr, LabelMatcher, MatchOp, PromQLValue, Sample, Series, UnaryOp,
};
use arrow::record_batch::RecordBatch;
use crate::Chronix;
pub(crate) type HistogramBuckets = BTreeMap<Vec<(String, String)>, Vec<(f64, f64)>>;
pub struct CompiledMatcher {
pub(crate) name: String,
pub(crate) op: MatchOp,
pub(crate) value: String,
pub(crate) regex: Option<Regex>,
}
impl CompiledMatcher {
pub fn compile(m: &LabelMatcher) -> Result<Self, EvalError> {
let regex = match m.op {
MatchOp::RegexMatch | MatchOp::RegexNotMatch => {
let pattern = format!("^(?s:{})$", m.value);
Some(
Regex::new(&pattern)
.map_err(|e| EvalError(format!("invalid regex '{}': {e}", m.value)))?,
)
}
_ => None,
};
Ok(Self {
name: m.name.clone(),
op: m.op,
value: m.value.clone(),
regex,
})
}
#[must_use]
pub fn matches(&self, labels: &[(String, String)]) -> bool {
let label_value = labels
.iter()
.find(|(k, _)| k == &self.name)
.map(|(_, v)| v.as_str())
.unwrap_or("");
match self.op {
MatchOp::Equal => label_value == self.value,
MatchOp::NotEqual => label_value != self.value,
MatchOp::RegexMatch => self
.regex
.as_ref()
.is_some_and(|re| re.is_match(label_value)),
MatchOp::RegexNotMatch => self
.regex
.as_ref()
.is_none_or(|re| !re.is_match(label_value)),
}
}
}
pub(crate) fn compile_post_filters(
matchers: &[LabelMatcher],
) -> Result<Vec<CompiledMatcher>, EvalError> {
matchers
.iter()
.filter(|m| m.name != "__name__" && m.op != MatchOp::Equal)
.map(CompiledMatcher::compile)
.collect()
}
pub fn compile_label_matchers(
matchers: &[LabelMatcher],
) -> Result<Vec<CompiledMatcher>, EvalError> {
matchers
.iter()
.filter(|m| m.name != "__name__")
.map(CompiledMatcher::compile)
.collect()
}
fn check_no_duplicate_labels(series: &[Series]) -> Result<(), EvalError> {
let mut seen: std::collections::HashSet<&[(String, String)]> =
std::collections::HashSet::with_capacity(series.len());
for s in series {
if !seen.insert(s.labels.as_slice()) {
return Err(EvalError(
"vector cannot contain metrics with the same labelset".into(),
));
}
}
Ok(())
}
#[must_use]
pub fn label_set_matches(compiled: &[CompiledMatcher], labels: &[(String, String)]) -> bool {
compiled.iter().all(|m| m.matches(labels))
}
#[derive(Clone, PartialEq, Eq, Hash)]
pub(crate) struct ScanKey {
pub(crate) measurement: String,
pub(crate) tags: Vec<(String, String)>,
pub(crate) start: i64,
pub(crate) end: i64,
}
const SCAN_CACHE_CAPACITY: usize = 8;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct ScanStats {
pub scans: u64,
pub cache_hits: u64,
}
pub struct PromQLEvaluator {
pub(crate) db: Arc<Chronix>,
pub(crate) namespace: Option<String>,
pub(crate) regex_cache: std::cell::RefCell<HashMap<String, regex::Regex>>,
pub(crate) scan_cache: std::cell::RefCell<HashMap<ScanKey, Arc<Vec<RecordBatch>>>>,
pub(crate) scan_stats: std::cell::Cell<ScanStats>,
}
const DEFAULT_SUBQUERY_STEP_NS: i64 = 60_000_000_000;
#[derive(Debug, Clone)]
pub struct QueryParams {
pub time: i64,
pub start: Option<i64>,
pub end: Option<i64>,
pub step: Option<i64>,
pub lookback_delta: i64,
pub range_fetch_start: Option<i64>,
pub range_fetch_end: Option<i64>,
pub max_series: usize,
pub max_memory_bytes: usize,
pub max_points: usize,
pub deadline: Option<Deadline>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Deadline {
expires_at: std::time::Instant,
budget: std::time::Duration,
}
impl Deadline {
#[must_use]
pub fn after(budget: std::time::Duration) -> Option<Self> {
(!budget.is_zero()).then(|| Self {
expires_at: std::time::Instant::now() + budget,
budget,
})
}
pub fn check(self) -> Result<(), EvalError> {
if std::time::Instant::now() >= self.expires_at {
return Err(EvalError(format!(
"query exceeded its time budget of {:?}",
self.budget
)));
}
Ok(())
}
}
impl Default for QueryParams {
fn default() -> Self {
Self {
time: 0,
start: None,
end: None,
step: None,
lookback_delta: 300_000_000_000, range_fetch_start: None,
range_fetch_end: None,
max_series: 0,
max_memory_bytes: 0,
max_points: 11_000,
deadline: None,
}
}
}
#[derive(Debug)]
pub struct EvalError(pub String);
impl std::fmt::Display for EvalError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "PromQL evaluation error: {}", self.0)
}
}
impl std::error::Error for EvalError {}
impl PromQLEvaluator {
pub fn new(db: Arc<Chronix>) -> Self {
Self {
db,
namespace: None,
regex_cache: std::cell::RefCell::new(HashMap::new()),
scan_cache: std::cell::RefCell::new(HashMap::new()),
scan_stats: std::cell::Cell::new(ScanStats::default()),
}
}
#[must_use]
pub fn with_namespace(mut self, namespace: Option<String>) -> Self {
self.namespace = namespace;
self
}
#[must_use]
pub fn scan_stats(&self) -> ScanStats {
self.scan_stats.get()
}
pub fn instant_query(
&self,
expr: &Expr,
params: &QueryParams,
) -> Result<PromQLValue, EvalError> {
self.eval(expr, params)
}
pub fn range_query(&self, expr: &Expr, params: &QueryParams) -> Result<PromQLValue, EvalError> {
let start = params
.start
.ok_or_else(|| EvalError("range query requires start".into()))?;
let end = params
.end
.ok_or_else(|| EvalError("range query requires end".into()))?;
let step = params
.step
.ok_or_else(|| EvalError("range query requires step".into()))?;
if step <= 0 {
return Err(EvalError("step must be positive".into()));
}
if start > end {
return Err(EvalError("start must not be after end".into()));
}
let points = (i128::from(end) - i128::from(start)) / i128::from(step) + 1;
if params.max_points > 0 && points > params.max_points as i128 {
return Err(EvalError(format!(
"query would produce {points} evaluation points, exceeding the limit of {}",
params.max_points
)));
}
let mut series_map: BTreeMap<Vec<(String, String)>, Vec<Sample>> = BTreeMap::new();
let max_series = params.max_series;
let max_memory = params.max_memory_bytes;
let mut estimated_bytes: usize = 0;
let mut t = start;
while t <= end {
if let Some(deadline) = params.deadline {
deadline.check()?;
}
let step_params = QueryParams {
time: t,
start: Some(start),
end: Some(end),
step: params.step,
lookback_delta: params.lookback_delta,
range_fetch_start: Some(start - params.lookback_delta),
range_fetch_end: Some(end),
max_series,
max_memory_bytes: max_memory,
max_points: params.max_points,
deadline: params.deadline,
};
let result = self.eval(expr, &step_params)?;
if let PromQLValue::Vector(series_list) = result {
for series in series_list {
if max_series > 0
&& !series_map.contains_key(&series.labels)
&& series_map.len() >= max_series
{
return Err(EvalError(format!(
"query exceeded maximum series limit ({max_series})"
)));
}
for sample in &series.samples {
series_map
.entry(series.labels.clone())
.or_default()
.push(Sample {
timestamp: t,
value: sample.value,
});
estimated_bytes += std::mem::size_of::<Sample>();
}
if max_memory > 0 && estimated_bytes > max_memory {
return Err(EvalError(format!(
"range query exceeded memory limit ({estimated_bytes} bytes > {max_memory} bytes)"
)));
}
}
}
let Some(next) = t.checked_add(step) else {
break;
};
t = next;
}
let matrix: Vec<Series> = series_map
.into_iter()
.map(|(labels, samples)| Series::floats(labels, samples))
.collect();
Ok(PromQLValue::Matrix(matrix))
}
pub(crate) fn eval(&self, expr: &Expr, params: &QueryParams) -> Result<PromQLValue, EvalError> {
let value = self.eval_node(expr, params)?;
if let PromQLValue::Vector(series) = &value {
check_no_duplicate_labels(series)?;
}
Ok(value)
}
fn eval_node(&self, expr: &Expr, params: &QueryParams) -> Result<PromQLValue, EvalError> {
match expr {
Expr::NumberLiteral(n) => Ok(PromQLValue::Scalar(*n)),
Expr::StringLiteral(s) => Ok(PromQLValue::String(s.clone())),
Expr::Paren(inner) => self.eval(inner, params),
Expr::UnaryExpr { op, expr } => {
let val = self.eval(expr, params)?;
match op {
UnaryOp::Neg => match val {
PromQLValue::Scalar(n) => Ok(PromQLValue::Scalar(-n)),
PromQLValue::Vector(series) => {
let negated = series
.into_iter()
.map(|mut s| {
s.labels.retain(|(k, _)| k != "__name__");
for sample in &mut s.samples {
sample.value = -sample.value;
}
s
})
.collect();
Ok(PromQLValue::Vector(negated))
}
_ => Err(EvalError("cannot negate matrix or string".into())),
},
}
}
Expr::VectorSelector {
name,
matchers,
offset,
at,
} => self.eval_vector_selector(name, matchers, offset, *at, params),
Expr::MatrixSelector { vector, range } => {
self.eval_matrix_selector(vector, range, params)
}
Expr::Call { func, args } => self.eval_function(func, args, params),
Expr::Aggregation {
op,
expr,
param,
modifier,
} => self.eval_aggregation(*op, expr, param.as_deref(), modifier, params),
Expr::BinaryExpr {
op,
lhs,
rhs,
bool_mod,
matching,
} => self.eval_binary(*op, lhs, rhs, *bool_mod, matching.as_ref(), params),
Expr::Subquery {
expr: inner,
range,
step,
offset,
at,
} => self.eval_subquery(inner, *range, *step, *offset, *at, params),
}
}
fn eval_subquery(
&self,
inner: &Expr,
range: Duration,
step: Option<Duration>,
offset: Option<Duration>,
at: Option<crate::promql::ast::AtModifier>,
params: &QueryParams,
) -> Result<PromQLValue, EvalError> {
let eval_time = match at {
None => params.time,
Some(at) => at.resolve(
params.start.unwrap_or(params.time),
params.end.unwrap_or(params.time),
),
};
let offset_ns = offset.map_or(0, |o| o.as_nanos());
let range_ns = range.as_nanos();
let step_ns = if let Some(s) = step {
let s_ns = s.as_nanos();
if s_ns <= 0 {
return Err(EvalError("subquery step must be positive".into()));
}
s_ns
} else {
DEFAULT_SUBQUERY_STEP_NS
};
let window_end = eval_time - offset_ns;
let window_start = window_end - range_ns;
let (prefetch_start, prefetch_end) =
match (params.range_fetch_start, params.range_fetch_end) {
(Some(rs), Some(re)) => (
rs.saturating_sub(range_ns).saturating_sub(offset_ns),
re.saturating_sub(offset_ns),
),
_ => (
window_start.saturating_sub(params.lookback_delta),
window_end,
),
};
let mut series_map: BTreeMap<Vec<(String, String)>, Vec<Sample>> = BTreeMap::new();
let mut t = window_start
.div_euclid(step_ns)
.checked_mul(step_ns)
.ok_or_else(|| EvalError("subquery window overflows the step grid".into()))?;
if t <= window_start {
t = t
.checked_add(step_ns)
.ok_or_else(|| EvalError("subquery window overflows the step grid".into()))?;
}
let points = (i128::from(window_end) - i128::from(t)) / i128::from(step_ns) + 1;
if params.max_points > 0 && points > params.max_points as i128 {
return Err(EvalError(format!(
"subquery would produce {points} evaluation points, exceeding the limit of {}",
params.max_points
)));
}
while t <= window_end {
if let Some(deadline) = params.deadline {
deadline.check()?;
}
let step_params = QueryParams {
time: t,
step: Some(step_ns),
lookback_delta: params.lookback_delta,
range_fetch_start: Some(prefetch_start),
range_fetch_end: Some(prefetch_end),
start: params.start,
end: params.end,
max_series: params.max_series,
max_memory_bytes: params.max_memory_bytes,
max_points: params.max_points,
deadline: params.deadline,
};
let result = self.eval(inner, &step_params)?;
if let PromQLValue::Vector(series_list) = result {
for series in series_list {
for sample in &series.samples {
series_map
.entry(series.labels.clone())
.or_default()
.push(Sample {
timestamp: t,
value: sample.value,
});
}
}
}
let Some(next) = t.checked_add(step_ns) else {
break;
};
t = next;
}
let matrix: Vec<Series> = series_map
.into_iter()
.map(|(labels, samples)| Series::floats(labels, samples))
.collect();
Ok(PromQLValue::Matrix(matrix))
}
}
#[cfg(test)]
mod deadline_tests {
use super::{Deadline, PromQLEvaluator, QueryParams};
use std::sync::Arc;
use std::time::Duration;
#[test]
fn a_zero_budget_is_no_deadline() {
assert!(Deadline::after(Duration::ZERO).is_none());
assert!(Deadline::after(Duration::from_secs(1)).is_some());
}
#[test]
fn a_spent_budget_says_what_it_was() {
let deadline = Deadline::after(Duration::from_nanos(1)).expect("non-zero");
std::thread::sleep(Duration::from_millis(2));
let err = deadline.check().expect_err("the budget is spent");
assert!(err.0.contains("time budget"), "{err}");
}
#[test]
fn an_expired_deadline_stops_a_range_query_before_its_first_step() {
let tmp = tempfile::tempdir().expect("tempdir");
let config = chronix_core::ChronixConfig::builder()
.data_dir(tmp.path().to_path_buf())
.build()
.expect("config");
let db = Arc::new(crate::Chronix::open(config).expect("open"));
let now = 1_700_000_000_000_000_000i64;
let points: Vec<chronix_core::Point> = (0..100)
.map(|i| {
chronix_core::Point::new(
chronix_core::SeriesKey::new("cpu", crate::tags! { "host" => "a" })
.expect("key"),
crate::fields! { "value" => f64::from(i) },
now + i64::from(i) * 1_000_000_000,
)
.expect("point")
})
.collect();
let _ = db.insert_batch(&points).expect("insert");
let expr = crate::promql::parse("cpu").expect("parse");
let evaluator = PromQLEvaluator::new(db);
let deadline = Deadline::after(Duration::from_nanos(1));
std::thread::sleep(Duration::from_millis(2));
let params = QueryParams {
start: Some(now),
end: Some(now + 100_000_000_000),
step: Some(1_000_000_000),
deadline,
..Default::default()
};
let err = evaluator
.range_query(&expr, ¶ms)
.expect_err("an expired deadline must stop the evaluation");
assert!(err.0.contains("time budget"), "{err}");
let ok = QueryParams {
deadline: None,
..params
};
evaluator
.range_query(&expr, &ok)
.expect("the same query without a budget answers");
}
}