From 485fae96123a640fbc63a087ffa0cb7a8239cd2c Mon Sep 17 00:00:00 2001 From: zz_y Date: Sun, 17 May 2026 21:07:57 -0600 Subject: [PATCH] feat(planner): (metric, role)-keyed WorkloadStore + PlanStore for multi-role metrics (B2 full restructure) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Closes the B2 thread (Sum-as-separate-aggregation alongside sketch). PR #282's bundle deferred this with note "needs (metric, role)-keyed restructure"; the B2-planning-agent's analysis concluded the restructure is the architecturally clean solution despite its size (vs the side-table alternative). What changed: * New AggRole enum (Quantile / Sum / Count / Topk / Other) in control_plane/src/workload.rs + derive_agg_role() classifier that resolves from (a) sketch_family_override (DDSketch/KLL → Quantile, HLL → Count, CountSketch → Topk, CMS → Other) then (b) outermost PromQL token of query_string. * WorkloadStore keyed by (metric: String, role: AggRole) — append-not-collapse on duplicate metric inserts. New `get_all_for_metric` + `keys()` accessors for emit-path walks. * PlanStore keyed by (metric, role) — one CollectionPlan per pair; rollback/diff/expired now per-pair. `metrics()` dedups across roles to preserve the metrics-exposer wire shape. * agent_to_metric → agent_to_metrics: Vec<(metric, role)> in replan.rs so a single agent can serve multiple (metric, role) plans. register_agent / unregister_agent updated; push_config_to_agent picks the first pair (the 5-sketch routing-connector edge YAML carries every metric's pipeline anyway, so one push covers them all). * Replanner::replan_metric stays as a wrapper that loops over every role registered for the metric; new replan_metric_role(metric, role) for targeted single-role replans. replan_expired loops over expired (metric, role) pairs; handle_violation re-plans every pair the agent serves. * main.rs pre-pop loop now derives the role per WorkloadEntry via derive_agg_role() before set — the three http_requests_total entries in mvp-workload.yaml (sum / sum+ rate / count) now persist as TWO distinct keys (http_requests_total, Sum) + (http_requests_total, Count) instead of collapsing onto one with only the last entry's plan surviving. * handle_plan derives the role from the QuerySpec's query_string + optional sketch_type override and threads it into store.set / workload_store.set / register_agent. * HTTP endpoints (/api/v1/plan/:metric, /diff, /rollback, /config) preserved at metric granularity; responses extended with a per-role `roles: [...]` array (additive, no breaking URL change). Rollback returns BAD_REQUEST when no role has a previous plan (matches pre-B2 single-role behaviour). * Emit helpers (collect_metric_to_family, collect_metric_to_grouping_labels) walk the new (metric, role) keyed store and take the FIRST sketch-binding role per metric for the routing-connector OTTL conditions — Sum-shaped roles correctly decline bind_workload_typed and flow through the metrics/raw_passthrough pipeline. * Bootstrap typed path (emit_bootstrap_typed) walks every role's workload until one binds, instead of the prior single-role lookup. * metrics_exposer's refresh_plan_ids iterates per-pair so each (metric, role) emits its own asap_active_plan_id row under the metric label. Role folded into the plan_id hash so distinct-role plans produce distinct ids. Test plan: * cargo test -p control_plane --lib: 740 pass (was 719; +21 new tests covering AggRole classification, WorkloadStore multi-role coexistence, PlanStore role-keyed rollback/diff/expired, Replanner replan_metric_role + multi-role agent registration + violation handling, and the synthetic http_requests_total 3-entry regression test). * cargo test -p control_plane --bins: 27 pass (unchanged). * cargo test -p data_plane --lib: 712 pass (no data_plane changes — backend's ingest_precompute_for_agg_config already handles AggKind::ExactAgg(Sum)). Smoke test (NOT run by this commit): The mvp-smoke-test/run_smoke.sh end-to-end check requires docker compose + the full ASAP stack + a backend-streaming.yaml that includes a Sum-shaped aggregationType entry for http_requests_total. The static backend-streaming.yaml in deploy/configs/ today only carries DDSketch entries — that's a deployment-config follow-up, NOT a controller change. Without the static YAML update the data_plane has no Sum capability for http_requests_total at startup and the `sum by (zone)` PromQL still surfaces capability-miss until the dynamic streaming-config POST path (Phase 5 controller → backend hot-reload) plumbs through to the live backend. Ambiguous-case decisions (documented inline): * count_over_time(...) classifies to Count (cardinality semantics). count(over_time(sum_over_time(...))) shape would be Sum via the outer aggregation. * rate / increase classify to Sum (both bind to ExactAgg(Sum) on the data plane). * CountMinSketch override classifies to Other (frequency estimator without an inherent topk shape — keep Topk pure for sketch-with-heap families). Closes B2. Co-Authored-By: Claude Opus 4.7 (1M context) --- control_plane/src/emit/mod.rs | 52 ++-- control_plane/src/main.rs | 325 ++++++++++++++++++----- control_plane/src/metrics_exposer.rs | 30 ++- control_plane/src/replan.rs | 328 +++++++++++++++++------ control_plane/src/store/mod.rs | 179 +++++++++---- control_plane/src/store/workload.rs | 254 ++++++++++++++++-- control_plane/src/workload.rs | 375 +++++++++++++++++++++++++++ 7 files changed, 1312 insertions(+), 231 deletions(-) diff --git a/control_plane/src/emit/mod.rs b/control_plane/src/emit/mod.rs index 602506a7..b9d6230f 100644 --- a/control_plane/src/emit/mod.rs +++ b/control_plane/src/emit/mod.rs @@ -287,18 +287,26 @@ pub fn collect_metric_to_family( ) -> std::collections::HashMap { let mut out = std::collections::HashMap::new(); for entry in registry.entries() { - let Some((workload, _wc)) = workload_store.get(&entry.metric_name) else { - continue; - }; - let Some(physical_expr) = crate::optimizer::rules::bind_workload_typed(&workload) else { - // `http_requests_total` and other raw-passthrough metrics - // land here — correctly excluded so they fall through to - // the routing connector's default `metrics/raw_passthrough` - // pipeline in the emitter. - continue; - }; - if let Some(kind) = extract_root_sketch_kind(&physical_expr) { - out.insert(entry.metric_name.clone(), kind); + // B2 (metric, role) restructure: walk every role registered + // for this metric. Sum-shaped roles (raw passthrough / ExactAgg) + // decline `bind_workload_typed` and correctly fall through to + // the routing-connector's default `metrics/raw_passthrough` + // pipeline. The FIRST sketch-shaped role that binds wins the + // entry — Quantile / Count / Topk all map to a single sketch + // family per metric in the current emitter's wire shape (one + // routing-connector OTTL condition per metric → one pipeline). + // Multi-sketch-per-metric coverage stays a follow-up; the + // streaming-config (`emit_backend_streaming_config_json`) + // already iterates per-aggregation so the backend serves all + // roles even when the edge emits only the first sketch shape. + for (_, workload, _wc) in workload_store.get_all_for_metric(&entry.metric_name) { + let Some(physical_expr) = crate::optimizer::rules::bind_workload_typed(&workload) else { + continue; + }; + if let Some(kind) = extract_root_sketch_kind(&physical_expr) { + out.entry(entry.metric_name.clone()).or_insert(kind); + break; + } } } out @@ -329,10 +337,20 @@ pub fn collect_metric_to_grouping_labels( ) -> std::collections::HashMap> { let mut out = std::collections::HashMap::new(); for entry in registry.entries() { - let Some((workload, _wc)) = workload_store.get(&entry.metric_name) else { - continue; - }; - out.insert(entry.metric_name.clone(), workload.group_by_labels.clone()); + // B2 (metric, role): the FIRST registered role's grouping + // labels win — in practice all roles for a metric share the + // same `grouping_labels` since the YAML field lives on the + // WorkloadEntry. Using `get_all_for_metric().first()` keeps + // pre-B2 semantics ("the entry the controller pre-popped first + // wins") in the common case AND lets a multi-role metric still + // emit a single keep_keys OTTL processor per metric. + if let Some((_, workload, _)) = workload_store + .get_all_for_metric(&entry.metric_name) + .into_iter() + .next() + { + out.insert(entry.metric_name.clone(), workload.group_by_labels.clone()); + } } out } @@ -525,8 +543,10 @@ mod runtime_tests { data: types_v2::DataShape::default(), }; if let Ok(wl) = analyzer.analyze(spec) { + let role = crate::workload::derive_agg_role(entry); store.set( &entry.metric_name, + role, wl, types::WorkloadCharacteristics::default(), ); diff --git a/control_plane/src/main.rs b/control_plane/src/main.rs index f3ccea66..09512bb7 100644 --- a/control_plane/src/main.rs +++ b/control_plane/src/main.rs @@ -209,10 +209,16 @@ async fn main() { let pushed = r.push_config_to_agent(&aid).await; // If the agent has no prior assignment, assign it a workload // from the registry (if available). + // + // B2 (metric, role): bind the on-connect default to + // the first agent-role registry entry's CLASSIFIED + // role (via `derive_agg_role`) so the workload-store + // lookup in `push_config_to_agent` resolves. if !pushed { if let Some(registry) = reg.read().await.as_ref() { if let Some(entry) = registry.first_for_role("agent") { - r.register_agent(&aid, &entry.metric_name).await; + let role = control_plane::workload::derive_agg_role(entry); + r.register_agent(&aid, &entry.metric_name, role).await; r.push_config_to_agent(&aid).await; } } @@ -319,16 +325,29 @@ async fn main() { shape: types_v2::QueryShape::default(), data: types_v2::DataShape::default(), }; + // B2 full restructure — derive the AggRole for this entry + // BEFORE store insertion so collisions on metric name don't + // overwrite a prior role's entry. The pre-B2 loop wrote + // `set(metric, ...)` and silently dropped every entry but + // the LAST one when a metric appeared multiple times in the + // YAML — that's the bug that caused `sum by (zone) + // (http_requests_total)` to return `ExactAgg(Sum) + // capability not satisfied` (entries 2 + 3 of + // mvp-workload.yaml were both Sum-shaped but only the + // count(...) entry 4 survived, with DDSketch from the + // unrelated `http_requests_total_latency_ms` quantile + // entry). + let role = control_plane::workload::derive_agg_role(entry); match analyzer.analyze(spec) { Ok(wl) => { let wc = types::WorkloadCharacteristics::default(); let plan = planner.plan(&wl, Some(&wc)); let metric_name = wl.metric_name.clone(); - plan_store.set(&metric_name, plan); - workload_store.set(&metric_name, wl, wc); + plan_store.set(&metric_name, role, plan); + workload_store.set(&metric_name, role, wl, wc); } Err(e) => { - warn!(metric = %entry.metric_name, error = %e, + warn!(metric = %entry.metric_name, role = %role, error = %e, "failed to pre-populate plan from workload registry"); } } @@ -528,9 +547,26 @@ async fn handle_plan( } plan.precompute = build_precompute_engine_jobs(&workload, "backend:4317"); - st.store.set(&workload.metric_name, plan.clone()); + // B2 (metric, role): derive the role from the request's + // query_string + optional `sketch_type` override so the + // store keys at (metric, role) granularity. Without the role + // a second POST for the same metric with a different shape + // (Quantile vs Sum) would silently overwrite the prior plan. + let role = { + let entry = control_plane::workload::WorkloadEntry { + metric_name: workload.metric_name.clone(), + query_string: query_string.clone(), + accuracy_sla: workload.accuracy_sla, + assign_to_role: String::from("agent"), + sketch_family_override: workload.sketch_type_override.clone(), + target_path: None, + grouping_labels: workload.group_by_labels.clone(), + }; + control_plane::workload::derive_agg_role(&entry) + }; + st.store.set(&workload.metric_name, role, plan.clone()); // Persist workload so the replanner can re-run plan() without the original spec. - st.workload_store.set(&workload.metric_name, workload.clone(), wc); + st.workload_store.set(&workload.metric_name, role, workload.clone(), wc); // ── Push agent config to agent-role collectors ──────────────────────────── if let Ok(agent_yaml) = generate_agent_collector_config(&plan.agent_config, &st.opamp_endpoint) { @@ -785,11 +821,11 @@ async fn handle_plan( } } - // ── Update scrape-endpoint sketch types and agent→metric mapping ────────── + // ── Update scrape-endpoint sketch types and agent→(metric, role) mapping ── let sketch_type = plan.agent_config.sketch_type.clone(); for agent_id in st.opamp.connected_agents().await { st.scraper.set_sketch_type(&agent_id, sketch_type.clone()).await; - st.replanner.register_agent(&agent_id, &workload.metric_name).await; + st.replanner.register_agent(&agent_id, &workload.metric_name, role).await; } // `plan_summary` was computed in the single algebra pipeline above. @@ -870,14 +906,38 @@ async fn handle_get_plan( State(st): State, Path(metric): Path, ) -> impl IntoResponse { - match st.store.get(&metric) { - Ok(plan) => (StatusCode::OK, Json(json!({ - "metric": metric, - "sketch_type": plan.agent_config.sketch_type.to_string(), - "valid_until": plan.valid_until, - }))).into_response(), - Err(e) => (StatusCode::NOT_FOUND, e.to_string()).into_response(), + // B2 (metric, role): return every role's plan for this metric. + // Wire shape (additive, no breaking change): when only one role is + // registered, the response still carries the pre-B2 top-level + // `sketch_type` / `valid_until` fields for backward compat. The + // new `roles` array is always present so clients can opt in to + // the multi-role view. + let plans = st.store.get_all_for_metric(&metric); + if plans.is_empty() { + return (StatusCode::NOT_FOUND, format!("plan not found for metric {metric:?}")) + .into_response(); } + let roles: Vec = plans + .iter() + .map(|(role, plan)| { + json!({ + "role": role.as_str(), + "sketch_type": plan.agent_config.sketch_type.to_string(), + "valid_until": plan.valid_until, + }) + }) + .collect(); + let first = &plans[0].1; + ( + StatusCode::OK, + Json(json!({ + "metric": metric, + "sketch_type": first.agent_config.sketch_type.to_string(), + "valid_until": first.valid_until, + "roles": roles, + })), + ) + .into_response() } async fn handle_rollback( @@ -886,18 +946,70 @@ async fn handle_rollback( ) -> impl IntoResponse { // Reset the baseline so the next POST /api/v1/plan re-runs the cost // model and establishes a fresh baseline plan for this metric. + // + // B2 (metric, role): rollback ALL roles for this metric. The + // response surfaces the per-role outcome so clients can see which + // roles had a previous plan and which were no-ops. Pre-B2 callers + // who fired a rollback on a metric got back `{rolled_back: true}` + // unconditionally for a single role; the new shape stays additive + // (carries `rolled_back: true` when ≥1 role rolled back). st.planner.reset(&metric); - match st.store.rollback(&metric) { - Ok(plan) => { - if let Ok(yaml) = generate_agent_collector_config(&plan.agent_config, &st.opamp_endpoint) { - st.opamp.push_to_role(AgentRole::Agent, RemoteConfig { - config_hash: short_hash(&yaml), yaml, - }).await; + let plans = st.store.get_all_for_metric(&metric); + if plans.is_empty() { + return ( + StatusCode::BAD_REQUEST, + format!("plan not found for metric {metric:?}"), + ) + .into_response(); + } + let mut per_role = Vec::with_capacity(plans.len()); + let mut any_rolled_back = false; + for (role, _) in plans { + match st.store.rollback(&metric, role) { + Ok(plan) => { + if let Ok(yaml) = + generate_agent_collector_config(&plan.agent_config, &st.opamp_endpoint) + { + st.opamp + .push_to_role( + AgentRole::Agent, + RemoteConfig { + config_hash: short_hash(&yaml), + yaml, + }, + ) + .await; + } + per_role.push(json!({ "role": role.as_str(), "rolled_back": true })); + any_rolled_back = true; + } + Err(e) => { + per_role.push(json!({ + "role": role.as_str(), + "rolled_back": false, + "reason": e.to_string(), + })); } - (StatusCode::OK, Json(json!({ "metric": metric, "rolled_back": true }))).into_response() } - Err(e) => (StatusCode::BAD_REQUEST, e.to_string()).into_response(), } + // Pre-B2 contract: return BAD_REQUEST when no role could roll + // back (e.g. every role's plan has no `previous` slot). The + // multi-role variants are surfaced in the `roles` array so + // callers can distinguish "rolled back N of K" cases. + let status = if any_rolled_back { + StatusCode::OK + } else { + StatusCode::BAD_REQUEST + }; + ( + status, + Json(json!({ + "metric": metric, + "rolled_back": any_rolled_back, + "roles": per_role, + })), + ) + .into_response() } async fn handle_agents(State(st): State) -> impl IntoResponse { @@ -911,16 +1023,28 @@ async fn handle_get_config( State(st): State, Path(metric): Path, ) -> impl IntoResponse { - match st.store.get(&metric) { - Ok(plan) => match generate_agent_collector_config(&plan.agent_config, &st.opamp_endpoint) { - Ok(yaml) => ( - StatusCode::OK, - [("content-type", "application/yaml")], - yaml, - ).into_response(), - Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), - }, - Err(e) => (StatusCode::NOT_FOUND, e.to_string()).into_response(), + // B2 (metric, role): the legacy `generate_agent_collector_config` + // emits a SINGLE-pipeline YAML — when a metric has multiple roles, + // pick the FIRST registered role's plan. The 5-sketch routing- + // connector emit path (the `USE_TYPED_STAGE_SPLIT` typed pipeline) + // is the supported multi-role wire shape; this endpoint stays + // legacy-compat by picking one role's plan. + let plan = st.store.get_all_for_metric(&metric).into_iter().next(); + let Some((_role, plan)) = plan else { + return ( + StatusCode::NOT_FOUND, + format!("plan not found for metric {metric:?}"), + ) + .into_response(); + }; + match generate_agent_collector_config(&plan.agent_config, &st.opamp_endpoint) { + Ok(yaml) => ( + StatusCode::OK, + [("content-type", "application/yaml")], + yaml, + ) + .into_response(), + Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()).into_response(), } } @@ -1049,14 +1173,20 @@ async fn emit_bootstrap_typed( // 1. Resolve the metric this bootstrap should target. // When the agent has a prior pinned assignment we honour it // (pre-existing on_connect contract). The replanner's - // `agent_to_metric()` is the source of truth for this mapping. + // `agent_to_metrics()` is the source of truth for this mapping. + // + // B2 (metric, role): an agent may pin multiple `(metric, role)` + // pairs. The bootstrap returns a single edge YAML, so we pick + // the FIRST pair's metric — the 5-sketch routing-connector + // pipeline emitted below covers every metric in the registry, + // not just this one. let pinned_metric: Option = if let Some(aid) = pinned_agent_id { st.replanner - .agent_to_metric() + .agent_to_metrics() .read() .await .get(aid) - .cloned() + .and_then(|v| v.first().map(|(m, _)| m.clone())) } else { None }; @@ -1095,12 +1225,20 @@ async fn emit_bootstrap_typed( // 2-3. Walk candidates: first metric that pre-populated the // workload store AND binds via the typed path provides the // base edge_cfg shape. + // + // B2 (metric, role): a metric may have multiple roles + // registered; we walk every role's workload entry until one + // binds. The Sum-shaped roles (raw passthrough) decline the + // typed bind, so for `http_requests_total` the + // Quantile-shaped role on `http_requests_total_latency_ms` + // stays the source of the edge config. let mut chosen: Option<(String, crate::sketch_algebra::PhysicalExpr)> = None; - for cand in &candidates { - let Some((wl, _wc)) = st.workload_store.get(cand) else { continue }; - if let Some(expr) = optimizer::rules::bind_workload_typed(&wl) { - chosen = Some((cand.clone(), expr)); - break; + 'outer: for cand in &candidates { + for (_, wl, _) in st.workload_store.get_all_for_metric(cand) { + if let Some(expr) = optimizer::rules::bind_workload_typed(&wl) { + chosen = Some((cand.clone(), expr)); + break 'outer; + } } } let (metric, physical_expr) = chosen.ok_or_else(|| { @@ -1202,18 +1340,63 @@ async fn handle_plan_diff( State(st): State, Path(metric): Path, ) -> impl IntoResponse { - match st.store.diff(&metric) { - Ok(Some(diff)) => (StatusCode::OK, Json(json!({ - "metric": metric, + // B2 (metric, role): a metric may carry multiple roles; surface a + // per-role `roles` array. The top-level `has_diff` is true iff at + // least one role has a diff. Pre-B2 single-role clients see + // `has_diff` and the `diff` field of the first role with one; + // the new shape stays additive (no URL change, response keys + // preserved). + let plans = st.store.get_all_for_metric(&metric); + if plans.is_empty() { + return ( + StatusCode::NOT_FOUND, + format!("plan not found for metric {metric:?}"), + ) + .into_response(); + } + let mut roles: Vec = Vec::with_capacity(plans.len()); + let mut first_diff: Option = None; + let mut any_has_diff = false; + for (role, _) in plans { + match st.store.diff(&metric, role) { + Ok(Some(diff)) => { + let diff_json = serde_json::to_value(&diff).unwrap_or(serde_json::Value::Null); + if first_diff.is_none() { + first_diff = Some(diff_json.clone()); + } + any_has_diff = true; + roles.push(json!({ + "role": role.as_str(), + "has_diff": true, + "diff": diff_json, + })); + } + Ok(None) => { + roles.push(json!({ "role": role.as_str(), "has_diff": false })); + } + Err(e) => { + roles.push(json!({ + "role": role.as_str(), + "error": e.to_string(), + })); + } + } + } + let body = if any_has_diff { + json!({ + "metric": metric, "has_diff": true, - "diff": diff, - }))).into_response(), - Ok(None) => (StatusCode::OK, Json(json!({ - "metric": metric, + "diff": first_diff, + "roles": roles, + }) + } else { + json!({ + "metric": metric, "has_diff": false, - }))).into_response(), - Err(e) => (StatusCode::NOT_FOUND, e.to_string()).into_response(), - } + "roles": roles, + }) + }; + (StatusCode::OK, Json(body)).into_response() } /// Returns the current EMA cost model state — blended benchmark + observed costs @@ -1468,7 +1651,7 @@ mod api_tests { repeat_every: None, accuracy_sla: 0.01, latency_sla: None, sketch_type_override: None, exact_required: false, quantiles: vec![], }; - st.store.set("m", RulesPlanner::new().plan(&wl)); + st.store.set("m", control_plane::workload::AggRole::Quantile, RulesPlanner::new().plan(&wl)); let req = Request::builder() .method("POST").uri("/api/v1/plan/m/rollback") .body(Body::empty()).unwrap(); @@ -1699,8 +1882,14 @@ mod api_tests { let wl = analyzer.analyze(spec).unwrap(); let wc = types::WorkloadCharacteristics::default(); let plan = planner.plan(&wl, Some(&wc)); - plan_store.set("http_latency", plan); - workload_store.set("http_latency", wl, wc); + // B2 (metric, role): pre-populate using the same role the + // on_connect callback's `derive_agg_role(entry)` will compute + // for this test's workloads.yaml entry (no query_string + no + // sketch_family_override → AggRole::Other). Without matching + // the role, `push_config_to_agent`'s workload_store.get + // returns None and the on_connect path silently bails. + plan_store.set("http_latency", control_plane::workload::AggRole::Other, plan); + workload_store.set("http_latency", control_plane::workload::AggRole::Other, wl, wc); // Build replanner and late-binding cells. let replanner_cell: Arc>>> = @@ -1730,7 +1919,8 @@ mod api_tests { if !pushed { if let Some(registry) = reg.read().await.as_ref() { if let Some(entry) = registry.first_for_role("agent") { - r.register_agent(&aid, &entry.metric_name).await; + let role = control_plane::workload::derive_agg_role(entry); + r.register_agent(&aid, &entry.metric_name, role).await; r.push_config_to_agent(&aid).await; } } @@ -1829,8 +2019,8 @@ mod api_tests { let wl = analyzer.analyze(spec).unwrap(); let wc = types::WorkloadCharacteristics::default(); let plan = planner.plan(&wl, Some(&wc)); - plan_store.set("metric_a", plan); - workload_store.set("metric_a", wl, wc); + plan_store.set("metric_a", control_plane::workload::AggRole::Quantile, plan); + workload_store.set("metric_a", control_plane::workload::AggRole::Quantile, wl, wc); let replanner = Arc::new(Replanner::new( Arc::clone(&planner), @@ -1849,8 +2039,12 @@ mod api_tests { tokio::time::sleep(Duration::from_millis(100)).await; // Register agent-a for metric_a, agent-b is NOT registered for metric_a. - replanner.register_agent("agent-a", "metric_a").await; - replanner.register_agent("agent-b", "metric_b").await; + replanner + .register_agent("agent-a", "metric_a", control_plane::workload::AggRole::Quantile) + .await; + replanner + .register_agent("agent-b", "metric_b", control_plane::workload::AggRole::Quantile) + .await; // Trigger replan for metric_a. let ok = replanner.replan_metric("metric_a").await; @@ -2058,8 +2252,8 @@ mod api_tests { let wl = analyzer.analyze(spec).expect("analyze"); let wc = types::WorkloadCharacteristics::default(); let plan = state.planner.plan(&wl, Some(&wc)); - state.store.set(metric, plan); - state.workload_store.set(metric, wl, wc); + state.store.set(metric, control_plane::workload::AggRole::Quantile, plan); + state.workload_store.set(metric, control_plane::workload::AggRole::Quantile, wl, wc); // 4. Swap in the populated registry. state.workload_registry = registry; @@ -2276,8 +2470,8 @@ mod api_tests { let wl = analyzer.analyze(spec).expect("analyze"); let wc = types::WorkloadCharacteristics::default(); let plan = state.planner.plan(&wl, Some(&wc)); - state.store.set(*m, plan); - state.workload_store.set(*m, wl, wc); + state.store.set(*m, control_plane::workload::AggRole::Quantile, plan); + state.workload_store.set(*m, control_plane::workload::AggRole::Quantile, wl, wc); } // 4. Swap in the populated registry. @@ -2471,8 +2665,9 @@ mod api_tests { let wc = types::WorkloadCharacteristics::default(); let plan = state.planner.plan(&wl, Some(&wc)); let metric_name = wl.metric_name.clone(); - state.store.set(&metric_name, plan); - state.workload_store.set(&metric_name, wl, wc); + let role = control_plane::workload::derive_agg_role(entry); + state.store.set(&metric_name, role, plan); + state.workload_store.set(&metric_name, role, wl, wc); } } diff --git a/control_plane/src/metrics_exposer.rs b/control_plane/src/metrics_exposer.rs index 72d33730..23786306 100644 --- a/control_plane/src/metrics_exposer.rs +++ b/control_plane/src/metrics_exposer.rs @@ -198,8 +198,13 @@ impl MetricsRegistry { // ever observed; without this a re-plan would leave the old // plan_id label permanently emitting a stale value. self.active_plan_id.reset(); - for metric in plan_store.metrics() { - let Ok(plan) = plan_store.get(&metric) else { + // B2 (metric, role): iterate per-pair so each role's plan + // emits its own `(metric, plan_id)` gauge value. The metric + // label retains the un-decorated metric name (pre-B2 wire + // shape) — a metric with multiple roles surfaces multiple + // active_plan_id rows under the same metric label. + for (metric, role) in plan_store.keys() { + let Ok(plan) = plan_store.get(&metric, role) else { continue; }; // Stable hash of the plan's debug repr — good enough for @@ -207,9 +212,12 @@ impl MetricsRegistry { let mut hasher = DefaultHasher::new(); // Cover the fields the planner actually changes per // re-plan: agent sketch+mode+delta+grouping, and valid_until - // (to catch refresh-only re-plans). + // (to catch refresh-only re-plans). Role is folded in so + // distinct-role plans produce distinct ids even when their + // agent_config fields happen to coincide. format!( - "{:?}|{:?}|{:?}|{:?}|{:?}|{:?}", + "{:?}|{:?}|{:?}|{:?}|{:?}|{:?}|{:?}", + role, plan.agent_config.sketch_type, plan.agent_config.mode, plan.agent_config.delta_transmission, @@ -457,10 +465,16 @@ mod tests { } } + use crate::workload::AggRole; let plan_store = Arc::new(PlanStore::new()); - plan_store.set("http_requests_total", make_plan(SketchType::DDSketch, 600)); + plan_store.set( + "http_requests_total", + AggRole::Quantile, + make_plan(SketchType::DDSketch, 600), + ); plan_store.set( "http_requests_total_latency_ms", + AggRole::Quantile, make_plan(SketchType::HLL, 600), ); @@ -483,7 +497,11 @@ mod tests { ); // Re-plan with a different sketch must change the plan_id label. let before = text.clone(); - plan_store.set("http_requests_total", make_plan(SketchType::KLL, 600)); + plan_store.set( + "http_requests_total", + AggRole::Quantile, + make_plan(SketchType::KLL, 600), + ); registry.refresh_plan_ids(&plan_store); let mfs = registry.registry.gather(); let mut buf = Vec::new(); diff --git a/control_plane/src/replan.rs b/control_plane/src/replan.rs index 7b81c710..5ecf5aa5 100644 --- a/control_plane/src/replan.rs +++ b/control_plane/src/replan.rs @@ -33,6 +33,7 @@ use crate::optimizer::{cost as cost_model, rules}; use crate::physical::stage_split; use crate::store::{PlanStore, WorkloadStore}; use crate::types::QueryWorkload; +use crate::workload::AggRole; fn short_hash(s: &str) -> String { use std::collections::hash_map::DefaultHasher; @@ -69,9 +70,13 @@ pub struct Replanner { /// path still works — it just skips the workload-registry archive /// extension and only adds the freshness probes. workload_registry: Option>, - /// Maps agent_id → metric_name so violation callbacks can look up which - /// metric a particular agent is serving. - agent_to_metric: Arc>>, + /// Maps `agent_id → Vec<(metric_name, role)>` so violation callbacks + /// can look up which `(metric, role)` pairs a particular agent is + /// serving. **B2 restructure**: one agent can serve multiple + /// `(metric, role)` pairs (e.g. an agent handling all three of + /// `http_requests_total`'s roles registered by `mvp-workload.yaml` + /// entries 2/3/4). The vector preserves insertion order. + agent_to_metrics: Arc>>>, } impl Replanner { @@ -92,7 +97,7 @@ impl Replanner { opamp_endpoint: opamp_endpoint.into(), backend_client: None, workload_registry: None, - agent_to_metric: Arc::new(RwLock::new(HashMap::new())), + agent_to_metrics: Arc::new(RwLock::new(HashMap::new())), } } @@ -119,25 +124,40 @@ impl Replanner { // ── Agent registry ──────────────────────────────────────────────────────── - /// Record that `agent_id` is serving `metric`. Called from `handle_plan` - /// after pushing configs so violations can be mapped back to a metric. - pub async fn register_agent(&self, agent_id: impl Into, metric: impl Into) { - self.agent_to_metric - .write() - .await - .insert(agent_id.into(), metric.into()); + /// Record that `agent_id` is serving `(metric, role)`. Called from + /// `handle_plan` after pushing configs so violations can be mapped + /// back. Idempotent: re-registering the same `(metric, role)` for + /// the same agent leaves the vector unchanged (dedup). + /// + /// **B2 contract**: an agent may serve MULTIPLE `(metric, role)` + /// pairs concurrently (the controller may push a single agent the + /// edge YAML for every role of every metric it owns). Each call + /// appends a new pair if not already present; the inverse + /// [`unregister_agent`] drops all of them at once. + pub async fn register_agent( + &self, + agent_id: impl Into, + metric: impl Into, + role: AggRole, + ) { + let key = (metric.into(), role); + let mut map = self.agent_to_metrics.write().await; + let entry = map.entry(agent_id.into()).or_default(); + if !entry.contains(&key) { + entry.push(key); + } } - /// Remove the mapping for a disconnected agent. + /// Remove every `(metric, role)` mapping for a disconnected agent. pub async fn unregister_agent(&self, agent_id: &str) { - self.agent_to_metric.write().await.remove(agent_id); + self.agent_to_metrics.write().await.remove(agent_id); } - /// Returns a read-only reference to the agent→metric mapping so that - /// callers (e.g. the on_connect callback) can check if an agent has a - /// prior assignment. - pub fn agent_to_metric(&self) -> &Arc>> { - &self.agent_to_metric + /// Returns a read-only reference to the agent→`(metric, role)` list + /// mapping so callers (e.g. the on_connect callback) can check if an + /// agent has a prior assignment. + pub fn agent_to_metrics(&self) -> &Arc>>> { + &self.agent_to_metrics } // ── Config push helpers ────────────────────────────────────────────────── @@ -164,8 +184,13 @@ impl Replanner { /// request (workload missing from store, `bind_workload_typed` /// declines the shape, no `Edge` entry, emit failure) — caller /// then falls back to the legacy emitter. - fn try_emit_typed_edge_yaml(&self, metric: &str, agent_id: &str) -> Option { - let (workload, _wc) = self.workload_store.get(metric)?; + fn try_emit_typed_edge_yaml( + &self, + metric: &str, + role: AggRole, + agent_id: &str, + ) -> Option { + let (workload, _wc) = self.workload_store.get(metric, role)?; self.try_emit_typed_edge_yaml_for_workload(&workload, agent_id) } @@ -255,32 +280,43 @@ impl Replanner { /// the OpAMP-on-connect side of the typed emit so reconnecting /// agents receive the same routed YAML as fresh-connect agents. pub async fn push_config_to_agent(&self, agent_id: &str) -> bool { - let metric = self.agent_to_metric.read().await.get(agent_id).cloned(); - let Some(metric) = metric else { return false }; + // B2: an agent may serve multiple `(metric, role)` pairs. Push + // the config for the FIRST pair on connect — same shape as the + // pre-B2 single-mapping path. (The 5-sketch routing-connector + // edge YAML, once `metric_to_family` is populated by + // `collect_metric_to_family`, carries pipelines for every + // metric+role anyway, so a single push covers all of them.) + let pair = self + .agent_to_metrics + .read() + .await + .get(agent_id) + .and_then(|v| v.first().cloned()); + let Some((metric, role)) = pair else { return false }; - let Ok(plan) = self.plan_store.get(&metric) else { + let Ok(plan) = self.plan_store.get(&metric, role) else { return false; }; let yaml = if stage_split::typed_stage_split_enabled() { - match self.try_emit_typed_edge_yaml(&metric, agent_id) { + match self.try_emit_typed_edge_yaml(&metric, role, agent_id) { Some(y) => { info!( - agent = agent_id, metric = %metric, bytes = y.len(), + agent = agent_id, metric = %metric, role = %role, bytes = y.len(), "[USE_TYPED_STAGE_SPLIT] pushed typed edge YAML on connect" ); y } None => { warn!( - agent = agent_id, metric = %metric, + agent = agent_id, metric = %metric, role = %role, "[USE_TYPED_STAGE_SPLIT] typed emit failed on connect; \ falling back to legacy generate_agent_collector_config" ); match generate_agent_collector_config(&plan.agent_config, &self.opamp_endpoint) { Ok(y) => y, Err(_) => { - warn!(agent = agent_id, metric = %metric, "failed to generate agent config on connect"); + warn!(agent = agent_id, metric = %metric, role = %role, "failed to generate agent config on connect"); return false; } } @@ -290,7 +326,7 @@ impl Replanner { match generate_agent_collector_config(&plan.agent_config, &self.opamp_endpoint) { Ok(y) => y, Err(_) => { - warn!(agent = agent_id, metric = %metric, "failed to generate agent config on connect"); + warn!(agent = agent_id, metric = %metric, role = %role, "failed to generate agent config on connect"); return false; } } @@ -305,21 +341,44 @@ impl Replanner { }, ) .await; - info!(agent = agent_id, metric = %metric, "pushed config to reconnecting agent"); + info!(agent = agent_id, metric = %metric, role = %role, "pushed config to reconnecting agent"); true } // ── Re-plan helpers ─────────────────────────────────────────────────────── - /// Re-plans a single metric and pushes updated configs. - /// Returns `true` if re-planning succeeded, `false` if the metric is unknown. + /// Re-plans every role registered for `metric` and pushes updated + /// configs. Returns `true` if at least one role was re-planned, + /// `false` if the metric has no roles registered at all. + /// + /// **B2 wrapper**: a single metric may carry multiple `(metric, role)` + /// pairs; this function loops over them and delegates per-role to + /// [`Self::replan_metric_role`]. Callers that only want to re-plan + /// a single role should call `replan_metric_role` directly. pub async fn replan_metric(&self, metric: &str) -> bool { - let Some((workload, wc)) = self.workload_store.get(metric) else { + let pairs = self.workload_store.get_all_for_metric(metric); + if pairs.is_empty() { warn!(metric, "replan requested but workload not found in store"); return false; + } + let mut any = false; + for (role, _, _) in pairs { + if self.replan_metric_role(metric, role).await { + any = true; + } + } + any + } + + /// Re-plans a single `(metric, role)` pair and pushes updated configs. + /// Returns `true` on success, `false` if the pair is unknown. + pub async fn replan_metric_role(&self, metric: &str, role: AggRole) -> bool { + let Some((workload, wc)) = self.workload_store.get(metric, role) else { + warn!(metric, role = %role, "replan requested but workload not found in store"); + return false; }; - info!(metric, "re-planning metric"); + info!(metric, role = %role, "re-planning metric+role"); // Reset the baseline so the cost model runs fresh rather than returning the // previously established baseline — the whole point of a re-plan is to @@ -327,21 +386,23 @@ impl Replanner { self.planner.reset(metric); let mut plan = self.planner.plan(&workload, Some(&wc)); plan.precompute = build_precompute_engine_jobs(&workload, "backend:4317"); - self.plan_store.set(metric, plan.clone()); + self.plan_store.set(metric, role, plan.clone()); - // Push agent config only to agents registered for this specific metric, - // rather than broadcasting to all agent-role collectors. Same gate - // as `push_config_to_agent` — typed path on, legacy fallback on - // emit failure or when the gate is off. + // Push agent config only to agents registered for this specific + // `(metric, role)` pair, rather than broadcasting to all + // agent-role collectors. Same gate as `push_config_to_agent` + // — typed path on, legacy fallback on emit failure or when the + // gate is off. // // Issue #2: emit per-agent inside the push loop so each agent's // opamp `X-Agent-ID` header carries its actual id (the agent // re-presents this header after the controller-pushed config // triggers a Docker restart). - let agents = self.agent_to_metric.read().await; + let key = (metric.to_string(), role); + let agents = self.agent_to_metrics.read().await; let target_agents: Vec = agents .iter() - .filter(|(_, m)| m.as_str() == metric) + .filter(|(_, pairs)| pairs.contains(&key)) .map(|(id, _)| id.clone()) .collect(); drop(agents); @@ -418,34 +479,41 @@ impl Replanner { true } - /// Re-plans all metrics whose `valid_until` has already passed. + /// Re-plans every `(metric, role)` pair whose `valid_until` has + /// already passed. pub async fn replan_expired(&self) { let expired = self.plan_store.expired(chrono::Utc::now()); if expired.is_empty() { return; } - info!(count = expired.len(), "re-planning expired metrics"); - for metric in expired { - self.replan_metric(&metric).await; + info!(count = expired.len(), "re-planning expired (metric, role) pairs"); + for (metric, role) in expired { + self.replan_metric_role(&metric, role).await; } } - /// Called from the violation callback. Looks up the metric served by - /// `agent_id` and triggers an immediate re-plan. + /// Called from the violation callback. Looks up the `(metric, role)` + /// pairs served by `agent_id` and triggers an immediate re-plan of + /// each one. pub async fn handle_violation(&self, agent_id: &str) { - let metric = self.agent_to_metric.read().await.get(agent_id).cloned(); - match metric { - Some(m) => { - info!(agent = agent_id, metric = %m, "SLA violation → triggering re-plan"); - self.replan_metric(&m).await; - } - None => { - warn!( - agent = agent_id, - "SLA violation but no metric mapping found; re-planning all expired" - ); - self.replan_expired().await; - } + let pairs = self + .agent_to_metrics + .read() + .await + .get(agent_id) + .cloned() + .unwrap_or_default(); + if pairs.is_empty() { + warn!( + agent = agent_id, + "SLA violation but no (metric, role) mapping found; re-planning all expired" + ); + self.replan_expired().await; + return; + } + for (metric, role) in pairs { + info!(agent = agent_id, metric = %metric, role = %role, "SLA violation → triggering re-plan"); + self.replan_metric_role(&metric, role).await; } } @@ -554,13 +622,13 @@ mod tests { async fn replan_known_metric_updates_plan_store() { let r = make_replanner(); let (wl, wc) = test_workload("latency"); - r.workload_store.set("latency", wl, wc); - r.plan_store.set("latency", make_plan()); + r.workload_store.set("latency", AggRole::Quantile, wl, wc); + r.plan_store.set("latency", AggRole::Quantile, make_plan()); let ok = r.replan_metric("latency").await; assert!(ok); // Plan store should now have a new entry (valid_until in the future). - let updated = r.plan_store.get("latency").unwrap(); + let updated = r.plan_store.get("latency", AggRole::Quantile).unwrap(); assert!(updated.valid_until > Utc::now()); } @@ -568,22 +636,22 @@ mod tests { async fn replan_expired_replans_only_expired() { let r = make_replanner(); let (wl, wc) = test_workload("old"); - r.workload_store.set("old", wl, wc); + r.workload_store.set("old", AggRole::Quantile, wl, wc); // Insert an already-expired plan. let mut expired_plan = make_plan(); expired_plan.valid_until = Utc::now() - chrono::Duration::seconds(60); - r.plan_store.set("old", expired_plan); + r.plan_store.set("old", AggRole::Quantile, expired_plan); // Insert a still-active plan for "active". let (awl, awc) = test_workload("active"); - r.workload_store.set("active", awl, awc); - r.plan_store.set("active", make_plan()); + r.workload_store.set("active", AggRole::Quantile, awl, awc); + r.plan_store.set("active", AggRole::Quantile, make_plan()); r.replan_expired().await; // "old" should now have a freshly computed plan. - let old_plan = r.plan_store.get("old").unwrap(); + let old_plan = r.plan_store.get("old", AggRole::Quantile).unwrap(); assert!(old_plan.valid_until > Utc::now()); } @@ -591,25 +659,129 @@ mod tests { async fn register_then_violation_replans_correct_metric() { let r = make_replanner(); let (wl, wc) = test_workload("req_rate"); - r.workload_store.set("req_rate", wl, wc); - r.plan_store.set("req_rate", make_plan()); + r.workload_store + .set("req_rate", AggRole::Quantile, wl, wc); + r.plan_store.set("req_rate", AggRole::Quantile, make_plan()); - r.register_agent("agent-1", "req_rate").await; + r.register_agent("agent-1", "req_rate", AggRole::Quantile).await; r.handle_violation("agent-1").await; // Plan should have been refreshed. - assert!(r.plan_store.get("req_rate").is_ok()); + assert!(r.plan_store.get("req_rate", AggRole::Quantile).is_ok()); } #[tokio::test] async fn unregister_removes_mapping() { let r = make_replanner(); - r.register_agent("a1", "m").await; + r.register_agent("a1", "m", AggRole::Quantile).await; r.unregister_agent("a1").await; // After unregister, handle_violation falls back to replan_expired (no-op). r.handle_violation("a1").await; // should not panic } + // ── B2 multi-role regression tests ──────────────────────────────────────── + + /// A metric with two different `(metric, role)` registrations + /// keeps both plans live after replan. Pre-B2 the store collapsed + /// them onto one key and the second replan would overwrite the + /// first; this regression test pins the new contract. + #[tokio::test] + async fn replan_multi_role_metric_updates_both_plans() { + let r = make_replanner(); + let (wl_q, wc_q) = test_workload("http_requests_total"); + let mut wl_s = wl_q.clone(); + wl_s.aggregations = vec![AggType::Quantile]; // analyzer-shaped (test fixture) + let wc_s = wc_q.clone(); + + r.workload_store + .set("http_requests_total", AggRole::Quantile, wl_q, wc_q); + r.workload_store + .set("http_requests_total", AggRole::Sum, wl_s, wc_s); + r.plan_store + .set("http_requests_total", AggRole::Quantile, make_plan()); + r.plan_store + .set("http_requests_total", AggRole::Sum, make_plan()); + + // `replan_metric` is the wrapper that loops over every role + // registered for the metric. + let ok = r.replan_metric("http_requests_total").await; + assert!(ok, "wrapper replan_metric should succeed for ≥1 role"); + + // Both roles' plans persist independently. + assert!(r + .plan_store + .get("http_requests_total", AggRole::Quantile) + .is_ok()); + assert!(r.plan_store.get("http_requests_total", AggRole::Sum).is_ok()); + } + + /// Targeted single-role replan via [`Replanner::replan_metric_role`] + /// only touches the specified role's plan and leaves the other + /// role's plan unchanged. + #[tokio::test] + async fn replan_metric_role_only_touches_target_role() { + let r = make_replanner(); + let (wl, wc) = test_workload("m"); + r.workload_store + .set("m", AggRole::Quantile, wl.clone(), wc.clone()); + r.workload_store.set("m", AggRole::Sum, wl, wc); + + // Make the Sum-role plan expired and Quantile plan fresh. + let mut sum_plan = make_plan(); + sum_plan.valid_until = Utc::now() - chrono::Duration::seconds(60); + r.plan_store.set("m", AggRole::Sum, sum_plan); + let fresh = make_plan(); + let fresh_ts = fresh.valid_until; + r.plan_store.set("m", AggRole::Quantile, fresh); + + let ok = r.replan_metric_role("m", AggRole::Sum).await; + assert!(ok); + + // The Quantile plan stays untouched (same valid_until as + // before the replan). + let q = r.plan_store.get("m", AggRole::Quantile).unwrap(); + assert_eq!(q.valid_until, fresh_ts); + // Sum has been re-planned (new valid_until in the future). + let s = r.plan_store.get("m", AggRole::Sum).unwrap(); + assert!(s.valid_until > Utc::now()); + } + + /// One agent serving multiple `(metric, role)` pairs receives a + /// re-plan for every one of them on violation. + #[tokio::test] + async fn agent_serving_multiple_roles_triggers_per_role_replan() { + let r = make_replanner(); + let (wl, wc) = test_workload("m"); + r.workload_store + .set("m", AggRole::Quantile, wl.clone(), wc.clone()); + r.workload_store.set("m", AggRole::Sum, wl, wc); + r.plan_store.set("m", AggRole::Quantile, make_plan()); + r.plan_store.set("m", AggRole::Sum, make_plan()); + + // One agent serves both roles. + r.register_agent("agent-1", "m", AggRole::Quantile).await; + r.register_agent("agent-1", "m", AggRole::Sum).await; + + // Sanity: agent_to_metrics() carries both pairs in order. + let pairs = r + .agent_to_metrics() + .read() + .await + .get("agent-1") + .cloned() + .unwrap(); + assert_eq!( + pairs, + vec![ + ("m".to_string(), AggRole::Quantile), + ("m".to_string(), AggRole::Sum) + ] + ); + + // Handle violation — both roles should re-plan without panic. + r.handle_violation("agent-1").await; + } + // ── Typed-emit path tests ───────────────────────────────────────────────── // // These tests exercise the `USE_TYPED_STAGE_SPLIT`-gated emit path @@ -662,11 +834,11 @@ mod tests { let r = make_replanner(); let (wl, wc) = test_workload("latency"); - r.workload_store.set("latency", wl, wc); - r.plan_store.set("latency", make_plan()); + r.workload_store.set("latency", AggRole::Quantile, wl, wc); + r.plan_store.set("latency", AggRole::Quantile, make_plan()); let yaml = r - .try_emit_typed_edge_yaml("latency", "test-agent") + .try_emit_typed_edge_yaml("latency", AggRole::Quantile, "test-agent") .expect("typed emit should succeed for a quantile workload"); // gorillas3 — archive-tier write to MinIO. Without this the @@ -717,12 +889,12 @@ mod tests { let r = make_replanner(); let (wl, wc) = test_workload("latency"); - r.workload_store.set("latency", wl, wc); - r.plan_store.set("latency", make_plan()); + r.workload_store.set("latency", AggRole::Quantile, wl, wc); + r.plan_store.set("latency", AggRole::Quantile, make_plan()); // Drive the legacy emitter directly — same code // `push_config_to_agent` runs when the gate is off. - let plan = r.plan_store.get("latency").unwrap(); + let plan = r.plan_store.get("latency", AggRole::Quantile).unwrap(); let yaml = generate_agent_collector_config(&plan.agent_config, &r.opamp_endpoint) .expect("legacy emit should succeed"); diff --git a/control_plane/src/store/mod.rs b/control_plane/src/store/mod.rs index bf292f18..33a84043 100644 --- a/control_plane/src/store/mod.rs +++ b/control_plane/src/store/mod.rs @@ -1,11 +1,17 @@ pub mod workload; -pub use workload::WorkloadStore; +pub use workload::{WorkloadKey, WorkloadStore}; use chrono::{DateTime, Utc}; use std::collections::HashMap; use std::sync::RwLock; use crate::types::CollectionPlan; +use crate::workload::AggRole; + +/// Composite key `(metric_name, role)` for the plan store — same +/// shape as [`WorkloadKey`]. See [`crate::workload::AggRole`] for +/// the B2 restructure rationale. +pub type PlanKey = (String, AggRole); #[derive(Debug)] pub struct PlanStore { @@ -14,7 +20,7 @@ pub struct PlanStore { #[derive(Debug, Default)] struct StoreInner { - entries: HashMap, + entries: HashMap, } #[derive(Debug, Clone)] @@ -26,10 +32,10 @@ struct Entry { #[derive(Debug, thiserror::Error)] pub enum StoreError { - #[error("plan not found for metric {0:?}")] - NotFound(String), - #[error("no previous plan for metric {0:?}")] - NoPrevious(String), + #[error("plan not found for metric {0:?} role {1:?}")] + NotFound(String, AggRole), + #[error("no previous plan for metric {0:?} role {1:?}")] + NoPrevious(String, AggRole), } impl PlanStore { @@ -39,13 +45,16 @@ impl PlanStore { } } - pub fn set(&self, metric: impl Into, plan: CollectionPlan) { - let metric = metric.into(); + /// Insert or replace the plan for `(metric, role)`. When a prior + /// plan exists, it is preserved as `previous` so [`Self::rollback`] + /// can restore it. + pub fn set(&self, metric: impl Into, role: AggRole, plan: CollectionPlan) { + let key: PlanKey = (metric.into(), role); let mut inner = self.inner.write().unwrap(); - match inner.entries.get_mut(&metric) { + match inner.entries.get_mut(&key) { None => { inner.entries.insert( - metric, + key, Entry { current: plan, previous: None, @@ -62,37 +71,70 @@ impl PlanStore { } } - pub fn get(&self, metric: &str) -> Result { + pub fn get(&self, metric: &str, role: AggRole) -> Result { self.inner .read() .unwrap() .entries - .get(metric) + .get(&(metric.to_string(), role)) .map(|e| e.current.clone()) - .ok_or_else(|| StoreError::NotFound(metric.to_string())) + .ok_or_else(|| StoreError::NotFound(metric.to_string(), role)) + } + + /// Returns all `(role, plan)` pairs for `metric` across every role. + /// Empty vec when no role has a plan registered for the metric. + pub fn get_all_for_metric(&self, metric: &str) -> Vec<(AggRole, CollectionPlan)> { + self.inner + .read() + .unwrap() + .entries + .iter() + .filter(|((m, _), _)| m == metric) + .map(|((_, role), e)| (*role, e.current.clone())) + .collect() } - pub fn rollback(&self, metric: &str) -> Result { + pub fn rollback(&self, metric: &str, role: AggRole) -> Result { + let key: PlanKey = (metric.to_string(), role); let mut inner = self.inner.write().unwrap(); let e = inner .entries - .get_mut(metric) - .ok_or_else(|| StoreError::NotFound(metric.to_string()))?; + .get_mut(&key) + .ok_or_else(|| StoreError::NotFound(metric.to_string(), role))?; let prev = e .previous .take() - .ok_or_else(|| StoreError::NoPrevious(metric.to_string()))?; + .ok_or_else(|| StoreError::NoPrevious(metric.to_string(), role))?; e.current = prev.clone(); e.updated_at = Utc::now(); Ok(prev) } - pub fn metrics(&self) -> Vec { + /// Returns every `(metric, role)` key currently in the store. + pub fn keys(&self) -> Vec { self.inner.read().unwrap().entries.keys().cloned().collect() } - pub fn expired(&self, now: DateTime) -> Vec { + /// Returns every distinct metric name currently in the store + /// (dedup'd across roles). Used by the metrics-exposer's plan-id + /// gauge which aggregates per metric. + pub fn metrics(&self) -> Vec { + let mut out: Vec = self + .inner + .read() + .unwrap() + .entries + .keys() + .map(|(m, _)| m.clone()) + .collect(); + out.sort(); + out.dedup(); + out + } + + /// Returns every `(metric, role)` whose `valid_until` is in the past. + pub fn expired(&self, now: DateTime) -> Vec { self.inner .read() .unwrap() @@ -103,14 +145,14 @@ impl PlanStore { .collect() } - /// Returns a diff between the current and previous plan for `metric`. - /// Returns `None` if no previous plan exists. - pub fn diff(&self, metric: &str) -> Result, StoreError> { + /// Returns a diff between the current and previous plan for the + /// `(metric, role)` pair. Returns `None` if no previous plan exists. + pub fn diff(&self, metric: &str, role: AggRole) -> Result, StoreError> { let inner = self.inner.read().unwrap(); let e = inner .entries - .get(metric) - .ok_or_else(|| StoreError::NotFound(metric.to_string()))?; + .get(&(metric.to_string(), role)) + .ok_or_else(|| StoreError::NotFound(metric.to_string(), role))?; let Some(prev) = &e.previous else { return Ok(None); }; @@ -153,7 +195,6 @@ pub struct PlanDiff { mod tests { use super::*; use crate::types::*; - use std::time::Duration; fn make_plan(valid_secs: i64) -> CollectionPlan { let valid_until = Utc::now() + chrono::Duration::seconds(valid_secs); @@ -188,15 +229,18 @@ mod tests { fn set_and_get() { let s = PlanStore::new(); let plan = make_plan(600); - s.set("latency", plan.clone()); - let got = s.get("latency").unwrap(); + s.set("latency", AggRole::Quantile, plan.clone()); + let got = s.get("latency", AggRole::Quantile).unwrap(); assert_eq!(got.valid_until, plan.valid_until); } #[test] fn get_not_found() { let s = PlanStore::new(); - assert!(matches!(s.get("missing"), Err(StoreError::NotFound(_)))); + assert!(matches!( + s.get("missing", AggRole::Quantile), + Err(StoreError::NotFound(..)) + )); } #[test] @@ -204,58 +248,103 @@ mod tests { let s = PlanStore::new(); let p1 = make_plan(100); let p2 = make_plan(200); - s.set("m", p1.clone()); - s.set("m", p2.clone()); - let rolled = s.rollback("m").unwrap(); + s.set("m", AggRole::Quantile, p1.clone()); + s.set("m", AggRole::Quantile, p2.clone()); + let rolled = s.rollback("m", AggRole::Quantile).unwrap(); assert_eq!(rolled.valid_until, p1.valid_until); // After rollback, Get should return p1. - assert_eq!(s.get("m").unwrap().valid_until, p1.valid_until); + assert_eq!( + s.get("m", AggRole::Quantile).unwrap().valid_until, + p1.valid_until + ); } #[test] fn rollback_no_previous() { let s = PlanStore::new(); - s.set("m", make_plan(600)); - assert!(matches!(s.rollback("m"), Err(StoreError::NoPrevious(_)))); + s.set("m", AggRole::Quantile, make_plan(600)); + assert!(matches!( + s.rollback("m", AggRole::Quantile), + Err(StoreError::NoPrevious(..)) + )); } #[test] fn rollback_not_found() { let s = PlanStore::new(); - assert!(matches!(s.rollback("x"), Err(StoreError::NotFound(_)))); + assert!(matches!( + s.rollback("x", AggRole::Quantile), + Err(StoreError::NotFound(..)) + )); } #[test] - fn metrics_list() { + fn metrics_dedups_across_roles() { let s = PlanStore::new(); - s.set("a", make_plan(600)); - s.set("b", make_plan(600)); - let mut m = s.metrics(); - m.sort(); + s.set("a", AggRole::Quantile, make_plan(600)); + s.set("a", AggRole::Sum, make_plan(600)); + s.set("b", AggRole::Sum, make_plan(600)); + let m = s.metrics(); assert_eq!(m, vec!["a", "b"]); } #[test] fn expired() { let s = PlanStore::new(); - s.set("old", make_plan(-1)); // already expired - s.set("active", make_plan(600)); + s.set("old", AggRole::Quantile, make_plan(-1)); // already expired + s.set("active", AggRole::Sum, make_plan(600)); let exp = s.expired(Utc::now()); - assert_eq!(exp, vec!["old"]); + assert_eq!(exp, vec![("old".to_string(), AggRole::Quantile)]); + } + + #[test] + fn different_roles_for_same_metric_coexist() { + // The PlanStore's role-keyed mirror of WorkloadStore's + // same-named test. Pins the B2 contract: per-role plans + // persist independently. + let s = PlanStore::new(); + let plan_q = make_plan(600); + let plan_s = make_plan(700); + let plan_c = make_plan(800); + s.set("http_requests_total", AggRole::Quantile, plan_q.clone()); + s.set("http_requests_total", AggRole::Sum, plan_s.clone()); + s.set("http_requests_total", AggRole::Count, plan_c.clone()); + + assert_eq!( + s.get("http_requests_total", AggRole::Quantile) + .unwrap() + .valid_until, + plan_q.valid_until + ); + assert_eq!( + s.get("http_requests_total", AggRole::Sum) + .unwrap() + .valid_until, + plan_s.valid_until + ); + assert_eq!( + s.get("http_requests_total", AggRole::Count) + .unwrap() + .valid_until, + plan_c.valid_until + ); + + let all = s.get_all_for_metric("http_requests_total"); + assert_eq!(all.len(), 3); } #[test] fn concurrent_access() { use std::sync::Arc; let s = Arc::new(PlanStore::new()); - s.set("m", make_plan(600)); + s.set("m", AggRole::Quantile, make_plan(600)); let handles: Vec<_> = (0..8) .map(|_| { let s = Arc::clone(&s); std::thread::spawn(move || { for _ in 0..100 { - let _ = s.get("m"); + let _ = s.get("m", AggRole::Quantile); } }) }) diff --git a/control_plane/src/store/workload.rs b/control_plane/src/store/workload.rs index 149bc564..36767d5f 100644 --- a/control_plane/src/store/workload.rs +++ b/control_plane/src/store/workload.rs @@ -1,13 +1,24 @@ //! Persists the `QueryWorkload` + `WorkloadCharacteristics` associated with -//! each planned metric so that the re-planner can re-run `plan()` without -//! needing the original `QuerySpec` HTTP payload. +//! each planned `(metric, AggRole)` pair so that the re-planner can re-run +//! `plan()` without needing the original `QuerySpec` HTTP payload. +//! +//! **B2 full restructure**: the store is keyed by `(metric, AggRole)` +//! rather than `metric` alone. A single metric (e.g. `http_requests_total`) +//! that carries multiple PromQL shapes (sum/quantile/count) registers +//! one entry per role, each with its own plan and downstream +//! `AggregationConfig`. See [`crate::workload::AggRole`] for the +//! classification rules and the B2 PR description. use std::collections::HashMap; use std::sync::RwLock; use crate::types::{QueryWorkload, WorkloadCharacteristics}; +use crate::workload::AggRole; + +/// Composite key `(metric_name, role)` for the store. +pub type WorkloadKey = (String, AggRole); pub struct WorkloadStore { - inner: RwLock>, + inner: RwLock>, } impl WorkloadStore { @@ -17,17 +28,68 @@ impl WorkloadStore { } } - pub fn set(&self, metric: impl Into, wl: QueryWorkload, wc: WorkloadCharacteristics) { - self.inner.write().unwrap().insert(metric.into(), (wl, wc)); + /// Insert (or replace) the entry for a `(metric, role)` pair. + /// + /// **B2 contract**: prior collisions on metric alone would overwrite + /// (silently dropping all but the last YAML entry). Now collisions + /// only happen when metric AND role coincide — re-registering the + /// same `(metric, role)` is the legitimate update path (controller + /// HTTP `POST /api/v1/plan` re-issuing the same shape). + pub fn set( + &self, + metric: impl Into, + role: AggRole, + wl: QueryWorkload, + wc: WorkloadCharacteristics, + ) { + self.inner + .write() + .unwrap() + .insert((metric.into(), role), (wl, wc)); + } + + /// Returns a clone of `(workload, characteristics)` for the + /// `(metric, role)` pair if known. + pub fn get( + &self, + metric: &str, + role: AggRole, + ) -> Option<(QueryWorkload, WorkloadCharacteristics)> { + self.inner + .read() + .unwrap() + .get(&(metric.to_string(), role)) + .cloned() } - /// Returns a clone of `(workload, characteristics)` if the metric is known. - pub fn get(&self, metric: &str) -> Option<(QueryWorkload, WorkloadCharacteristics)> { - self.inner.read().unwrap().get(metric).cloned() + /// Returns every `(workload, characteristics)` pair registered for + /// `metric`, across all roles. Empty vec when nothing is registered + /// for the metric. Order is unspecified — sort if determinism matters. + pub fn get_all_for_metric( + &self, + metric: &str, + ) -> Vec<(AggRole, QueryWorkload, WorkloadCharacteristics)> { + self.inner + .read() + .unwrap() + .iter() + .filter(|((m, _), _)| m == metric) + .map(|((_, role), (wl, wc))| (*role, wl.clone(), wc.clone())) + .collect() } - pub fn remove(&self, metric: &str) { - self.inner.write().unwrap().remove(metric); + /// Returns every `(metric, role)` key currently registered. Used by + /// emit paths that walk the registry to build per-pair routing / + /// aggregation tables. + pub fn keys(&self) -> Vec { + self.inner.read().unwrap().keys().cloned().collect() + } + + pub fn remove(&self, metric: &str, role: AggRole) { + self.inner + .write() + .unwrap() + .remove(&(metric.to_string(), role)); } } @@ -57,33 +119,183 @@ mod tests { #[test] fn set_and_get() { let s = WorkloadStore::new(); - s.set("latency", wl("latency"), WorkloadCharacteristics::default()); - let (got, _) = s.get("latency").unwrap(); + s.set( + "latency", + AggRole::Quantile, + wl("latency"), + WorkloadCharacteristics::default(), + ); + let (got, _) = s.get("latency", AggRole::Quantile).unwrap(); assert_eq!(got.metric_name, "latency"); } #[test] fn unknown_metric_returns_none() { let s = WorkloadStore::new(); - assert!(s.get("nope").is_none()); + assert!(s.get("nope", AggRole::Quantile).is_none()); + } + + #[test] + fn unknown_role_for_known_metric_returns_none() { + let s = WorkloadStore::new(); + s.set("m", AggRole::Quantile, wl("m"), WorkloadCharacteristics::default()); + assert!(s.get("m", AggRole::Sum).is_none()); } #[test] - fn overwrite_replaces() { + fn overwrite_same_role_replaces() { let s = WorkloadStore::new(); - s.set("m", wl("m"), WorkloadCharacteristics::default()); + s.set("m", AggRole::Quantile, wl("m"), WorkloadCharacteristics::default()); let mut updated = wl("m"); updated.accuracy_sla = 0.05; - s.set("m", updated, WorkloadCharacteristics::default()); - let (got, _) = s.get("m").unwrap(); + s.set("m", AggRole::Quantile, updated, WorkloadCharacteristics::default()); + let (got, _) = s.get("m", AggRole::Quantile).unwrap(); assert_eq!(got.accuracy_sla, 0.05); } #[test] - fn remove_clears_entry() { + fn different_roles_for_same_metric_coexist() { + // The core B2 contract: a single metric can carry multiple roles + // and each entry persists independently of the others. + let s = WorkloadStore::new(); + let mut wl_q = wl("http_requests_total"); + wl_q.accuracy_sla = 0.01; + let mut wl_s = wl("http_requests_total"); + wl_s.accuracy_sla = 0.02; + let mut wl_c = wl("http_requests_total"); + wl_c.accuracy_sla = 0.03; + s.set( + "http_requests_total", + AggRole::Quantile, + wl_q, + WorkloadCharacteristics::default(), + ); + s.set( + "http_requests_total", + AggRole::Sum, + wl_s, + WorkloadCharacteristics::default(), + ); + s.set( + "http_requests_total", + AggRole::Count, + wl_c, + WorkloadCharacteristics::default(), + ); + + // All three persist (the pre-B2 store would have collapsed + // them onto one key, only the last survives). + assert_eq!( + s.get("http_requests_total", AggRole::Quantile) + .unwrap() + .0 + .accuracy_sla, + 0.01 + ); + assert_eq!( + s.get("http_requests_total", AggRole::Sum) + .unwrap() + .0 + .accuracy_sla, + 0.02 + ); + assert_eq!( + s.get("http_requests_total", AggRole::Count) + .unwrap() + .0 + .accuracy_sla, + 0.03 + ); + // `get_all_for_metric` surfaces all three. + let all = s.get_all_for_metric("http_requests_total"); + assert_eq!(all.len(), 3); + } + + #[test] + fn remove_clears_only_target_role() { + let s = WorkloadStore::new(); + s.set("m", AggRole::Quantile, wl("m"), WorkloadCharacteristics::default()); + s.set("m", AggRole::Sum, wl("m"), WorkloadCharacteristics::default()); + s.remove("m", AggRole::Quantile); + assert!(s.get("m", AggRole::Quantile).is_none()); + assert!(s.get("m", AggRole::Sum).is_some()); + } + + #[test] + fn three_http_requests_total_entries_persist_after_pre_pop_loop_mirror() { + // B2 regression: mirror the (analyzer → planner → + // workload_store.set) pre-pop loop from main.rs for the + // three http_requests_total entries in mvp-workload.yaml. + // Pre-B2 only the LAST entry survived; the role-keyed store + // keeps Sum + Count distinct. (The two Sum-shaped entries + // legitimately overwrite each other under the same key — + // that's the operator-update path.) + use crate::workload::{derive_agg_role, WorkloadEntry}; + let entries = vec![ + WorkloadEntry { + metric_name: "http_requests_total".into(), + query_string: Some("sum by (zone) (http_requests_total)".into()), + accuracy_sla: 0.0, + assign_to_role: "gateway".into(), + sketch_family_override: None, + target_path: None, + grouping_labels: vec!["zone".into()], + }, + WorkloadEntry { + metric_name: "http_requests_total".into(), + query_string: Some("sum by (zone) (rate(http_requests_total[5m]))".into()), + accuracy_sla: 0.01, + assign_to_role: "agent".into(), + sketch_family_override: None, + target_path: None, + grouping_labels: vec!["zone".into()], + }, + WorkloadEntry { + metric_name: "http_requests_total".into(), + query_string: Some(r#"count(http_requests_total{zone="z0"})"#.into()), + accuracy_sla: 0.0, + assign_to_role: "archive".into(), + sketch_family_override: None, + target_path: None, + grouping_labels: vec!["zone".into()], + }, + ]; + + let store = WorkloadStore::new(); + for entry in &entries { + let role = derive_agg_role(entry); + store.set( + &entry.metric_name, + role, + wl(&entry.metric_name), + WorkloadCharacteristics::default(), + ); + } + + // Both distinct roles persist after the loop (vs pre-B2: only + // the last `set` survives because the key was metric only). + let all = store.get_all_for_metric("http_requests_total"); + let roles: std::collections::HashSet<_> = all.iter().map(|(r, _, _)| *r).collect(); + assert!( + roles.contains(&crate::workload::AggRole::Sum), + "Sum-role plan must survive after the pre-pop loop; got {roles:?}" + ); + assert!( + roles.contains(&crate::workload::AggRole::Count), + "Count-role plan must survive after the pre-pop loop; got {roles:?}" + ); + } + + #[test] + fn keys_returns_all_pairs() { let s = WorkloadStore::new(); - s.set("m", wl("m"), WorkloadCharacteristics::default()); - s.remove("m"); - assert!(s.get("m").is_none()); + s.set("a", AggRole::Quantile, wl("a"), WorkloadCharacteristics::default()); + s.set("b", AggRole::Sum, wl("b"), WorkloadCharacteristics::default()); + let mut keys = s.keys(); + keys.sort(); + assert_eq!( + keys, + vec![("a".to_string(), AggRole::Quantile), ("b".to_string(), AggRole::Sum)] + ); } } diff --git a/control_plane/src/workload.rs b/control_plane/src/workload.rs index 3492a2bb..94d373bd 100644 --- a/control_plane/src/workload.rs +++ b/control_plane/src/workload.rs @@ -9,6 +9,156 @@ use tracing::{info, warn}; use crate::types::SketchType; +/// Aggregation role a single (metric, query-shape) pair plays in the planner. +/// +/// **Why this exists** (B2 full restructure): a single metric can carry +/// MULTIPLE aggregation roles when the workload YAML registers more than +/// one PromQL shape for it. The canonical case is +/// `http_requests_total`, which the MVP demo's `mvp-workload.yaml` +/// registers three times (entries 2/3/4 of [`deploy/configs/mvp-workload.yaml`]): +/// * `sum by (zone) (http_requests_total)` → [`AggRole::Sum`] +/// * `sum by (zone) (rate(http_requests_total[5m]))` → [`AggRole::Sum`] +/// (rate binds to ExactAgg(Sum)-shaped capability) +/// * `count(http_requests_total{zone="z0"})` → [`AggRole::Count`] +/// +/// Before this enum: the `WorkloadStore` was keyed by metric name alone +/// and `set(metric, …)` overwrote on collision — only the LAST entry +/// survived, so the `sum by (zone)` query (entries 2 + 3) lost its plan +/// and the data-plane refused it with `ExactAgg(Sum) capability not +/// satisfied` (the surviving plan was DDSketch from entry 1's quantile +/// shape). +/// +/// After: the store is keyed by `(metric, role)` so each shape gets its +/// own plan, its own `AggregationConfig` on the backend's streaming +/// config, and its own routing-connector pipeline. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum AggRole { + /// `quantile_over_time`, `quantile(...)`, `histogram_quantile(...)`, + /// or workload entries with `sketch_family_override: DDSketch | KLL`. + /// Routes to a quantile-shaped sketch (DDSketch / KLL). + Quantile, + /// Bare counter selector, `sum(...)`, `sum_over_time(...)`, + /// `rate(...)`, `increase(...)`. All bind to ExactAgg(Sum)-shaped + /// capability on the data plane; the streaming-config emits an + /// `aggregation_type: Sum` rather than a sketch. + Sum, + /// `count(...)`, `count_over_time(...)`, `count_distinct_over_time(...)`, + /// or workload entries with `sketch_family_override: HLL`. Routes + /// to HLL when a sketch is appropriate, otherwise to a Sum-as-count + /// exact-aggregation. + Count, + /// `topk(...)`, `topk_over_time(...)`, or workload entries with + /// `sketch_family_override: CountSketch | CountMinSketch`. Routes + /// to CountSketch / CMS-with-heap. + Topk, + /// Fallback bucket — specialized sketch families that don't fit the + /// four shapes above (e.g. CountMinSketch frequency without a topk + /// outer), or PromQL shapes the role-classifier can't recognise + /// today. Preserves "unique" semantics for the (metric, role) key + /// so multiple unrecognised entries still don't collide. + Other, +} + +impl AggRole { + /// Stable lowercase tag for routing-connector pipeline names, log + /// fields, and HTTP path discriminators. + pub fn as_str(&self) -> &'static str { + match self { + AggRole::Quantile => "quantile", + AggRole::Sum => "sum", + AggRole::Count => "count", + AggRole::Topk => "topk", + AggRole::Other => "other", + } + } +} + +impl std::fmt::Display for AggRole { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(self.as_str()) + } +} + +/// Classify a single workload entry into its `AggRole`. +/// +/// Resolution order: +/// 1. **`sketch_family_override`** wins when present: +/// * `DDSketch` / `KLL` → [`AggRole::Quantile`] +/// * `HLL` → [`AggRole::Count`] +/// * `CountSketch` → [`AggRole::Topk`] +/// * `CountMinSketch` → [`AggRole::Other`] (frequency — no +/// single canonical shape; we keep it out of `Topk` so the topk +/// variant stays semantically pure for sketch-with-heap families) +/// 2. **PromQL AST classification** by outermost-function name in +/// `query_string`. Recognised function tokens: +/// * `quantile_over_time` / `quantile` / `histogram_quantile` → [`AggRole::Quantile`] +/// * `sum` / `sum_over_time` / `rate` / `increase` → [`AggRole::Sum`] +/// * `count` / `count_over_time` / `count_distinct_over_time` → [`AggRole::Count`] +/// * `topk` / `topk_over_time` → [`AggRole::Topk`] +/// * Bare metric selector (no outer function) → [`AggRole::Sum`] +/// (matches PromQL's instant-vector semantics — a bare counter +/// sums values across time-aligned samples). +/// 3. Fallback: [`AggRole::Other`]. +/// +/// Ambiguous case decisions (documented for the B2 PR): +/// * `count_over_time(metric)` — Count (cardinality semantics). +/// The Sum-shaped alternative is rare in practice; users who want +/// it write `sum_over_time(count(...))` which classifies as Sum. +/// * `rate` / `increase` — Sum. Both bind to ExactAgg(Sum) on the +/// data plane (see `data_plane/src/precompute_engine/ingest_handler.rs`'s +/// handling of `AggKind::ExactAgg { Sum }`). +pub fn derive_agg_role(entry: &WorkloadEntry) -> AggRole { + // 1. `sketch_family_override` wins. + if let Some(family) = entry.sketch_family_override.as_ref() { + return match family { + SketchType::DDSketch | SketchType::KLL => AggRole::Quantile, + SketchType::HLL => AggRole::Count, + SketchType::CountSketch => AggRole::Topk, + SketchType::CountMinSketch => AggRole::Other, + }; + } + + // 2. PromQL AST classification — outermost-token sniff. We don't + // need a full AST walk: PromQL function calls always lead with + // `(`, so the leading identifier carries the shape. For + // nested aggregations the OUTERMOST one drives the role (it's + // the one bound to the data-plane capability). + if let Some(qs) = entry.query_string.as_ref() { + let trimmed = qs.trim_start(); + // Find the leading identifier — letters / underscores up to + // the first non-identifier char (`(`, space, `{`, etc.). + let token_end = trimmed + .find(|c: char| !(c.is_ascii_alphanumeric() || c == '_')) + .unwrap_or(trimmed.len()); + let leading = &trimmed[..token_end]; + if !leading.is_empty() { + match leading { + "quantile_over_time" | "quantile" | "histogram_quantile" => { + return AggRole::Quantile + } + "sum" | "sum_over_time" | "rate" | "irate" | "increase" => return AggRole::Sum, + "count" | "count_over_time" | "count_distinct_over_time" => return AggRole::Count, + "topk" | "topk_over_time" => return AggRole::Topk, + _ => {} + } + } + // Bare metric selector — no outer function call. The leading + // token is a metric name (or empty if the query starts with a + // brace). Treat as Sum (PromQL's default instant-vector + // interpretation aligns with Sum-shaped capability). + if !trimmed.is_empty() && !trimmed.starts_with('{') { + return AggRole::Sum; + } + } + + // 3. No query_string and no override — default to Sum (Mode-3 + // raw-passthrough entries in the workload registry typically + // declare a bare metric with `assign_to_role: archive` and no + // `query_string`). + AggRole::Other +} + /// A single workload entry from the workloads YAML file. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct WorkloadEntry { @@ -276,6 +426,141 @@ mod tests { assert_eq!(entries[5].sketch_family_override, None); } + // ── AggRole tests (B2 full restructure) ────────────────────────────── + + fn entry(metric: &str, q: Option<&str>, override_: Option) -> WorkloadEntry { + WorkloadEntry { + metric_name: metric.into(), + query_string: q.map(|s| s.to_string()), + accuracy_sla: 0.01, + assign_to_role: "agent".into(), + sketch_family_override: override_, + target_path: None, + grouping_labels: vec![], + } + } + + #[test] + fn agg_role_override_pins_quantile_for_dd_and_kll() { + assert_eq!( + derive_agg_role(&entry("m", None, Some(SketchType::DDSketch))), + AggRole::Quantile + ); + assert_eq!( + derive_agg_role(&entry("m", None, Some(SketchType::KLL))), + AggRole::Quantile + ); + } + + #[test] + fn agg_role_override_pins_count_for_hll() { + assert_eq!( + derive_agg_role(&entry("m", None, Some(SketchType::HLL))), + AggRole::Count + ); + } + + #[test] + fn agg_role_override_pins_topk_for_countsketch() { + assert_eq!( + derive_agg_role(&entry("m", None, Some(SketchType::CountSketch))), + AggRole::Topk + ); + } + + #[test] + fn agg_role_override_pins_other_for_cms() { + // CountMinSketch is a frequency estimator without an inherent + // topk shape — keep it in `Other` so the topk variant stays + // pure for sketch-with-heap families. + assert_eq!( + derive_agg_role(&entry("m", None, Some(SketchType::CountMinSketch))), + AggRole::Other + ); + } + + #[test] + fn agg_role_quantile_query_strings() { + for q in [ + "quantile_over_time(0.99, m[5m])", + "quantile(0.5, m)", + "histogram_quantile(0.99, rate(m_bucket[5m]))", + ] { + assert_eq!( + derive_agg_role(&entry("m", Some(q), None)), + AggRole::Quantile, + "query `{q}` should classify as Quantile" + ); + } + } + + #[test] + fn agg_role_sum_query_strings() { + for q in [ + "sum by (zone) (m)", + "sum_over_time(m[5m])", + "rate(m[5m])", + "increase(m[5m])", + "sum by (zone) (rate(m[5m]))", + ] { + assert_eq!( + derive_agg_role(&entry("m", Some(q), None)), + AggRole::Sum, + "query `{q}` should classify as Sum" + ); + } + } + + #[test] + fn agg_role_count_query_strings() { + for q in [ + "count(m)", + "count_over_time(m[5m])", + r#"count(m{zone="z0"})"#, + ] { + assert_eq!( + derive_agg_role(&entry("m", Some(q), None)), + AggRole::Count, + "query `{q}` should classify as Count" + ); + } + } + + #[test] + fn agg_role_topk_query_strings() { + for q in ["topk(5, m)", "topk_over_time(3, m[5m])"] { + assert_eq!( + derive_agg_role(&entry("m", Some(q), None)), + AggRole::Topk, + "query `{q}` should classify as Topk" + ); + } + } + + #[test] + fn agg_role_bare_metric_selector_is_sum() { + assert_eq!( + derive_agg_role(&entry("http_requests_total", Some("http_requests_total"), None)), + AggRole::Sum + ); + } + + #[test] + fn agg_role_no_query_string_no_override_is_other() { + assert_eq!(derive_agg_role(&entry("m", None, None)), AggRole::Other); + } + + #[test] + fn agg_role_override_beats_query_string() { + // An explicit `sketch_family_override: HLL` paired with a + // `count(...)` query — both happen to classify as Count, but + // we exercise the override priority with a deliberately + // mismatched pair (override KLL on a count query) to pin the + // override-first rule. + let e = entry("m", Some("count(m)"), Some(SketchType::KLL)); + assert_eq!(derive_agg_role(&e), AggRole::Quantile); + } + #[test] fn deserialize_sketch_family_override_lowercase_aliases() { // Lowercase / kebab-case spellings also accepted, plus the two @@ -303,6 +588,96 @@ mod tests { ); } + #[test] + fn three_synthetic_http_requests_total_entries_classify_to_two_distinct_roles() { + // Synthetic mirror of `deploy/configs/mvp-workload.yaml` + // entries 2/3/4 — proves `derive_agg_role` produces distinct + // roles for the three http_requests_total shapes. Pre-B2 the + // workload store collapsed these onto one key and only the + // last entry's plan survived; the (metric, role) keyed store + // + per-entry role classification fixes that. + let entries = vec![ + entry( + "http_requests_total", + Some("sum by (zone) (http_requests_total)"), + None, + ), + entry( + "http_requests_total", + Some("sum by (zone) (rate(http_requests_total[5m]))"), + None, + ), + entry( + "http_requests_total", + Some(r#"count(http_requests_total{zone="z0"})"#), + None, + ), + ]; + let roles: Vec = entries.iter().map(derive_agg_role).collect(); + assert_eq!(roles, vec![AggRole::Sum, AggRole::Sum, AggRole::Count]); + // The store distinguishes Sum vs Count keys, so two of the + // three entries (the two Sum-shaped ones) still collide + // under (metric, role). That's the documented behaviour — + // two YAML entries with the SAME (metric, role) overwrite, + // which is the legitimate "operator updated their workload" + // path. The fix scope is collisions across DIFFERENT shapes, + // not idempotent re-registers. + let distinct: std::collections::HashSet<_> = roles.iter().copied().collect(); + assert_eq!(distinct.len(), 2, "Sum + Count = 2 distinct roles"); + } + + #[test] + fn live_mvp_workload_yaml_assigns_three_roles_to_http_requests_total() { + // B2 full restructure regression: the live + // `deploy/configs/mvp-workload.yaml` carries THREE entries for + // `http_requests_total` (entries 2/3/4 — sum/sum+rate/count). + // Pre-B2 these collapsed onto one workload-store key and + // dropped two of the three plans, so the `sum by (zone)` + // query returned `ExactAgg(Sum) capability not satisfied`. + // + // The fix is the `(metric, role)` key + the per-entry role + // classification via `derive_agg_role`. This test pins that + // the three entries classify to two distinct roles (`Sum` for + // entries 2 + 3, `Count` for entry 4) — the multi-row keyed + // store can persist them all simultaneously. + use std::path::PathBuf; + let mut path = PathBuf::from(env!("CARGO_MANIFEST_DIR")); + path.pop(); + path.push("deploy/configs/mvp-workload.yaml"); + if !path.exists() { + // Live file not in this checkout; skip silently (matches + // the sibling override test below). + return; + } + let registry = WorkloadRegistry::load(path.to_str().unwrap()); + let http_requests_entries: Vec<&WorkloadEntry> = registry + .entries() + .iter() + .filter(|e| e.metric_name == "http_requests_total") + .collect(); + assert!( + http_requests_entries.len() >= 3, + "mvp-workload.yaml is expected to carry ≥3 entries for \ + http_requests_total (sum, sum(rate), count); got {}", + http_requests_entries.len() + ); + let roles: Vec = http_requests_entries + .iter() + .map(|e| derive_agg_role(e)) + .collect(); + // At least one Sum and at least one Count among the entries. + assert!( + roles.contains(&AggRole::Sum), + "expected ≥1 Sum-role entry among http_requests_total in \ + mvp-workload.yaml; got {roles:?}" + ); + assert!( + roles.contains(&AggRole::Count), + "expected ≥1 Count-role entry among http_requests_total in \ + mvp-workload.yaml; got {roles:?}" + ); + } + #[test] fn live_mvp_workload_yaml_loads_with_overrides() { // Smoke-test the live deploy file. Confirms entries 5–8 carry