From 2501389cb4699d132976ca42c6001029c3ca05e5 Mon Sep 17 00:00:00 2001 From: Zeying Zhu Date: Wed, 15 Apr 2026 15:27:20 -0400 Subject: [PATCH] feat: SimpleEngine re-snapshots StreamingConfig per query (PR E phase 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Makes `POST /api/v1/streaming-config` actually useful for query-time behavior. Before this PR, SimpleEngine snapshotted the startup config into `Arc` once at construction and ignored every subsequent swap — meaning a controller push via PR #10's endpoint was observable only by the GET debug endpoint, not by queries. Phase 2 threads the shared `HotReloadStreamingConfig` handle all the way into SimpleEngine and re-snapshots per query: main.rs builds one HotReloadStreamingConfig ├─> HttpServer::with_hot_reload_config(handle.clone()) │ ↑ POST /api/v1/streaming-config swaps via this handle │ └─> SimpleEngine::new_with_hot_reload(handle.clone()) ↑ each query calls streaming_config_snapshot() which load_full()s from the same underlying ArcSwap Because `HotReloadStreamingConfig` clones share the same inner `Arc>`, a swap on one clone is immediately visible through the other. That's the entire mechanism. ## Changes ### `SimpleEngine` * Field `streaming_config: Arc` → `streaming_config_source: HotReloadStreamingConfig` * New constructor `SimpleEngine::new_with_hot_reload(...)` takes the shared handle. `main.rs` uses this path. * Existing `SimpleEngine::new(...)` signature preserved for tests, binaries, and legacy callers. Internally wraps the provided `Arc` in a FRESH `HotReloadStreamingConfig` so external unrelated handles don't leak in — the `legacy_new_constructor_is_independent_of_external_handle` test pins this behavior. * New accessor `streaming_config_snapshot(&self) -> Arc`. Each call observes whatever was most recently pushed via PR #10. Doc comment explains the per-query-binding discipline internal callers must follow. ### 13 internal call sites in `simple_engine.rs` Each site that reads from `streaming_config` is converted to use either: * Inlined `self.streaming_config_snapshot().get_aggregation_config(...).map(|c| c.X.clone())...` when the access is a single-expression chain whose result copies out of the borrowed `AggregationConfig` (e.g. `window_size * 1000` or a `.clone()` on an owned field). * A per-helper `let streaming_config = self.streaming_config_snapshot();` binding when a borrowed `&AggregationConfig` outlives the initial expression (e.g. `create_store_query_plan`, which binds the config once and reads multiple fields from it). The binding ensures a consistent view for that helper's entire execution and keeps the borrowed reference valid through the local Arc's lifetime. Per-query entry points do NOT bind a single top-level snapshot today. Instead each helper re-snapshots independently. Under a concurrent swap that races with a running query, two helpers in the same query could see different configs — but both views are internally consistent, and the resulting query answer remains correct under EITHER view (old or new). Swaps are rare; the torn-state window is microseconds. Binding a single per-query snapshot and plumbing it through helper signatures would be a much larger refactor and is tracked as phase 3. ### `main.rs` Switches from `SimpleEngine::new(streaming_config.clone(), ...)` to `SimpleEngine::new_with_hot_reload(hot_reload_config.clone(), ...)`. The same `HotReloadStreamingConfig` handle now flows to both the HTTP server (which handles POST) and the engine (which reads), so a POST is observable by the next query. ## Tests Three new unit tests in a `hot_reload_phase2_tests` module at the end of `simple_engine.rs`: 1. `streaming_config_snapshot_starts_at_initial_config` — baseline sanity check that the initial config is observable through the new accessor. 2. `streaming_config_snapshot_observes_post_construction_swap` — **the core PR E phase 2 guarantee**: build the engine with a shared handle, swap the handle from outside, call the accessor again, and verify the new config is visible. Also asserts the previous snapshot remains internally consistent (per-query stability). 3. `legacy_new_constructor_is_independent_of_external_handle` — verifies that the backwards-compat `SimpleEngine::new` path snapshots into a FRESH internal `HotReloadStreamingConfig`, so a swap on an unrelated external handle does NOT leak into test/bin-only constructions. ## Validation * cargo check --all-targets: clean * cargo clippy --all-targets -- -D warnings: clean * cargo fmt --check: clean * cargo test -p query_engine_rust --lib: 495 passed (up 3 from main — the phase 2 unit tests) ## What closes now Paired with PR #10 (endpoint) and PR #11 (capability-miss notify), phase 2 closes the query-side of the control loop end-to-end: query hits SimpleEngine miss → PR #11 fires fire-and-forget POST to controller /api/v1/plan → §5.2 fallback serves the current query → controller generates new plan → controller POSTs to backend /api/v1/streaming-config (PR #10) → swap lands in the shared HotReloadStreamingConfig → next query re-snapshots via PR E phase 2 and finds a match ## What's still NOT hot-reloaded * **Ingest router / OTLP receiver** — still routes by startup `AggregationConfig` clones. Adding a new agg_id via POST does not yet cause incoming metrics to be routed into that agg_id on the ingest side. Tracked as phase 3. * **In-flight precompute worker GroupState** — unchanged from phase 1. Existing groups complete with construction-time config; new groups pick up the new config. Intended semantics. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../src/engines/simple_engine.rs | 250 ++++++++++++++++-- asap-query-engine/src/main.rs | 12 +- 2 files changed, 231 insertions(+), 31 deletions(-) diff --git a/asap-query-engine/src/engines/simple_engine.rs b/asap-query-engine/src/engines/simple_engine.rs index 3a711fef..9378d336 100644 --- a/asap-query-engine/src/engines/simple_engine.rs +++ b/asap-query-engine/src/engines/simple_engine.rs @@ -139,7 +139,16 @@ pub struct SimpleEngine { store: Arc, // promsketch_store: Option>, inference_config: InferenceConfig, - streaming_config: Arc, + /// Hot-reloadable `StreamingConfig` handle. Internal read sites + /// call `Self::streaming_config_snapshot()` which re-snapshots + /// from this handle, so runtime swaps pushed through PR #10's + /// `POST /api/v1/streaming-config` endpoint take effect on the + /// **next** query without restarting the binary (PR E phase 2). + /// Clones of `HotReloadStreamingConfig` share the same + /// underlying `ArcSwap`, so when `main.rs` hands the same handle + /// to both `SimpleEngine` and `HttpServer::with_hot_reload_config`, + /// a POST is immediately visible to the next query. + streaming_config_source: crate::data_model::HotReloadStreamingConfig, prometheus_scrape_interval: u64, controller_patterns: HashMap>, query_language: QueryLanguage, @@ -152,6 +161,13 @@ pub struct SimpleEngine { } impl SimpleEngine { + /// Construct a `SimpleEngine` with a static `Arc`. + /// Wraps the config in a fresh `HotReloadStreamingConfig` internally + /// — callers that need to share the hot-reload handle with the HTTP + /// server should use `new_with_hot_reload` instead so a POST to + /// `/api/v1/streaming-config` is visible to both. The `_static` + /// variant stays as the simple entry point for tests, binaries, + /// and legacy callers that don't own a `HotReloadStreamingConfig`. pub fn new( store: Arc, // promsketch_store: Option>, @@ -159,6 +175,27 @@ impl SimpleEngine { streaming_config: Arc, prometheus_scrape_interval: u64, query_language: QueryLanguage, + ) -> Self { + let hot_reload = crate::data_model::HotReloadStreamingConfig::from_arc(streaming_config); + Self::new_with_hot_reload( + store, + inference_config, + hot_reload, + prometheus_scrape_interval, + query_language, + ) + } + + /// Construct a `SimpleEngine` that shares a `HotReloadStreamingConfig` + /// handle with another holder (typically the HTTP server). This is + /// the constructor `main.rs` should call so `POST /api/v1/streaming-config` + /// is observable by the next query. + pub fn new_with_hot_reload( + store: Arc, + inference_config: InferenceConfig, + streaming_config_source: crate::data_model::HotReloadStreamingConfig, + prometheus_scrape_interval: u64, + query_language: QueryLanguage, ) -> Self { // Create temporal pattern blocks let mut temporal_pattern_blocks = HashMap::new(); @@ -296,7 +333,7 @@ impl SimpleEngine { store, // promsketch_store, inference_config, - streaming_config, + streaming_config_source, prometheus_scrape_interval, controller_patterns, query_language, @@ -304,6 +341,22 @@ impl SimpleEngine { } } + /// Take a fresh snapshot of the current `StreamingConfig`. Each + /// call observes whatever was most recently pushed through PR #10's + /// `POST /api/v1/streaming-config` endpoint. The returned `Arc` + /// is stable for the caller's lifetime — a concurrent swap + /// produces a new `Arc` and leaves the one returned here alone. + /// + /// Internal read sites inside `SimpleEngine` bind this once per + /// logical unit of work (typically per query-handler invocation + /// or per helper call) and use the local Arc for the duration, + /// so references into the underlying `StreamingConfig` stay + /// valid and a single query sees internally-consistent config + /// fields even if a concurrent swap lands mid-query. + pub fn streaming_config_snapshot(&self) -> Arc { + self.streaming_config_source.snapshot() + } + /// Attach a `ControllerClient` so capability misses fire a /// fire-and-forget notification to the DataCollector controller. /// Builder-style method — takes self by value and returns it so @@ -327,9 +380,8 @@ impl SimpleEngine { &self, requirements: &QueryRequirements, ) -> Option { - let result = self - .streaming_config - .find_compatible_aggregation(requirements); + let streaming_config = self.streaming_config_snapshot(); + let result = streaming_config.find_compatible_aggregation(requirements); if result.is_none() { crate::drivers::query::controller_client::spawn_capability_miss_notify( &self.controller_client, @@ -619,9 +671,11 @@ impl SimpleEngine { (0, end_timestamp) } AggregationType::SetAggregator => { - // Latest window only + // Latest window only. `.map(|c| c.window_size * 1000)` + // copies out a u64 so the snapshot only needs to live + // for the duration of the expression. let window_size = self - .streaming_config + .streaming_config_snapshot() .get_aggregation_config(agg_info.aggregation_id_for_key) .map(|config| config.window_size * 1000) .ok_or_else(|| { @@ -653,9 +707,13 @@ impl SimpleEngine { timestamps: &QueryTimestamps, agg_info: &AggregationIdInfo, ) -> Result { + // Bind a single snapshot of the streaming config for this + // helper's entire execution. `aggregation_config_for_value` + // is a borrow that outlives the initial expression, so the + // Arc it borrows from must outlive this scope. + let streaming_config = self.streaming_config_snapshot(); // Get aggregation config for value to determine window type - let aggregation_config_for_value = self - .streaming_config + let aggregation_config_for_value = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_value) .ok_or_else(|| { format!( @@ -1246,14 +1304,13 @@ impl SimpleEngine { let do_merge = query_pattern_type == QueryPatternType::OnlyTemporal || query_pattern_type == QueryPatternType::OneTemporalOneSpatial; - let grouping_labels = self - .streaming_config + let streaming_config = self.streaming_config_snapshot(); + let grouping_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_value) .map(|config| config.grouping_labels.clone()) .unwrap_or_else(|| query_output_labels.clone()); - let aggregated_labels = self - .streaming_config + let aggregated_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_key) .map(|config| config.aggregated_labels.clone()) .unwrap_or_else(KeyByLabelNames::empty); @@ -1418,7 +1475,7 @@ impl SimpleEngine { let step_ms = (step * 1000.0) as u64; let tumbling_window_ms = self - .streaming_config + .streaming_config_snapshot() .get_aggregation_config(base_context.agg_info.aggregation_id_for_value) .map(|c| c.window_size * 1000)?; @@ -1742,10 +1799,10 @@ impl SimpleEngine { let mut aggregation_type_for_key: Option = None; let mut aggregation_type_for_value: Option = None; + let streaming_config = self.streaming_config_snapshot(); if query_config_aggregations.len() == 2 { for aggregation in query_config_aggregations { - let aggregation_type = self - .streaming_config + let aggregation_type = streaming_config .get_aggregation_config(aggregation.aggregation_id) .map(|config| config.aggregation_type) .ok_or_else(|| { @@ -1781,8 +1838,7 @@ impl SimpleEngine { } else { // Single aggregation: key and value share the same aggregation let id = query_config_aggregations[0].aggregation_id; - let agg_type = self - .streaming_config + let agg_type = streaming_config .get_aggregation_config(id) .map(|config| config.aggregation_type) .ok_or_else(|| format!("No streaming config for aggregation_id {id}"))?; @@ -2075,14 +2131,13 @@ impl SimpleEngine { }) .ok()?; - let grouping_labels = self - .streaming_config + let streaming_config = self.streaming_config_snapshot(); + let grouping_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_value) .map(|config| config.grouping_labels.clone()) .unwrap_or_else(|| metadata.query_output_labels.clone()); - let aggregated_labels = self - .streaming_config + let aggregated_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_key) .map(|config| config.aggregated_labels.clone()) .unwrap_or_else(KeyByLabelNames::empty); @@ -2258,14 +2313,13 @@ impl SimpleEngine { }) .ok()?; - let grouping_labels = self - .streaming_config + let streaming_config = self.streaming_config_snapshot(); + let grouping_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_value) .map(|config| config.grouping_labels.clone()) .unwrap_or_else(|| query_metadata.query_output_labels.clone()); - let aggregated_labels = self - .streaming_config + let aggregated_labels = streaming_config .get_aggregation_config(agg_info.aggregation_id_for_key) .map(|config| config.aggregated_labels.clone()) .unwrap_or_else(KeyByLabelNames::empty); @@ -3107,7 +3161,7 @@ impl SimpleEngine { // Get window size let tumbling_window_ms = self - .streaming_config + .streaming_config_snapshot() .get_aggregation_config(base_context.agg_info.aggregation_id_for_value) .map(|config| config.window_size * 1000)?; @@ -4542,3 +4596,145 @@ mod sketch_query_tests { // assert!(result.is_none()); // } } + +// ─── PR E phase 2: per-query re-snapshot tests ───────────────────────── +#[cfg(test)] +mod hot_reload_phase2_tests { + use super::*; + use crate::data_model::{ + AggregationType, CleanupPolicy, HotReloadStreamingConfig, InferenceConfig, QueryLanguage, + StreamingConfig, WindowType, + }; + use crate::stores::simple_map_store::SimpleMapStore; + use promql_utilities::data_model::key_by_label_names::KeyByLabelNames; + + fn dummy_agg(id: u64, metric: &str) -> crate::data_model::AggregationConfig { + crate::data_model::AggregationConfig::new( + id, + AggregationType::Sum, + String::new(), + std::collections::HashMap::new(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + KeyByLabelNames::empty(), + String::new(), + 60, + 60, + WindowType::Tumbling, + String::new(), + metric.to_string(), + None, + None, + None, + None, + ) + } + + fn cfg_with_agg(id: u64, metric: &str) -> StreamingConfig { + let mut map = std::collections::HashMap::new(); + map.insert(id, dummy_agg(id, metric)); + StreamingConfig::new(map) + } + + fn build_engine(handle: HotReloadStreamingConfig) -> SimpleEngine { + let streaming_config = Arc::new(StreamingConfig::default()); + let store = Arc::new(SimpleMapStore::new( + streaming_config, + CleanupPolicy::NoCleanup, + )); + let inference_config = + InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::NoCleanup); + SimpleEngine::new_with_hot_reload( + store, + inference_config, + handle, + 15000, + QueryLanguage::promql, + ) + } + + #[test] + fn streaming_config_snapshot_starts_at_initial_config() { + let handle = HotReloadStreamingConfig::new(cfg_with_agg(101, "metric_a")); + let engine = build_engine(handle); + let snap = engine.streaming_config_snapshot(); + assert_eq!(snap.aggregation_configs.len(), 1); + assert!(snap.aggregation_configs.contains_key(&101)); + } + + #[test] + fn streaming_config_snapshot_observes_post_construction_swap() { + // Core of PR E phase 2: once the engine is built, swapping + // the shared HotReloadStreamingConfig handle must take + // effect on the next snapshot call — this is the guarantee + // that makes POST /api/v1/streaming-config actually useful + // for query-time behavior. + let handle = HotReloadStreamingConfig::new(cfg_with_agg(101, "metric_a")); + let engine = build_engine(handle.clone()); + + // Initial snapshot: id 101 only. + let snap_before = engine.streaming_config_snapshot(); + assert_eq!(snap_before.aggregation_configs.len(), 1); + assert!(snap_before.aggregation_configs.contains_key(&101)); + assert!(!snap_before.aggregation_configs.contains_key(&202)); + + // Simulate a controller push via `HotReloadStreamingConfig::swap`. + // Clones of the handle share the same underlying ArcSwap, so a + // swap on `handle` is observable through the engine's stored + // clone. + handle.swap(cfg_with_agg(202, "metric_b")); + + // Next snapshot: id 202, id 101 gone. This is exactly what a + // POST-push-then-query sequence must produce. + let snap_after = engine.streaming_config_snapshot(); + assert_eq!(snap_after.aggregation_configs.len(), 1); + assert!(snap_after.aggregation_configs.contains_key(&202)); + assert!(!snap_after.aggregation_configs.contains_key(&101)); + + // The old snapshot is still internally consistent — it's a + // separate Arc that was cheap-cloned before the swap and + // continues to reflect the pre-swap state. This matches the + // per-query-entry-snapshot contract: a query that started + // before the swap sees old config for its entire execution. + assert!(snap_before.aggregation_configs.contains_key(&101)); + } + + #[test] + fn legacy_new_constructor_is_independent_of_external_handle() { + // The legacy `SimpleEngine::new` path wraps the provided + // Arc in a FRESH HotReloadStreamingConfig + // internally, so external swaps must NOT leak in. This is + // the behavior tests and binaries that don't own a shared + // handle depend on. + let external_handle = HotReloadStreamingConfig::new(cfg_with_agg(101, "metric_a")); + let streaming_config = external_handle.snapshot(); + + let store = Arc::new(SimpleMapStore::new( + Arc::clone(&streaming_config), + CleanupPolicy::NoCleanup, + )); + let inference_config = + InferenceConfig::new(QueryLanguage::promql, CleanupPolicy::NoCleanup); + let engine = SimpleEngine::new( + store, + inference_config, + streaming_config, + 15000, + QueryLanguage::promql, + ); + + // External swap should NOT be visible inside the engine — the + // legacy constructor snapshotted the initial Arc into its own + // fresh hot-reload wrapper. + external_handle.swap(cfg_with_agg(999, "metric_swapped")); + + let engine_snap = engine.streaming_config_snapshot(); + assert_eq!(engine_snap.aggregation_configs.len(), 1); + assert!( + engine_snap.aggregation_configs.contains_key(&101), + "legacy `new` constructor should pin the initial config, \ + external swaps to unrelated handles must not leak in" + ); + assert!(!engine_snap.aggregation_configs.contains_key(&999)); + } +} diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 8b8c36fa..eabf28e4 100644 --- a/asap-query-engine/src/main.rs +++ b/asap-query-engine/src/main.rs @@ -350,13 +350,17 @@ async fn main() -> Result<()> { // None // }; - // Setup query engine + // Setup query engine. SimpleEngine shares the same + // HotReloadStreamingConfig handle as the HTTP server, so a POST + // to /api/v1/streaming-config is observable by the next query + // (PR E phase 2). Without sharing the handle, SimpleEngine + // would take a one-time snapshot at construction and ignore + // subsequent swaps. let engine = { - let mut engine = SimpleEngine::new( + let mut engine = SimpleEngine::new_with_hot_reload( store.clone(), - // promsketch_store.clone(), inference_config, - streaming_config.clone(), + hot_reload_config.clone(), args.prometheus_scrape_interval, args.query_language, );