From 163c1e51235817e88816916407c17539914dc2ab Mon Sep 17 00:00:00 2001 From: mctursh Date: Thu, 13 Aug 2026 12:53:54 +0100 Subject: [PATCH 1/2] feat(query-tracker): validate weighted-mode config at load --- crates/core/src/config.rs | 92 ++++++++++++++++++++++++++++++++++++--- 1 file changed, 87 insertions(+), 5 deletions(-) diff --git a/crates/core/src/config.rs b/crates/core/src/config.rs index c6d865e..af3de6a 100644 --- a/crates/core/src/config.rs +++ b/crates/core/src/config.rs @@ -734,17 +734,29 @@ pub enum PriorityMode { /// cost-per-hit (and needs no rate roll). Weighted { /// Weight on demand (request count in the window). Default `0.0`. - #[serde(rename = "demand-weight", default)] + #[serde( + rename = "demand-weight", + default, + deserialize_with = "deserialize_nonneg_finite_f64" + )] demand_weight: f64, /// Weight on supply (`idx_scan` in the window, halved via /// `SCANS_PER_REQUEST` so it is comparable to demand). Only contributes /// for created indexes — candidates have no supply yet. Default `0.0`. - #[serde(rename = "supply-weight", default)] + #[serde( + rename = "supply-weight", + default, + deserialize_with = "deserialize_nonneg_finite_f64" + )] supply_weight: f64, /// Weight on failed/timed-out requests in the window, to prioritize /// patterns that currently *cannot* be served without an index. /// Default `0.0`. - #[serde(rename = "failure-weight", default)] + #[serde( + rename = "failure-weight", + default, + deserialize_with = "deserialize_nonneg_finite_f64" + )] failure_weight: f64, /// Weight on the measured latency **gain** from the index — the ratio of /// the (compensated) without-index average cost to the with-index @@ -756,7 +768,11 @@ pub enum PriorityMode { /// regression guard applies it too). Stays neutral (`gain = 1`) until the /// pattern has served requests both with and without the index (so a /// ratio can be formed), or when the weight is `0`. Default `0.0`. - #[serde(rename = "latency-weight", default)] + #[serde( + rename = "latency-weight", + default, + deserialize_with = "deserialize_nonneg_finite_f64" + )] latency_weight: f64, /// Window over which the demand/supply/failure counts are measured. A /// background task snapshots the counters at this cadence and stores the @@ -1051,7 +1067,8 @@ pub struct QueryTrackerConfig { /// a no-op — the raw wall-clock averages are compared as-is. #[serde( rename = "without-index-compensation-factor", - default = "QueryTrackerConfig::default_without_index_compensation_factor" + default = "QueryTrackerConfig::default_without_index_compensation_factor", + deserialize_with = "deserialize_pos_finite_f64" )] pub without_index_compensation_factor: f64, @@ -1476,6 +1493,39 @@ where }) } +/// Validate a `weighted` priority weight: it must be finite and non-negative. +/// The weights are Display-formatted into the score SQL, so a non-finite value +/// (`inf`/`nan`) becomes a bare token Postgres reads as an unknown column and +/// fails every create/evict pass; a negative weight silently inverts the ranking. +fn deserialize_nonneg_finite_f64<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + let value = f64::deserialize(deserializer)?; + if !value.is_finite() || value < 0.0 { + return Err(serde::de::Error::custom(format!( + "a `weighted` priority weight must be finite and non-negative, got {value}" + ))); + } + Ok(value) +} + +/// Validate `without-index-compensation-factor`: finite and strictly positive. +/// It multiplies the without-index cost, so `0` or negative breaks the latency +/// gain comparison and the regression guard. +fn deserialize_pos_finite_f64<'de, D>(deserializer: D) -> Result +where + D: Deserializer<'de>, +{ + let value = f64::deserialize(deserializer)?; + if !value.is_finite() || value <= 0.0 { + return Err(serde::de::Error::custom(format!( + "`without-index-compensation-factor` must be finite and positive, got {value}" + ))); + } + Ok(value) +} + #[cfg(test)] mod tests { use super::*; @@ -1501,4 +1551,36 @@ mod tests { // Neutral value guard by default: candidate and incumbent scores compared as-is. assert_eq!(c.value_guard_creation_bias, 1.0); } + + // A non-finite weight would be formatted into the score SQL as a bare `inf`/`NaN` + // token and error every scoring pass; a negative weight inverts the ranking. Both + // must be rejected at load, not at runtime. + #[test] + fn weighted_weights_must_be_finite_and_nonneg() { + let base = "indexer-metrics = \"http://localhost:9100/metrics\"\n"; + let load = |pm: &str| { + from_str::(&format!( + "{base}priority-mode = {{ weighted = {{ {pm} }} }}" + )) + }; + assert!(load("demand-weight = 1.0").is_ok()); + assert!(load("demand-weight = -1.0").is_err()); + assert!(load("latency-weight = inf").is_err()); + assert!(load("supply-weight = nan").is_err()); + } + + // The compensation factor multiplies the without-index cost, so it must be finite + // and strictly positive. + #[test] + fn compensation_factor_must_be_finite_and_positive() { + let base = "indexer-metrics = \"http://localhost:9100/metrics\"\n"; + let load = |value: &str| { + from_str::(&format!( + "{base}without-index-compensation-factor = {value}" + )) + }; + assert!(load("1.5").is_ok()); + assert!(load("0.0").is_err()); + assert!(load("-1.0").is_err()); + } } From 740de2b547106a5e15ce46e34f61f260845e997d Mon Sep 17 00:00:00 2001 From: mctursh Date: Thu, 13 Aug 2026 14:46:11 +0100 Subject: [PATCH 2/2] feat(query-tracker): make apply_batch atomic to prevent demand double-count --- crates/query-tracker/src/modules/ingest.rs | 32 ++++++++++++++----- crates/query-tracker/src/modules/store/mod.rs | 23 +++++++------ 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/crates/query-tracker/src/modules/ingest.rs b/crates/query-tracker/src/modules/ingest.rs index fc8f144..b3fc119 100644 --- a/crates/query-tracker/src/modules/ingest.rs +++ b/crates/query-tracker/src/modules/ingest.rs @@ -19,15 +19,15 @@ use crate::modules::store::Store; use crate::stats::metrics; use cloudbreak_core::modules::index_identity::IndexIdentity; use cloudbreak_core::modules::query_tracker_api::{TrackBatch, TrackResponse}; -use sea_orm::DbErr; +use sea_orm::{DbErr, TransactionTrait}; use solana_pubkey::Pubkey; use std::str::FromStr; use tracing::error; pub async fn apply_batch(store: &Store, batch: TrackBatch) -> Result { - let mut accepted = 0usize; let mut skipped = 0usize; + let mut resolved = Vec::with_capacity(batch.observations.len()); for obs in batch.observations { let program = match Pubkey::from_str(&obs.program) { Ok(p) => p, @@ -55,19 +55,35 @@ pub async fn apply_batch(store: &Store, batch: TrackBatch) -> Result( &self, + conn: &C, identity: &IndexIdentity, count: u32, cost_us: u64, @@ -153,11 +154,9 @@ impl Store { fingerprints: &HashSet, example_request: Option, ) -> Result<(), DbErr> { - let backend = self.db.get_database_backend(); + let backend = conn.get_database_backend(); let pattern_id = identity.pattern_id(); - let txn = self.db.begin().await?; - // New rows are always `candidate` (no index yet), so their first cost // goes to the without-index bucket. On update we route by the row's // current status: `created` → with-index bucket, anything else → @@ -203,7 +202,7 @@ impl Store { ], ); - let existing_hll: Option> = txn + let existing_hll: Option> = conn .query_one(insert) .await? .and_then(|row| row.try_get::>>("", "variety_hll").ok()) @@ -214,7 +213,7 @@ impl Store { sketch.insert_many(fingerprints.iter().copied()); let hll_bytes = sketch.to_bytes(); let estimate = sketch.estimate() as i64; - txn.execute(Statement::from_sql_and_values( + conn.execute(Statement::from_sql_and_values( backend, "UPDATE index_patterns SET variety_hll = $2, variety_estimate = $3 \ WHERE pattern_id = $1", @@ -223,7 +222,7 @@ impl Store { .await?; } - txn.commit().await + Ok(()) } // ---- creation ---------------------------------------------------------