diff --git a/control_plane/src/backend_plan/from_stage_config.rs b/control_plane/src/backend_plan/from_stage_config.rs index 02f4f81d1..a840420a8 100644 --- a/control_plane/src/backend_plan/from_stage_config.rs +++ b/control_plane/src/backend_plan/from_stage_config.rs @@ -77,7 +77,36 @@ pub fn from_stage_config( ); } - let mut routing = Vec::with_capacity(cfg.readouts.len()); + // Exact aggregates are already finalized by their accumulator and do not + // have a SketchQuery readout node. Their warm route therefore comes from + // the physical aggregation itself; approximate routes remain readout- + // driven below. + let mut routing = cfg + .aggregations + .iter() + .filter_map(|agg| { + let planner_types::post_asap::SummaryFamilyType::ExactAggregate(kind, _) = &agg.family + else { + return None; + }; + let fingerprint = *fingerprint_by_agg_id.get(agg.aggregation_id.as_str())?; + let agg_type = match kind { + planner_types::post_asap::ExactKind::Sum + | planner_types::post_asap::ExactKind::Count => asap_types::AggregationType::Sum, + planner_types::post_asap::ExactKind::MinMax => asap_types::AggregationType::MinMax, + planner_types::post_asap::ExactKind::Increase + | planner_types::post_asap::ExactKind::Rate => { + asap_types::AggregationType::Increase + } + }; + Some(RoutingEntry { + satisfies: Capability::ExactAgg(agg_type), + materialization: fingerprint, + storage_backend: StorageBackend::SketchStore, + }) + }) + .collect::>(); + routing.reserve(cfg.readouts.len()); for readout in &cfg.readouts { let Some(&fingerprint) = fingerprint_by_agg_id.get(readout.aggregation_id.as_str()) else { // Orphan readout (no matching aggregation in this cycle's @@ -94,11 +123,14 @@ pub fn from_stage_config( continue; }; let satisfies = capability_for_readout(agg, &readout.op)?; - routing.push(RoutingEntry { + let route = RoutingEntry { satisfies, materialization: fingerprint, storage_backend: StorageBackend::SketchStore, - }); + }; + if !routing.contains(&route) { + routing.push(route); + } } let plan_monitors = monitors diff --git a/control_plane/src/physical/compiler.rs b/control_plane/src/physical/compiler.rs index 0090f2fab..6dd6b7f91 100644 --- a/control_plane/src/physical/compiler.rs +++ b/control_plane/src/physical/compiler.rs @@ -1713,13 +1713,14 @@ impl PhysicalCompiler { for (ordinal, selected) in selected.into_iter().enumerate() { let metric = selected.metric.clone(); let aggregation_id = format!("{}:{ordinal}:{}", query.query_id, metric); + // Rate is a readout over the same reset-aware counter state + // as Increase. Keep that semantic distinction in QueryPlan, + // while the physical store binds both to Increase state. + let physical_family = physical_materialization_family(&selected.family); let aggregation = BackendAggregation { aggregation_id: aggregation_id.clone(), metric_name: metric.clone(), - family: SummaryFamilyType::Sketch( - selected.kind.clone(), - planner_types::post_asap::GroupingStrategy::PerSubpopulationInstance, - ), + family: physical_family, window_secs: query.window_secs, spatial_filter: String::new(), grouping: query.group_by.clone(), @@ -1756,16 +1757,18 @@ impl PhysicalCompiler { } } aggregations.push(aggregation); - readouts.push(BackendReadout { - aggregation_id, - op: selected.readout.clone(), - }); + if let Some(readout) = selected.readout.clone() { + readouts.push(BackendReadout { + aggregation_id, + op: readout, + }); + } collector_materializations.push(CollectorMaterialization { query_id: query.query_id.clone(), materialization, metric: metric.clone(), - algorithm: format!("{:?}", selected.kind.algorithm()).to_ascii_lowercase(), - parameters: sketch_params_json(&selected.params), + algorithm: selected.algorithm, + parameters: selected.parameters, group_by: query.group_by.clone(), window_secs: query.window_secs, abstract_window_framework: planner_selection.window_framework.clone(), @@ -1909,7 +1912,7 @@ impl PhysicalCompiler { fingerprint.0 )) })?; - if &materialization.family != node_family + if materialization.family != physical_materialization_family(node_family) || materialization.window.size_ms != query.window_secs.saturating_mul(1_000) || materialization.group_by != query.group_by { @@ -2264,12 +2267,13 @@ fn select_lifecycle( }) } -struct SelectedSketch { +struct SelectedMaterialization { node_identity: usize, metric: String, - kind: planner_types::post_asap::SketchKind, - params: planner_types::post_asap::SketchParams, - readout: SketchQuery, + family: SummaryFamilyType, + readout: Option, + algorithm: String, + parameters: Value, } /// Collect every executable materialization leaf in the selected post-ASAP @@ -2280,11 +2284,11 @@ struct SelectedSketch { /// fallback node and no unused warm state is provisioned. fn collect_selected_materializations( node: &Rc, -) -> Result, String> { +) -> Result, String> { fn walk( node: &Rc, readout: Option<&SketchQuery>, - selected: &mut Vec, + selected: &mut Vec, ) -> Result<(), String> { match &node.expr { SummaryExpr::SummaryEstimate { @@ -2304,15 +2308,35 @@ fn collect_selected_materializations( let metric = summary_agg_metric(node).ok_or_else(|| { "SummaryAgg has no unique time-series source in post-ASAP IR".to_string() })?; - selected.push(SelectedSketch { + selected.push(SelectedMaterialization { node_identity: Rc::as_ptr(node) as usize, metric, - kind: kind.clone(), - params: kind.params().clone(), - readout: readout.clone(), + family: SummaryFamilyType::Sketch( + kind.clone(), + planner_types::post_asap::GroupingStrategy::PerSubpopulationInstance, + ), + readout: Some(readout.clone()), + algorithm: format!("{:?}", kind.algorithm()).to_ascii_lowercase(), + parameters: sketch_params_json(kind.params()), }); } } + SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(kind, params), + .. + } => { + let metric = summary_agg_metric(node).ok_or_else(|| { + "SummaryAgg has no unique time-series source in post-ASAP IR".to_string() + })?; + selected.push(SelectedMaterialization { + node_identity: Rc::as_ptr(node) as usize, + metric, + family: SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), + readout: None, + algorithm: format!("{kind:?}").to_ascii_lowercase(), + parameters: Value::Object(Default::default()), + }); + } SummaryExpr::KeepPreAsap(_) | SummaryExpr::SummaryAgg { .. } | SummaryExpr::SummaryJoin { .. } @@ -2327,6 +2351,18 @@ fn collect_selected_materializations( Ok(selected) } +fn physical_materialization_family(family: &SummaryFamilyType) -> SummaryFamilyType { + match family { + SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Rate, _) => { + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Increase, + planner_types::post_asap::ExactParams::Increase, + ) + } + _ => family.clone(), + } +} + fn sketch_params_json(params: &planner_types::post_asap::SketchParams) -> Value { use planner_types::post_asap::SketchParams as P; match params { @@ -2688,6 +2724,53 @@ mod tests { assert_eq!(bundle.envelope.planner_revision, PLANNER_REVISION); } + #[test] + fn backend_local_compiler_materializes_declared_exact_promql_cases() { + for (query_id, promql, expected_readout) in [ + ( + "q-rate", + "rate(m[1m])", + crate::query_plan::ExactReadout::Rate, + ), + ( + "q-increase", + "increase(m[1m])", + crate::query_plan::ExactReadout::Increase, + ), + ( + "q-sum", + "sum_over_time(m[1m])", + crate::query_plan::ExactReadout::Sum, + ), + ] { + let mut deployment = environment(10_000); + deployment.target = PhysicalDeploymentTarget::BackendLocalRemoteWrite; + deployment.collector_ids.clear(); + let plan = PhysicalCompiler + .compile(request(query_id, promql), deployment) + .unwrap_or_else(|error| panic!("{promql} must compile: {error}")); + assert_eq!(plan.backend_plan.materializations.len(), 1, "{promql}"); + assert_eq!(plan.query_plan.entries.len(), 1, "{promql}"); + assert!(plan.collector_plans.is_empty(), "{promql}"); + let entry = plan.query_plan.entries.values().next().unwrap(); + assert!(matches!( + entry.nodes.get(&entry.root), + Some(crate::query_plan::QueryPlanNode::ExactReadout { readout, .. }) + if *readout == expected_readout + )); + if expected_readout == crate::query_plan::ExactReadout::Rate { + let materialization = plan.backend_plan.materializations.values().next().unwrap(); + assert_eq!( + materialization.family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Increase, + planner_types::post_asap::ExactParams::Increase, + ) + ); + } + } + } + #[test] fn checked_in_backend_local_snapshot_is_canonical_and_compilable() { let source = include_str!("../../../docs/examples/asapquery-planning-snapshot.json"); diff --git a/control_plane/src/query_plan.rs b/control_plane/src/query_plan.rs index a0c4c003d..61e5796ec 100644 --- a/control_plane/src/query_plan.rs +++ b/control_plane/src/query_plan.rs @@ -241,6 +241,10 @@ pub enum QueryPlanNode { input: QueryNodeId, query: QueryReadout, }, + ExactReadout { + input: QueryNodeId, + readout: ExactReadout, + }, SummaryMerge { inputs: Vec, }, @@ -253,12 +257,22 @@ impl QueryPlanNode { pub fn inputs(&self) -> &[QueryNodeId] { match self { Self::ReadMaterialization { .. } | Self::ExactFallback { .. } => &[], - Self::SummaryEstimate { input, .. } => std::slice::from_ref(input), + Self::SummaryEstimate { input, .. } | Self::ExactReadout { input, .. } => { + std::slice::from_ref(input) + } Self::SummaryMerge { inputs } => inputs, } } } +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum ExactReadout { + Sum, + Increase, + Rate, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] pub enum QueryReadout { @@ -325,19 +339,24 @@ where reduction, child, .. - } => { - if !matches!( - family, - SummaryFamilyType::ExactAggregate(..) | SummaryFamilyType::Sketch(..) - ) { - return Err(QueryPlanError::UnsupportedNode(format!( - "summary family {family:?}" - ))); + } => match family { + SummaryFamilyType::ExactAggregate(..) | SummaryFamilyType::Sketch(..) => { + let mut binding = (self.bind)(node, family)?; + binding.output_grouping = physical_grouping(reduction, child)?; + if let Some(readout) = exact_readout(family) { + let input = QueryNodeId(self.next_id); + self.next_id += 1; + self.nodes + .insert(input, QueryPlanNode::ReadMaterialization { binding }); + QueryPlanNode::ExactReadout { input, readout } + } else { + QueryPlanNode::ReadMaterialization { binding } + } } - let mut binding = (self.bind)(node, family)?; - binding.output_grouping = physical_grouping(reduction, child)?; - QueryPlanNode::ReadMaterialization { binding } - } + other => QueryPlanNode::ExactFallback { + reason: format!("summary family {other:?} is not executable by the warm tier"), + }, + }, SummaryExpr::SummaryEstimate { summary_input, query, @@ -374,6 +393,16 @@ where } } +fn exact_readout(family: &SummaryFamilyType) -> Option { + use planner_types::post_asap::ExactKind; + match family { + SummaryFamilyType::ExactAggregate(ExactKind::Sum, _) => Some(ExactReadout::Sum), + SummaryFamilyType::ExactAggregate(ExactKind::Increase, _) => Some(ExactReadout::Increase), + SummaryFamilyType::ExactAggregate(ExactKind::Rate, _) => Some(ExactReadout::Rate), + _ => None, + } +} + fn physical_grouping( reduction: &Reduction, child: &SummaryNode, diff --git a/data_plane/src/drivers/query/adapters/prometheus_http.rs b/data_plane/src/drivers/query/adapters/prometheus_http.rs index 8b5083bf8..329b93045 100644 --- a/data_plane/src/drivers/query/adapters/prometheus_http.rs +++ b/data_plane/src/drivers/query/adapters/prometheus_http.rs @@ -158,7 +158,11 @@ impl PrometheusHttpAdapter { .as_secs_f64() }; - Ok(ParsedQueryRequest { query, time }) + Ok(ParsedQueryRequest { + query, + time, + timeout: params.get("timeout").cloned(), + }) } /// Helper to parse range query parameters @@ -206,6 +210,7 @@ impl PrometheusHttpAdapter { start, end, step, + timeout: params.get("timeout").cloned(), }) } } diff --git a/data_plane/src/drivers/query/adapters/traits.rs b/data_plane/src/drivers/query/adapters/traits.rs index bbeb5997b..29446bab6 100644 --- a/data_plane/src/drivers/query/adapters/traits.rs +++ b/data_plane/src/drivers/query/adapters/traits.rs @@ -16,6 +16,8 @@ use crate::query_engines::QueryResult; pub struct ParsedQueryRequest { pub query: String, pub time: f64, + /// Prometheus timeout syntax, preserved verbatim for exact fallback. + pub timeout: Option, } /// Parsed range query request with validated parameters @@ -25,6 +27,8 @@ pub struct ParsedRangeQueryRequest { pub start: f64, // epoch seconds pub end: f64, // epoch seconds pub step: f64, // seconds, must be multiple of tumbling window + /// Prometheus timeout syntax, preserved verbatim for exact fallback. + pub timeout: Option, } /// Result of query execution (before formatting for protocol) diff --git a/data_plane/src/drivers/query/fallback/mod.rs b/data_plane/src/drivers/query/fallback/mod.rs index bfdcca46e..5a15b4542 100644 --- a/data_plane/src/drivers/query/fallback/mod.rs +++ b/data_plane/src/drivers/query/fallback/mod.rs @@ -6,7 +6,7 @@ use axum::{ use serde_json::Value; use std::collections::HashMap; -use crate::drivers::query::adapters::ParsedQueryRequest; +use crate::drivers::query::adapters::{ParsedQueryRequest, ParsedRangeQueryRequest}; /// Response format from fallback backend #[derive(Debug, Clone)] @@ -60,6 +60,22 @@ pub trait FallbackClient: Send + Sync { self.execute_query(request).await } + /// Execute a Prometheus range query against the fallback backend. + async fn execute_range_query( + &self, + _request: &ParsedRangeQueryRequest, + ) -> Result { + Err(StatusCode::NOT_IMPLEMENTED) + } + + async fn execute_range_query_with_headers( + &self, + request: &ParsedRangeQueryRequest, + _headers: HashMap, + ) -> Result { + self.execute_range_query(request).await + } + /// Get runtime info from the fallback backend (optional) /// /// # Returns diff --git a/data_plane/src/drivers/query/fallback/prometheus.rs b/data_plane/src/drivers/query/fallback/prometheus.rs index 45f040e26..f32e5e8d9 100644 --- a/data_plane/src/drivers/query/fallback/prometheus.rs +++ b/data_plane/src/drivers/query/fallback/prometheus.rs @@ -1,5 +1,5 @@ use super::{FallbackClient, FallbackResponse}; -use crate::drivers::query::adapters::ParsedQueryRequest; +use crate::drivers::query::adapters::{ParsedQueryRequest, ParsedRangeQueryRequest}; use async_trait::async_trait; use axum::http::StatusCode; use reqwest::Client; @@ -19,6 +19,16 @@ impl PrometheusHttpFallback { base_url, } } + + fn apply_headers( + mut request: reqwest::RequestBuilder, + headers: std::collections::HashMap, + ) -> reqwest::RequestBuilder { + for (name, value) in headers { + request = request.header(name, value); + } + request + } } #[async_trait] @@ -39,10 +49,13 @@ impl FallbackClient for PrometheusHttpFallback { debug!("Full forwarding URL: {}", full_url); // Prepare query parameters for forwarding - let query_params = vec![ + let mut query_params = vec![ ("query", request.query.clone()), ("time", request.time.to_string()), ]; + if let Some(timeout) = &request.timeout { + query_params.push(("timeout", timeout.clone())); + } debug!("Final query parameters for forwarding: {:?}", query_params); @@ -95,6 +108,95 @@ impl FallbackClient for PrometheusHttpFallback { } } + async fn execute_query_with_headers( + &self, + request: &ParsedQueryRequest, + headers: std::collections::HashMap, + ) -> Result { + let full_url = format!("{}/api/v1/query", self.base_url.trim_end_matches('/')); + let mut query_params = vec![ + ("query", request.query.clone()), + ("time", request.time.to_string()), + ]; + if let Some(timeout) = &request.timeout { + query_params.push(("timeout", timeout.clone())); + } + let request = Self::apply_headers(self.client.get(full_url).query(&query_params), headers); + let response = match request + .timeout(std::time::Duration::from_secs(30)) + .send() + .await + { + Ok(response) => response, + Err(error) => { + use crate::drivers::query::adapters::PrometheusResponse; + return Ok(FallbackResponse::Json( + serde_json::to_value(PrometheusResponse::error( + "internal", + &format!("Failed to forward query to Prometheus: {error}"), + )) + .expect("Prometheus error response serializes"), + )); + } + }; + let status = response.status(); + let payload = match response.json::().await { + Ok(payload) => payload, + Err(_) => { + use crate::drivers::query::adapters::PrometheusResponse; + serde_json::to_value(PrometheusResponse::error( + "internal", + "Failed to parse Prometheus response", + )) + .expect("Prometheus error response serializes") + } + }; + if !status.is_success() { + debug!(status = %status, "Prometheus instant fallback returned an error response"); + } + Ok(FallbackResponse::Json(payload)) + } + + async fn execute_range_query( + &self, + request: &ParsedRangeQueryRequest, + ) -> Result { + self.execute_range_query_with_headers(request, std::collections::HashMap::new()) + .await + } + + async fn execute_range_query_with_headers( + &self, + request: &ParsedRangeQueryRequest, + headers: std::collections::HashMap, + ) -> Result { + let full_url = format!("{}/api/v1/query_range", self.base_url.trim_end_matches('/')); + let mut query_params = vec![ + ("query", request.query.clone()), + ("start", request.start.to_string()), + ("end", request.end.to_string()), + ("step", request.step.to_string()), + ]; + if let Some(timeout) = &request.timeout { + query_params.push(("timeout", timeout.clone())); + } + let request = Self::apply_headers(self.client.get(full_url).query(&query_params), headers); + let response = request + .timeout(std::time::Duration::from_secs(30)) + .send() + .await + .map_err(|_| StatusCode::BAD_GATEWAY)?; + let status = response.status(); + let payload = response + .json::() + .await + .map_err(|_| StatusCode::BAD_GATEWAY)?; + if !status.is_success() { + debug!(status = %status, "Prometheus range fallback returned an error response"); + } + Ok(FallbackResponse::Json(payload)) + } + async fn get_runtime_info(&self) -> Result { debug!("Fetching runtime info from Prometheus fallback"); @@ -154,3 +256,63 @@ impl FallbackClient for PrometheusHttpFallback { } } } + +#[cfg(test)] +mod tests { + use super::*; + use axum::{extract::Query, http::HeaderMap, routing::get, Json, Router}; + use std::{collections::HashMap, sync::Arc}; + use tokio::{net::TcpListener, sync::Mutex}; + + #[tokio::test] + async fn range_fallback_preserves_parameters_and_request_context() { + let captured = Arc::new(Mutex::new(None)); + let handler_capture = Arc::clone(&captured); + let app = Router::new().route( + "/api/v1/query_range", + get( + move |Query(params): Query>, headers: HeaderMap| { + let captured = Arc::clone(&handler_capture); + async move { + *captured.lock().await = Some((params, headers)); + Json(serde_json::json!({ + "status": "success", + "data": {"resultType": "matrix", "result": []} + })) + } + }, + ), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + tokio::spawn(async move { axum::serve(listener, app).await.unwrap() }); + + let fallback = PrometheusHttpFallback::new(format!("http://{address}")); + let request = ParsedRangeQueryRequest { + query: "rate(requests_total[1m])".into(), + start: 100.25, + end: 200.75, + step: 15.5, + timeout: Some("7s".into()), + }; + fallback + .execute_range_query_with_headers( + &request, + HashMap::from([ + ("authorization".into(), "Bearer secret".into()), + ("x-scope-orgid".into(), "tenant-a".into()), + ]), + ) + .await + .unwrap(); + + let (params, headers) = captured.lock().await.take().unwrap(); + assert_eq!(params.get("query"), Some(&request.query)); + assert_eq!(params.get("start").map(String::as_str), Some("100.25")); + assert_eq!(params.get("end").map(String::as_str), Some("200.75")); + assert_eq!(params.get("step").map(String::as_str), Some("15.5")); + assert_eq!(params.get("timeout").map(String::as_str), Some("7s")); + assert_eq!(headers["authorization"], "Bearer secret"); + assert_eq!(headers["x-scope-orgid"], "tenant-a"); + } +} diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index 8eb17ea0f..fac8b854c 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -72,6 +72,22 @@ fn extract_tenant(headers: &axum::http::HeaderMap) -> String { .unwrap_or_else(|| crate::query_engines::routing::DEFAULT_TENANT.to_string()) } +/// Copy only request context that is safe and meaningful for the configured +/// Prometheus fallback. Hop-by-hop and client-controlled transport headers are +/// intentionally excluded. +fn extract_fallback_headers(headers: &HeaderMap) -> HashMap { + const FORWARDED: [&str; 3] = ["authorization", "x-scope-orgid", "x-asap-tenant"]; + FORWARDED + .into_iter() + .filter_map(|name| { + headers + .get(name) + .and_then(|value| value.to_str().ok()) + .map(|value| (name.to_string(), value.to_string())) + }) + .collect() +} + #[derive(Debug, Clone)] pub struct HttpServerConfig { pub port: u16, @@ -1143,7 +1159,16 @@ async fn process_via_simple_engine( ); use crate::drivers::query::adapters::QueryExecutionResult; use crate::query_engines::routing::query_engine_routing::QueryEngine; - match state.query_engine.execute(&parsed_request.query).await { + let evaluation_ms = if parsed_request.time > 0.0 { + (parsed_request.time * 1_000.0) as u64 + } else { + unix_time_ms() + }; + match state + .query_engine + .execute_at(&parsed_request.query, evaluation_ms) + .await + { Ok(query_result) => { let query_duration = query_start_time.elapsed(); debug!( @@ -1529,6 +1554,7 @@ async fn handle_instant_query( // `X-ASAP-Tenant` header (default `"default"`). Captured before // `parse_get_request` consumes `query_params`. let tenant = extract_tenant(&headers); + let forwarding_headers = extract_fallback_headers(&headers); let parsed_request = match state.adapter.parse_get_request(query_params).await { Ok(req) => { @@ -1555,7 +1581,7 @@ async fn handle_instant_query( &state, &parsed_request, start_time, - HashMap::new(), + forwarding_headers, engine_override, &tenant, ) @@ -1599,13 +1625,7 @@ async fn handle_instant_query_post( let start_time = Instant::now(); debug!("=== INCOMING POST REQUEST ==="); - // Extract headers we want to forward - let mut forwarding_headers = HashMap::new(); - if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) { - if let Ok(auth_str) = auth.to_str() { - forwarding_headers.insert("Authorization".to_string(), auth_str.to_string()); - } - } + let forwarding_headers = extract_fallback_headers(&headers); // Check content type to determine how to parse the body let content_type = headers @@ -1871,6 +1891,8 @@ async fn process_range_query_request( state: &AppState, parsed_request: &ParsedRangeQueryRequest, start_time: Instant, + forwarding_headers: HashMap, + tenant: &str, ) -> Response { // Check if handling is enabled if !state.config.handle_http_requests { @@ -1913,11 +1935,7 @@ async fn process_range_query_request( let end_ms = (parsed_request.end * 1000.0) as u64; let step_ms = (parsed_request.step * 1000.0) as u64; - // Range handlers don't read the `X-ASAP-Tenant` header, so resolve - // the metric's storage axis against the `default` tenant's routing - // table — the same resolution the instant path applies when no - // tenant header is present. - let metric_storage = resolve_metric_storage(state, &parsed_request.query, "default"); + let metric_storage = resolve_metric_storage(state, &parsed_request.query, tenant); // Match the instant path's hardcoding of `(Sum, Approximate)`. // TODO: derive `accuracy` (and `stat`) from the request rather than @@ -1995,16 +2013,26 @@ async fn process_range_query_request( registered = ?registered, "EngineRouter (range): no engine registered for any compatible backend", ); - ( - StatusCode::SERVICE_UNAVAILABLE, - Json(serde_json::json!({ - "status": "error", - "errorType": "internal", - "error": format!( - "no engine registered for any compatible backend; tried {tried:?}, registered={registered:?}" - )})), - ) - .into_response() + if let Some(fallback) = &state.fallback { + match fallback + .execute_range_query_with_headers(parsed_request, forwarding_headers) + .await + { + Ok(response) => response.into_response(), + Err(status) => status.into_response(), + } + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "status": "error", + "errorType": "internal", + "error": format!( + "no engine registered for any compatible backend; tried {tried:?}, registered={registered:?}" + )})), + ) + .into_response() + } } Err(EngineRouterError::AllFailed { last }) => { use crate::query_engines::EngineError; @@ -2021,9 +2049,19 @@ async fn process_range_query_request( falling through to unsupported", parsed_request.query ); - match state.adapter.format_unsupported_query_response().await { - Ok(json) => json.into_response(), - Err(status) => status.into_response(), + if let Some(fallback) = &state.fallback { + match fallback + .execute_range_query_with_headers(parsed_request, forwarding_headers) + .await + { + Ok(response) => response.into_response(), + Err(status) => status.into_response(), + } + } else { + match state.adapter.format_unsupported_query_response().await { + Ok(json) => json.into_response(), + Err(status) => status.into_response(), + } } } EngineError::Backend { .. } => ( @@ -2041,6 +2079,7 @@ async fn process_range_query_request( async fn handle_range_query( query_params: Query>, + headers: HeaderMap, State(state): State, ) -> Response { let _timer = srv_metrics::start_query_timer(srv_metrics::QUERY_TYPE_RANGE); @@ -2068,13 +2107,25 @@ async fn handle_range_query( }; } }; + let tenant = extract_tenant(&headers); - let response = process_range_query_request(&state, &parsed_request, start_time).await; + let response = process_range_query_request( + &state, + &parsed_request, + start_time, + extract_fallback_headers(&headers), + &tenant, + ) + .await; srv_metrics::record_query_outcome(srv_metrics::QUERY_TYPE_RANGE, query_status_label(&response)); response } -async fn handle_range_query_post(State(state): State, body: Bytes) -> Response { +async fn handle_range_query_post( + State(state): State, + headers: HeaderMap, + body: Bytes, +) -> Response { let _timer = srv_metrics::start_query_timer(srv_metrics::QUERY_TYPE_RANGE); let start_time = Instant::now(); debug!("=== INCOMING RANGE QUERY POST REQUEST ==="); @@ -2129,8 +2180,16 @@ async fn handle_range_query_post(State(state): State, body: Bytes) -> }; } }; + let tenant = extract_tenant(&headers); - let response = process_range_query_request(&state, &parsed_request, start_time).await; + let response = process_range_query_request( + &state, + &parsed_request, + start_time, + extract_fallback_headers(&headers), + &tenant, + ) + .await; srv_metrics::record_query_outcome(srv_metrics::QUERY_TYPE_RANGE, query_status_label(&response)); response } diff --git a/data_plane/src/precompute_engine/operators/increase_accumulator.rs b/data_plane/src/precompute_engine/operators/increase_accumulator.rs index 9d55a5ad9..bbbc67b89 100644 --- a/data_plane/src/precompute_engine/operators/increase_accumulator.rs +++ b/data_plane/src/precompute_engine/operators/increase_accumulator.rs @@ -8,6 +8,9 @@ use std::collections::HashMap; use asap_types::Statistic; +const RESET_AWARE_WIRE_MAGIC: &[u8; 8] = b"ASAPINC2"; +const RESET_AWARE_WIRE_EXTENSION_LEN: usize = 8 + 8 + 8; + /// Accumulator for tracking increases in counter metrics /// Stores the starting and last seen measurements with timestamps #[derive(Debug, Clone, Serialize, Deserialize)] @@ -16,26 +19,97 @@ pub struct IncreaseAccumulator { pub starting_timestamp: i64, pub last_seen_measurement: Measurement, pub last_seen_timestamp: i64, + /// Sum of monotonic deltas, adding the post-reset value whenever the + /// counter decreases. This is the reset correction Prometheus applies. + #[serde(default)] + pub total_increase: f64, + #[serde(default)] + pub sample_count: u64, } impl IncreaseAccumulator { + /// Return the number of bytes occupied by one accumulator at the start of + /// `buffer`. Old persisted values end after `last_seen_timestamp`; reset- + /// aware values carry a magic-prefixed extension. The magic makes this + /// safe when the buffer also contains the next keyed entry. + pub(crate) fn serialized_len_from_prefix( + buffer: &[u8], + ) -> Result> { + if buffer.len() < 4 { + return Err("Buffer too short for starting measurement length".into()); + } + let starting_len = u32::from_le_bytes(buffer[0..4].try_into()?) as usize; + let last_len_offset = 4usize + .checked_add(starting_len) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < last_len_offset + 4 { + return Err("Buffer too short for last seen measurement length".into()); + } + let last_len = + u32::from_le_bytes(buffer[last_len_offset..last_len_offset + 4].try_into()?) as usize; + let legacy_len = last_len_offset + .checked_add(4) + .and_then(|offset| offset.checked_add(last_len)) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < legacy_len { + return Err("Buffer too short for last seen timestamp".into()); + } + let has_extension = buffer.len() >= legacy_len + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[legacy_len..legacy_len + RESET_AWARE_WIRE_MAGIC.len()] + == RESET_AWARE_WIRE_MAGIC; + Ok(legacy_len + + if has_extension { + RESET_AWARE_WIRE_EXTENSION_LEN + } else { + 0 + }) + } + pub fn new( starting_measurement: Measurement, starting_timestamp: i64, last_seen_measurement: Measurement, last_seen_timestamp: i64, ) -> Self { + let total_increase = if last_seen_timestamp <= starting_timestamp { + 0.0 + } else if last_seen_measurement.value >= starting_measurement.value { + last_seen_measurement.value - starting_measurement.value + } else { + last_seen_measurement.value + }; + let sample_count = if last_seen_timestamp > starting_timestamp { + 2 + } else { + 1 + }; Self { starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, + total_increase, + sample_count, } } pub fn update(&mut self, measurement: Measurement, timestamp: i64) { + if timestamp < self.last_seen_timestamp { + return; + } + if timestamp == self.last_seen_timestamp { + return; + } + if measurement.value >= self.last_seen_measurement.value { + self.total_increase += measurement.value - self.last_seen_measurement.value; + } else { + self.total_increase += measurement.value; + } self.last_seen_measurement = measurement; self.last_seen_timestamp = timestamp; + self.sample_count = self.sample_count.saturating_add(1); } pub fn deserialize_from_json(data: &Value) -> Result> { @@ -50,12 +124,19 @@ impl IncreaseAccumulator { .as_i64() .ok_or("Missing or invalid 'last_seen_timestamp' field")?; - Ok(Self::new( + let mut accumulator = Self::new( starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, - )) + ); + accumulator.total_increase = data["total_increase"] + .as_f64() + .unwrap_or(accumulator.total_increase); + accumulator.sample_count = data["sample_count"] + .as_u64() + .unwrap_or(accumulator.sample_count); + Ok(accumulator) } pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { @@ -132,12 +213,30 @@ impl IncreaseAccumulator { buffer[offset + 7], ]); - Ok(Self::new( + let mut accumulator = Self::new( starting_measurement, starting_timestamp, last_seen_measurement, last_seen_timestamp, - )) + ); + offset += 8; + if buffer.len() >= offset + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[offset..offset + RESET_AWARE_WIRE_MAGIC.len()] == RESET_AWARE_WIRE_MAGIC + { + offset += RESET_AWARE_WIRE_MAGIC.len(); + accumulator.total_increase = f64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked total-increase bytes"), + ); + offset += 8; + accumulator.sample_count = u64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked sample-count bytes"), + ); + } + Ok(accumulator) } } @@ -148,6 +247,8 @@ impl SerializableToSink for IncreaseAccumulator { "starting_timestamp": self.starting_timestamp, "last_seen_measurement": self.last_seen_measurement.serialize_to_json(), "last_seen_timestamp": self.last_seen_timestamp, + "total_increase": self.total_increase, + "sample_count": self.sample_count, }) } @@ -170,6 +271,9 @@ impl SerializableToSink for IncreaseAccumulator { // Last seen timestamp buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); + buffer.extend_from_slice(RESET_AWARE_WIRE_MAGIC); + buffer.extend_from_slice(&self.total_increase.to_le_bytes()); + buffer.extend_from_slice(&self.sample_count.to_le_bytes()); buffer } @@ -183,16 +287,21 @@ impl MergeableAccumulator for IncreaseAccumulator { return Err("No accumulators to merge".into()); } + let mut accumulators = accumulators; + accumulators.sort_by_key(|accumulator| accumulator.starting_timestamp); let mut result = accumulators[0].clone(); for acc in &accumulators[1..] { - // Use the earlier starting point - if acc.starting_timestamp < result.starting_timestamp { - result.starting_measurement = acc.starting_measurement.clone(); - result.starting_timestamp = acc.starting_timestamp; + if acc.starting_timestamp > result.last_seen_timestamp { + result.total_increase += + if acc.starting_measurement.value >= result.last_seen_measurement.value { + acc.starting_measurement.value - result.last_seen_measurement.value + } else { + acc.starting_measurement.value + }; } - - // Use the later last seen point + result.total_increase += acc.total_increase; + result.sample_count = result.sample_count.saturating_add(acc.sample_count); if acc.last_seen_timestamp > result.last_seen_timestamp { result.last_seen_measurement = acc.last_seen_measurement.clone(); result.last_seen_timestamp = acc.last_seen_timestamp; @@ -262,10 +371,13 @@ impl AggregateCore for IncreaseAccumulator { &self, statistic: asap_types::Statistic, _key: &Option, - _query_kwargs: &std::collections::HashMap, + query_kwargs: &std::collections::HashMap, ) -> Result> { use crate::storage_engines::types::SingleSubpopulationAggregate; - self.query(statistic, None) + self.query( + statistic, + (!query_kwargs.is_empty()).then_some(query_kwargs), + ) } } @@ -275,24 +387,9 @@ impl SingleSubpopulationAggregate for IncreaseAccumulator { statistic: Statistic, query_kwargs: Option<&HashMap>, ) -> Result> { - // IncreaseAccumulator doesn't use query_kwargs, assert it's None - if query_kwargs.is_some() { - return Err("IncreaseAccumulator does not support query parameters".into()); - } - match statistic { - Statistic::Increase => { - Ok(self.last_seen_measurement.value - self.starting_measurement.value) - } - Statistic::Rate => { - // Convert to per second; timestamps are in milliseconds - let time_diff = (self.last_seen_timestamp - self.starting_timestamp) as f64; - if time_diff <= 0.0 { - return Err("Invalid time difference for rate calculation".into()); - } - let value_diff = self.last_seen_measurement.value - self.starting_measurement.value; - Ok(value_diff / time_diff * 1000.0) - } + Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), + Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), // For instant `sum [by (...)] (counter_metric)` Prometheus // sums the latest cumulative value of each matching series. // The IncreaseAccumulator already tracks that latest value @@ -314,6 +411,65 @@ impl SingleSubpopulationAggregate for IncreaseAccumulator { } } +impl IncreaseAccumulator { + fn extrapolated_value( + &self, + query_kwargs: Option<&HashMap>, + is_rate: bool, + ) -> Result> { + if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { + return Err("at least two ordered counter samples are required".into()); + } + let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; + let Some(kwargs) = query_kwargs else { + return Ok(if is_rate { + self.total_increase / sampled_interval + } else { + self.total_increase + }); + }; + let range_start = kwargs + .get("range_start_ms") + .ok_or("missing range_start_ms")? + .parse::()?; + let range_end = kwargs + .get("range_end_ms") + .ok_or("missing range_end_ms")? + .parse::()?; + if range_end <= range_start { + return Err("invalid counter evaluation range".into()); + } + + let mut duration_to_start = + (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; + let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; + let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; + let extrapolation_threshold = average_sample_interval * 1.1; + + if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { + let duration_to_zero = + sampled_interval * (self.starting_measurement.value / self.total_increase); + duration_to_start = duration_to_start.min(duration_to_zero); + } + let mut extrapolate_to = sampled_interval; + extrapolate_to += if duration_to_start < extrapolation_threshold { + duration_to_start.max(0.0) + } else { + average_sample_interval / 2.0 + }; + extrapolate_to += if duration_to_end < extrapolation_threshold { + duration_to_end.max(0.0) + } else { + average_sample_interval / 2.0 + }; + let mut factor = extrapolate_to / sampled_interval; + if is_rate { + factor /= (range_end - range_start) as f64 / 1000.0; + } + Ok(self.total_increase * factor) + } +} + pub struct IncreaseAccumulatorFactory; impl SingleSubpopulationAggregateFactory for IncreaseAccumulatorFactory { @@ -428,6 +584,33 @@ mod tests { assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); } + #[test] + fn prometheus_counter_reset_and_boundary_extrapolation() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + acc.update(Measurement::new(20.0), 20_000); + acc.update(Measurement::new(3.0), 30_000); + acc.update(Measurement::new(13.0), 50_000); + assert_eq!(acc.total_increase, 23.0); + assert_eq!(acc.sample_count, 4); + + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + let increase = + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) + .unwrap(); + let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) + .unwrap(); + assert!((increase - 34.5).abs() < 1e-12); + assert!((rate - 0.575).abs() < 1e-12); + } + #[test] fn test_increase_accumulator_sum_is_latest_cumulative_value() { // Instant `sum ()` semantics: the per-series summand is @@ -517,6 +700,13 @@ mod tests { acc.last_seen_timestamp, deserialized_bytes.last_seen_timestamp ); + assert_eq!(acc.total_increase, deserialized_bytes.total_increase); + assert_eq!(acc.sample_count, deserialized_bytes.sample_count); + + let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; + let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); + assert_eq!(legacy_value.total_increase, 15.0); + assert_eq!(legacy_value.sample_count, 2); } #[test] diff --git a/data_plane/src/precompute_engine/operators/multiple_increase_accumulator.rs b/data_plane/src/precompute_engine/operators/multiple_increase_accumulator.rs index 3cbad3337..7e2268a75 100644 --- a/data_plane/src/precompute_engine/operators/multiple_increase_accumulator.rs +++ b/data_plane/src/precompute_engine/operators/multiple_increase_accumulator.rs @@ -89,25 +89,11 @@ impl MultipleIncreaseAccumulator { if offset >= buffer.len() { return Err("Buffer too short for increase accumulator data".into()); } - let increase_data = IncreaseAccumulator::deserialize_from_bytes(&buffer[offset..])?; - - // Calculate consumed bytes for IncreaseAccumulator - // Structure: starting_measurement_len(4) + starting_measurement + starting_timestamp(8) + - // last_seen_measurement_len(4) + last_seen_measurement + last_seen_timestamp(8) - let starting_measurement_len = u32::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - ]) as usize; - let last_seen_measurement_len = u32::from_le_bytes([ - buffer[offset + 4 + starting_measurement_len + 8], - buffer[offset + 4 + starting_measurement_len + 8 + 1], - buffer[offset + 4 + starting_measurement_len + 8 + 2], - buffer[offset + 4 + starting_measurement_len + 8 + 3], - ]) as usize; let consumed_bytes = - 4 + starting_measurement_len + 8 + 4 + last_seen_measurement_len + 8; + IncreaseAccumulator::serialized_len_from_prefix(&buffer[offset..])?; + let increase_data = IncreaseAccumulator::deserialize_from_bytes( + &buffer[offset..offset + consumed_bytes], + )?; offset += consumed_bytes; accumulator.increases.insert(key, increase_data); @@ -197,19 +183,15 @@ impl AggregateCore for MultipleIncreaseAccumulator { .downcast_ref::() .ok_or("Failed to downcast to MultipleIncreaseAccumulator")?; - // Clone self once, then merge other's data in-place + // Clone self once, then merge each matching counter with the same + // reset-aware, boundary-aware implementation used by the unkeyed path. let mut merged = self.clone(); for (key, data) in &other_multiple_increase.increases { if let Some(existing_data) = merged.increases.get_mut(key) { - // Merge in-place: take earliest start, latest end - if data.starting_timestamp < existing_data.starting_timestamp { - existing_data.starting_measurement = data.starting_measurement.clone(); - existing_data.starting_timestamp = data.starting_timestamp; - } - if data.last_seen_timestamp > existing_data.last_seen_timestamp { - existing_data.last_seen_measurement = data.last_seen_measurement.clone(); - existing_data.last_seen_timestamp = data.last_seen_timestamp; - } + *existing_data = IncreaseAccumulator::merge_accumulators(vec![ + existing_data.clone(), + data.clone(), + ])?; } else { merged.increases.insert(key.clone(), data.clone()); } @@ -252,14 +234,14 @@ impl MultipleSubpopulationAggregate for MultipleIncreaseAccumulator { &self, statistic: Statistic, key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, + query_kwargs: Option<&HashMap>, ) -> Result> { let data = self .increases .get(key) .ok_or_else(|| format!("Key {key} not found in MultipleIncreaseAccumulator"))?; - data.query(statistic, None) + data.query(statistic, query_kwargs) } fn clone_boxed(&self) -> Box { @@ -280,16 +262,8 @@ impl MergeableAccumulator for MultipleIncreaseAccum for accumulator in accumulators { for (key, data) in accumulator.increases { if let Some(existing_data) = result.increases.get_mut(&key) { - // Merge in-place without cloning existing_data - // Take the earliest start time and latest end time - if data.starting_timestamp < existing_data.starting_timestamp { - existing_data.starting_measurement = data.starting_measurement; - existing_data.starting_timestamp = data.starting_timestamp; - } - if data.last_seen_timestamp > existing_data.last_seen_timestamp { - existing_data.last_seen_measurement = data.last_seen_measurement; - existing_data.last_seen_timestamp = data.last_seen_timestamp; - } + *existing_data = + IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; } else { result.increases.insert(key, data); } @@ -442,27 +416,44 @@ mod tests { let mut acc = MultipleIncreaseAccumulator::new(); let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); + let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); + reset_aware.update(Measurement::new(3.0), 3000); + acc.update(key.clone(), reset_aware); + acc.update( + second_key.clone(), + create_test_increase_accumulator(4.0, 9.0), + ); // Test JSON serialization let json_value = acc.serialize_to_json(); let deserialized = MultipleIncreaseAccumulator::deserialize_from_json(&json_value).unwrap(); - assert_eq!(deserialized.increases.len(), 1); + assert_eq!(deserialized.increases.len(), 2); let deserialized_acc = deserialized.increases.get(&key).unwrap(); assert_eq!(deserialized_acc.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc.last_seen_measurement.value, 25.0); + assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc.total_increase, 18.0); // Test binary serialization let bytes = acc.serialize_to_bytes(); let deserialized_bytes = MultipleIncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(deserialized_bytes.increases.len(), 1); + assert_eq!(deserialized_bytes.increases.len(), 2); let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 25.0); + assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc_bytes.total_increase, 18.0); + assert_eq!( + deserialized_bytes + .increases + .get(&second_key) + .unwrap() + .last_seen_measurement + .value, + 9.0 + ); } #[test] diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index a26878b05..042cc2855 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -339,21 +339,10 @@ impl ASAPQueryEngine { /// `handle_range_query_promql` returns `None`. /// /// Time semantics follow Prometheus's - /// `/api/v1/query_range?start&end&step` spec: the result is a - /// `matrix` (one row per series, each row carrying multiple - /// (timestamp, value) samples). The warm-tier reducer naturally - /// produces one sample per window_close in `[start, end]`, so - /// the matrix is sampled at the underlying aggregation's window - /// boundaries — typically a finer grid than the user's `step` - /// when window_size < step. (The Prometheus spec says - /// evaluate at each step `t = start, start+step, …, end`; the - /// warm tier returns at native window-close granularity instead. - /// This is more data, not less — clients that expect exact step - /// timestamps can downsample, or route step-precise queries to - /// the cold tier via the EngineRouter.) - /// - /// `step` is currently accepted for API compatibility but unused - /// — see the granularity-mismatch note above. + /// `/api/v1/query_range?start&end&step` contract: the immutable + /// QueryPlan is evaluated independently at `start + n*step`, and + /// the result contains exactly those timestamps. Native summary + /// pane boundaries are an internal detail and never become API steps. pub async fn execute_range_promql_modern( &self, query: &str, @@ -379,12 +368,12 @@ impl ASAPQueryEngine { physical_plan.backend_plan.plan_version, readiness_requirement(query_entry), )); - crate::query_engines::asap_query_engine::live_serve::serve_from_query_plan( + crate::query_engines::asap_query_engine::live_serve::serve_range_steps_from_query_plan( idx, query_entry, start_ms, end_ms, - false, + step_ms, ) } Err(reason) => Err(crate::query_engines::asap_query_engine::post_asap_planner::LoweringSkip::QueryNotPlanned(reason.to_string())), @@ -657,16 +646,25 @@ impl crate::query_engines::routing::query_engine_routing::QueryEngine for ASAPQu &self, query: &str, ) -> Result + { + let now_ms = std::time::SystemTime::now() + .duration_since(std::time::SystemTime::UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0); + self.execute_at(query, now_ms).await + } + + async fn execute_at( + &self, + query: &str, + now_ms: u64, + ) -> Result { // One authoritative warm path: ASAPPlanner post-ASAP DAG → // BackendPlan/materialization resolver → SID lookup → DAG executor. // A typed resolver/executor error becomes CapabilityMiss, which lets // EngineRouter continue to the archive backend. if let Some(idx) = self.sketch_index.as_ref() { - let now_ms = std::time::SystemTime::now() - .duration_since(std::time::SystemTime::UNIX_EPOCH) - .map(|d| d.as_millis() as u64) - .unwrap_or(0); let physical_plan = self.physical_plan_snapshot(); let mut readiness = None; let planned = match physical_plan.as_ref() { diff --git a/data_plane/src/query_engines/asap_query_engine/live_serve.rs b/data_plane/src/query_engines/asap_query_engine/live_serve.rs index c5f78c028..9efc044ec 100644 --- a/data_plane/src/query_engines/asap_query_engine/live_serve.rs +++ b/data_plane/src/query_engines/asap_query_engine/live_serve.rs @@ -18,12 +18,14 @@ //! registered sid) means the query fails over to archive — there is no //! other sketch-tier path left to try. +use std::collections::BTreeMap; + use control_plane::types_v2::AccuracyTarget; use crate::query_engines::asap_query_engine::post_asap_planner::LoweringSkip; use crate::query_engines::asap_query_engine::post_asap_readout::{ execute_post_asap_instant, execute_post_asap_readout, execute_query_plan_instant, - execute_query_plan_readout, + execute_query_plan_readout, fold_coverage, }; use crate::storage_engines::sketch_db::index::SketchStore; use crate::storage_engines::sketch_db::query::ASAPTierResult; @@ -193,6 +195,92 @@ pub fn serve_from_query_plan( }) } +/// Evaluate a planned range query at the exact Prometheus timestamps +/// `start, start+step, ... <= end`. Each point executes the same immutable +/// QueryPlan DAG with its own lookback interval; native summary window-close +/// timestamps never leak into the public range response. +pub fn serve_range_steps_from_query_plan( + index: &SketchStore, + entry: &control_plane::query_plan::QueryPlanEntry, + start_ms: u64, + end_ms: u64, + step_ms: u64, +) -> Result { + const MAX_RANGE_STEPS: usize = 11_000; + if step_ms == 0 || start_ms > end_ms { + return Err(LoweringSkip::InvalidQueryPlan( + "range query requires start <= end and step > 0".into(), + )); + } + let step_count = end_ms.saturating_sub(start_ms) / step_ms + 1; + if step_count > MAX_RANGE_STEPS as u64 { + return Err(LoweringSkip::MaterializationNotReady(format!( + "range query has {step_count} steps; warm execution limit is {MAX_RANGE_STEPS}" + ))); + } + let window_ms = entry + .materialization_bindings() + .into_iter() + .map(|binding| binding.window_ms) + .max() + .unwrap_or(0); + let mut rows = BTreeMap::, Vec<(i64, f64)>>::new(); + let mut total_coverage = None; + let mut evaluation_ms = start_ms; + loop { + let (outcome, t0_ms) = execute_query_plan_instant(index, entry, evaluation_ms)?; + if !coverage_covers_closed_windows(outcome.coverage, t0_ms, evaluation_ms, window_ms) { + return Err(LoweringSkip::MaterializationNotReady(format!( + "coverage {:?} is incomplete at range evaluation {evaluation_ms}", + outcome.coverage + ))); + } + if outcome.series.is_empty() { + return Err(LoweringSkip::MaterializationNotReady(format!( + "planned summary produced no series at range evaluation {evaluation_ms}" + ))); + } + fold_coverage(&mut total_coverage, outcome.coverage); + for (labels, samples) in outcome.series { + let Some((_, value)) = samples.last() else { + return Err(LoweringSkip::MaterializationNotReady(format!( + "planned summary produced no value at range evaluation {evaluation_ms}" + ))); + }; + rows.entry(labels) + .or_default() + .push((evaluation_ms as i64, *value)); + } + let Some(next) = evaluation_ms.checked_add(step_ms) else { + break; + }; + if next > end_ms { + break; + } + evaluation_ms = next; + } + Ok(ASAPTierResult { + series: rows.into_iter().collect(), + coverage: total_coverage, + }) +} + +fn coverage_covers_closed_windows( + coverage: Option<(u64, u64)>, + t0_ms: u64, + t1_ms: u64, + window_ms: u64, +) -> bool { + let Some((start, end)) = coverage else { + return false; + }; + window_ms > 0 + && start <= end + && start.saturating_sub(window_ms) <= t0_ms + && end >= t1_ms + && end.saturating_sub(start).saturating_add(window_ms) >= t1_ms.saturating_sub(t0_ms) +} + pub fn serve_instant_from_query_plan( index: &SketchStore, entry: &control_plane::query_plan::QueryPlanEntry, @@ -353,6 +441,76 @@ mod tests { assert!(!result.is_empty()); } + #[test] + fn formal_range_plan_returns_exact_requested_steps() { + let idx = SketchStore::new(); + let policy = asap_types::PolicyFingerprint(901); + idx.register(SketchInstanceMetadata { + sid: 9, + metric_name: "bytes".into(), + group_by_keys: std::collections::BTreeSet::new(), + capability: Some(Capability::ExactAgg(asap_types::AggregationType::Sum)), + agg_kind: AggKind::ExactAgg { + agg_type: asap_types::AggregationType::Sum, + parameters_canonical: String::new(), + spatial_filter_canonical: String::new(), + }, + accuracy: None, + first_seen_unix_ms: 0, + retired_at_ms: None, + expires_at_ms: None, + policy_fp: policy, + }); + for (start, end, value) in [(0, 1_000, 1.0), (1_000, 2_000, 2.0), (2_000, 3_000, 3.0)] { + idx.append_precompute( + 9, + BTreeMap::new(), + (start, end), + Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(value)), + ); + } + let entry = control_plane::query_plan::QueryPlanEntry { + query_id: "q-sum".into(), + canonical_promql: "sum_over_time(bytes[1s])".into(), + root: control_plane::query_plan::QueryNodeId(0), + nodes: BTreeMap::from([ + ( + control_plane::query_plan::QueryNodeId(0), + control_plane::query_plan::QueryPlanNode::ExactReadout { + input: control_plane::query_plan::QueryNodeId(1), + readout: control_plane::query_plan::ExactReadout::Sum, + }, + ), + ( + control_plane::query_plan::QueryNodeId(1), + control_plane::query_plan::QueryPlanNode::ReadMaterialization { + binding: control_plane::query_plan::MaterializationBinding { + materialization: policy, + metric: "bytes".into(), + sid_grouping: vec![], + output_grouping: control_plane::query_plan::PhysicalGrouping::PerEntity, + window_ms: 1_000, + }, + }, + ), + ]), + instant: control_plane::query_plan::InstantExecution { + lookback_ms: 1_000, + full_history: false, + cumulative_readout: true, + }, + fallback: control_plane::query_plan::FallbackPolicy::ExactBackend, + }; + + let result = serve_range_steps_from_query_plan(&idx, &entry, 1_000, 3_000, 1_000) + .expect("step-precise range execution"); + assert_eq!(result.series.len(), 1); + assert_eq!( + result.series[0].1, + vec![(1_000, 1.0), (2_000, 2.0), (3_000, 3.0)] + ); + } + #[test] fn flag_on_global_merge_shape_is_served_merged_not_declined() { // Previously `flag_on_ambiguous_shape_falls_back`, asserting diff --git a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs index a62c8b977..cdaa7a4fb 100644 --- a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs +++ b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs @@ -174,6 +174,34 @@ impl QueryNodeRuntime for PhysicalQueryRuntime<'_> { .collect::, _>>() .map(PhysicalQueryOutput::Value) } + QueryPlanNode::ExactReadout { readout, .. } => { + let [PhysicalQueryOutput::State(groups)] = inputs else { + return Err(PhysicalNodeError::ExpectedState); + }; + groups + .iter() + .map(|(key, state)| { + state + .exact_value_for( + *readout, + &None, + self.context.t0_ms, + self.context.t1_ms, + ) + .map(|value| { + ( + key.clone(), + SummaryValue::Points( + vec![(self.context.t1_ms as i64, value)], + state.exact_coverage(), + ), + ) + }) + .ok_or(PhysicalNodeError::ExpectedState) + }) + .collect::, _>>() + .map(PhysicalQueryOutput::Value) + } QueryPlanNode::SummaryMerge { .. } => { let mut by_group: BTreeMap, Vec> = BTreeMap::new(); @@ -604,4 +632,74 @@ mod tests { // `exact_coverage` by this module's A0 test in `summary_executor.rs`. assert_eq!(outcome.coverage, Some((2_000, 2_000))); } + + #[test] + fn exact_query_plan_rate_uses_reset_aware_readout() { + let idx = SketchStore::new(); + let policy = asap_types::PolicyFingerprint(777); + idx.register(SketchInstanceMetadata { + sid: 7, + metric_name: "requests_total".into(), + group_by_keys: std::collections::BTreeSet::new(), + capability: Some(Capability::ExactAgg(asap_types::AggregationType::Increase)), + agg_kind: AggKind::ExactAgg { + agg_type: asap_types::AggregationType::Increase, + parameters_canonical: String::new(), + spatial_filter_canonical: String::new(), + }, + accuracy: None, + first_seen_unix_ms: 0, + retired_at_ms: None, + expires_at_ms: None, + policy_fp: policy, + }); + use crate::storage_engines::types::Measurement; + let mut accumulator = crate::precompute_engine::operators::IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + accumulator.update(Measurement::new(20.0), 20_000); + accumulator.update(Measurement::new(3.0), 30_000); + accumulator.update(Measurement::new(13.0), 50_000); + idx.append_precompute(7, BTreeMap::new(), (0, 60_000), Box::new(accumulator)); + + let entry = control_plane::query_plan::QueryPlanEntry { + query_id: "q-rate".into(), + canonical_promql: "rate(requests_total[1m])".into(), + root: control_plane::query_plan::QueryNodeId(0), + nodes: BTreeMap::from([ + ( + control_plane::query_plan::QueryNodeId(0), + QueryPlanNode::ExactReadout { + input: control_plane::query_plan::QueryNodeId(1), + readout: control_plane::query_plan::ExactReadout::Rate, + }, + ), + ( + control_plane::query_plan::QueryNodeId(1), + QueryPlanNode::ReadMaterialization { + binding: control_plane::query_plan::MaterializationBinding { + materialization: policy, + metric: "requests_total".into(), + sid_grouping: vec![], + output_grouping: control_plane::query_plan::PhysicalGrouping::PerEntity, + window_ms: 60_000, + }, + }, + ), + ]), + instant: control_plane::query_plan::InstantExecution { + lookback_ms: 60_000, + full_history: false, + cumulative_readout: true, + }, + fallback: control_plane::query_plan::FallbackPolicy::ExactBackend, + }; + let outcome = execute_query_plan_readout(&idx, &entry, 0, 60_000, true) + .expect("execute exact rate DAG"); + let value = outcome.series[0].1[0].1; + assert!((value - 0.575).abs() < 1e-12, "reset-aware rate={value}"); + } } diff --git a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs index cc7efc88f..3510bc9ff 100644 --- a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs +++ b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs @@ -210,7 +210,7 @@ impl GroupState { for acc in windows.values() { merged = Some(match merged.take() { None => acc.clone_boxed_core(), - Some(m) => m.merge_with(acc.as_ref()).ok()?, + Some(current) => current.merge_with(acc.as_ref()).ok()?, }); } } @@ -219,6 +219,50 @@ impl GroupState { .ok() } + /// Finalize a compiler-declared exact readout. Rate and increase share + /// reset-aware Increase state physically, but remain distinct operations + /// in QueryPlan so serving never infers semantics from PromQL text. + pub fn exact_value_for( + &self, + readout: control_plane::query_plan::ExactReadout, + key: &Option, + range_start_ms: u64, + range_end_ms: u64, + ) -> Option { + let GroupState::ExactAgg { entries, agg_type } = self else { + return None; + }; + let stat = match (readout, agg_type) { + ( + control_plane::query_plan::ExactReadout::Sum, + AggregationType::Sum | AggregationType::MultipleSum, + ) => asap_types::Statistic::Sum, + ( + control_plane::query_plan::ExactReadout::Increase, + AggregationType::Increase | AggregationType::MultipleIncrease, + ) => asap_types::Statistic::Increase, + ( + control_plane::query_plan::ExactReadout::Rate, + AggregationType::Increase | AggregationType::MultipleIncrease, + ) => asap_types::Statistic::Rate, + _ => return None, + }; + let mut merged: Option> = None; + for windows in entries { + for acc in windows.values() { + merged = Some(match merged.take() { + None => acc.clone_boxed_core(), + Some(m) => m.merge_with(acc.as_ref()).ok()?, + }); + } + } + let query_kwargs = std::collections::HashMap::from([ + ("range_start_ms".to_string(), range_start_ms.to_string()), + ("range_end_ms".to_string(), range_end_ms.to_string()), + ]); + merged?.query_statistic(stat, key, &query_kwargs).ok() + } + /// Coverage analog of `exact_value` — folds `(min_window_end_ms, /// max_window_end_ms)` from every window observed across the group's /// entries, via the same `fold_coverage` helper diff --git a/data_plane/src/query_engines/routing/query_engine_routing.rs b/data_plane/src/query_engines/routing/query_engine_routing.rs index af88f54fe..59a6ab955 100644 --- a/data_plane/src/query_engines/routing/query_engine_routing.rs +++ b/data_plane/src/query_engines/routing/query_engine_routing.rs @@ -87,6 +87,18 @@ pub trait QueryEngine: Send + Sync { /// Answer `query` against this engine's storage tier. async fn execute(&self, query: &str) -> Result; + /// Evaluate an instant query at an explicit Unix-millisecond timestamp. + /// Engines without a time-aware implementation retain their existing + /// behavior; the ASAP QueryPlan runtime overrides this method. + async fn execute_at( + &self, + query: &str, + evaluation_ms: u64, + ) -> Result { + let _ = evaluation_ms; + self.execute(query).await + } + /// Forward a range query to this engine's storage tier. /// /// Params are milliseconds since epoch. The default returns `CapabilityMiss`