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, );