diff --git a/asap-query-engine/examples/engine_config.yaml b/asap-query-engine/examples/engine_config.yaml index ec934c2..bc8f359 100644 --- a/asap-query-engine/examples/engine_config.yaml +++ b/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. diff --git a/asap-query-engine/src/engine_config.rs b/asap-query-engine/src/engine_config.rs index 8ef1ef6..5591d6d 100644 --- a/asap-query-engine/src/engine_config.rs +++ b/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#" diff --git a/asap-query-engine/src/main.rs b/asap-query-engine/src/main.rs index 95ec12e..3a5e4f7 100644 --- a/asap-query-engine/src/main.rs +++ b/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, @@ -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")); + } +} diff --git a/asap-query-engine/src/tests/prometheus_forwarding_tests.rs b/asap-query-engine/src/tests/prometheus_forwarding_tests.rs index 27d09cc..03a7b4c 100644 --- a/asap-query-engine/src/tests/prometheus_forwarding_tests.rs +++ b/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> { +async fn start_mock_prometheus_server() -> Result> { use axum::{extract::Query, response::Json, routing::get, Router}; use serde_json::json; use std::collections::HashMap; @@ -88,7 +88,8 @@ async fn start_mock_prometheus_server(port: u16) -> Result<(), Box Result<(), Box (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; diff --git a/asap-tools/experiments/config/config.yaml b/asap-tools/experiments/config/config.yaml index 15ef039..61f09e5 100644 --- a/asap-tools/experiments/config/config.yaml +++ b/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-*") diff --git a/asap-tools/experiments/experiment_run_e2e.py b/asap-tools/experiments/experiment_run_e2e.py index 9c931ce..343a861 100644 --- a/asap-tools/experiments/experiment_run_e2e.py +++ b/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( diff --git a/asap-tools/experiments/experiment_run_grafana_demo.py b/asap-tools/experiments/experiment_run_grafana_demo.py index 0ecd3dd..fa10298 100644 --- a/asap-tools/experiments/experiment_run_grafana_demo.py +++ b/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( diff --git a/asap-tools/experiments/experiment_utils/services/__init__.py b/asap-tools/experiments/experiment_utils/services/__init__.py index bc399b6..0b8dd02 100644 --- a/asap-tools/experiments/experiment_utils/services/__init__.py +++ b/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", diff --git a/asap-tools/experiments/experiment_utils/services/prometheus.py b/asap-tools/experiments/experiment_utils/services/prometheus.py index fb86d79..07d5413 100644 --- a/asap-tools/experiments/experiment_utils/services/prometheus.py +++ b/asap-tools/experiments/experiment_utils/services/prometheus.py @@ -33,6 +33,10 @@ def get_query_endpoint_port(self) -> int: """Get the query endpoint port for Prometheus.""" return 9090 + def get_health_endpoint(self) -> str: + """Get Prometheus health check endpoint.""" + return "/-/ready" + def start(self, experiment_output_dir: str, **kwargs) -> None: """ Start Prometheus service. diff --git a/asap-tools/experiments/experiment_utils/services/query_engine.py b/asap-tools/experiments/experiment_utils/services/query_engine.py index 1629285..7450df1 100644 --- a/asap-tools/experiments/experiment_utils/services/query_engine.py +++ b/asap-tools/experiments/experiment_utils/services/query_engine.py @@ -14,6 +14,24 @@ from experiment_utils.providers.base import InfrastructureProvider +def resolve_backend_config( + backend: dict, + prometheus_service: BaseService, + provider: InfrastructureProvider, + node_offset: int, + forward_unsupported_queries: bool, +) -> dict: + """Resolve runtime backend URL and health settings for the query engine.""" + backend_config = dict(backend) + if backend_config["type"] == "prometheus": + prometheus_host = provider.get_node_ip(node_offset) + prometheus_port = prometheus_service.get_query_endpoint_port() + backend_config["server"] = f"http://{prometheus_host}:{prometheus_port}" + backend_config["health_endpoint"] = prometheus_service.get_health_endpoint() + backend_config["forward_unsupported_queries"] = forward_unsupported_queries + return backend_config + + class BaseQueryEngineService(BaseService): """Base class for query engine services.""" @@ -238,9 +256,10 @@ def start( dump_precomputes: Whether to dump precomputed values lock_strategy: Lock strategy for SimpleMapStore (global or per-key) backend_config: Fully resolved BackendConfig dict with type tag and all - backend-specific fields (url/server/database/index as needed) - plus forward_unsupported_queries. Matches the BackendConfig - tagged union in asap-query-engine/src/engine_config.rs. + backend-specific fields (url/server/database/index as needed), + optional health_endpoint, and forward_unsupported_queries. + Matches the BackendConfig tagged union in + asap-query-engine/src/engine_config.rs. http_port: Port for the query engine's HTTP API server remote_write_port: Port the precompute engine listens on for Prometheus remote write; should match streaming.remote_write.base_port (default 8080)