diff --git a/Cargo.lock b/Cargo.lock index fde1bfb27f..5e87ae51ce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6770,6 +6770,7 @@ dependencies = [ "o2_enterprise", "o2_openfga", "object_store", + "openobserve-vrl", "opentelemetry", "opentelemetry-proto", "parking_lot 0.12.5", @@ -6842,6 +6843,7 @@ dependencies = [ "o2_enterprise", "o2_openfga", "openobserve-core", + "openobserve-vrl", "prometheus", "proto", "regex", @@ -6853,6 +6855,18 @@ dependencies = [ "wal", ] +[[package]] +name = "openobserve-vrl" +version = "0.1.0" +dependencies = [ + "async-trait", + "config", + "enrichment", + "serde_json", + "tokio", + "vrl", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -9211,14 +9225,20 @@ version = "0.1.0" dependencies = [ "arrow", "arrow-schema", - "async-trait", + "chrono", "config", - "enrichment", + "datafusion", + "hashbrown 0.16.1", + "infra", + "log", "o2_enterprise", + "openobserve-vrl", + "parking_lot 0.12.5", "serde", "serde_json", + "tokio", + "tracing", "utoipa", - "vrl", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index e0154e68cf..19cc4242e2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -239,6 +239,7 @@ members = [ "src/api_query", "src/api_common", "src/core", + "src/vrl", "src/search", "src/jobs", "src/config", @@ -269,6 +270,7 @@ flight = { path = "src/flight" } tantivy_utils = { path = "src/tantivy_utils" } web = { path = "src/web" } openobserve-core = { path = "src/core" } +openobserve-vrl = { path = "src/vrl" } search = { path = "src/search" } openobserve-api = { path = "src/api" } openobserve-api-management = { path = "src/api_management" } diff --git a/deny.toml b/deny.toml index d91a703db8..64502b56e4 100644 --- a/deny.toml +++ b/deny.toml @@ -28,6 +28,7 @@ exceptions = [ { name = "openobserve-api-query", allow = ["AGPL-3.0"] }, { name = "openobserve-core", allow = ["AGPL-3.0"] }, { name = "openobserve-jobs", allow = ["AGPL-3.0"] }, + { name = "openobserve-vrl", allow = ["AGPL-3.0"] }, { name = "search", allow = ["AGPL-3.0"] }, # Empty dependency-resolution stubs for the private enterprise crates. { name = "o2_dex", allow = ["AGPL-3.0"] }, diff --git a/src/config/src/config.rs b/src/config/src/config.rs index abf1267583..57668ca9da 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -61,6 +61,7 @@ pub static COMMIT_HASH: &str = env!("GIT_COMMIT_HASH"); pub static BUILD_DATE: &str = env!("GIT_BUILD_DATE"); pub const META_ORG_ID: &str = "_meta"; +pub const DEFAULT_ORG: &str = "default"; pub const MMDB_CITY_FILE_NAME: &str = "GeoLite2-City.mmdb"; pub const MMDB_ASN_FILE_NAME: &str = "GeoLite2-ASN.mmdb"; diff --git a/src/core/Cargo.toml b/src/core/Cargo.toml index 0c85416348..9bbeac11d0 100644 --- a/src/core/Cargo.toml +++ b/src/core/Cargo.toml @@ -20,15 +20,17 @@ enterprise = [ "search/enterprise", ] cloud = ["infra/cloud", "enterprise", "o2_enterprise/cloud"] -vectorscan = ["o2_enterprise/vectorscan"] +vectorscan = ["o2_enterprise/vectorscan", "search/vectorscan"] vortex = [ "dep:vortex", "dep:vortex-datafusion", "config/vortex", "o2_enterprise?/vortex", + "search/vortex", ] [dependencies] +openobserve-vrl.workspace = true search.workspace = true axum.workspace = true http.workspace = true diff --git a/src/core/src/cipher/mod.rs b/src/core/src/cipher/mod.rs index 1db2d28bcb..c3192865c3 100644 --- a/src/core/src/cipher/mod.rs +++ b/src/core/src/cipher/mod.rs @@ -13,12 +13,11 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . +pub use ::search::cipher::registry; use o2_enterprise::enterprise::cipher::http_repr::HttpKey; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; -pub mod registry; - #[derive(Clone, Debug, Deserialize, ToSchema)] pub struct KeyAddRequest { pub name: String, diff --git a/src/core/src/common/infra/config.rs b/src/core/src/common/infra/config.rs index be7563ad74..edd9b1aa5e 100644 --- a/src/core/src/common/infra/config.rs +++ b/src/core/src/common/infra/config.rs @@ -21,7 +21,6 @@ use config::{ alerts::alert::Alert, destinations::{Destination, Template}, folder::Folder, - function::Transform, pipeline::Pipeline, promql::ClusterLeader, ratelimit::CachedUserRoles, @@ -33,24 +32,18 @@ use config::{ use dashmap::DashMap; use hashbrown::HashMap; use infra::table::short_urls::ShortUrlRecord; -use parking_lot::RwLock; -pub use search::enrichment::ENRICHMENT_TABLES; -use vector_enrichment::TableRegistry; +pub use openobserve_vrl::{QUERY_FUNCTIONS, enrichment::ENRICHMENT_TABLES}; use crate::{ common::meta::{ maxmind::MaxmindClient, organization::{OrgStatus, Organization, OrganizationSetting}, }, - service::{ - db::scheduler as db_scheduler, enrichment_table::geoip::Geoip, - pipeline::batch_execution::ExecutablePipeline, - }, + service::{db::scheduler as db_scheduler, pipeline::batch_execution::ExecutablePipeline}, }; // global cache variables pub static KVS: Lazy> = Lazy::new(Default::default); -pub static QUERY_FUNCTIONS: Lazy> = Lazy::new(DashMap::default); pub static USERS: Lazy> = Lazy::new(DashMap::default); pub static ORG_USERS: Lazy> = @@ -74,22 +67,9 @@ pub static REALTIME_ALERT_TRIGGERS: Lazy> = Lazy::new(Default::default); pub static DESTINATIONS: Lazy> = Lazy::new(Default::default); -pub static ENRICHMENT_REGISTRY: Lazy> = - Lazy::new(|| Arc::new(TableRegistry::default())); - pub static MAXMIND_DB_CLIENT: Lazy>>> = Lazy::new(|| Arc::new(tokio::sync::RwLock::new(None))); -pub static GEOIP_CITY_TABLE: Lazy>>> = - Lazy::new(|| Arc::new(RwLock::new(None))); - -pub static GEOIP_ASN_TABLE: Lazy>>> = - Lazy::new(|| Arc::new(RwLock::new(None))); - -#[cfg(feature = "enterprise")] -pub static GEOIP_ENT_TABLE: Lazy>>> = - Lazy::new(|| Arc::new(RwLock::new(None))); - pub static STREAM_EXECUTABLE_PIPELINES: Lazy>> = Lazy::new(Default::default); pub static PIPELINE_STREAM_MAPPING: Lazy> = @@ -139,10 +119,6 @@ mod tests { // Test Arc variables let users_rum_token = USERS_RUM_TOKEN.clone(); assert_eq!(users_rum_token.len(), 0); - - // Test enrichment registry - let registry = ENRICHMENT_REGISTRY.clone(); - assert!(!std::ptr::eq(registry.as_ref(), std::ptr::null())); } #[tokio::test] @@ -258,29 +234,6 @@ mod tests { assert!(read_guard.is_none()); } - #[test] - fn test_geoip_tables() { - // Test GEOIP_CITY_TABLE - let city_table = GEOIP_CITY_TABLE.clone(); - let read_guard = city_table.read(); - assert!(read_guard.is_none()); - drop(read_guard); - - // Test GEOIP_ASN_TABLE - let asn_table = GEOIP_ASN_TABLE.clone(); - let read_guard = asn_table.read(); - assert!(read_guard.is_none()); - drop(read_guard); - - #[cfg(feature = "enterprise")] - { - // Test GEOIP_ENT_TABLE (only available with enterprise feature) - let ent_table = GEOIP_ENT_TABLE.clone(); - let read_guard = ent_table.read(); - assert!(read_guard.is_none()); - } - } - #[test] fn test_dashmap_operations() { // Test basic DashMap operations on KVS @@ -303,6 +256,8 @@ mod tests { #[test] fn test_query_functions_cache() { + use config::meta::function::Transform; + let test_key = "test_function".to_string(); let test_transform = Transform { function: "test_function".to_string(), diff --git a/src/core/src/common/meta/organization.rs b/src/core/src/common/meta/organization.rs index 8dc431128a..7de9bc19c4 100644 --- a/src/core/src/common/meta/organization.rs +++ b/src/core/src/common/meta/organization.rs @@ -13,14 +13,13 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . +pub use config::DEFAULT_ORG; use config::{meta::user::UserRole, stats::MemorySize}; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; #[cfg(feature = "cloud")] use super::user::InviteStatus; - -pub const DEFAULT_ORG: &str = "default"; pub const CUSTOM: &str = "custom"; pub const USER_DEFAULT: &str = "user_default"; pub const THRESHOLD: i64 = 9383939382; diff --git a/src/core/src/common/utils/functions.rs b/src/core/src/common/utils/functions.rs index 41e79bc064..ce436476b2 100644 --- a/src/core/src/common/utils/functions.rs +++ b/src/core/src/common/utils/functions.rs @@ -5,7 +5,7 @@ // the Free Software Foundation, either version 3 of the License, or // (at your option) any later version. // -// This program is distributed in the hope that it will be useful +// This program is distributed in the hope that it will be useful, // but WITHOUT ANY WARRANTY; without even the implied warranty of // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the // GNU Affero General Public License for more details. @@ -13,160 +13,6 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::collections::HashMap; - -use config::{ - GEO_IP_ASN_ENRICHMENT_TABLE, GEO_IP_CITY_ENRICHMENT_TABLE, meta::function::VRLCompilerConfig, +pub use openobserve_vrl::{ + get_all_transform_keys, get_enrichment_tables, get_vrl_compiler_config, init_vrl_runtime, }; -use vector_enrichment::{Table, TableRegistry}; - -use crate::common::{ - infra::config::{ENRICHMENT_TABLES, GEOIP_ASN_TABLE, GEOIP_CITY_TABLE}, - meta::organization::DEFAULT_ORG, -}; - -pub async fn get_all_transform_keys(org_id: &str) -> Vec { - let org_key = &format!("{org_id}/"); - - crate::common::infra::config::QUERY_FUNCTIONS - .clone() - .iter() - .filter_map(|transform| { - let key = transform.key(); - key.strip_prefix(org_key).map(|x| x.to_string()) - }) - .collect() -} - -pub fn init_vrl_runtime() -> vrl::compiler::runtime::Runtime { - vrl::compiler::runtime::Runtime::new(vrl::prelude::state::RuntimeState::default()) -} - -pub fn get_vrl_compiler_config(org_id: &str) -> VRLCompilerConfig { - let en_tables = ENRICHMENT_TABLES.clone(); - let mut functions = vrl::stdlib::all(); - functions.append(&mut vector_enrichment::vrl_functions()); - let registry = TableRegistry::default(); - let mut tables: HashMap> = HashMap::new(); - - for table in en_tables.iter() { - if table.org_id == org_id || table.org_id == DEFAULT_ORG { - tables.insert( - table.stream_name.to_owned(), - Box::new(table.value().clone()), - ); - } - } - drop(en_tables); - - if let Some(v) = GEOIP_CITY_TABLE.read().as_ref() { - tables.insert(GEO_IP_CITY_ENRICHMENT_TABLE.to_owned(), Box::new(v.clone())); - } - - if let Some(v) = GEOIP_ASN_TABLE.read().as_ref() { - tables.insert(GEO_IP_ASN_ENRICHMENT_TABLE.to_owned(), Box::new(v.clone())); - } - - #[cfg(feature = "enterprise")] - if o2_enterprise::enterprise::common::config::get_config() - .common - .enable_enterprise_mmdb - { - let geoip_ent = crate::common::infra::config::GEOIP_ENT_TABLE.read(); - if let Some(table) = geoip_ent.as_ref() { - tables.insert( - o2_enterprise::enterprise::common::config::GEO_IP_ENTERPRISE_ENRICHMENT_TABLE - .to_owned(), - Box::new(table.clone()), - ); - } - }; - - registry.load(tables); - let mut config = vrl::compiler::CompileConfig::default(); - config.set_custom(registry); - VRLCompilerConfig { config, functions } -} - -#[cfg(test)] -mod tests { - - use std::sync::Arc; - - use super::*; - - #[test] - fn test_init_vrl_runtime() { - let runtime = init_vrl_runtime(); - assert!(runtime.is_empty()); - } - - #[test] - fn test_get_vrl_compiler_config() { - let config = get_vrl_compiler_config("default"); - assert!(!config.functions.is_empty()); - assert!(config.config.get_custom::().is_some()); - } - - #[tokio::test] - async fn test_get_all_transform_keys() { - // Setup test data - let test_org = "test_org"; - let test_key = format!("{test_org}/test_transform"); - crate::common::infra::config::QUERY_FUNCTIONS.insert( - test_key.clone(), - config::meta::function::Transform { - name: test_key.clone(), - function: ".".to_string(), - params: "row".to_string(), - num_args: 1, - trans_type: Some(0), - streams: Some(vec![config::meta::function::StreamOrder { - stream: "test".to_string(), - order: 1, - stream_type: config::meta::stream::StreamType::Logs, - is_removed: false, - apply_before_flattening: false, - }]), - }, - ); - - // Test the function - let keys = get_all_transform_keys(test_org).await; - assert!(keys.contains(&"test_transform".to_string())); - } - - #[test] - fn test_get_vrl_compiler_config_with_enrichment_tables() { - // Setup test data - let test_org = "test_org"; - let test_stream = "test_stream"; - let table_name = "test_table"; - let en_table = crate::service::enrichment::StreamTable { - org_id: test_org.to_string(), - stream_name: test_stream.to_string(), - data: Arc::new(vec![]), - }; - ENRICHMENT_TABLES.insert(table_name.to_string(), en_table); - - // Test the function - let config = get_vrl_compiler_config(test_org); - let registry = config.config.get_custom::().unwrap(); - - // Verify enrichment table is loaded - assert!(registry.table_ids().contains(&test_stream.to_string())); - } - - #[tokio::test] - async fn test_get_all_transform_keys_empty() { - let keys = get_all_transform_keys("nonexistent_org").await; - assert!(keys.is_empty()); - } - - #[test] - fn test_get_vrl_compiler_config_empty() { - let config = get_vrl_compiler_config("nonexistent_org"); - let registry = config.config.get_custom::().unwrap(); - assert!(registry.table_ids().is_empty()); - } -} diff --git a/src/core/src/service/file_downloader.rs b/src/core/src/service/file_downloader.rs index 0a7fad02d2..ce055621b7 100644 --- a/src/core/src/service/file_downloader.rs +++ b/src/core/src/service/file_downloader.rs @@ -13,720 +13,4 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{ - collections::VecDeque, - sync::{Arc, LazyLock as Lazy}, -}; - -use config::{ - cluster::LOCAL_NODE, - get_config, - meta::cluster::{Role, RoleGroup, get_internal_grpc_token}, - metrics, - utils::time::{day_micros, now_micros}, -}; -use futures_util::StreamExt; -use hashbrown::{HashMap, HashSet}; -use infra::{cache::file_data, cluster}; -use proto::cluster_rpc::{SimpleFileList, event_client::EventClient}; -use tokio::sync::{ - Mutex, - mpsc::{Receiver, Sender}, -}; -use tonic::{codec::CompressionEncoding, metadata::MetadataValue}; - -/// (trace_id, file_id, account, file, size, cache_type) -type FileInfo = (String, i64, String, String, usize, file_data::CacheType); - -mod processing_files { - use hashbrown::HashSet; - use parking_lot::RwLock; - - use super::*; - - static PROCESSING_FILES: Lazy>> = - Lazy::new(|| RwLock::new(HashSet::new())); - - pub fn is_processing(file_name: &str) -> bool { - PROCESSING_FILES.read().contains(file_name) - } - - pub fn add(file_name: &str) { - PROCESSING_FILES.write().insert(file_name.to_string()); - } - - pub fn remove(file_name: &str) { - PROCESSING_FILES.write().remove(file_name); - } -} - -mod queued_files { - use hashbrown::HashSet; - use parking_lot::RwLock; - - use super::*; - - static QUEUED_FILES: Lazy>> = Lazy::new(|| RwLock::new(HashSet::new())); - - /// Returns true if newly inserted (caller should enqueue), false if already present (skip). - pub fn try_add(file_name: &str) -> bool { - QUEUED_FILES.write().insert(file_name.to_string()) - } - - pub fn remove(file_name: &str) { - QUEUED_FILES.write().remove(file_name); - } -} - -struct DownloadQueue { - sender: Sender, - receiver: Arc>>, -} - -impl DownloadQueue { - fn new(sender: Sender, receiver: Arc>>) -> Self { - Self { sender, receiver } - } -} - -struct PriorityDownloadQueue { - stack: Mutex>, - notify: tokio::sync::Notify, - max_size: usize, -} - -impl PriorityDownloadQueue { - fn new(max_size: usize) -> Self { - Self { - stack: Mutex::new(VecDeque::with_capacity(max_size)), - notify: tokio::sync::Notify::new(), - max_size, - } - } - - // Returns false if at cap — caller must remove from queued_files and skip metric inc. - async fn push(&self, file_info: FileInfo) -> bool { - let mut stack = self.stack.lock().await; - if stack.len() >= self.max_size { - return false; - } - stack.push_back(file_info); - drop(stack); - self.notify.notify_one(); - true - } - - // Blocks until an item is available. LIFO via pop_back. - // notified() is created before the lock check to avoid missed-wakeup race. - // Chains notify_one() so all workers drain the queue under burst load. - async fn pop(&self) -> FileInfo { - loop { - let notified = self.notify.notified(); - { - let mut stack = self.stack.lock().await; - if let Some(item) = stack.pop_back() { - if !stack.is_empty() { - self.notify.notify_one(); - } - return item; - } - } - notified.await; - } - } -} - -const FILE_DOWNLOAD_QUEUE_SIZE: usize = 10000; -static FILE_DOWNLOAD_CHANNEL: Lazy = Lazy::new(|| { - let (tx, rx) = tokio::sync::mpsc::channel::(FILE_DOWNLOAD_QUEUE_SIZE); - DownloadQueue::new(tx, Arc::new(Mutex::new(rx))) -}); - -static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy = - Lazy::new(|| PriorityDownloadQueue::new(FILE_DOWNLOAD_QUEUE_SIZE)); - -pub async fn run() -> Result<(), anyhow::Error> { - let cfg = get_config(); - // worker tasks: handle normal queue download (FIFO via mpsc channel) - for thread in 0..cfg.limit.file_download_thread_num { - let rx = FILE_DOWNLOAD_CHANNEL.receiver.clone(); - tokio::spawn(async move { - loop { - let ret = rx.lock().await.recv().await; - match ret { - None => { - log::debug!("[FILE_CACHE_DOWNLOAD:JOB:NORMAL] Receiving channel is closed"); - break; - } - Some((trace_id, id, account, file, file_size, cache)) => { - // transition: queued → not-queued (must run before any skip path) - queued_files::remove(&file); - - // check if the file is already being downloaded - if processing_files::is_processing(&file) { - log::warn!( - "[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it" - ); - // update metrics - metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .dec(); - continue; - } - - // add the file to processing set - processing_files::add(&file); - - // download the file - match download_file( - thread, &trace_id, id, &account, &file, file_size, cache, - ) - .await - { - Ok(data_len) => { - if data_len > 0 && data_len != file_size { - log::warn!( - "[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it", - ); - } - } - Err(e) => { - log::error!( - "[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} to cache {cache:?} err: {e}", - ); - } - } - - // remove the file from processing set - processing_files::remove(&file); - - // update metrics - metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .dec(); - } - } - } - }); - } - - // worker tasks: handle priority queue download (LIFO via VecDeque::pop_back) - for thread in 0..cfg.limit.file_download_priority_queue_thread_num { - tokio::spawn(async move { - loop { - let (trace_id, id, account, file, file_size, cache) = - PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await; - - // transition: queued → not-queued (must run before any skip path) - queued_files::remove(&file); - - // check if the file is already being downloaded - if processing_files::is_processing(&file) { - log::warn!( - "[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it" - ); - metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .dec(); - continue; - } - - // add the file to processing set - processing_files::add(&file); - - // download the file - match download_file(thread, &trace_id, id, &account, &file, file_size, cache).await - { - Ok(data_len) => { - if data_len > 0 && data_len != file_size { - log::warn!( - "[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it", - ); - } - } - Err(e) => { - log::error!( - "[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} to cache {cache:?} err: {e}", - ); - } - } - - // remove the file from processing set - processing_files::remove(&file); - - // update metrics - metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .dec(); - } - }); - } - Ok(()) -} - -async fn download_file( - thread: usize, - trace_id: &str, - file_id: i64, - account: &str, - file_name: &str, - file_size: usize, - cache_type: file_data::CacheType, -) -> Result { - let cfg = get_config(); - - // download file from node - if cfg.cache_latest_files.download_from_node - && let Ok(ok) = - download_file_with_consistent_hash(file_id, account, file_name, file_size).await - && ok - { - return Ok(file_size); - } - - // download from object store - let size = if file_size > 0 { Some(file_size) } else { None }; - let start = std::time::Instant::now(); - let ret = match cache_type { - file_data::CacheType::Memory => { - let mut disk_exists = false; - let mem_exists = file_data::memory::exist(file_name).await; - if !mem_exists && !cfg.memory_cache.skip_disk_check { - disk_exists = file_data::disk::exist(file_name).await; - } - if !mem_exists && (cfg.memory_cache.skip_disk_check || !disk_exists) { - file_data::memory::download(account, file_name, size).await - } else { - Ok(0) - } - } - file_data::CacheType::Disk => { - if !file_data::disk::exist(file_name).await { - file_data::disk::download(account, file_name, size).await - } else { - Ok(0) - } - } - _ => Ok(0), - }; - log::debug!( - "[FILE_CACHE_DOWNLOAD:JOB:{thread}] [trace_id {trace_id}] download file: {file_name}, ret: {:?}, took: {} ms", - ret, - start.elapsed().as_millis() - ); - ret -} - -async fn download_file_with_consistent_hash( - file_id: i64, - account: &str, - file_name: &str, - file_size: usize, -) -> Result { - let role_group = if LOCAL_NODE.is_interactive_querier() { - RoleGroup::Interactive - } else { - RoleGroup::Background - }; - let Some(node_name) = cluster::get_node_from_consistent_hash( - &file_id.to_string(), - &Role::Querier, - Some(role_group), - ) - .await - else { - return Ok(false); - }; - // get node by file_id - let Some(node) = cluster::get_cached_node_by_name(&node_name).await else { - return Ok(false); - }; - // download file from node - let Ok(failed) = download_from_node( - &node.grpc_addr, - &[( - file_id, - account.to_string(), - file_name.to_string(), - file_size as i64, - 0, - )], - ) - .await - else { - return Ok(false); - }; - // failed is empty means download success - Ok(failed.is_empty()) -} - -// download files from node and return download failed files -// file: (account, file, size, ts) -pub async fn download_from_node( - addr: &str, - files: &[(i64, String, String, i64, i64)], -) -> Result, anyhow::Error> { - let start = std::time::Instant::now(); - let cfg = get_config(); - log::debug!( - "[FILE_CACHE_DOWNLOAD:gRPC] Download files from node start, files: {:?}", - files.len() - ); - - let token: MetadataValue<_> = get_internal_grpc_token() - .parse() - .map_err(|_| anyhow::anyhow!("Invalid token"))?; - - let channel = infra::client::grpc::get_cached_channel(addr).await?; - let client = EventClient::with_interceptor(channel, move |mut req: tonic::Request<()>| { - req.metadata_mut().insert("authorization", token.clone()); - Ok(req) - }); - - let file_size_map = files - .iter() - .filter_map(|(_, _, f, s, _)| { - if *s > cfg.cache_latest_files.download_node_size * 1024 * 1024 { - None - } else { - Some((f, *s as usize)) - } - }) - .collect::>(); - if file_size_map.is_empty() { - return Ok(files.to_vec()); - } - let request = tonic::Request::new(SimpleFileList { - files: files.iter().map(|(_, _, f, ..)| f.to_string()).collect(), - }); - - let resp = client - .send_compressed(CompressionEncoding::Gzip) - .accept_compressed(CompressionEncoding::Gzip) - .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) - .get_files(request) - .await - .map_err(|e| anyhow::anyhow!("Failed to get files from {addr}, {e}"))?; - - let mut file_contents = HashMap::new(); - let mut downloaded_files = HashSet::new(); - let mut resp_stream = resp.into_inner(); - while let Some(resp) = resp_stream.next().await { - let resp = match resp { - Ok(resp) => resp, - Err(err) => { - if err.code() == tonic::Code::NotFound { - log::debug!( - "[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: file not found", - err.message(), - addr - ); - continue; - } - return Err(anyhow::anyhow!( - "Failed to download file from {addr}, {err}" - )); - } - }; - for content in resp.entries { - let entry = file_contents - .entry(content.filename.clone()) - .or_insert(bytes::BytesMut::new()); - entry.extend_from_slice(&content.content); - downloaded_files.insert(content.filename); - } - } - - log::debug!( - "[FILE_CACHE_DOWNLOAD:gRPC] Successfully retrieved {} files from {} in {} ms", - downloaded_files.len(), - addr, - start.elapsed().as_millis() - ); - - // Cache the file contents - for (file, content) in file_contents { - let data = content.freeze(); - if let Some(size) = file_size_map.get(&file) - && *size != data.len() - { - log::warn!( - "[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: size mismatch, expected {} but got {}", - file, - addr, - size, - data.len() - ); - downloaded_files.remove(&file); - continue; - } - if let Err(e) = infra::cache::file_data::set(&file, data).await { - log::error!("[FILE_CACHE_DOWNLOAD:gRPC] Failed to cache file {file}: {e}"); - downloaded_files.remove(&file); - } - } - - // Return list of failed files - let failed_files: Vec<_> = files - .iter() - .filter(|(_, f, ..)| !downloaded_files.contains(f)) - .cloned() - .collect(); - - log::debug!( - "[FILE_CACHE_DOWNLOAD:gRPC] Failed to retrieve {} files from {} in {} ms", - failed_files.len(), - addr, - start.elapsed().as_millis() - ); - - Ok(failed_files) -} - -pub async fn queue_download( - trace_id: String, - id: i64, - account: String, - file: String, - size: i64, - ts: i64, - cache_type: file_data::CacheType, -) -> Result<(), anyhow::Error> { - log::debug!( - "[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] enqueue file: {file}, size: {size}, ts: {ts}" - ); - - // files with data older than the cache max age are not worth caching: with - // time_lru they become the oldest bucket and are evicted first, so the - // download would be wasted work. Skip the queue and let the query read - // them directly from object storage. - if exceeds_cache_max_age(ts, cache_type) { - log::debug!( - "[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] skip file: {file}, data is older than the cache max age" - ); - return Ok(()); - } - - let cfg = get_config(); - - // skip if already queued or already being downloaded - if !queued_files::try_add(&file) { - return Ok(()); - } - if processing_files::is_processing(&file) { - queued_files::remove(&file); - return Ok(()); - } - - if cfg.limit.file_download_enable_priority_queue - && should_prioritize_file(ts, cfg.limit.file_download_priority_queue_window_secs) - { - if PRIORITY_FILE_DOWNLOAD_CHANNEL - .push(( - trace_id, - id, - account, - file.clone(), - size as usize, - cache_type, - )) - .await - { - metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .inc(); - } else { - queued_files::remove(&file); - log::warn!( - "[FILE_CACHE_DOWNLOAD:JOB] priority queue full ({FILE_DOWNLOAD_QUEUE_SIZE}), dropping file: {file}" - ); - } - } else { - if let Err(e) = FILE_DOWNLOAD_CHANNEL - .sender - .send(( - trace_id, - id, - account, - file.clone(), - size as usize, - cache_type, - )) - .await - { - queued_files::remove(&file); - return Err(e.into()); - } - metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE - .with_label_values::<&str>(&[]) - .inc(); - } - Ok(()) -} - -// if the file timestamp is in the past window, it should be prioritized -fn should_prioritize_file(ts: i64, window_secs: i64) -> bool { - let window_micros = window_secs * 1_000_000; - let now = now_micros(); - ts > now - window_micros -} - -/// Returns true if the file's data is older than the cache max age and should -/// not be downloaded into the cache. -pub fn exceeds_cache_max_age(ts: i64, cache_type: file_data::CacheType) -> bool { - let cfg = get_config(); - let max_age_days = match cache_type { - file_data::CacheType::Memory => cfg.memory_cache.max_age_days, - file_data::CacheType::Disk => cfg.disk_cache.max_age_days, - file_data::CacheType::None => 0, - }; - exceeds_max_age(ts, max_age_days) -} - -// if the file data is older than max_age_days, it should not be cached. -// max_age_days == 0 means no limit, ts <= 0 means the timestamp is unknown. -fn exceeds_max_age(ts: i64, max_age_days: i64) -> bool { - if max_age_days <= 0 || ts <= 0 { - return false; - } - ts < now_micros() - day_micros(max_age_days) -} - -#[cfg(test)] -mod tests { - use config::utils::time::{day_micros, hour_micros, now_micros}; - - use super::{ - FileInfo, PriorityDownloadQueue, exceeds_max_age, file_data, processing_files, queued_files, - }; - - #[test] - fn test_exceeds_max_age_disabled() { - // max_age_days == 0 means no limit, nothing is too old - let one_year_ago = now_micros() - day_micros(365); - assert!(!exceeds_max_age(one_year_ago, 0)); - assert!(!exceeds_max_age(one_year_ago, -1)); - } - - #[test] - fn test_exceeds_max_age_unknown_ts() { - // unknown timestamp should not be skipped - assert!(!exceeds_max_age(0, 3)); - assert!(!exceeds_max_age(-1, 3)); - } - - #[test] - fn test_exceeds_max_age_recent_file() { - let one_hour_ago = now_micros() - hour_micros(1); - assert!(!exceeds_max_age(one_hour_ago, 3)); - } - - #[test] - fn test_exceeds_max_age_old_file() { - let thirty_days_ago = now_micros() - day_micros(30); - assert!(exceeds_max_age(thirty_days_ago, 3)); - } - - #[test] - fn test_is_processing_false_initially() { - assert!(!processing_files::is_processing( - "unique_not_added_file_abc123.parquet" - )); - } - - #[test] - fn test_add_and_is_processing() { - let name = "test_add_file_xyz999.parquet"; - processing_files::add(name); - assert!(processing_files::is_processing(name)); - processing_files::remove(name); - } - - #[test] - fn test_remove_clears_processing() { - let name = "test_remove_file_qqq888.parquet"; - processing_files::add(name); - processing_files::remove(name); - assert!(!processing_files::is_processing(name)); - } - - #[test] - fn test_normal_queue_cap_honored() { - let cap = 3; - let (tx, _rx) = tokio::sync::mpsc::channel::(cap); - - for i in 0..cap { - let result = tx.try_send(( - format!("trace-{i}"), - i as i64, - "org".into(), - format!("file-{i}.parquet"), - 1024, - file_data::CacheType::Disk, - )); - assert!(result.is_ok(), "send {i} should succeed under cap"); - } - - let overflow = tx.try_send(( - "trace-overflow".into(), - 99, - "org".into(), - "overflow.parquet".into(), - 1024, - file_data::CacheType::Disk, - )); - assert!(overflow.is_err(), "send beyond cap must fail"); - } - - #[test] - fn test_queued_files_no_duplicates() { - let file = "dedup_test_unique_zz1234.parquet"; - - // first add succeeds - assert!(queued_files::try_add(file)); - // same file rejected while already queued - assert!(!queued_files::try_add(file)); - - // after remove, can be queued again - queued_files::remove(file); - assert!(queued_files::try_add(file)); - - // cleanup - queued_files::remove(file); - } - - #[tokio::test] - async fn test_priority_queue_cap_honored() { - let cap = 3; - let queue = PriorityDownloadQueue::new(cap); - - // fill to cap — all succeed - for i in 0..cap { - let ok = queue - .push(( - format!("trace-{i}"), - i as i64, - "org".into(), - format!("file-{i}.parquet"), - 1024, - file_data::CacheType::Disk, - )) - .await; - assert!(ok, "push {i} should succeed under cap"); - } - - // one more must be rejected - let overflow = queue - .push(( - "trace-overflow".into(), - 99, - "org".into(), - "overflow.parquet".into(), - 1024, - file_data::CacheType::Disk, - )) - .await; - assert!(!overflow, "push beyond cap must return false"); - } -} +pub use infra::cache::file_downloader::*; diff --git a/src/core/src/service/ingestion/mod.rs b/src/core/src/service/ingestion/mod.rs index da1319d611..8c5b8589b3 100644 --- a/src/core/src/service/ingestion/mod.rs +++ b/src/core/src/service/ingestion/mod.rs @@ -29,7 +29,7 @@ use config::{ ider::SnowflakeIdGenerator, meta::{ alerts::alert::Alert, - function::{VRLResultResolver, VRLRuntimeConfig}, + function::VRLResultResolver, self_reporting::usage::{RequestStats, TriggerData, TriggerDataStatus, TriggerDataType}, stream::{PartitionTimeLevel, StreamParams, StreamPartition, StreamType}, }, @@ -41,10 +41,7 @@ use infra::{ schema::STREAM_RECORD_ID_GENERATOR, }; use proto::cluster_rpc::IngestionType; -use vrl::{ - compiler::{CompilationResult, TargetValueRef, runtime::Runtime}, - prelude::state, -}; +use vrl::compiler::{TargetValueRef, runtime::Runtime}; use super::{ db::{alerts::alert, pipeline}, @@ -55,11 +52,8 @@ use crate::{ common::{ infra::config::{REALTIME_ALERT_TRIGGERS, STREAM_ALERTS}, meta::{ingestion::IngestionRequest, stream::SchemaRecords}, - utils::{ - functions::get_vrl_compiler_config, - js::{ - JSRuntimeConfig, apply_js_fn as apply_js, compile_js_function as compile_js_func, - }, + utils::js::{ + JSRuntimeConfig, apply_js_fn as apply_js, compile_js_function as compile_js_func, }, }, service::{ @@ -72,6 +66,8 @@ use crate::{ pub mod grpc; pub mod ingestion_service; +pub use openobserve_vrl::compile_vrl_function; + pub type TriggerAlertData = Vec<(Alert, Vec>)>; /// Global atomic counter for round-robin distribution of requests across memory table buckets. @@ -90,34 +86,6 @@ pub fn get_thread_id() -> usize { REQUEST_COUNTER.fetch_add(1, Ordering::Relaxed) % cfg.limit.http_worker_num } -pub fn compile_vrl_function(func: &str, org_id: &str) -> Result { - if func.contains("get_env_var") { - return Err(std::io::Error::other("get_env_var is not supported")); - } - - let external = state::ExternalEnv::default(); - let vrl_config = get_vrl_compiler_config(org_id); - match vrl::compiler::compile_with_external( - func, - &vrl_config.functions, - &external, - vrl_config.config, - ) { - Ok(CompilationResult { - program, - warnings: _, - config, - }) => Ok(VRLRuntimeConfig { - program, - config, - fields: vec![], - }), - Err(e) => Err(std::io::Error::other( - vrl::diagnostic::Formatter::new(func, e).to_string(), - )), - } -} - /// Compile a JS function and return configuration /// This wraps the common::utils::js::compile_js_function pub fn compile_js_function(func: &str, org_id: &str) -> Result { diff --git a/src/core/src/service/search/cache/cacher.rs b/src/core/src/service/search/cache/cacher.rs index 4386bb7f88..24cfac2cb6 100644 --- a/src/core/src/service/search/cache/cacher.rs +++ b/src/core/src/service/search/cache/cacher.rs @@ -35,7 +35,7 @@ use crate::{ MultiCachedQueryResponse, result_utils::{get_ts_value, has_non_timestamp_ordering, is_timestamp_field}, }, - sql::{RE_HISTOGRAM, Sql, visitor::histogram_interval::generate_histogram_interval}, + sql::Sql, }, }; @@ -907,46 +907,6 @@ pub async fn delete_cache( Ok(true) } -pub fn handle_histogram( - origin_sql: &mut String, - q_time_range: (i64, i64), - histogram_interval: i64, -) { - let caps = if let Some(caps) = RE_HISTOGRAM.captures(origin_sql.as_str()) { - caps - } else { - return; - }; - - // 0th capture is the whole histogram(...) , - // 1st capture is the comma-delimited list of args - // ideally there should be at least one arg, otherwise df with anyways complain, - // so we we return from here if capture[1] is None - let args = match caps.get(1) { - Some(v) => v - .as_str() - .split(',') - .map(|v| v.trim().trim_matches(|v| v == '\'' || v == '"')) - .collect::>(), - None => return, - }; - - let interval = if histogram_interval > 0 { - format!("{histogram_interval} second") - } else { - args.get(1) - .map_or_else(|| generate_histogram_interval(q_time_range), |v| *v) - .to_string() - }; - - let field = args.first().unwrap_or(&"_timestamp"); - - *origin_sql = origin_sql.replace( - caps.get(0).unwrap().as_str(), - &format!("histogram({field},'{interval}')"), - ); -} - fn calculate_deltas_multi( results: &[CachedQueryResponse], start_time: i64, @@ -1164,16 +1124,6 @@ mod tests { assert_eq!(deltas, expected_deltas); } - #[test] - fn test_handle_histogram() { - // Test case 1: Basic histogram with numeric interval - let mut sql = "SELECT histogram(_timestamp, '10 seconds') FROM logs".to_string(); - let time_range = (1640995200000000, 1641081600000000); // 2022-01-01 to 2022-01-02 - handle_histogram(&mut sql, time_range, 10); - assert!(sql.contains("histogram(_timestamp,")); - assert!(sql.contains("second")); - } - #[test] fn test_get_ts_col_order_by() { let sql = Sql { diff --git a/src/core/src/service/search/cache/mod.rs b/src/core/src/service/search/cache/mod.rs index 06567492bd..545d151516 100644 --- a/src/core/src/service/search/cache/mod.rs +++ b/src/core/src/service/search/cache/mod.rs @@ -568,7 +568,11 @@ pub async fn prepare_cache_response( } else { sql.time_range }; - cacher::handle_histogram(&mut origin_sql, q_time_range, req.query.histogram_interval); + crate::service::search::sql::histogram::handle_histogram( + &mut origin_sql, + q_time_range, + req.query.histogram_interval, + ); } // calculate hash for the query with version (after normalizing histogram interval) diff --git a/src/core/src/service/search/cluster/flight.rs b/src/core/src/service/search/cluster/flight.rs index d6525ac952..a5f3cf5828 100644 --- a/src/core/src/service/search/cluster/flight.rs +++ b/src/core/src/service/search/cluster/flight.rs @@ -30,9 +30,7 @@ use config::{ metrics, utils::{json, time::now_micros, took_watcher::TookWatcher}, }; -use datafusion::{ - common::TableReference, physical_plan::visit_execution_plan, prelude::SessionContext, -}; +use datafusion::{common::TableReference, physical_plan::visit_execution_plan}; use hashbrown::{HashMap, HashSet}; use infra::{ cluster, @@ -56,16 +54,14 @@ use crate::service::{ search::{ SearchResult, datafusion::{ - exec::{DataFusionContextBuilder, register_udf}, + context::{SearchContextBuilder, register_table}, optimizer::{ context::{ PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext, }, - create_physical_plan, generate_analyzer_rules, generate_optimizer_rules, - generate_physical_optimizer_rules, + create_physical_plan, }, plan_metrics::get_peak_memory_from_ctx, - table_provider::{catalog::StreamTypeProvider, empty_table::NewEmptyTable}, }, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, sql::Sql, @@ -664,93 +660,6 @@ pub async fn partition_file_by_hash( partitions } -pub struct SearchContextBuilder { - pub target_partitions: usize, - pub contexts: Vec, -} - -impl Default for SearchContextBuilder { - fn default() -> Self { - Self::new() - } -} - -impl SearchContextBuilder { - pub fn new() -> Self { - Self { - target_partitions: 0, - contexts: vec![], - } - } - - pub fn target_partitions(mut self, target_partitions: usize) -> Self { - self.target_partitions = target_partitions; - self - } - - pub fn add_context(mut self, context: PhysicalOptimizerContext) -> Self { - self.contexts.push(context); - self - } - - pub async fn build(self, req: &Request, sql: &Arc) -> Result { - let analyzer_rules = generate_analyzer_rules(sql); - let optimizer_rules = generate_optimizer_rules(sql); - let physical_optimizer_rules = generate_physical_optimizer_rules(req, sql, self.contexts); - let mut ctx = DataFusionContextBuilder::new() - .trace_id(&req.trace_id) - .work_group(req.work_group.clone()) - .analyzer_rules(analyzer_rules) - .optimizer_rules(optimizer_rules) - .physical_optimizer_rules(physical_optimizer_rules) - .sorted_by_time(sql.sorted_by_time) - .build(self.target_partitions) - .await?; - - // register udf - register_udf(&ctx, &req.org_id)?; - datafusion_functions_json::register_all(&mut ctx)?; - - Ok(ctx) - } -} - -pub async fn register_table(ctx: &SessionContext, sql: &Sql) -> Result<()> { - // register schema provider - let mut registed_schema = HashSet::new(); - for (stream, _) in &sql.schemas { - let stream_type = stream.stream_type(); - if !stream.has_stream_type() || registed_schema.contains(&stream_type) { - continue; - } - registed_schema.insert(stream_type.clone()); - let schema_provider = StreamTypeProvider::create(&stream_type).await?; - let _ = ctx - .catalog("datafusion") - .unwrap() - .as_ref() - .register_schema(&stream_type, schema_provider); - } - - // register table - for (stream, schema) in &sql.schemas { - let schema = schema - .schema() - .as_ref() - .clone() - .with_metadata(Default::default()); - let stream_name = stream.to_quoted_string(); - let table = Arc::new( - NewEmptyTable::new(&stream_name, Arc::new(schema)) - .with_partitions(ctx.state().config().target_partitions()) - .with_sorted_by_time(sql.sorted_by_time), - ); - ctx.register_table(&stream_name, table)?; - } - - Ok(()) -} - #[tracing::instrument(name = "service:search:cluster:flight:get_file_id_lists", skip_all)] pub async fn get_file_id_lists( trace_id: &str, diff --git a/src/core/src/service/search/datafusion/context.rs b/src/core/src/service/search/datafusion/context.rs new file mode 100644 index 0000000000..a31584cbc0 --- /dev/null +++ b/src/core/src/service/search/datafusion/context.rs @@ -0,0 +1,114 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{collections::HashSet, sync::Arc}; + +use config::{datafusion::request::Request, meta::sql::TableReferenceExt}; +use datafusion::prelude::SessionContext; +use infra::errors::Result; + +use super::{ + exec::{DataFusionContextBuilder, register_udf}, + optimizer::{ + context::PhysicalOptimizerContext, generate_analyzer_rules, generate_optimizer_rules, + generate_physical_optimizer_rules, + }, + table_provider::{catalog::StreamTypeProvider, empty_table::NewEmptyTable}, +}; +use crate::service::search::sql::Sql; + +pub struct SearchContextBuilder { + pub target_partitions: usize, + pub contexts: Vec, +} + +impl Default for SearchContextBuilder { + fn default() -> Self { + Self::new() + } +} + +impl SearchContextBuilder { + pub fn new() -> Self { + Self { + target_partitions: 0, + contexts: vec![], + } + } + + pub fn target_partitions(mut self, target_partitions: usize) -> Self { + self.target_partitions = target_partitions; + self + } + + pub fn add_context(mut self, context: PhysicalOptimizerContext) -> Self { + self.contexts.push(context); + self + } + + pub async fn build(self, req: &Request, sql: &Arc) -> Result { + let analyzer_rules = generate_analyzer_rules(sql); + let optimizer_rules = generate_optimizer_rules(sql); + let physical_optimizer_rules = generate_physical_optimizer_rules(req, sql, self.contexts); + let mut ctx = DataFusionContextBuilder::new() + .trace_id(&req.trace_id) + .work_group(req.work_group.clone()) + .analyzer_rules(analyzer_rules) + .optimizer_rules(optimizer_rules) + .physical_optimizer_rules(physical_optimizer_rules) + .sorted_by_time(sql.sorted_by_time) + .build(self.target_partitions) + .await?; + + register_udf(&ctx, &sql.org_id)?; + datafusion_functions_json::register_all(&mut ctx)?; + + Ok(ctx) + } +} + +pub async fn register_table(ctx: &SessionContext, sql: &Sql) -> Result<()> { + let mut registered_schemas = HashSet::new(); + for (stream, _) in &sql.schemas { + let stream_type = stream.stream_type(); + if !stream.has_stream_type() || registered_schemas.contains(&stream_type) { + continue; + } + registered_schemas.insert(stream_type.clone()); + let schema_provider = StreamTypeProvider::create(&stream_type).await?; + let _ = ctx + .catalog("datafusion") + .unwrap() + .as_ref() + .register_schema(&stream_type, schema_provider); + } + + for (stream, schema) in &sql.schemas { + let schema = schema + .schema() + .as_ref() + .clone() + .with_metadata(Default::default()); + let stream_name = stream.to_quoted_string(); + let table = Arc::new( + NewEmptyTable::new(&stream_name, Arc::new(schema)) + .with_partitions(ctx.state().config().target_partitions()) + .with_sorted_by_time(sql.sorted_by_time), + ); + ctx.register_table(&stream_name, table)?; + } + + Ok(()) +} diff --git a/src/core/src/service/search/datafusion/distributed_plan/rewrite.rs b/src/core/src/service/search/datafusion/distributed_plan/rewrite.rs index 8e002f8815..c0cc83d80a 100644 --- a/src/core/src/service/search/datafusion/distributed_plan/rewrite.rs +++ b/src/core/src/service/search/datafusion/distributed_plan/rewrite.rs @@ -15,6 +15,7 @@ use std::sync::Arc; +use ::search::types::QueryParams; use config::meta::{inverted_index::IndexOptimizeMode, stream::FileKey}; use datafusion::{ arrow::datatypes::SchemaRef, @@ -32,8 +33,7 @@ use datafusion::{ use o2_enterprise::enterprise::search::datafusion::distributed_plan::metadata_count_exec::MetadataCountExec; use crate::service::search::{ - datafusion::plan::tantivy_optimize_exec::TantivyOptimizeExec, grpc::QueryParams, - index::IndexCondition, + datafusion::plan::tantivy_optimize_exec::TantivyOptimizeExec, index::IndexCondition, }; pub fn aggregate_optimize_rewrite( diff --git a/src/core/src/service/search/datafusion/mod.rs b/src/core/src/service/search/datafusion/mod.rs index d16d74f37f..307a28a5b2 100644 --- a/src/core/src/service/search/datafusion/mod.rs +++ b/src/core/src/service/search/datafusion/mod.rs @@ -15,6 +15,7 @@ use std::str::FromStr; +pub mod context; pub mod distributed_plan; pub mod exec; pub mod merge; diff --git a/src/core/src/service/search/datafusion/plan/projections.rs b/src/core/src/service/search/datafusion/plan/projections.rs index d8b985a7f4..fa204ab8ec 100644 --- a/src/core/src/service/search/datafusion/plan/projections.rs +++ b/src/core/src/service/search/datafusion/plan/projections.rs @@ -27,9 +27,8 @@ use datafusion::{ use hashbrown::HashSet; use crate::service::search::{ - cache::cacher::handle_histogram, - cluster::flight::{SearchContextBuilder, register_table}, - sql::Sql, + datafusion::context::{SearchContextBuilder, register_table}, + sql::{Sql, histogram::handle_histogram}, }; /// Structure to store the result schema info diff --git a/src/core/src/service/search/datafusion/plan/regex_projections.rs b/src/core/src/service/search/datafusion/plan/regex_projections.rs index 8889ce2934..edb8d60363 100644 --- a/src/core/src/service/search/datafusion/plan/regex_projections.rs +++ b/src/core/src/service/search/datafusion/plan/regex_projections.rs @@ -19,7 +19,7 @@ use config::{datafusion::request::Request, meta::projections::ProjectionColumnMa use datafusion::common::tree_node::TreeNode; use crate::service::search::{ - cluster::flight::{SearchContextBuilder, register_table}, + datafusion::context::{SearchContextBuilder, register_table}, sql::Sql, }; diff --git a/src/core/src/service/search/datafusion/plan/tantivy_optimize_exec.rs b/src/core/src/service/search/datafusion/plan/tantivy_optimize_exec.rs index fb6f765aaf..f54317802d 100644 --- a/src/core/src/service/search/datafusion/plan/tantivy_optimize_exec.rs +++ b/src/core/src/service/search/datafusion/plan/tantivy_optimize_exec.rs @@ -15,6 +15,7 @@ use std::{collections::HashSet, sync::Arc}; +use ::search::types::QueryParams; use arrow::array::{ Array, Int64Array, TimestampMicrosecondArray, TimestampNanosecondArray, UInt64Array, }; @@ -41,10 +42,7 @@ use datafusion::{ }; use futures::TryStreamExt; -use crate::service::search::{ - grpc::{QueryParams, storage::tantivy_search}, - index::IndexCondition, -}; +use crate::service::search::{index::IndexCondition, tantivy::tantivy_search}; #[derive(Debug)] pub struct TantivyOptimizeExec { diff --git a/src/core/src/service/search/datafusion/udf/cipher_udf.rs b/src/core/src/service/search/datafusion/udf/cipher_udf.rs index 93e325b568..850d6b8426 100644 --- a/src/core/src/service/search/datafusion/udf/cipher_udf.rs +++ b/src/core/src/service/search/datafusion/udf/cipher_udf.rs @@ -18,6 +18,7 @@ const BASE64_CHARS: &str = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ use std::sync::{Arc, LazyLock as Lazy}; +use ::search::cipher::registry::REGISTRY; use arrow::array::{ArrayRef, StringArray}; use config::utils::str; use datafusion::{ @@ -32,8 +33,6 @@ use datafusion::{ use o2_enterprise::enterprise::cipher::Cipher; use serde_json::Value; -use crate::cipher::registry::REGISTRY; - /// The name of the decrypt UDF given to DataFusion. pub const DECRYPT_UDF_NAME: &str = "decrypt_path"; /// The name of the decrypt UDF given to DataFusion. diff --git a/src/core/src/service/search/datafusion/udf/transform_udf.rs b/src/core/src/service/search/datafusion/udf/transform_udf.rs index 83a2f08d5e..269a3ff276 100644 --- a/src/core/src/service/search/datafusion/udf/transform_udf.rs +++ b/src/core/src/service/search/datafusion/udf/transform_udf.rs @@ -26,11 +26,10 @@ use datafusion::{ logical_expr::{ColumnarValue, ScalarUDF, Volatility}, prelude::create_udf, }; +use openobserve_vrl::{QUERY_FUNCTIONS, compile_vrl_function}; use vector_enrichment::TableRegistry; use vrl::compiler::{TargetValueRef, VrlRuntime, runtime::Runtime}; -use crate::{common::infra::config::QUERY_FUNCTIONS, service::ingestion::compile_vrl_function}; - type FnType = Arc Result + Sync + Send>; fn create_user_df(fn_name: &str, num_args: u8, pow_scalar: FnType) -> ScalarUDF { @@ -49,13 +48,11 @@ fn create_user_df(fn_name: &str, num_args: u8, pow_scalar: FnType) -> ScalarUDF pub fn get_all_transform(org_id: &str) -> Result> { let mut udf_list = Vec::new(); - for transform in QUERY_FUNCTIONS.clone().iter() { - let key = transform.key(); - // do not register ingest_time transforms - if key.contains(org_id) { + for transform in QUERY_FUNCTIONS.iter() { + if transform.key().contains(org_id) { udf_list.push(get_udf_vrl( - transform.name.to_owned(), - transform.function.to_owned().as_str(), + transform.name.clone(), + transform.function.as_str(), &transform.params, transform.num_args, org_id, @@ -67,12 +64,12 @@ pub fn get_all_transform(org_id: &str) -> Result> { fn get_udf_vrl( fn_name: String, - func: &str, + source: &str, params: &str, num_args: u8, org_id: &str, ) -> Result { - let local_func = func.trim().to_owned(); + let local_func = source.trim().to_owned(); let local_fn_params = params.to_owned(); let local_org_id = org_id.to_owned(); @@ -95,10 +92,10 @@ fn get_udf_vrl( } obj_str.push_str(&format!(" \n {}", local_func)); match compile_vrl_function(&obj_str, &local_org_id) { - Ok(res) => { - let registry = res.config.get_custom::().unwrap(); + Ok(result) => { + let registry = result.config.get_custom::().unwrap(); registry.finish_load(); - let result = apply_vrl_fn(&mut runtime, res.program); + let result = apply_vrl_fn(&mut runtime, result.program); if result != json::Value::Null { res_data_vec.insert(i, json::get_string_value(&result)); } else { diff --git a/src/core/src/service/search/grpc/mod.rs b/src/core/src/service/search/grpc/mod.rs index 8b3b862c2f..1f8d332030 100644 --- a/src/core/src/service/search/grpc/mod.rs +++ b/src/core/src/service/search/grpc/mod.rs @@ -16,14 +16,8 @@ use std::sync::Arc; use arrow_schema::Schema; -use config::meta::{ - search::ScanStats, - stream::{FileKey, StreamType}, -}; -use datafusion::{ - datasource::TableProvider, execution::cache::cache_manager::FileStatisticsCache, - sql::TableReference, -}; +use config::meta::{search::ScanStats, stream::FileKey}; +use datafusion::{datasource::TableProvider, execution::cache::cache_manager::FileStatisticsCache}; use hashbrown::HashSet; use infra::errors::Result; @@ -32,29 +26,12 @@ use super::{datafusion::exec::TableBuilder, index::IndexCondition}; pub mod flight; pub mod storage; pub mod wal; +pub use ::search::{file_cache::calc_target_partitions, types::QueryParams}; + pub use super::tantivy::cache as tantivy_result_cache; pub type SearchTable = Result<(Vec>, ScanStats, HashSet)>; -#[derive(Debug)] -pub struct QueryParams { - pub trace_id: String, - pub org_id: String, - pub stream: TableReference, - pub stream_type: StreamType, - pub stream_name: String, - pub time_range: (i64, i64), - pub work_group: Option, - pub use_inverted_index: bool, -} - -/// Linear interpolation: cached_ratio=0 -> query_thread_num, cached_ratio=1 -> cpu_num. -pub fn calc_target_partitions(cpu_num: usize, query_thread_num: usize, cached_ratio: f64) -> usize { - (cpu_num as i64 - + ((query_thread_num as i64 - cpu_num as i64) as f64 * (1.0 - cached_ratio)) as i64) - as usize -} - /// Create tables from files, automatically splitting them based on time range overlap: /// - Files completely within the query time range: no timestamp filter applied /// - Files partially overlapping with the query time range: timestamp filter applied diff --git a/src/core/src/service/search/grpc/storage.rs b/src/core/src/service/search/grpc/storage.rs index 248a1fd554..d66f8f2e41 100644 --- a/src/core/src/service/search/grpc/storage.rs +++ b/src/core/src/service/search/grpc/storage.rs @@ -15,10 +15,11 @@ use std::sync::Arc; +pub use ::search::file_cache::{cache_files, calc_target_partitions}; use arrow_schema::Schema; use config::{ cluster::LOCAL_NODE, - get_config, is_local_disk_storage, + get_config, meta::{ inverted_index::IndexOptimizeMode, search::{ScanStats, StorageType}, @@ -36,7 +37,6 @@ use infra::{ use itertools::Itertools; use tracing::Instrument; -pub use super::calc_target_partitions; pub use crate::service::search::tantivy::tantivy_search; use crate::service::{ file_list, @@ -297,134 +297,6 @@ pub async fn search( Ok((tables, scan_stats, HashSet::new())) } -#[tracing::instrument(name = "service:search:grpc:storage:cache_files", skip_all)] -pub async fn cache_files( - trace_id: &str, - files: &[(i64, &String, &String, i64, i64)], - scan_stats: &mut ScanStats, - file_type: &str, -) -> (file_data::CacheType, u64, u64) { - // check how many files already cached - let mut cached_files = HashSet::with_capacity(files.len()); - let (mut cache_hits, mut cache_misses) = (0, 0); - - let start = std::time::Instant::now(); - for (_id, _account, file, _size, max_ts) in files.iter() { - if file_data::memory::exist(file).await { - scan_stats.querier_memory_cached_files += 1; - cached_files.insert(file); - cache_hits += 1; - } else if file_data::disk::exist(file).await { - scan_stats.querier_disk_cached_files += 1; - cached_files.insert(file); - cache_hits += 1; - } else { - cache_misses += 1; - }; - - // Record file access metrics - let stream_type = if file_type == "index" { - config::meta::stream::StreamType::Index - } else { - // Determine stream type from the file path - if file.contains("/logs/") { - config::meta::stream::StreamType::Logs - } else if file.contains("/metrics/") { - config::meta::stream::StreamType::Metrics - } else if file.contains("/traces/") { - config::meta::stream::StreamType::Traces - } else { - config::meta::stream::StreamType::Logs // Default - } - }; - - let current_time = chrono::Utc::now().timestamp_micros(); - let file_age_seconds = (current_time - max_ts) / 1_000_000; - let file_age_hours = file_age_seconds as f64 / 3600.0; - - if file_age_hours > 0.0 { - config::metrics::FILE_ACCESS_TIME - .with_label_values(&[&stream_type.to_string()]) - .observe(file_age_hours); - } - } - - let check_cache_took = start.elapsed().as_millis() as usize; - if check_cache_took > 1000 { - log::warn!( - "[trace_id {trace_id}] search->storage: check file cache took: {check_cache_took} ms", - ); - } - - let files_num = files.len() as i64; - if files_num == scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files { - // all files are cached - return (file_data::CacheType::Disk, cache_hits, cache_misses); - } - - // check cache size - let cfg = get_config(); - let cache_type = if cfg.memory_cache.enabled - && scan_stats.compressed_size < cfg.memory_cache.skip_size as i64 - { - // if scan_compressed_size < ZO_MEMORY_CACHE_SKIP_SIZE, use memory cache - file_data::CacheType::Memory - } else if !is_local_disk_storage() - && cfg.disk_cache.enabled - && scan_stats.compressed_size < cfg.disk_cache.skip_size as i64 - { - // if scan_compressed_size < ZO_DISK_CACHE_SKIP_SIZE, use disk cache - file_data::CacheType::Disk - } else { - // no cache, the files are too big than cache size - return (file_data::CacheType::None, cache_hits, cache_misses); - }; - - let trace_id = trace_id.to_string(); - let files = files - .iter() - .filter_map(|(id, account, file, size, ts)| { - if cached_files.contains(&file) { - None - } else { - Some((*id, account.to_string(), file.to_string(), *size, *ts)) - } - }) - .collect_vec(); - let file_type = file_type.to_string(); - tokio::spawn(async move { - let files_num = files.len(); - for (id, account, file, size, ts) in files { - if let Err(e) = crate::service::file_downloader::queue_download( - trace_id.clone(), - id, - account, - file.clone(), - size, - ts, - cache_type, - ) - .await - { - log::error!( - "[trace_id {trace_id}] error in queuing file {file} for background download: {e}" - ); - } - } - log::info!( - "[trace_id {trace_id}] search->storage: successfully enqueued {files_num} files of {file_type} for background download into {cache_type:?}", - ); - }); - - // if cached file less than 50% of the total files, return None - if scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files < files_num / 2 - { - (file_data::CacheType::None, cache_hits, cache_misses) - } else { - (cache_type, cache_hits, cache_misses) - } -} - /// Check bloom filter for the file list #[tracing::instrument(name = "service:search:grpc:storage:check_bloom_filter", skip_all)] pub async fn check_bloom_filter( diff --git a/src/core/src/service/search/sql/histogram.rs b/src/core/src/service/search/sql/histogram.rs index 6e6d8ec0c7..672d957cb7 100644 --- a/src/core/src/service/search/sql/histogram.rs +++ b/src/core/src/service/search/sql/histogram.rs @@ -23,6 +23,45 @@ use sqlparser::{ parser::Parser, }; +use super::{RE_HISTOGRAM, visitor::histogram_interval::generate_histogram_interval}; + +pub fn handle_histogram( + origin_sql: &mut String, + q_time_range: (i64, i64), + histogram_interval: i64, +) { + let Some(captures) = RE_HISTOGRAM.captures(origin_sql.as_str()) else { + return; + }; + + let args = match captures.get(1) { + Some(args) => args + .as_str() + .split(',') + .map(|value| { + value + .trim() + .trim_matches(|value| value == '\'' || value == '"') + }) + .collect::>(), + None => return, + }; + + let interval = if histogram_interval > 0 { + format!("{histogram_interval} second") + } else { + args.get(1) + .map_or_else(|| generate_histogram_interval(q_time_range), |value| *value) + .to_string() + }; + let field = args.first().unwrap_or(&"_timestamp"); + + *origin_sql = origin_sql.replace( + captures.get(0).unwrap().as_str(), + &format!("histogram({field},'{interval}')"), + ); +} + /// Converts an original query to a histogram query /// Extracts WHERE clause and builds histogram query with provided stream name pub fn convert_to_histogram_query( @@ -201,6 +240,16 @@ pub fn histogram_bucket_start(timestamp_us: i64, interval_us: i64) -> i64 { mod tests { use super::*; + #[test] + fn test_handle_histogram() { + // Test case 1: Basic histogram with numeric interval + let mut sql = "SELECT histogram(_timestamp, '10 seconds') FROM logs".to_string(); + let time_range = (1640995200000000, 1641081600000000); // 2022-01-01 to 2022-01-02 + handle_histogram(&mut sql, time_range, 10); + assert!(sql.contains("histogram(_timestamp,")); + assert!(sql.contains("second")); + } + #[test] fn test_convert_simple_query() { let original_query = "SELECT * FROM \"logs\" WHERE status = 500"; diff --git a/src/core/src/service/search/sql/visitor/group_by.rs b/src/core/src/service/search/sql/visitor/group_by.rs index 7524713821..f10426e3ad 100644 --- a/src/core/src/service/search/sql/visitor/group_by.rs +++ b/src/core/src/service/search/sql/visitor/group_by.rs @@ -28,7 +28,7 @@ use hashbrown::HashSet; use infra::errors::Error; use crate::service::search::{ - cluster::flight::{SearchContextBuilder, register_table}, + datafusion::context::{SearchContextBuilder, register_table}, sql::Sql, }; diff --git a/src/core/src/service/search/super_cluster/leader.rs b/src/core/src/service/search/super_cluster/leader.rs index c324e54902..317c4cdace 100644 --- a/src/core/src/service/search/super_cluster/leader.rs +++ b/src/core/src/service/search/super_cluster/leader.rs @@ -44,10 +44,12 @@ use tracing_opentelemetry::OpenTelemetrySpanExt; use crate::service::search::{ SEARCH_SERVER, SearchResult, - cluster::flight::{SearchContextBuilder, register_table}, - datafusion::optimizer::{ - context::{PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext}, - create_physical_plan, + datafusion::{ + context::{SearchContextBuilder, register_table}, + optimizer::{ + context::{PhysicalOptimizerContext, RemoteScanContext, StreamingAggregationContext}, + create_physical_plan, + }, }, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, sql::Sql, diff --git a/src/core/src/service/search/tantivy/mod.rs b/src/core/src/service/search/tantivy/mod.rs index 2b6086b0ce..2720a909e6 100644 --- a/src/core/src/service/search/tantivy/mod.rs +++ b/src/core/src/service/search/tantivy/mod.rs @@ -22,6 +22,10 @@ mod warm; use std::{ops::Bound, sync::Arc}; +use ::search::{ + file_cache::{cache_files, calc_target_partitions}, + types::QueryParams, +}; use arrow::{ buffer::{BooleanBuffer, MutableBuffer}, util::bit_util, @@ -62,7 +66,6 @@ use self::{ warm::WarmPlan, }; use crate::service::search::{ - grpc::{QueryParams, calc_target_partitions, storage::cache_files}, index::IndexCondition, inspector::{SearchInspectorFieldsBuilder, search_inspector_fields}, }; diff --git a/src/core/src/service/search/utils.rs b/src/core/src/service/search/utils.rs index e4a615cf1d..ad014dfb95 100644 --- a/src/core/src/service/search/utils.rs +++ b/src/core/src/service/search/utils.rs @@ -19,6 +19,7 @@ use std::{ sync::{Arc, atomic::Ordering}, }; +use ::search::CAPPED_RESULTS_MSG; use config::meta::search::{PARTIAL_ERROR_RESPONSE_MESSAGE, ScanStats}; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanVisitor}; use infra::runtime::DATAFUSION_RUNTIME; @@ -26,7 +27,7 @@ use sqlparser::ast::{BinaryOperator, Expr}; use tokio::sync::Mutex; use super::datafusion::distributed_plan::remote_scan_exec::RemoteScanExec; -use crate::{common::meta::search::CAPPED_RESULTS_MSG, service::search::sql::Sql}; +use crate::service::search::sql::Sql; type Cleanup = Pin + Send>>; diff --git a/src/infra/src/cache/file_downloader.rs b/src/infra/src/cache/file_downloader.rs new file mode 100644 index 0000000000..69452025f2 --- /dev/null +++ b/src/infra/src/cache/file_downloader.rs @@ -0,0 +1,733 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{ + collections::VecDeque, + sync::{Arc, LazyLock as Lazy}, +}; + +use config::{ + cluster::LOCAL_NODE, + get_config, + meta::cluster::{Role, RoleGroup, get_internal_grpc_token}, + metrics, + utils::time::{day_micros, now_micros}, +}; +use futures::StreamExt; +use hashbrown::{HashMap, HashSet}; +use proto::cluster_rpc::{SimpleFileList, event_client::EventClient}; +use tokio::sync::{ + Mutex, + mpsc::{Receiver, Sender}, +}; +use tonic::{codec::CompressionEncoding, metadata::MetadataValue}; + +use crate::{cache::file_data, cluster}; + +/// (trace_id, file_id, account, file, size, cache_type) +type FileInfo = (String, i64, String, String, usize, file_data::CacheType); + +mod processing_files { + use hashbrown::HashSet; + use parking_lot::RwLock; + + use super::*; + + static PROCESSING_FILES: Lazy>> = + Lazy::new(|| RwLock::new(HashSet::new())); + + pub fn is_processing(file_name: &str) -> bool { + PROCESSING_FILES.read().contains(file_name) + } + + pub fn add(file_name: &str) { + PROCESSING_FILES.write().insert(file_name.to_string()); + } + + pub fn remove(file_name: &str) { + PROCESSING_FILES.write().remove(file_name); + } +} + +mod queued_files { + use hashbrown::HashSet; + use parking_lot::RwLock; + + use super::*; + + static QUEUED_FILES: Lazy>> = Lazy::new(|| RwLock::new(HashSet::new())); + + /// Returns true if newly inserted (caller should enqueue), false if already present (skip). + pub fn try_add(file_name: &str) -> bool { + QUEUED_FILES.write().insert(file_name.to_string()) + } + + pub fn remove(file_name: &str) { + QUEUED_FILES.write().remove(file_name); + } +} + +struct DownloadQueue { + sender: Sender, + receiver: Arc>>, +} + +impl DownloadQueue { + fn new(sender: Sender, receiver: Arc>>) -> Self { + Self { sender, receiver } + } +} + +struct PriorityDownloadQueue { + stack: Mutex>, + notify: tokio::sync::Notify, + max_size: usize, +} + +impl PriorityDownloadQueue { + fn new(max_size: usize) -> Self { + Self { + stack: Mutex::new(VecDeque::with_capacity(max_size)), + notify: tokio::sync::Notify::new(), + max_size, + } + } + + // Returns false if at cap — caller must remove from queued_files and skip metric inc. + async fn push(&self, file_info: FileInfo) -> bool { + let mut stack = self.stack.lock().await; + if stack.len() >= self.max_size { + return false; + } + stack.push_back(file_info); + drop(stack); + self.notify.notify_one(); + true + } + + // Blocks until an item is available. LIFO via pop_back. + // notified() is created before the lock check to avoid missed-wakeup race. + // Chains notify_one() so all workers drain the queue under burst load. + async fn pop(&self) -> FileInfo { + loop { + let notified = self.notify.notified(); + { + let mut stack = self.stack.lock().await; + if let Some(item) = stack.pop_back() { + if !stack.is_empty() { + self.notify.notify_one(); + } + return item; + } + } + notified.await; + } + } +} + +const FILE_DOWNLOAD_QUEUE_SIZE: usize = 10000; +static FILE_DOWNLOAD_CHANNEL: Lazy = Lazy::new(|| { + let (tx, rx) = tokio::sync::mpsc::channel::(FILE_DOWNLOAD_QUEUE_SIZE); + DownloadQueue::new(tx, Arc::new(Mutex::new(rx))) +}); + +static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy = + Lazy::new(|| PriorityDownloadQueue::new(FILE_DOWNLOAD_QUEUE_SIZE)); + +pub async fn run() -> Result<(), anyhow::Error> { + let cfg = get_config(); + // worker tasks: handle normal queue download (FIFO via mpsc channel) + for thread in 0..cfg.limit.file_download_thread_num { + let rx = FILE_DOWNLOAD_CHANNEL.receiver.clone(); + tokio::spawn(async move { + loop { + let ret = rx.lock().await.recv().await; + match ret { + None => { + log::debug!("[FILE_CACHE_DOWNLOAD:JOB:NORMAL] Receiving channel is closed"); + break; + } + Some((trace_id, id, account, file, file_size, cache)) => { + // transition: queued → not-queued (must run before any skip path) + queued_files::remove(&file); + + // check if the file is already being downloaded + if processing_files::is_processing(&file) { + log::warn!( + "[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it" + ); + // update metrics + metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .dec(); + continue; + } + + // add the file to processing set + processing_files::add(&file); + + // download the file + match download_file( + thread, &trace_id, id, &account, &file, file_size, cache, + ) + .await + { + Ok(data_len) => { + if data_len > 0 && data_len != file_size { + log::warn!( + "[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it", + ); + } + } + Err(e) => { + log::error!( + "[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {file} to cache {cache:?} err: {e}", + ); + } + } + + // remove the file from processing set + processing_files::remove(&file); + + // update metrics + metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .dec(); + } + } + } + }); + } + + // worker tasks: handle priority queue download (LIFO via VecDeque::pop_back) + for thread in 0..cfg.limit.file_download_priority_queue_thread_num { + tokio::spawn(async move { + loop { + let (trace_id, id, account, file, file_size, cache) = + PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await; + + // transition: queued → not-queued (must run before any skip path) + queued_files::remove(&file); + + // check if the file is already being downloaded + if processing_files::is_processing(&file) { + log::warn!( + "[trace_id {trace_id}] [thread {thread}] search->storage: file {file} is already being downloaded, will skip it" + ); + metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .dec(); + continue; + } + + // add the file to processing set + processing_files::add(&file); + + // download the file + match download_file(thread, &trace_id, id, &account, &file, file_size, cache).await + { + Ok(data_len) => { + if data_len > 0 && data_len != file_size { + log::warn!( + "[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} found size mismatch, expected: {file_size}, actual: {data_len}, will skip it", + ); + } + } + Err(e) => { + log::error!( + "[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {file} to cache {cache:?} err: {e}", + ); + } + } + + // remove the file from processing set + processing_files::remove(&file); + + // update metrics + metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .dec(); + } + }); + } + Ok(()) +} + +async fn download_file( + thread: usize, + trace_id: &str, + file_id: i64, + account: &str, + file_name: &str, + file_size: usize, + cache_type: file_data::CacheType, +) -> Result { + let cfg = get_config(); + + // download file from node + if cfg.cache_latest_files.download_from_node + && let Ok(ok) = + download_file_with_consistent_hash(file_id, account, file_name, file_size).await + && ok + { + return Ok(file_size); + } + + // download from object store + let size = if file_size > 0 { Some(file_size) } else { None }; + let start = std::time::Instant::now(); + let ret = match cache_type { + file_data::CacheType::Memory => { + let mut disk_exists = false; + let mem_exists = file_data::memory::exist(file_name).await; + if !mem_exists && !cfg.memory_cache.skip_disk_check { + disk_exists = file_data::disk::exist(file_name).await; + } + if !mem_exists && (cfg.memory_cache.skip_disk_check || !disk_exists) { + file_data::memory::download(account, file_name, size).await + } else { + Ok(0) + } + } + file_data::CacheType::Disk => { + if !file_data::disk::exist(file_name).await { + file_data::disk::download(account, file_name, size).await + } else { + Ok(0) + } + } + _ => Ok(0), + }; + log::debug!( + "[FILE_CACHE_DOWNLOAD:JOB:{thread}] [trace_id {trace_id}] download file: {file_name}, ret: {:?}, took: {} ms", + ret, + start.elapsed().as_millis() + ); + ret +} + +async fn download_file_with_consistent_hash( + file_id: i64, + account: &str, + file_name: &str, + file_size: usize, +) -> Result { + let role_group = if LOCAL_NODE.is_interactive_querier() { + RoleGroup::Interactive + } else { + RoleGroup::Background + }; + let Some(node_name) = cluster::get_node_from_consistent_hash( + &file_id.to_string(), + &Role::Querier, + Some(role_group), + ) + .await + else { + return Ok(false); + }; + // get node by file_id + let Some(node) = cluster::get_cached_node_by_name(&node_name).await else { + return Ok(false); + }; + // download file from node + let Ok(failed) = download_from_node( + &node.grpc_addr, + &[( + file_id, + account.to_string(), + file_name.to_string(), + file_size as i64, + 0, + )], + ) + .await + else { + return Ok(false); + }; + // failed is empty means download success + Ok(failed.is_empty()) +} + +// download files from node and return download failed files +// file: (account, file, size, ts) +pub async fn download_from_node( + addr: &str, + files: &[(i64, String, String, i64, i64)], +) -> Result, anyhow::Error> { + let start = std::time::Instant::now(); + let cfg = get_config(); + log::debug!( + "[FILE_CACHE_DOWNLOAD:gRPC] Download files from node start, files: {:?}", + files.len() + ); + + let token: MetadataValue<_> = get_internal_grpc_token() + .parse() + .map_err(|_| anyhow::anyhow!("Invalid token"))?; + + let channel = crate::client::grpc::get_cached_channel(addr).await?; + let client = EventClient::with_interceptor(channel, move |mut req: tonic::Request<()>| { + req.metadata_mut().insert("authorization", token.clone()); + Ok(req) + }); + + let file_size_map = files + .iter() + .filter_map(|(_, _, f, s, _)| { + if *s > cfg.cache_latest_files.download_node_size * 1024 * 1024 { + None + } else { + Some((f, *s as usize)) + } + }) + .collect::>(); + if file_size_map.is_empty() { + return Ok(files.to_vec()); + } + let request = tonic::Request::new(SimpleFileList { + files: files.iter().map(|(_, _, f, ..)| f.to_string()).collect(), + }); + + let resp = client + .send_compressed(CompressionEncoding::Gzip) + .accept_compressed(CompressionEncoding::Gzip) + .max_decoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) + .max_encoding_message_size(cfg.grpc.max_message_size * 1024 * 1024) + .get_files(request) + .await + .map_err(|e| anyhow::anyhow!("Failed to get files from {addr}, {e}"))?; + + let mut file_contents = HashMap::new(); + let mut downloaded_files = HashSet::new(); + let mut resp_stream = resp.into_inner(); + while let Some(resp) = resp_stream.next().await { + let resp = match resp { + Ok(resp) => resp, + Err(err) => { + if err.code() == tonic::Code::NotFound { + log::debug!( + "[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: file not found", + err.message(), + addr + ); + continue; + } + return Err(anyhow::anyhow!( + "Failed to download file from {addr}, {err}" + )); + } + }; + for content in resp.entries { + let entry = file_contents + .entry(content.filename.clone()) + .or_insert(bytes::BytesMut::new()); + entry.extend_from_slice(&content.content); + downloaded_files.insert(content.filename); + } + } + + log::debug!( + "[FILE_CACHE_DOWNLOAD:gRPC] Successfully retrieved {} files from {} in {} ms", + downloaded_files.len(), + addr, + start.elapsed().as_millis() + ); + + // Cache the file contents + for (file, content) in file_contents { + let data = content.freeze(); + if let Some(size) = file_size_map.get(&file) + && *size != data.len() + { + log::warn!( + "[FILE_CACHE_DOWNLOAD:gRPC] Failed to download file {} from {}: size mismatch, expected {} but got {}", + file, + addr, + size, + data.len() + ); + downloaded_files.remove(&file); + continue; + } + if let Err(e) = file_data::set(&file, data).await { + log::error!("[FILE_CACHE_DOWNLOAD:gRPC] Failed to cache file {file}: {e}"); + downloaded_files.remove(&file); + } + } + + // Return list of failed files + let failed_files: Vec<_> = files + .iter() + .filter(|(_, f, ..)| !downloaded_files.contains(f)) + .cloned() + .collect(); + + log::debug!( + "[FILE_CACHE_DOWNLOAD:gRPC] Failed to retrieve {} files from {} in {} ms", + failed_files.len(), + addr, + start.elapsed().as_millis() + ); + + Ok(failed_files) +} + +pub async fn queue_download( + trace_id: String, + id: i64, + account: String, + file: String, + size: i64, + ts: i64, + cache_type: file_data::CacheType, +) -> Result<(), anyhow::Error> { + log::debug!( + "[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] enqueue file: {file}, size: {size}, ts: {ts}" + ); + + // files with data older than the cache max age are not worth caching: with + // time_lru they become the oldest bucket and are evicted first, so the + // download would be wasted work. Skip the queue and let the query read + // them directly from object storage. + if exceeds_cache_max_age(ts, cache_type) { + log::debug!( + "[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] skip file: {file}, data is older than the cache max age" + ); + return Ok(()); + } + + let cfg = get_config(); + + // skip if already queued or already being downloaded + if !queued_files::try_add(&file) { + return Ok(()); + } + if processing_files::is_processing(&file) { + queued_files::remove(&file); + return Ok(()); + } + + if cfg.limit.file_download_enable_priority_queue + && should_prioritize_file(ts, cfg.limit.file_download_priority_queue_window_secs) + { + if PRIORITY_FILE_DOWNLOAD_CHANNEL + .push(( + trace_id, + id, + account, + file.clone(), + size as usize, + cache_type, + )) + .await + { + metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .inc(); + } else { + queued_files::remove(&file); + log::warn!( + "[FILE_CACHE_DOWNLOAD:JOB] priority queue full ({FILE_DOWNLOAD_QUEUE_SIZE}), dropping file: {file}" + ); + } + } else { + if let Err(e) = FILE_DOWNLOAD_CHANNEL + .sender + .send(( + trace_id, + id, + account, + file.clone(), + size as usize, + cache_type, + )) + .await + { + queued_files::remove(&file); + return Err(e.into()); + } + metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE + .with_label_values::<&str>(&[]) + .inc(); + } + Ok(()) +} + +// if the file timestamp is in the past window, it should be prioritized +fn should_prioritize_file(ts: i64, window_secs: i64) -> bool { + let window_micros = window_secs * 1_000_000; + let now = now_micros(); + ts > now - window_micros +} + +/// Returns true if the file's data is older than the cache max age and should +/// not be downloaded into the cache. +pub fn exceeds_cache_max_age(ts: i64, cache_type: file_data::CacheType) -> bool { + let cfg = get_config(); + let max_age_days = match cache_type { + file_data::CacheType::Memory => cfg.memory_cache.max_age_days, + file_data::CacheType::Disk => cfg.disk_cache.max_age_days, + file_data::CacheType::None => 0, + }; + exceeds_max_age(ts, max_age_days) +} + +// if the file data is older than max_age_days, it should not be cached. +// max_age_days == 0 means no limit, ts <= 0 means the timestamp is unknown. +fn exceeds_max_age(ts: i64, max_age_days: i64) -> bool { + if max_age_days <= 0 || ts <= 0 { + return false; + } + ts < now_micros() - day_micros(max_age_days) +} + +#[cfg(test)] +mod tests { + use config::utils::time::{day_micros, hour_micros, now_micros}; + + use super::{ + FileInfo, PriorityDownloadQueue, exceeds_max_age, file_data, processing_files, queued_files, + }; + + #[test] + fn test_exceeds_max_age_disabled() { + // max_age_days == 0 means no limit, nothing is too old + let one_year_ago = now_micros() - day_micros(365); + assert!(!exceeds_max_age(one_year_ago, 0)); + assert!(!exceeds_max_age(one_year_ago, -1)); + } + + #[test] + fn test_exceeds_max_age_unknown_ts() { + // unknown timestamp should not be skipped + assert!(!exceeds_max_age(0, 3)); + assert!(!exceeds_max_age(-1, 3)); + } + + #[test] + fn test_exceeds_max_age_recent_file() { + let one_hour_ago = now_micros() - hour_micros(1); + assert!(!exceeds_max_age(one_hour_ago, 3)); + } + + #[test] + fn test_exceeds_max_age_old_file() { + let thirty_days_ago = now_micros() - day_micros(30); + assert!(exceeds_max_age(thirty_days_ago, 3)); + } + + #[test] + fn test_is_processing_false_initially() { + assert!(!processing_files::is_processing( + "unique_not_added_file_abc123.parquet" + )); + } + + #[test] + fn test_add_and_is_processing() { + let name = "test_add_file_xyz999.parquet"; + processing_files::add(name); + assert!(processing_files::is_processing(name)); + processing_files::remove(name); + } + + #[test] + fn test_remove_clears_processing() { + let name = "test_remove_file_qqq888.parquet"; + processing_files::add(name); + processing_files::remove(name); + assert!(!processing_files::is_processing(name)); + } + + #[test] + fn test_normal_queue_cap_honored() { + let cap = 3; + let (tx, _rx) = tokio::sync::mpsc::channel::(cap); + + for i in 0..cap { + let result = tx.try_send(( + format!("trace-{i}"), + i as i64, + "org".into(), + format!("file-{i}.parquet"), + 1024, + file_data::CacheType::Disk, + )); + assert!(result.is_ok(), "send {i} should succeed under cap"); + } + + let overflow = tx.try_send(( + "trace-overflow".into(), + 99, + "org".into(), + "overflow.parquet".into(), + 1024, + file_data::CacheType::Disk, + )); + assert!(overflow.is_err(), "send beyond cap must fail"); + } + + #[test] + fn test_queued_files_no_duplicates() { + let file = "dedup_test_unique_zz1234.parquet"; + + // first add succeeds + assert!(queued_files::try_add(file)); + // same file rejected while already queued + assert!(!queued_files::try_add(file)); + + // after remove, can be queued again + queued_files::remove(file); + assert!(queued_files::try_add(file)); + + // cleanup + queued_files::remove(file); + } + + #[tokio::test] + async fn test_priority_queue_cap_honored() { + let cap = 3; + let queue = PriorityDownloadQueue::new(cap); + + // fill to cap — all succeed + for i in 0..cap { + let ok = queue + .push(( + format!("trace-{i}"), + i as i64, + "org".into(), + format!("file-{i}.parquet"), + 1024, + file_data::CacheType::Disk, + )) + .await; + assert!(ok, "push {i} should succeed under cap"); + } + + // one more must be rejected + let overflow = queue + .push(( + "trace-overflow".into(), + 99, + "org".into(), + "overflow.parquet".into(), + 1024, + file_data::CacheType::Disk, + )) + .await; + assert!(!overflow, "push beyond cap must return false"); + } +} diff --git a/src/infra/src/cache/mod.rs b/src/infra/src/cache/mod.rs index b6f49595df..156b95fa46 100644 --- a/src/infra/src/cache/mod.rs +++ b/src/infra/src/cache/mod.rs @@ -15,6 +15,7 @@ pub mod bytes_cache; pub mod file_data; +pub mod file_downloader; pub mod meta; pub mod stats; pub mod storage; diff --git a/src/jobs/Cargo.toml b/src/jobs/Cargo.toml index 54bb39a601..589ac3adb2 100644 --- a/src/jobs/Cargo.toml +++ b/src/jobs/Cargo.toml @@ -33,6 +33,7 @@ hashbrown.workspace = true infra.workspace = true ingester.workspace = true log.workspace = true +openobserve-vrl.workspace = true o2_dex = { workspace = true, optional = true } o2_enterprise = { workspace = true, optional = true, default-features = false } o2_openfga = { workspace = true, optional = true } diff --git a/src/jobs/src/job/mmdb_downloader.rs b/src/jobs/src/job/mmdb_downloader.rs index 4a8d7c19f1..30c60c259e 100644 --- a/src/jobs/src/job/mmdb_downloader.rs +++ b/src/jobs/src/job/mmdb_downloader.rs @@ -16,7 +16,8 @@ use std::sync::LazyLock; use config::{ - MMDB_ASN_FILE_NAME, MMDB_CITY_FILE_NAME, + GEO_IP_ASN_ENRICHMENT_TABLE, GEO_IP_CITY_ENRICHMENT_TABLE, MMDB_ASN_FILE_NAME, + MMDB_CITY_FILE_NAME, utils::download_utils::{download_file, is_digest_different}, }; #[cfg(feature = "enterprise")] @@ -24,10 +25,7 @@ use o2_enterprise::enterprise::common::config::get_config as get_o2_config; use tokio::time; use crate::{ - common::{ - infra::config::{GEOIP_ASN_TABLE, GEOIP_CITY_TABLE, MAXMIND_DB_CLIENT}, - meta::maxmind::MaxmindClient, - }, + common::{infra::config::MAXMIND_DB_CLIENT, meta::maxmind::MaxmindClient}, service::enrichment_table::geoip::{Geoip, GeoipConfig, MMDB_INIT_NOTIFIER}, }; @@ -156,8 +154,8 @@ async fn update_maxmind_table(fname: &str) { Ok(_) => { #[cfg(feature = "enterprise")] if get_o2_config().common.enable_enterprise_mmdb { - let mut geoip = crate::common::infra::config::GEOIP_ENT_TABLE.write(); - *geoip = Some( + openobserve_vrl::register_global_enrichment_table( + o2_enterprise::enterprise::common::config::GEO_IP_ENTERPRISE_ENRICHMENT_TABLE, Geoip::new(GeoipConfig::new( &get_o2_config().common.mmdb_enterprise_file_name, )) @@ -166,11 +164,15 @@ async fn update_maxmind_table(fname: &str) { } if fname.ends_with(MMDB_CITY_FILE_NAME) { - let mut geoip_city = GEOIP_CITY_TABLE.write(); - *geoip_city = Some(Geoip::new(GeoipConfig::new(MMDB_CITY_FILE_NAME)).unwrap()); + openobserve_vrl::register_global_enrichment_table( + GEO_IP_CITY_ENRICHMENT_TABLE, + Geoip::new(GeoipConfig::new(MMDB_CITY_FILE_NAME)).unwrap(), + ); } else if fname.ends_with(MMDB_ASN_FILE_NAME) { - let mut geoip_asn = GEOIP_ASN_TABLE.write(); - *geoip_asn = Some(Geoip::new(GeoipConfig::new(MMDB_ASN_FILE_NAME)).unwrap()); + openobserve_vrl::register_global_enrichment_table( + GEO_IP_ASN_ENRICHMENT_TABLE, + Geoip::new(GeoipConfig::new(MMDB_ASN_FILE_NAME)).unwrap(), + ); }; } Err(e) => log::warn!("Failed to update maxmind table with path: {fname}, {e}"), diff --git a/src/search/Cargo.toml b/src/search/Cargo.toml index 8673ec6ae7..52330896d5 100644 --- a/src/search/Cargo.toml +++ b/src/search/Cargo.toml @@ -7,18 +7,26 @@ publish = false [features] default = [] -enterprise = [] +enterprise = ["dep:hashbrown", "dep:o2_enterprise", "dep:parking_lot"] +vectorscan = [] +vortex = [] [dependencies] -async-trait.workspace = true +chrono.workspace = true config.workspace = true +datafusion.workspace = true +hashbrown = { workspace = true, optional = true } +infra.workspace = true +log.workspace = true +openobserve-vrl.workspace = true +o2_enterprise = { workspace = true, optional = true, default-features = false } +parking_lot = { workspace = true, optional = true } serde.workspace = true +tokio.workspace = true +tracing.workspace = true utoipa.workspace = true -vector-enrichment.workspace = true -vrl.workspace = true [dev-dependencies] arrow.workspace = true arrow-schema.workspace = true -o2_enterprise.workspace = true serde_json.workspace = true diff --git a/src/search/src/cipher/mod.rs b/src/search/src/cipher/mod.rs new file mode 100644 index 0000000000..b814e70253 --- /dev/null +++ b/src/search/src/cipher/mod.rs @@ -0,0 +1,16 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +pub mod registry; diff --git a/src/core/src/cipher/registry.rs b/src/search/src/cipher/registry.rs similarity index 100% rename from src/core/src/cipher/registry.rs rename to src/search/src/cipher/registry.rs diff --git a/src/search/src/enrichment.rs b/src/search/src/enrichment.rs index 2bdd49aa41..9c6bfa6468 100644 --- a/src/search/src/enrichment.rs +++ b/src/search/src/enrichment.rs @@ -5,7 +5,7 @@ // the Free Software Foundation, either version 3 of the License, or // (at your option) any later version. // -// This program is distributed in the hope that it will be useful +// This program is distributed in the hope that it will be useful, // but WITHOUT ANY WARRANTY; without even the implied warranty of // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the // GNU Affero General Public License for more details. @@ -13,209 +13,4 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{ - collections::BTreeMap, - sync::{Arc, LazyLock}, -}; - -use async_trait::async_trait; -use config::{RwHashMap, stats::MemorySize, utils::time::parse_str_to_time}; -use vector_enrichment::{Case, IndexHandle, Table}; -use vrl::value::{KeyString, ObjectMap, Value}; - -pub static ENRICHMENT_TABLES: LazyLock> = - LazyLock::new(Default::default); - -#[derive(Debug, Clone)] -pub struct StreamTable { - pub org_id: String, - pub stream_name: String, - pub data: Arc>, -} - -impl MemorySize for StreamTable { - fn mem_size(&self) -> usize { - std::mem::size_of::() - + self.org_id.mem_size() - + self.stream_name.mem_size() - + self.data.iter().map(|v| v.to_string().len()).sum::() - } -} - -#[async_trait] -impl Table for StreamTable { - fn find_table_row( - &self, - case: Case, - conditions: &[vector_enrichment::Condition], - select: Option<&[String]>, - _wildcard: Option<&Value>, - _index: Option, - ) -> Result { - let resp = get_data(self, conditions, select, case); - Ok(resp.first().cloned().unwrap_or_default()) - } - - fn find_table_rows( - &self, - case: Case, - conditions: &[vector_enrichment::Condition], - select: Option<&[String]>, - _wildcard: Option<&Value>, - _index: Option, - ) -> Result, vector_enrichment::Error> { - Ok(get_data(self, conditions, select, case)) - } - - fn add_index( - &mut self, - _case: Case, - _fields: &[&str], - ) -> Result { - Ok(IndexHandle(1)) - } - - fn index_fields(&self) -> Vec<(Case, Vec)> { - Vec::new() - } - - fn needs_reload(&self) -> bool { - false - } -} - -fn get_data( - table: &StreamTable, - conditions: &[vector_enrichment::Condition], - select: Option<&[String]>, - case: Case, -) -> Vec { - if conditions.is_empty() { - return vec![]; - } - - let filtered = table.data.iter().filter(|value| { - let Value::Object(map) = value else { - return false; - }; - - conditions.iter().all(|condition| match condition { - vector_enrichment::Condition::Equals { field, value } => match case { - Case::Insensitive => match (map.get(*field), value) { - (Some(Value::Bytes(left)), Value::Bytes(right)) => { - match (std::str::from_utf8(left), std::str::from_utf8(right)) { - (Ok(left), Ok(right)) => left.eq_ignore_ascii_case(right), - (Err(_), Err(_)) => left == right, - _ => false, - } - } - _ => false, - }, - Case::Sensitive => map.get(*field).is_some_and(|item| item == value), - }, - vector_enrichment::Condition::FromDate { field, from } => map - .get(*field) - .and_then(|value| value.as_str()) - .and_then(|value| parse_str_to_time(value.as_ref()).ok()) - .is_some_and(|date| date >= *from), - vector_enrichment::Condition::ToDate { field, to } => map - .get(*field) - .and_then(|value| value.as_str()) - .and_then(|value| parse_str_to_time(value.as_ref()).ok()) - .is_some_and(|date| date <= *to), - vector_enrichment::Condition::BetweenDates { field, from, to } => map - .get(*field) - .and_then(|value| value.as_str()) - .and_then(|value| parse_str_to_time(value.as_ref()).ok()) - .is_some_and(|date| date >= *from && date <= *to), - }) - }); - - filtered - .filter_map(Value::as_object) - .map(|map| { - if let Some(fields) = select { - fields - .iter() - .filter_map(|field| map.get_key_value(field.as_str())) - .map(|(key, value)| (key.clone(), value.clone())) - .collect::>() - } else { - map.clone() - } - }) - .collect() -} - -#[cfg(test)] -mod tests { - use serde_json::json; - use vector_enrichment::{Condition, Table}; - - use super::*; - - fn table() -> StreamTable { - StreamTable { - org_id: "default".to_string(), - stream_name: "users".to_string(), - data: Arc::new(vec![ - json!({"name": "Alice", "country": "US", "role": "admin"}).into(), - json!({"name": "Bob", "country": "GB", "role": "user"}).into(), - Value::from("not an object"), - ]), - } - } - - #[test] - fn matches_rows_using_requested_case_sensitivity() { - let table = table(); - let conditions = [Condition::Equals { - field: "name", - value: Value::from("ALICE"), - }]; - - let sensitive = table - .find_table_rows(Case::Sensitive, &conditions, None, None, None) - .unwrap(); - let insensitive = table - .find_table_rows(Case::Insensitive, &conditions, None, None, None) - .unwrap(); - - assert!(sensitive.is_empty()); - assert_eq!(insensitive.len(), 1); - assert_eq!(insensitive[0].get("name"), Some(&Value::from("Alice"))); - } - - #[test] - fn combines_conditions_and_projects_selected_fields() { - let table = table(); - let conditions = [ - Condition::Equals { - field: "country", - value: Value::from("US"), - }, - Condition::Equals { - field: "role", - value: Value::from("admin"), - }, - ]; - let select = ["name".to_string()]; - - let rows = table - .find_table_rows(Case::Sensitive, &conditions, Some(&select), None, None) - .unwrap(); - - assert_eq!(rows.len(), 1); - assert_eq!(rows[0].len(), 1); - assert_eq!(rows[0].get("name"), Some(&Value::from("Alice"))); - } - - #[test] - fn returns_no_rows_for_empty_conditions() { - let rows = table() - .find_table_rows(Case::Sensitive, &[], None, None, None) - .unwrap(); - - assert!(rows.is_empty()); - } -} +pub use openobserve_vrl::enrichment::*; diff --git a/src/search/src/file_cache.rs b/src/search/src/file_cache.rs new file mode 100644 index 0000000000..805665bf83 --- /dev/null +++ b/src/search/src/file_cache.rs @@ -0,0 +1,155 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::collections::HashSet; + +use config::{get_config, is_local_disk_storage, meta::search::ScanStats}; +use infra::cache::{file_data, file_downloader}; + +/// Linear interpolation: cached_ratio=0 -> query_thread_num, cached_ratio=1 -> cpu_num. +pub fn calc_target_partitions(cpu_num: usize, query_thread_num: usize, cached_ratio: f64) -> usize { + (cpu_num as i64 + + ((query_thread_num as i64 - cpu_num as i64) as f64 * (1.0 - cached_ratio)) as i64) + as usize +} + +#[tracing::instrument(name = "service:search:grpc:storage:cache_files", skip_all)] +pub async fn cache_files( + trace_id: &str, + files: &[(i64, &String, &String, i64, i64)], + scan_stats: &mut ScanStats, + file_type: &str, +) -> (file_data::CacheType, u64, u64) { + let mut cached_files = HashSet::with_capacity(files.len()); + let (mut cache_hits, mut cache_misses) = (0, 0); + + let start = std::time::Instant::now(); + for (_id, _account, file, _size, max_ts) in files.iter() { + if file_data::memory::exist(file).await { + scan_stats.querier_memory_cached_files += 1; + cached_files.insert(file); + cache_hits += 1; + } else if file_data::disk::exist(file).await { + scan_stats.querier_disk_cached_files += 1; + cached_files.insert(file); + cache_hits += 1; + } else { + cache_misses += 1; + } + + let stream_type = if file_type == "index" { + config::meta::stream::StreamType::Index + } else if file.contains("/logs/") { + config::meta::stream::StreamType::Logs + } else if file.contains("/metrics/") { + config::meta::stream::StreamType::Metrics + } else if file.contains("/traces/") { + config::meta::stream::StreamType::Traces + } else { + config::meta::stream::StreamType::Logs + }; + + let current_time = chrono::Utc::now().timestamp_micros(); + let file_age_seconds = (current_time - max_ts) / 1_000_000; + let file_age_hours = file_age_seconds as f64 / 3600.0; + + if file_age_hours > 0.0 { + config::metrics::FILE_ACCESS_TIME + .with_label_values(&[&stream_type.to_string()]) + .observe(file_age_hours); + } + } + + let check_cache_took = start.elapsed().as_millis() as usize; + if check_cache_took > 1000 { + log::warn!( + "[trace_id {trace_id}] search->storage: check file cache took: {check_cache_took} ms", + ); + } + + let files_num = files.len() as i64; + if files_num == scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files { + return (file_data::CacheType::Disk, cache_hits, cache_misses); + } + + let cfg = get_config(); + let cache_type = if cfg.memory_cache.enabled + && scan_stats.compressed_size < cfg.memory_cache.skip_size as i64 + { + file_data::CacheType::Memory + } else if !is_local_disk_storage() + && cfg.disk_cache.enabled + && scan_stats.compressed_size < cfg.disk_cache.skip_size as i64 + { + file_data::CacheType::Disk + } else { + return (file_data::CacheType::None, cache_hits, cache_misses); + }; + + let trace_id = trace_id.to_string(); + let files = files + .iter() + .filter_map(|(id, account, file, size, ts)| { + if cached_files.contains(file) { + None + } else { + Some((*id, account.to_string(), file.to_string(), *size, *ts)) + } + }) + .collect::>(); + let file_type = file_type.to_string(); + tokio::spawn(async move { + let files_num = files.len(); + for (id, account, file, size, ts) in files { + if let Err(e) = file_downloader::queue_download( + trace_id.clone(), + id, + account, + file.clone(), + size, + ts, + cache_type, + ) + .await + { + log::error!( + "[trace_id {trace_id}] error in queuing file {file} for background download: {e}" + ); + } + } + log::info!( + "[trace_id {trace_id}] search->storage: successfully enqueued {files_num} files of {file_type} for background download into {cache_type:?}", + ); + }); + + if scan_stats.querier_memory_cached_files + scan_stats.querier_disk_cached_files < files_num / 2 + { + (file_data::CacheType::None, cache_hits, cache_misses) + } else { + (cache_type, cache_hits, cache_misses) + } +} + +#[cfg(test)] +mod tests { + use super::calc_target_partitions; + + #[test] + fn target_partitions_interpolates_cache_ratio() { + assert_eq!(calc_target_partitions(8, 32, 0.0), 32); + assert_eq!(calc_target_partitions(8, 32, 0.5), 20); + assert_eq!(calc_target_partitions(8, 32, 1.0), 8); + } +} diff --git a/src/search/src/lib.rs b/src/search/src/lib.rs index 098af4b4b0..e6fff9bd20 100644 --- a/src/search/src/lib.rs +++ b/src/search/src/lib.rs @@ -13,7 +13,11 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . +#[cfg(feature = "enterprise")] +pub mod cipher; pub mod enrichment; +pub mod file_cache; +pub mod types; use config::meta::{ search::{Response, ScanStats}, diff --git a/src/search/src/types.rs b/src/search/src/types.rs new file mode 100644 index 0000000000..171d5aa5fb --- /dev/null +++ b/src/search/src/types.rs @@ -0,0 +1,29 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use config::meta::stream::StreamType; +use datafusion::sql::TableReference; + +#[derive(Debug)] +pub struct QueryParams { + pub trace_id: String, + pub org_id: String, + pub stream: TableReference, + pub stream_type: StreamType, + pub stream_name: String, + pub time_range: (i64, i64), + pub work_group: Option, + pub use_inverted_index: bool, +} diff --git a/src/vrl/Cargo.toml b/src/vrl/Cargo.toml new file mode 100644 index 0000000000..23f6149ac7 --- /dev/null +++ b/src/vrl/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "openobserve-vrl" +version.workspace = true +edition.workspace = true +license.workspace = true +publish = false + +[dependencies] +async-trait.workspace = true +config.workspace = true +vector-enrichment.workspace = true +vrl.workspace = true + +[dev-dependencies] +serde_json.workspace = true +tokio.workspace = true diff --git a/src/vrl/src/enrichment.rs b/src/vrl/src/enrichment.rs new file mode 100644 index 0000000000..459e36e13c --- /dev/null +++ b/src/vrl/src/enrichment.rs @@ -0,0 +1,219 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +use std::{ + collections::BTreeMap, + sync::{Arc, LazyLock}, +}; + +use async_trait::async_trait; +use config::{RwHashMap, stats::MemorySize, utils::time::parse_str_to_time}; +use vector_enrichment::{Case, IndexHandle, Table}; +use vrl::value::{KeyString, ObjectMap, Value}; + +pub static ENRICHMENT_TABLES: LazyLock> = + LazyLock::new(Default::default); + +#[derive(Debug, Clone)] +pub struct StreamTable { + pub org_id: String, + pub stream_name: String, + pub data: Arc>, +} + +impl MemorySize for StreamTable { + fn mem_size(&self) -> usize { + std::mem::size_of::() + + self.org_id.mem_size() + + self.stream_name.mem_size() + + self.data.iter().map(|v| v.to_string().len()).sum::() + } +} + +#[async_trait] +impl Table for StreamTable { + fn find_table_row( + &self, + case: Case, + conditions: &[vector_enrichment::Condition], + select: Option<&[String]>, + _wildcard: Option<&Value>, + _index: Option, + ) -> Result { + let resp = get_data(self, conditions, select, case); + Ok(resp.first().cloned().unwrap_or_default()) + } + + fn find_table_rows( + &self, + case: Case, + conditions: &[vector_enrichment::Condition], + select: Option<&[String]>, + _wildcard: Option<&Value>, + _index: Option, + ) -> Result, vector_enrichment::Error> { + Ok(get_data(self, conditions, select, case)) + } + + fn add_index( + &mut self, + _case: Case, + _fields: &[&str], + ) -> Result { + Ok(IndexHandle(1)) + } + + fn index_fields(&self) -> Vec<(Case, Vec)> { + Vec::new() + } + + fn needs_reload(&self) -> bool { + false + } +} + +fn get_data( + table: &StreamTable, + conditions: &[vector_enrichment::Condition], + select: Option<&[String]>, + case: Case, +) -> Vec { + if conditions.is_empty() { + return vec![]; + } + + let filtered = table.data.iter().filter(|value| { + let Value::Object(map) = value else { + return false; + }; + + conditions.iter().all(|condition| match condition { + vector_enrichment::Condition::Equals { field, value } => match case { + Case::Insensitive => match (map.get(*field), value) { + (Some(Value::Bytes(left)), Value::Bytes(right)) => { + match (std::str::from_utf8(left), std::str::from_utf8(right)) { + (Ok(left), Ok(right)) => left.eq_ignore_ascii_case(right), + (Err(_), Err(_)) => left == right, + _ => false, + } + } + _ => false, + }, + Case::Sensitive => map.get(*field).is_some_and(|item| item == value), + }, + vector_enrichment::Condition::FromDate { field, from } => map + .get(*field) + .and_then(|value| value.as_str()) + .and_then(|value| parse_str_to_time(value.as_ref()).ok()) + .is_some_and(|date| date >= *from), + vector_enrichment::Condition::ToDate { field, to } => map + .get(*field) + .and_then(|value| value.as_str()) + .and_then(|value| parse_str_to_time(value.as_ref()).ok()) + .is_some_and(|date| date <= *to), + vector_enrichment::Condition::BetweenDates { field, from, to } => map + .get(*field) + .and_then(|value| value.as_str()) + .and_then(|value| parse_str_to_time(value.as_ref()).ok()) + .is_some_and(|date| date >= *from && date <= *to), + }) + }); + + filtered + .filter_map(Value::as_object) + .map(|map| { + if let Some(fields) = select { + fields + .iter() + .filter_map(|field| map.get_key_value(field.as_str())) + .map(|(key, value)| (key.clone(), value.clone())) + .collect::>() + } else { + map.clone() + } + }) + .collect() +} + +#[cfg(test)] +mod tests { + use serde_json::json; + use vector_enrichment::{Condition, Table}; + + use super::*; + + fn table() -> StreamTable { + StreamTable { + org_id: "default".to_string(), + stream_name: "users".to_string(), + data: Arc::new(vec![ + json!({"name": "Alice", "country": "US", "role": "admin"}).into(), + json!({"name": "Bob", "country": "GB", "role": "user"}).into(), + Value::from("not an object"), + ]), + } + } + + #[test] + fn matches_rows_using_requested_case_sensitivity() { + let conditions = [Condition::Equals { + field: "name", + value: Value::from("ALICE"), + }]; + + let sensitive = table() + .find_table_rows(Case::Sensitive, &conditions, None, None, None) + .unwrap(); + let insensitive = table() + .find_table_rows(Case::Insensitive, &conditions, None, None, None) + .unwrap(); + + assert!(sensitive.is_empty()); + assert_eq!(insensitive.len(), 1); + assert_eq!(insensitive[0].get("name"), Some(&Value::from("Alice"))); + } + + #[test] + fn combines_conditions_and_projects_selected_fields() { + let conditions = [ + Condition::Equals { + field: "country", + value: Value::from("US"), + }, + Condition::Equals { + field: "role", + value: Value::from("admin"), + }, + ]; + let select = ["name".to_string()]; + + let rows = table() + .find_table_rows(Case::Sensitive, &conditions, Some(&select), None, None) + .unwrap(); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].len(), 1); + assert_eq!(rows[0].get("name"), Some(&Value::from("Alice"))); + } + + #[test] + fn returns_no_rows_for_empty_conditions() { + let rows = table() + .find_table_rows(Case::Sensitive, &[], None, None, None) + .unwrap(); + + assert!(rows.is_empty()); + } +} diff --git a/src/vrl/src/lib.rs b/src/vrl/src/lib.rs new file mode 100644 index 0000000000..4d6e5bfb2d --- /dev/null +++ b/src/vrl/src/lib.rs @@ -0,0 +1,209 @@ +// Copyright 2026 OpenObserve Inc. +// +// This program is free software: you can redistribute it and/or modify +// it under the terms of the GNU Affero General Public License as published by +// the Free Software Foundation, either version 3 of the License, or +// (at your option) any later version. +// +// This program is distributed in the hope that it will be useful, +// but WITHOUT ANY WARRANTY; without even the implied warranty of +// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +// GNU Affero General Public License for more details. +// +// You should have received a copy of the GNU Affero General Public License +// along with this program. If not, see . + +//! Shared ownership for OpenObserve transform definitions, enrichment tables, +//! and VRL compilation. This crate deliberately has no dependency on core or +//! the search engine so both can consume the same runtime state. + +use std::{collections::HashMap, sync::LazyLock}; + +use config::{ + DEFAULT_ORG, RwHashMap, + meta::function::{Transform, VRLCompilerConfig, VRLRuntimeConfig}, +}; +use vector_enrichment::{Table, TableRegistry}; +use vrl::compiler::{CompilationResult, runtime::Runtime}; + +pub mod enrichment; + +use enrichment::ENRICHMENT_TABLES; + +/// Query and ingestion transform definitions, keyed by `org_id/function_name`. +pub static QUERY_FUNCTIONS: LazyLock> = + LazyLock::new(Default::default); + +/// Process-wide enrichment tables such as MaxMind databases. +/// +/// Organization-owned stream enrichment tables live in [`ENRICHMENT_TABLES`]. +static GLOBAL_ENRICHMENT_TABLES: LazyLock>> = + LazyLock::new(Default::default); + +pub fn register_global_enrichment_table(name: impl Into, table: T) +where + T: Table + Send + Sync + 'static, +{ + GLOBAL_ENRICHMENT_TABLES.insert(name.into(), Box::new(table)); +} + +pub fn remove_global_enrichment_table(name: &str) { + GLOBAL_ENRICHMENT_TABLES.remove(name); +} + +pub async fn get_all_transform_keys(org_id: &str) -> Vec { + let org_key = format!("{org_id}/"); + + QUERY_FUNCTIONS + .iter() + .filter_map(|transform| { + let key = transform.key(); + key.strip_prefix(&org_key).map(ToOwned::to_owned) + }) + .collect() +} + +pub fn init_vrl_runtime() -> Runtime { + Runtime::new(vrl::prelude::state::RuntimeState::default()) +} + +pub fn get_enrichment_tables(org_id: &str) -> HashMap> { + let mut tables = HashMap::new(); + + for table in ENRICHMENT_TABLES.iter() { + if table.org_id == org_id || table.org_id == DEFAULT_ORG { + tables.insert( + table.stream_name.clone(), + Box::new(table.value().clone()) as Box, + ); + } + } + + for table in GLOBAL_ENRICHMENT_TABLES.iter() { + tables.insert(table.key().clone(), table.value().clone()); + } + + tables +} + +pub fn get_vrl_compiler_config(org_id: &str) -> VRLCompilerConfig { + let mut functions = vrl::stdlib::all(); + functions.append(&mut vector_enrichment::vrl_functions()); + let registry = TableRegistry::default(); + registry.load(get_enrichment_tables(org_id)); + let mut config = vrl::compiler::CompileConfig::default(); + config.set_custom(registry); + VRLCompilerConfig { config, functions } +} + +pub fn compile_vrl_function( + source: &str, + org_id: &str, +) -> Result { + if source.contains("get_env_var") { + return Err(std::io::Error::other("get_env_var is not supported")); + } + + let external = vrl::prelude::state::ExternalEnv::default(); + let vrl_config = get_vrl_compiler_config(org_id); + match vrl::compiler::compile_with_external( + source, + &vrl_config.functions, + &external, + vrl_config.config, + ) { + Ok(CompilationResult { + program, + warnings: _, + config, + }) => Ok(VRLRuntimeConfig { + program, + config, + fields: vec![], + }), + Err(error) => Err(std::io::Error::other( + vrl::diagnostic::Formatter::new(source, error).to_string(), + )), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use config::meta::function::Transform; + use vector_enrichment::TableRegistry; + + use super::{ + enrichment::{ENRICHMENT_TABLES, StreamTable}, + *, + }; + + #[test] + fn creates_vrl_compiler_config_with_org_enrichment_tables() { + let key = "vrl_crate_test/default/users"; + ENRICHMENT_TABLES.insert( + key.to_string(), + StreamTable { + org_id: "vrl_crate_test".to_string(), + stream_name: "users".to_string(), + data: Arc::new(vec![]), + }, + ); + + let config = get_vrl_compiler_config("vrl_crate_test"); + let registry = config.config.get_custom::().unwrap(); + assert!(registry.table_ids().contains(&"users".to_string())); + + ENRICHMENT_TABLES.remove(key); + } + + #[test] + fn includes_registered_global_enrichment_tables() { + let name = "vrl_crate_test_global"; + register_global_enrichment_table( + name, + StreamTable { + org_id: "default".to_string(), + stream_name: name.to_string(), + data: Arc::new(vec![]), + }, + ); + + let config = get_vrl_compiler_config("vrl_crate_test"); + let registry = config.config.get_custom::().unwrap(); + assert!(registry.table_ids().contains(&name.to_string())); + + remove_global_enrichment_table(name); + } + + #[tokio::test] + async fn returns_transform_keys_for_requested_org() { + let key = "vrl_crate_test/test_transform"; + QUERY_FUNCTIONS.insert( + key.to_string(), + Transform { + name: "test_transform".to_string(), + function: ".".to_string(), + params: "row".to_string(), + num_args: 1, + trans_type: Some(0), + streams: None, + }, + ); + + let keys = get_all_transform_keys("vrl_crate_test").await; + assert_eq!(keys, vec!["test_transform"]); + + QUERY_FUNCTIONS.remove(key); + } + + #[test] + fn rejects_environment_access() { + let error = match compile_vrl_function("get_env_var!(\"HOME\")", "default") { + Ok(_) => panic!("get_env_var must be rejected"), + Err(error) => error, + }; + assert_eq!(error.to_string(), "get_env_var is not supported"); + } +}