diff --git a/src/cli/basic/cli.rs b/src/cli/basic/cli.rs index 381208561a..c5eb50a52b 100644 --- a/src/cli/basic/cli.rs +++ b/src/cli/basic/cli.rs @@ -247,7 +247,7 @@ pub async fn cli() -> Result { if app.subcommand().is_none() { return Ok(false); } - let cfg = config::get_config(); + #[cfg(not(feature = "tokio-console"))] env_logger::init_from_env(env_logger::Env::new().default_filter_or("INFO")); @@ -266,6 +266,7 @@ pub async fn cli() -> Result { } // init infra, create data dir & tables + let cfg = config::get_config(); infra::init().await.expect("infra init failed"); match name { "reset" => { diff --git a/src/config/src/router.rs b/src/config/src/router.rs index d3fd792d27..8b2f11c658 100644 --- a/src/config/src/router.rs +++ b/src/config/src/router.rs @@ -66,17 +66,18 @@ pub const INGESTER_ROUTES: [&str; 11] = [ #[inline] pub fn is_querier_route(path: &str) -> bool { + let path = remove_base_uri(path); QUERIER_ROUTES.iter().any(|(route, skip_segments)| { if path.contains(route) { - let segments = path - .split('/') - .filter(|s| !s.is_empty()) - .collect::>(); - // check if we have enough segments - if segments.len() <= *skip_segments { - return false; + let mut segments = path.split('/').filter(|s| !s.is_empty()); + // Skip the required number of segments + for _ in 0..*skip_segments { + if segments.next().is_none() { + return false; + } } - let route_part = segments[*skip_segments..].join("/"); + // Join remaining segments without collecting into a Vec + let route_part = segments.collect::>().join("/"); route_part.starts_with(route) && INGESTER_ROUTES .iter() @@ -89,10 +90,41 @@ pub fn is_querier_route(path: &str) -> bool { #[inline] pub fn is_querier_route_by_body(path: &str) -> bool { + let path = remove_base_uri(path); QUERIER_ROUTES_BY_BODY.iter().any(|x| path.contains(x)) } #[inline] pub fn is_fixed_querier_route(path: &str) -> bool { + let path = remove_base_uri(path); FIXED_QUERIER_ROUTES.iter().any(|x| path.contains(x)) } + +#[inline] +pub fn is_ws_route(path: &str) -> bool { + let path = remove_base_uri(path); + let mut segments = path.split('/').filter(|s| !s.is_empty()); + // Skip first 3 segments + for _ in 0..3 { + if segments.next().is_none() { + return false; + } + } + segments.next() == Some("ws") + && INGESTER_ROUTES + .iter() + .all(|ingest_route| !path.ends_with(ingest_route)) +} + +#[inline] +fn remove_base_uri(path: &str) -> &str { + let base_uri = &crate::get_config().common.base_uri; + if base_uri.is_empty() { + return path; + } + if let Some(stripped) = path.strip_prefix(base_uri) { + stripped + } else { + path + } +} diff --git a/src/main.rs b/src/main.rs index cfdb8cdf25..9d8a7621a4 100644 --- a/src/main.rs +++ b/src/main.rs @@ -112,18 +112,16 @@ use tracing_subscriber::{ #[tokio::main] async fn main() -> Result<(), anyhow::Error> { - let cfg = get_config(); - #[cfg(feature = "tokio-console")] console_subscriber::ConsoleLayer::builder() .retention(Duration::from_secs( - cfg.tokio_console.tokio_console_retention, + get_config().tokio_console.tokio_console_retention, )) .server_addr( format!( "{}:{}", - cfg.tokio_console.tokio_console_server_addr, - cfg.tokio_console.tokio_console_server_port + get_config().tokio_console.tokio_console_server_addr, + get_config().tokio_console.tokio_console_server_port ) .as_str() .parse::()?, @@ -132,28 +130,29 @@ async fn main() -> Result<(), anyhow::Error> { // setup profiling #[cfg(feature = "profiling")] - let pprof_guard = if cfg.profiling.pprof_enabled || cfg.profiling.pprof_protobuf_enabled { - let guard = pprof::ProfilerGuardBuilder::default() - .frequency(1000) - .blocklist(&["libc", "libgcc", "pthread", "vdso"]) - .build() - .unwrap(); - Some(guard) - } else { - None - }; + let pprof_guard = + if get_config().profiling.pprof_enabled || get_config().profiling.pprof_protobuf_enabled { + let guard = pprof::ProfilerGuardBuilder::default() + .frequency(1000) + .blocklist(&["libc", "libgcc", "pthread", "vdso"]) + .build() + .unwrap(); + Some(guard) + } else { + None + }; // setup pyroscope #[cfg(feature = "pyroscope")] - let pyroscope_agent = if cfg.profiling.pyroscope_enabled { + let pyroscope_agent = if get_config().profiling.pyroscope_enabled { let agent = PyroscopeAgent::builder( - &cfg.profiling.pyroscope_server_url, - &cfg.profiling.pyroscope_project_name, + &get_config().profiling.pyroscope_server_url, + &get_config().profiling.pyroscope_project_name, ) .tags( [ - ("role", cfg.common.node_role.as_str()), - ("instance", cfg.common.instance_name.as_str()), + ("role", get_config().common.node_role.as_str()), + ("instance", get_config().common.instance_name.as_str()), ("version", config::VERSION), ] .to_vec(), @@ -172,13 +171,14 @@ async fn main() -> Result<(), anyhow::Error> { return Ok(()); } - let mut tracer_provider = None; + let cfg = get_config(); // setup logs #[cfg(feature = "tokio-console")] let enable_tokio_console = true; #[cfg(not(feature = "tokio-console"))] let enable_tokio_console = false; + let mut tracer_provider = None; let _guard: Option = if enable_tokio_console { None } else if cfg.log.events_enabled { diff --git a/src/router/http/mod.rs b/src/router/http/mod.rs index c20b1624f5..b599b80e60 100644 --- a/src/router/http/mod.rs +++ b/src/router/http/mod.rs @@ -20,7 +20,7 @@ use ::config::{ promql::RequestRangeQuery, search::{Request as SearchRequest, SearchPartitionRequest, ValuesRequest}, }, - router::{INGESTER_ROUTES, is_fixed_querier_route, is_querier_route, is_querier_route_by_body}, + router::{is_fixed_querier_route, is_querier_route, is_querier_route_by_body, is_ws_route}, utils::{json, rand::get_rand_element}, }; use actix_web::{ @@ -155,12 +155,7 @@ async fn dispatch( } // check if the request is a websocket request - let path_columns: Vec<&str> = path.split('/').collect(); - if *path_columns.get(3).unwrap_or(&"") == "ws" - && INGESTER_ROUTES - .iter() - .all(|ingest_route| !path.ends_with(ingest_route)) - { + if is_ws_route(&path) { return proxy_ws(req, payload, start).await; }