fix(query-engine): decouple backend health endpoint (#687)

* fix(query-engine): decouple backend health endpoint

* test(query-engine): allocate dynamic Prometheus mock ports

Author: milindsrivastava1997

asap-query-engine/examples/engine_config.yaml

@@ -36,6 +36,8 @@ http_server:
 backend:
   type: "prometheus"
   server: "http://localhost:9090"         # used for forwarding and planner context
+  # Backend-only startup health check; ASAPQuery keeps /api/v1/status/runtimeinfo.
+  health_endpoint: "/-/ready"
   forward_unsupported_queries: false      # when true, server must be reachable at startup
 
 # ClickHouse — exposes an SQL-over-HTTP API.

asap-query-engine/src/engine_config.rs

@@ -100,6 +100,10 @@ pub enum BackendConfig {
         /// Prometheus server URL used for query forwarding and planner context.
         #[serde(default = "default_prometheus_server")]
         server: String,
+        /// Backend health endpoint used only for the startup reachability check.
+        /// This is independent from ASAPQuery's own runtime-info endpoint.
+        #[serde(default = "default_prometheus_health_endpoint")]
+        health_endpoint: String,
         /// When true, queries not answerable from sketches are forwarded to `server`.
         /// The server must be reachable at startup.
         #[serde(default)]
@@ -145,6 +149,7 @@ impl Default for BackendConfig {
     fn default() -> Self {
         BackendConfig::Prometheus {
             server: default_prometheus_server(),
+            health_endpoint: default_prometheus_health_endpoint(),
             forward_unsupported_queries: false,
             fallback_timeout_secs: default_fallback_timeout_secs(),
         }
@@ -181,12 +186,48 @@ impl BackendConfig {
             } => *forward_unsupported_queries,
         }
     }
+
+    /// Return the URL used to check that the configured backend is reachable.
+    pub fn health_check_url(&self) -> String {
+        match self {
+            BackendConfig::Prometheus {
+                server,
+                health_endpoint,
+                ..
+            } => join_endpoint(server, health_endpoint),
+            BackendConfig::Clickhouse { url, .. } => join_endpoint(url, "/ping"),
+            BackendConfig::ElasticQuerydsl { url, .. } | BackendConfig::ElasticSql { url, .. } => {
+                join_endpoint(url, "/_cluster/health")
+            }
+        }
+    }
+
+    pub fn server_url(&self) -> &str {
+        match self {
+            BackendConfig::Prometheus { server, .. } => server,
+            BackendConfig::Clickhouse { url, .. }
+            | BackendConfig::ElasticQuerydsl { url, .. }
+            | BackendConfig::ElasticSql { url, .. } => url,
+        }
+    }
+}
+
+fn join_endpoint(server: &str, endpoint: &str) -> String {
+    format!(
+        "{}/{}",
+        server.trim_end_matches('/'),
+        endpoint.trim_start_matches('/')
+    )
 }
 
 fn default_prometheus_server() -> String {
     "http://localhost:9090".to_string()
 }
 
+fn default_prometheus_health_endpoint() -> String {
+    "/-/ready".to_string()
+}
+
 fn default_fallback_timeout_secs() -> u64 {
     30
 }
@@ -538,6 +579,10 @@ output_dir: "./output"
         assert!(matches!(config.backend, BackendConfig::Prometheus { .. }));
         assert_eq!(config.backend.query_language(), QueryLanguage::promql);
         assert!(!config.backend.forward_unsupported_queries());
+        assert_eq!(
+            config.backend.health_check_url(),
+            "http://localhost:9090/-/ready"
+        );
     }
 
     #[test]
@@ -649,6 +694,30 @@ backend:
         assert!(config.backend.forward_unsupported_queries());
     }
 
+    #[test]
+    fn backend_prometheus_uses_configured_health_endpoint() {
+        // VictoriaMetrics does not implement Prometheus's runtime-info path;
+        // its /health endpoint must be used for the backend startup check.
+        let yaml = r#"
+streaming_engine: "precompute"
+ingest:
+  type: "http_remote_write"
+  port: 9090
+output_dir: "./output"
+backend:
+  type: "prometheus"
+  server: "http://victoriametrics:8428/"
+  health_endpoint: "/health"
+  forward_unsupported_queries: true
+"#;
+        let config: EngineConfig = Figment::new().merge(Yaml::string(yaml)).extract().unwrap();
+
+        assert_eq!(
+            config.backend.health_check_url(),
+            "http://victoriametrics:8428/health"
+        );
+    }
+
     #[test]
     fn check_config_rejects_query_tracker_with_non_prometheus_backend() {
         let yaml = r#"

asap-query-engine/src/main.rs

@@ -315,6 +315,7 @@ async fn main() -> Result<()> {
             server,
             forward_unsupported_queries,
             fallback_timeout_secs,
+            ..
         } => AdapterConfig::prometheus_promql(
             server.clone(),
             *forward_unsupported_queries,
@@ -353,41 +354,11 @@ async fn main() -> Result<()> {
 
     if config.backend.forward_unsupported_queries() {
         let client = reqwest::Client::new();
-        let (health_url, backend_label) = match &config.backend {
-            BackendConfig::Prometheus { server, .. } => (
-                format!("{}/api/v1/status/runtimeinfo", server.trim_end_matches('/')),
-                server.clone(),
-            ),
-            BackendConfig::Clickhouse { url, .. } => {
-                (format!("{}/ping", url.trim_end_matches('/')), url.clone())
-            }
-            BackendConfig::ElasticQuerydsl { url, .. } | BackendConfig::ElasticSql { url, .. } => (
-                format!("{}/_cluster/health", url.trim_end_matches('/')),
-                url.clone(),
-            ),
-        };
-        match client
-            .get(&health_url)
-            .timeout(std::time::Duration::from_secs(5))
-            .send()
-            .await
-        {
-            Ok(resp) if resp.status().is_success() => {
-                info!("Backend reachable at {}", backend_label);
-            }
-            Ok(resp) => {
-                error!(
-                    "Backend at {} returned HTTP {} — cannot start",
-                    backend_label,
-                    resp.status()
-                );
-                std::process::exit(1);
-            }
-            Err(e) => {
-                error!("Cannot reach backend at {}: {}", backend_label, e);
-                std::process::exit(1);
-            }
+        if let Err(message) = check_backend_health(&client, &config.backend).await {
+            error!("{}", message);
+            std::process::exit(1);
         }
+        info!("Backend reachable at {}", config.backend.server_url());
     }
 
     let query_tracker = if config.query_tracker.enabled {
@@ -503,6 +474,29 @@ async fn main() -> Result<()> {
     Ok(())
 }
 
+async fn check_backend_health(
+    client: &reqwest::Client,
+    backend: &BackendConfig,
+) -> std::result::Result<(), String> {
+    let health_url = backend.health_check_url();
+    let backend_label = backend.server_url();
+
+    match client
+        .get(&health_url)
+        .timeout(std::time::Duration::from_secs(5))
+        .send()
+        .await
+    {
+        Ok(resp) if resp.status().is_success() => Ok(()),
+        Ok(resp) => Err(format!(
+            "Backend at {} returned HTTP {} — cannot start",
+            backend_label,
+            resp.status()
+        )),
+        Err(e) => Err(format!("Cannot reach backend at {}: {}", backend_label, e)),
+    }
+}
+
 /// Periodic memory diagnostics logger — runs every 30 seconds.
 async fn spawn_memory_diagnostics(
     store: Arc<SimpleMapStore>,
@@ -596,3 +590,35 @@ fn setup_logging(
     info!("Logs will be written to: {}/query_engine.log", output_dir);
     Ok(guard)
 }
+
+#[cfg(test)]
+mod tests {
+    use super::{check_backend_health, BackendConfig};
+    use axum::{routing::get, Router};
+    use tokio::net::TcpListener;
+
+    #[tokio::test]
+    async fn backend_health_check_uses_configured_endpoint_and_reports_failure() {
+        let app = Router::new().route(
+            "/health",
+            get(|| async { axum::http::StatusCode::SERVICE_UNAVAILABLE }),
+        );
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        tokio::spawn(async move {
+            axum::serve(listener, app).await.unwrap();
+        });
+
+        let backend = BackendConfig::Prometheus {
+            server: format!("http://{}", address),
+            health_endpoint: "/health".to_string(),
+            forward_unsupported_queries: true,
+            fallback_timeout_secs: 30,
+        };
+        let error = check_backend_health(&reqwest::Client::new(), &backend)
+            .await
+            .unwrap_err();
+
+        assert!(error.contains("returned HTTP 503"));
+    }
+}

asap-query-engine/src/tests/prometheus_forwarding_tests.rs

@@ -11,7 +11,7 @@ use tokio::net::TcpListener;
 use tokio::time::{sleep, Duration};
 
 /// Mock Prometheus server for testing
-async fn start_mock_prometheus_server(port: u16) -> Result<(), Box<dyn std::error::Error>> {
+async fn start_mock_prometheus_server() -> Result<u16, Box<dyn std::error::Error>> {
     use axum::{extract::Query, response::Json, routing::get, Router};
     use serde_json::json;
     use std::collections::HashMap;
@@ -88,15 +88,16 @@ async fn start_mock_prometheus_server(port: u16) -> Result<(), Box<dyn std::erro
         .route("/api/v1/query", get(mock_query_handler))
         .route("/api/v1/query_range", get(mock_range_query_handler));
 
-    let listener = TcpListener::bind(format!("127.0.0.1:{port}")).await?;
+    let listener = TcpListener::bind("127.0.0.1:0").await?;
+    let port = listener.local_addr()?.port();
 
     tokio::spawn(async move {
         axum::serve(listener, app).await.unwrap();
     });
 
     // Give the server time to start
     sleep(Duration::from_millis(100)).await;
-    Ok(())
+    Ok(port)
 }
 
 async fn setup_test_server(prometheus_port: u16) -> (HttpServer, u16) {
@@ -137,8 +138,7 @@ async fn setup_test_server(prometheus_port: u16) -> (HttpServer, u16) {
 #[tokio::test]
 async fn test_prometheus_forwarding_instant_query() {
     // Start mock Prometheus server
-    let prometheus_port = 19090;
-    start_mock_prometheus_server(prometheus_port).await.unwrap();
+    let prometheus_port = start_mock_prometheus_server().await.unwrap();
 
     // Start our HTTP server with forwarding enabled
     let (_server, server_port) = setup_test_server(prometheus_port).await;
@@ -168,8 +168,7 @@ async fn test_prometheus_forwarding_instant_query() {
 #[tokio::test]
 async fn test_prometheus_forwarding_error_handling() {
     // Start mock Prometheus server
-    let prometheus_port = 19092;
-    start_mock_prometheus_server(prometheus_port).await.unwrap();
+    let prometheus_port = start_mock_prometheus_server().await.unwrap();
 
     // Start our HTTP server with forwarding enabled
     let (_server, server_port) = setup_test_server(prometheus_port).await;
@@ -311,8 +310,7 @@ async fn test_prometheus_server_unreachable() {
 
 #[tokio::test]
 async fn test_prometheus_forwarding_range_query() {
-    let prometheus_port = 19094;
-    start_mock_prometheus_server(prometheus_port).await.unwrap();
+    let prometheus_port = start_mock_prometheus_server().await.unwrap();
 
     let (_server, server_port) = setup_test_server(prometheus_port).await;
 

asap-tools/experiments/config/config.yaml

@@ -87,7 +87,7 @@ cluster_data_directory: "/data/cluster_traces"  # Path to directory containing G
 # Backend configuration for the query engine (aligned with BackendConfig in asap-query-engine/src/engine_config.rs)
 backend:
   type: "prometheus"  # choices: ["prometheus", "clickhouse", "elastic_querydsl", "elastic_sql"]
-  # prometheus: server URL is built at runtime from the Prometheus service (no extra fields needed here)
+  # prometheus: server and health_endpoint are built at runtime from the monitoring service
   # clickhouse: url (e.g. "http://ch-host:8123"), database (e.g. "default")
   # elastic_querydsl / elastic_sql: url (e.g. "http://es-host:9200"), index (e.g. "metrics-*")
 

asap-tools/experiments/experiment_run_e2e.py

@@ -12,6 +12,7 @@
     KafkaService,
     FlinkService,
     QueryEngineRustService,
+    resolve_backend_config,
     ExporterServiceFactory,
     PrometheusKafkaAdapterService,
     ArroyoService,
@@ -484,22 +485,15 @@ def main(cfg: DictConfig):
 
             # in case we want to run query engine manually
             if not cfg.flow.replace_query_engine_with_dumb_consumer:
-                # Get prometheus port from prometheus service
-                prometheus_port = prometheus_service.get_query_endpoint_port()
                 # Get http port from query engine service
                 http_port = query_engine_service.get_http_port()
 
-                # Build a fully resolved BackendConfig dict.  For the prometheus
-                # backend the server URL depends on the runtime node IP, so we
-                # fill it in here rather than in config.yaml.
-                backend_config = dict(args.backend)
-                if backend_config["type"] == "prometheus":
-                    prometheus_host = provider.get_node_ip(args.node_offset)
-                    backend_config["server"] = (
-                        f"http://{prometheus_host}:{prometheus_port}"
-                    )
-                backend_config["forward_unsupported_queries"] = (
-                    args.forward_unsupported_queries
+                backend_config = resolve_backend_config(
+                    args.backend,
+                    prometheus_service,
+                    provider,
+                    args.node_offset,
+                    args.forward_unsupported_queries,
                 )
 
                 query_engine_service.start(

asap-tools/experiments/experiment_run_grafana_demo.py

@@ -11,6 +11,7 @@
 from experiment_utils.services import (
     KafkaService,
     QueryEngineRustService,
+    resolve_backend_config,
     ExporterServiceFactory,
     ArroyoService,
     ArroyoThroughputMonitor,
@@ -447,21 +448,16 @@ def main(cfg: DictConfig):
             )
         # in case we want to run query engine manually
         if not cfg.flow.replace_query_engine_with_dumb_consumer:
-            # Get prometheus port from prometheus service
-            prometheus_port = prometheus_service.get_query_endpoint_port()
             # Get http port from query engine service
             http_port = query_engine_service.get_http_port()
 
-            # Build a fully resolved BackendConfig dict.  For the prometheus
-            # backend the server URL depends on the runtime node IP, so we
-            # fill it in here rather than in config.yaml.
-            backend_config = dict(args.backend)
-            if backend_config["type"] == "prometheus":
-                prometheus_host = provider.get_node_ip(args.node_offset)
-                backend_config["server"] = f"http://{prometheus_host}:{prometheus_port}"
             # forward_unsupported_queries is forced True for the Grafana demo (line 63)
-            backend_config["forward_unsupported_queries"] = (
-                args.forward_unsupported_queries
+            backend_config = resolve_backend_config(
+                args.backend,
+                prometheus_service,
+                provider,
+                args.node_offset,
+                args.forward_unsupported_queries,
             )
 
             query_engine_service.start(

asap-tools/experiments/experiment_utils/services/__init__.py

@@ -10,6 +10,7 @@
 from .flink import FlinkService
 from .query_engine import (
     QueryEngineRustService,
+    resolve_backend_config,
 )
 from .monitoring import MonitoringService
 from .fake_exporters import (
@@ -131,6 +132,7 @@ def create_prometheus_service(cfg, provider, num_nodes: int, node_offset: int):
     "KafkaService",
     "FlinkService",
     "QueryEngineRustService",
+    "resolve_backend_config",
     "MonitoringService",
     "ExporterServiceFactory",
     "PythonExporterService",