Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions asap-query-engine/examples/engine_config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
69 changes: 69 additions & 0 deletions asap-query-engine/src/engine_config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)]
Expand Down Expand Up @@ -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(),
}
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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]
Expand Down Expand Up @@ -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#"
Expand Down
94 changes: 60 additions & 34 deletions asap-query-engine/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -315,6 +315,7 @@ async fn main() -> Result<()> {
server,
forward_unsupported_queries,
fallback_timeout_secs,
..
} => AdapterConfig::prometheus_promql(
server.clone(),
*forward_unsupported_queries,
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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>,
Expand Down Expand Up @@ -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"));
}
}
16 changes: 7 additions & 9 deletions asap-query-engine/src/tests/prometheus_forwarding_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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;

Expand Down
2 changes: 1 addition & 1 deletion asap-tools/experiments/config/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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-*")

Expand Down
20 changes: 7 additions & 13 deletions asap-tools/experiments/experiment_run_e2e.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
KafkaService,
FlinkService,
QueryEngineRustService,
resolve_backend_config,
ExporterServiceFactory,
PrometheusKafkaAdapterService,
ArroyoService,
Expand Down Expand Up @@ -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(
Expand Down
18 changes: 7 additions & 11 deletions asap-tools/experiments/experiment_run_grafana_demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from experiment_utils.services import (
KafkaService,
QueryEngineRustService,
resolve_backend_config,
ExporterServiceFactory,
ArroyoService,
ArroyoThroughputMonitor,
Expand Down Expand Up @@ -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(
Expand Down
2 changes: 2 additions & 0 deletions asap-tools/experiments/experiment_utils/services/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from .flink import FlinkService
from .query_engine import (
QueryEngineRustService,
resolve_backend_config,
)
from .monitoring import MonitoringService
from .fake_exporters import (
Expand Down Expand Up @@ -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",
Expand Down
Loading
Loading