feat(query-tracker): validate weighted-mode config and make apply_batch atomic

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<f64, D::Error>
+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<f64, D::Error>
+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::<QueryTrackerConfig>(&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::<QueryTrackerConfig>(&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());
+    }
 }

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<TrackResponse, DbErr> {
-    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<TrackRespon
             .as_ref()
             .and_then(|c| serde_json::to_value(c).ok());
 
+        resolved.push((
+            identity,
+            obs.count,
+            obs.total_cost_us,
+            obs.failed_count,
+            obs.value_fingerprints,
+            example_request,
+        ));
+    }
+
+    resolved.sort_by_cached_key(|(identity, ..)| identity.pattern_id());
+
+    let accepted = resolved.len();
+
+    let txn = store.db().begin().await?;
+    for (identity, count, cost_us, failed, fingerprints, example_request) in resolved {
         store
             .record_demand(
+                &txn,
                 &identity,
-                obs.count,
-                obs.total_cost_us,
-                obs.failed_count,
-                &obs.value_fingerprints,
+                count,
+                cost_us,
+                failed,
+                &fingerprints,
                 example_request,
             )
             .await?;
-
-        accepted += 1;
     }
+    txn.commit().await?;
 
     metrics::OBSERVATIONS_TOTAL.inc_by(accepted as u64);
     Ok(TrackResponse { accepted, skipped })

crates/query-tracker/src/modules/store/mod.rs

@@ -23,9 +23,7 @@ use crate::stats::variety::VarietySketch;
 use cloudbreak_core::PriorityMode;
 use cloudbreak_core::modules::index_identity::IndexIdentity;
 use patterns::{PatternRow, offsets_from_json, offsets_to_json, status};
-use sea_orm::{
-    ConnectionTrait, DatabaseConnection, DbErr, QueryResult, Statement, TransactionTrait, Value,
-};
+use sea_orm::{ConnectionTrait, DatabaseConnection, DbErr, QueryResult, Statement, Value};
 use std::collections::HashSet;
 
 /// Columns selected into a [`PatternRow`]; kept in one place so every read
@@ -142,22 +140,23 @@ impl Store {
     /// Fold one aggregated observation into the pattern's row: bump demand,
     /// cost and failure counters, refresh `last_demand_at`, resurrect an evicted
     /// pattern back to `candidate`, and merge value fingerprints into the
-    /// variety sketch. Runs in a transaction so the sketch read-modify-write is
-    /// consistent under concurrent `/track` requests.
-    pub async fn record_demand(
+    /// variety sketch. Runs against the caller's `conn`, which owns the
+    /// transaction, so the sketch read-modify-write stays consistent under
+    /// concurrent `/track` requests.
+    #[allow(clippy::too_many_arguments)]
+    pub async fn record_demand<C: ConnectionTrait>(
         &self,
+        conn: &C,
         identity: &IndexIdentity,
         count: u32,
         cost_us: u64,
         failed: u32,
         fingerprints: &HashSet<u64>,
         example_request: Option<serde_json::Value>,
     ) -> 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<Vec<u8>> = txn
+        let existing_hll: Option<Vec<u8>> = conn
             .query_one(insert)
             .await?
             .and_then(|row| row.try_get::<Option<Vec<u8>>>("", "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 ---------------------------------------------------------