refactor: make search engine sources crate-independent (#13234)

## Summary

- add `src/vrl` as a leaf package named `openobserve-vrl` with Rust path
`openobserve_vrl`
- move `QUERY_FUNCTIONS`, enrichment-table ownership, GeoIP/MaxMind
table registration, VRL compiler configuration, and shared runtime
helpers into `openobserve-vrl`
- let ingestion, jobs, core, and search use the shared VRL integration
without depending on one another
- keep DataFusion UDF registration direct through `Sql.org_id`; no
`QueryExecutionContext`, `PreparedQueryTransform`, or
`prepare_query_transforms`
- move the cipher registry, query parameters, and file-cache helpers to
`search`
- move the generic file downloader implementation to `infra`
- add feature forwarding for the `search` crate
- keep compatibility re-exports for existing API, ingestion, gRPC, jobs,
and Enterprise call sites

Enterprise companion:
https://github.com/openobserve/o2-enterprise/pull/2196
This commit is contained in:
Huaijin
2026-07-16 13:50:35 +00:00
committed by GitHub
parent 9b4ec202c2
commit 059058519e
43 changed files with 1667 additions and 1527 deletions
Generated
+23 -3
View File
@@ -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]]
+2
View File
@@ -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" }
+1
View File
@@ -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"] },
+1
View File
@@ -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";
+3 -1
View File
@@ -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
+1 -2
View File
@@ -13,12 +13,11 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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,
+4 -49
View File
@@ -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<RwHashMap<String, bytes::Bytes>> = Lazy::new(Default::default);
pub static QUERY_FUNCTIONS: Lazy<RwHashMap<String, Transform>> = Lazy::new(DashMap::default);
pub static USERS: Lazy<RwHashMap<String, infra::table::users::UserRecord>> =
Lazy::new(DashMap::default);
pub static ORG_USERS: Lazy<RwHashMap<String, infra::table::org_users::OrgUserRecord>> =
@@ -74,22 +67,9 @@ pub static REALTIME_ALERT_TRIGGERS: Lazy<RwAHashMap<String, db_scheduler::Trigge
Lazy::new(Default::default);
pub static ALERTS_TEMPLATES: Lazy<RwHashMap<String, Template>> = Lazy::new(Default::default);
pub static DESTINATIONS: Lazy<RwHashMap<String, Destination>> = Lazy::new(Default::default);
pub static ENRICHMENT_REGISTRY: Lazy<Arc<TableRegistry>> =
Lazy::new(|| Arc::new(TableRegistry::default()));
pub static MAXMIND_DB_CLIENT: Lazy<Arc<tokio::sync::RwLock<Option<MaxmindClient>>>> =
Lazy::new(|| Arc::new(tokio::sync::RwLock::new(None)));
pub static GEOIP_CITY_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
Lazy::new(|| Arc::new(RwLock::new(None)));
pub static GEOIP_ASN_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
Lazy::new(|| Arc::new(RwLock::new(None)));
#[cfg(feature = "enterprise")]
pub static GEOIP_ENT_TABLE: Lazy<Arc<RwLock<Option<Geoip>>>> =
Lazy::new(|| Arc::new(RwLock::new(None)));
pub static STREAM_EXECUTABLE_PIPELINES: Lazy<RwAHashMap<StreamParams, Vec<ExecutablePipeline>>> =
Lazy::new(Default::default);
pub static PIPELINE_STREAM_MAPPING: Lazy<RwAHashMap<String, StreamParams>> =
@@ -139,10 +119,6 @@ mod tests {
// Test Arc<RwLock> 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(),
+1 -2
View File
@@ -13,14 +13,13 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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;
+3 -157
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<String> {
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<String, Box<dyn Table + Send + Sync>> = 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::<TableRegistry>().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::<TableRegistry>().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::<TableRegistry>().unwrap();
assert!(registry.table_ids().is_empty());
}
}
+1 -717
View File
@@ -13,720 +13,4 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
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<RwLock<HashSet<String>>> =
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<RwLock<HashSet<String>>> = 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<FileInfo>,
receiver: Arc<Mutex<Receiver<FileInfo>>>,
}
impl DownloadQueue {
fn new(sender: Sender<FileInfo>, receiver: Arc<Mutex<Receiver<FileInfo>>>) -> Self {
Self { sender, receiver }
}
}
struct PriorityDownloadQueue {
stack: Mutex<VecDeque<FileInfo>>,
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<DownloadQueue> = Lazy::new(|| {
let (tx, rx) = tokio::sync::mpsc::channel::<FileInfo>(FILE_DOWNLOAD_QUEUE_SIZE);
DownloadQueue::new(tx, Arc::new(Mutex::new(rx)))
});
static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy<PriorityDownloadQueue> =
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<usize, anyhow::Error> {
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<bool, anyhow::Error> {
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<Vec<(i64, String, String, i64, i64)>, 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::<HashMap<_, _>>();
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::<FileInfo>(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::*;
+6 -38
View File
@@ -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<Map<String, Value>>)>;
/// 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<VRLRuntimeConfig, std::io::Error> {
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<JSRuntimeConfig, std::io::Error> {
+1 -51
View File
@@ -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::<Vec<&str>>(),
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 {
+5 -1
View File
@@ -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)
+3 -94
View File
@@ -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<PhysicalOptimizerContext>,
}
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<Sql>) -> Result<SessionContext> {
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,
@@ -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 <http://www.gnu.org/licenses/>.
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<PhysicalOptimizerContext>,
}
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<Sql>) -> Result<SessionContext> {
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(())
}
@@ -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(
@@ -15,6 +15,7 @@
use std::str::FromStr;
pub mod context;
pub mod distributed_plan;
pub mod exec;
pub mod merge;
@@ -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
@@ -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,
};
@@ -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 {
@@ -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.
@@ -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<dyn Fn(&[ColumnarValue]) -> Result<ColumnarValue> + 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<Vec<ScalarUDF>> {
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<Vec<ScalarUDF>> {
fn get_udf_vrl(
fn_name: String,
func: &str,
source: &str,
params: &str,
num_args: u8,
org_id: &str,
) -> Result<ScalarUDF> {
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::<TableRegistry>().unwrap();
Ok(result) => {
let registry = result.config.get_custom::<TableRegistry>().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 {
+4 -27
View File
@@ -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<Arc<dyn TableProvider>>, ScanStats, HashSet<u64>)>;
#[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<String>,
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
+2 -130
View File
@@ -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(
@@ -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::<Vec<_>>(),
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";
@@ -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,
};
@@ -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,
+4 -1
View File
@@ -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},
};
+2 -1
View File
@@ -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<Box<dyn Future<Output = ()> + Send>>;
+733
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<RwLock<HashSet<String>>> =
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<RwLock<HashSet<String>>> = 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<FileInfo>,
receiver: Arc<Mutex<Receiver<FileInfo>>>,
}
impl DownloadQueue {
fn new(sender: Sender<FileInfo>, receiver: Arc<Mutex<Receiver<FileInfo>>>) -> Self {
Self { sender, receiver }
}
}
struct PriorityDownloadQueue {
stack: Mutex<VecDeque<FileInfo>>,
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<DownloadQueue> = Lazy::new(|| {
let (tx, rx) = tokio::sync::mpsc::channel::<FileInfo>(FILE_DOWNLOAD_QUEUE_SIZE);
DownloadQueue::new(tx, Arc::new(Mutex::new(rx)))
});
static PRIORITY_FILE_DOWNLOAD_CHANNEL: Lazy<PriorityDownloadQueue> =
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<usize, anyhow::Error> {
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<bool, anyhow::Error> {
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<Vec<(i64, String, String, i64, i64)>, 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::<HashMap<_, _>>();
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::<FileInfo>(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");
}
}
+1
View File
@@ -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;
+1
View File
@@ -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 }
+13 -11
View File
@@ -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}"),
+13 -5
View File
@@ -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
+16
View File
@@ -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 <http://www.gnu.org/licenses/>.
pub mod registry;
+2 -207
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<RwHashMap<String, StreamTable>> =
LazyLock::new(Default::default);
#[derive(Debug, Clone)]
pub struct StreamTable {
pub org_id: String,
pub stream_name: String,
pub data: Arc<Vec<Value>>,
}
impl MemorySize for StreamTable {
fn mem_size(&self) -> usize {
std::mem::size_of::<StreamTable>()
+ self.org_id.mem_size()
+ self.stream_name.mem_size()
+ self.data.iter().map(|v| v.to_string().len()).sum::<usize>()
}
}
#[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<IndexHandle>,
) -> Result<ObjectMap, vector_enrichment::Error> {
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<IndexHandle>,
) -> Result<Vec<ObjectMap>, vector_enrichment::Error> {
Ok(get_data(self, conditions, select, case))
}
fn add_index(
&mut self,
_case: Case,
_fields: &[&str],
) -> Result<IndexHandle, vector_enrichment::Error> {
Ok(IndexHandle(1))
}
fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
Vec::new()
}
fn needs_reload(&self) -> bool {
false
}
}
fn get_data(
table: &StreamTable,
conditions: &[vector_enrichment::Condition],
select: Option<&[String]>,
case: Case,
) -> Vec<ObjectMap> {
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::<BTreeMap<KeyString, Value>>()
} 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::*;
+155
View File
@@ -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 <http://www.gnu.org/licenses/>.
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::<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) = 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);
}
}
+4
View File
@@ -13,7 +13,11 @@
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <http://www.gnu.org/licenses/>.
#[cfg(feature = "enterprise")]
pub mod cipher;
pub mod enrichment;
pub mod file_cache;
pub mod types;
use config::meta::{
search::{Response, ScanStats},
+29
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<String>,
pub use_inverted_index: bool,
}
+16
View File
@@ -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
+219
View File
@@ -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 <http://www.gnu.org/licenses/>.
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<RwHashMap<String, StreamTable>> =
LazyLock::new(Default::default);
#[derive(Debug, Clone)]
pub struct StreamTable {
pub org_id: String,
pub stream_name: String,
pub data: Arc<Vec<Value>>,
}
impl MemorySize for StreamTable {
fn mem_size(&self) -> usize {
std::mem::size_of::<StreamTable>()
+ self.org_id.mem_size()
+ self.stream_name.mem_size()
+ self.data.iter().map(|v| v.to_string().len()).sum::<usize>()
}
}
#[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<IndexHandle>,
) -> Result<ObjectMap, vector_enrichment::Error> {
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<IndexHandle>,
) -> Result<Vec<ObjectMap>, vector_enrichment::Error> {
Ok(get_data(self, conditions, select, case))
}
fn add_index(
&mut self,
_case: Case,
_fields: &[&str],
) -> Result<IndexHandle, vector_enrichment::Error> {
Ok(IndexHandle(1))
}
fn index_fields(&self) -> Vec<(Case, Vec<String>)> {
Vec::new()
}
fn needs_reload(&self) -> bool {
false
}
}
fn get_data(
table: &StreamTable,
conditions: &[vector_enrichment::Condition],
select: Option<&[String]>,
case: Case,
) -> Vec<ObjectMap> {
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::<BTreeMap<KeyString, Value>>()
} 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());
}
}
+209
View File
@@ -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 <http://www.gnu.org/licenses/>.
//! 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<RwHashMap<String, Transform>> =
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<RwHashMap<String, Box<dyn Table + Send + Sync>>> =
LazyLock::new(Default::default);
pub fn register_global_enrichment_table<T>(name: impl Into<String>, 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<String> {
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<String, Box<dyn Table + Send + Sync>> {
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<dyn Table + Send + Sync>,
);
}
}
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<VRLRuntimeConfig, std::io::Error> {
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::<TableRegistry>().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::<TableRegistry>().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");
}
}