fix: move event download to download manager (#6873)

This commit is contained in:
Hengfei Yang
2025-05-22 12:46:58 +08:00
committed by GitHub
parent 9211b15ef8
commit b105a9ef99
7 changed files with 143 additions and 140 deletions
+2 -2
View File
@@ -1194,7 +1194,7 @@ pub struct Limit {
#[env_config(name = "ZO_FILE_DOWNLOAD_PRIORITY_QUEUE_THREAD_NUM", default = 0)]
pub file_download_priority_queue_thread_num: usize,
#[env_config(name = "ZO_FILE_DOWNLOAD_PRIORITY_QUEUE_WINDOW_SECS", default = 3600)]
pub file_download_priority_queue_window_secs: u64,
pub file_download_priority_queue_window_secs: i64,
#[env_config(name = "ZO_FILE_DOWNLOAD_ENABLE_PRIORITY_QUEUE", default = true)]
pub file_download_enable_priority_queue: bool,
#[env_config(name = "ZO_QUERY_TIMEOUT", default = 600)]
@@ -1532,7 +1532,7 @@ pub struct CacheLatestFiles {
#[env_config(name = "ZO_CACHE_LATEST_FILES_DOWNLOAD_FROM_NODE", default = false)]
pub download_from_node: bool,
#[env_config(name = "ZO_CACHE_LATEST_FILES_DOWNLOAD_NODE_SIZE", default = 100)] // MB
pub download_node_size: usize,
pub download_node_size: i64,
}
#[derive(EnvConfig)]
+34 -16
View File
@@ -25,7 +25,7 @@ use config::{
};
use futures_util::StreamExt;
use hashbrown::{HashMap, HashSet};
use infra::cache::file_data::disk;
use infra::cache::file_data::{CacheType, TRACE_ID_FOR_CACHE_LATEST_FILE, disk};
use opentelemetry::global;
use proto::cluster_rpc::{
EmptyResponse, FileContent, FileContentResponse, FileList, SimpleFileList,
@@ -78,7 +78,8 @@ impl Event for Eventer {
files_to_download.push((
item.account.clone(),
item.key.clone(),
item.meta.compressed_size as usize,
item.meta.compressed_size,
item.meta.max_ts,
));
}
@@ -89,7 +90,8 @@ impl Event for Eventer {
files_to_download.push((
item.account.clone(),
ttv_file,
item.meta.index_size as usize,
item.meta.index_size,
item.meta.max_ts,
));
}
}
@@ -111,17 +113,33 @@ impl Event for Eventer {
}
// Fallback to individual downloads for failed files
for (account, file, size) in failed_files {
let size = if size > 0 { Some(size) } else { None };
if let Err(e) = infra::cache::file_data::download(&account, &file, size).await {
for (account, file, size, ts) in failed_files {
if let Err(e) = crate::job::queue_download(
TRACE_ID_FOR_CACHE_LATEST_FILE.to_string(),
account,
file,
size,
ts,
CacheType::Disk,
)
.await
{
log::error!("[gRPC:Event] Failed to cache file data: {}", e);
}
}
} else {
// Direct download when download_from_node_enabled is false
for (account, file, size) in files_to_download {
let size = if size > 0 { Some(size) } else { None };
if let Err(e) = infra::cache::file_data::download(&account, &file, size).await {
for (account, file, size, ts) in files_to_download {
if let Err(e) = crate::job::queue_download(
TRACE_ID_FOR_CACHE_LATEST_FILE.to_string(),
account,
file,
size,
ts,
CacheType::Disk,
)
.await
{
log::error!("[gRPC:Event] Failed to cache file data: {}", e);
}
}
@@ -249,8 +267,8 @@ async fn handle_file_chunked(
async fn download_from_node(
addr: &str,
files: &[(String, String, usize)],
) -> Result<Vec<(String, String, usize)>> {
files: &[(String, String, i64, i64)],
) -> Result<Vec<(String, String, i64, i64)>> {
let start = std::time::Instant::now();
let cfg = get_config();
log::debug!(
@@ -270,22 +288,22 @@ async fn download_from_node(
let file_size_map = files
.iter()
.filter_map(|(_, f, s)| {
.filter_map(|(_, f, s, _)| {
if *s > cfg.cache_latest_files.download_node_size * 1024 * 1024 {
None
} else {
Some((f, *s))
Some((f, *s as usize))
}
})
.collect::<HashMap<_, _>>();
if file_size_map.is_empty() {
return Ok(files
.iter()
.map(|(a, f, s)| (a.clone(), f.clone(), *s))
.map(|(a, f, s, ts)| (a.clone(), f.clone(), *s, *ts))
.collect());
}
let request = Request::new(SimpleFileList {
files: files.iter().map(|(_, f, _)| f.to_string()).collect(),
files: files.iter().map(|(_, f, ..)| f.to_string()).collect(),
});
let resp = client
@@ -358,7 +376,7 @@ async fn download_from_node(
// Return list of failed files
let failed_files: Vec<_> = files
.iter()
.filter(|(_, f, _)| !downloaded_files.contains(f))
.filter(|(_, f, ..)| !downloaded_files.contains(f))
.cloned()
.collect();
+1
View File
@@ -29,6 +29,7 @@ use hashlink::lru_cache::LruCache;
const DOWNLOAD_RETRY_TIMES: usize = 3;
const INITIAL_CACHE_SIZE: usize = 128;
pub const TRACE_ID_FOR_CACHE_LATEST_FILE: &str = "cache_latest_file";
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub enum CacheType {
+66 -56
View File
@@ -15,7 +15,7 @@
use std::{collections::VecDeque, sync::Arc};
use config::{get_config, meta::stream::FileMeta, metrics};
use config::{get_config, metrics, utils::time::now_micros};
use infra::cache::file_data;
use once_cell::sync::Lazy;
use tokio::sync::{
@@ -23,8 +23,8 @@ use tokio::sync::{
mpsc::{Receiver, Sender},
};
/// (account, file, size, cache_type)
type FileInfo = (String, String, usize, file_data::CacheType);
/// (trace_id, account, file, size, cache_type)
type FileInfo = (String, String, String, usize, file_data::CacheType);
struct DownloadQueue {
sender: Sender<FileInfo>,
@@ -76,41 +76,33 @@ pub async fn run() -> Result<(), anyhow::Error> {
let cfg = get_config();
// handle normal queue download
// move files
for _ in 0..cfg.limit.file_download_thread_num {
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] Receiving channel is closed");
log::debug!("[FILE_CACHE_DOWNLOAD:JOB:NORMAL] Receiving channel is closed");
break;
}
Some((account, file, file_size, cache)) => {
match download_file(&account, &file, file_size, cache).await {
Some((trace_id, account, file, file_size, cache)) => {
match download_file(thread, &trace_id, &account, &file, file_size, cache)
.await
{
Ok(data_len) => {
if data_len > 0 && data_len != file_size {
log::warn!(
"[FileDownloader] download file {} found size mismatch, expected: {}, actual: {}, will skip it",
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {} found size mismatch, expected: {}, actual: {}, will skip it",
file,
file_size,
data_len,
);
// update database
// if let Err(e) = file_list::update_compressed_size(&file,
// data_len) .await
// {
// log::error!(
// "[FileDownloader] update file size for file {} err:
// {}", file,
// e,
// );
// }
}
}
Err(e) => {
log::error!(
"[FileDownloader] download file {} to cache {:?} err: {}",
"[FILE_CACHE_DOWNLOAD:JOB:NORMAL] download file {} to cache {:?} err: {}",
file,
cache,
e,
@@ -137,19 +129,17 @@ pub async fn run() -> Result<(), anyhow::Error> {
let ret = rx.lock().await.recv().await;
match ret {
None => {
log::debug!(
"[FILE_CACHE_DOWNLOAD:PRIORITY_QUEUE:JOB] Receiving channel is closed"
);
log::debug!("[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] Receiving channel is closed");
if shutdown_tx.send(true).is_err() {
log::error!(
"[FILE_CACHE_DOWNLOAD:PRIORITY_QUEUE:JOB] Failed to send disconnect signal"
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] Failed to send disconnect signal"
);
}
break;
}
Some((account, file, file_size, cache)) => {
Some((trace_id, account, file, file_size, cache)) => {
PRIORITY_FILE_DOWNLOAD_CHANNEL
.push((account, file, file_size, cache))
.push((trace_id, account, file, file_size, cache))
.await;
}
}
@@ -157,26 +147,26 @@ pub async fn run() -> Result<(), anyhow::Error> {
});
// worker tasks: handle priority queue download
for _ in 0..cfg.limit.file_download_priority_queue_thread_num {
for thread in 0..cfg.limit.file_download_priority_queue_thread_num {
let mut shutdown_rx = shutdown_rx.clone();
tokio::spawn(async move {
loop {
tokio::select! {
_ = shutdown_rx.changed() => {
log::warn!(
"[FILE_CACHE_DOWNLOAD:PRIORITY_QUEUE:JOB] Received shutdown signal, exiting"
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] Received shutdown signal, exiting"
);
break;
}
_ = async {
let file_info = PRIORITY_FILE_DOWNLOAD_CHANNEL.pop().await;
match file_info {
Some((account, file, file_size, cache)) => {
match download_file(&account, &file, file_size, cache).await {
Some((trace_id, account, file, file_size, cache)) => {
match download_file(thread, &trace_id, &account, &file, file_size, cache).await {
Ok(data_len) => {
if data_len > 0 && data_len != file_size {
log::warn!(
"[FileDownloader::PriorityQueue] download file {} found size mismatch, expected: {}, actual: {}, will skip it",
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {} found size mismatch, expected: {}, actual: {}, will skip it",
file,
file_size,
data_len,
@@ -185,7 +175,7 @@ pub async fn run() -> Result<(), anyhow::Error> {
}
Err(e) => {
log::error!(
"[FileDownloader::PriorityQueue] download file {} to cache {:?} err: {}",
"[FILE_CACHE_DOWNLOAD:JOB:PRIORITY] download file {} to cache {:?} err: {}",
file,
cache,
e,
@@ -211,13 +201,17 @@ pub async fn run() -> Result<(), anyhow::Error> {
}
async fn download_file(
thread: usize,
trace_id: &str,
account: &str,
file_name: &str,
file_size: usize,
cache_type: file_data::CacheType,
) -> Result<usize, anyhow::Error> {
let cfg = get_config();
match cache_type {
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;
@@ -225,59 +219,75 @@ async fn download_file(
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, Some(file_size)).await
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, Some(file_size)).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
}
pub async fn queue_download(
trace_id: String,
account: String,
file: String,
size: i64,
ts: i64,
cache_type: file_data::CacheType,
to_priority_queue: bool,
) -> Result<(), anyhow::Error> {
if !to_priority_queue || !get_config().limit.file_download_enable_priority_queue {
FILE_DOWNLOAD_CHANNEL
.sender
.send((account, file, size as usize, cache_type))
.await?;
// update metrics
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
.with_label_values(&[])
.inc();
} else {
log::debug!(
"[FILE_CACHE_DOWNLOAD:JOB] [trace_id {trace_id}] enqueue file: {file}, size: {size}, ts: {ts}"
);
let cfg = get_config();
if cfg.limit.file_download_enable_priority_queue
&& should_prioritize_file(ts, cfg.limit.file_download_priority_queue_window_secs)
{
PRIORITY_FILE_DOWNLOAD_CHANNEL
.sender
.send((account, file.to_owned(), size as usize, cache_type))
.send((trace_id, account, file, size as usize, cache_type))
.await?;
// update metrics
metrics::FILE_DOWNLOADER_PRIORITY_QUEUE_SIZE
.with_label_values(&[])
.inc();
} else {
FILE_DOWNLOAD_CHANNEL
.sender
.send((
trace_id,
account,
file.to_owned(),
size as usize,
cache_type,
))
.await?;
// update metrics
metrics::FILE_DOWNLOADER_NORMAL_QUEUE_SIZE
.with_label_values(&[])
.inc();
}
Ok(())
}
pub fn should_prioritize_file(file_meta: &FileMeta) -> bool {
let cfg = get_config();
let window_micros = (cfg.limit.file_download_priority_queue_window_secs * 1_000_000) as i64;
let now = chrono::Utc::now().timestamp_micros();
// Check if the file's timestamp range overlaps with the current time window.
// A file is prioritized if its data is "fresh" - containing events from the recent past
// through the near future relative to current time, within the configured window.
file_meta.min_ts > now - window_micros && file_meta.max_ts < now + window_micros
// 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
}
+1 -1
View File
@@ -43,7 +43,7 @@ mod stats;
pub(crate) mod syslog_server;
mod telemetry;
pub use file_downloader::{queue_download, should_prioritize_file};
pub use file_downloader::queue_download;
pub use file_list_dump::FILE_LIST_SCHEMA;
pub use mmdb_downloader::MMDB_INIT_NOTIFIER;
+10 -20
View File
@@ -40,19 +40,16 @@ use itertools::Itertools;
use promql_parser::label::{MatchOp, Matchers};
use tracing::Instrument;
use crate::{
job::should_prioritize_file,
service::{
db, file_list,
search::{
datafusion::exec::register_table,
grpc::{
QueryParams,
storage::{cache_files, filter_file_list_by_tantivy_index},
},
index::{Condition, IndexCondition},
match_source,
use crate::service::{
db, file_list,
search::{
datafusion::exec::register_table,
grpc::{
QueryParams,
storage::{cache_files, filter_file_list_by_tantivy_index},
},
index::{Condition, IndexCondition},
match_source,
},
};
@@ -153,14 +150,7 @@ pub(crate) async fn create_context(
trace_id,
&files
.iter()
.map(|f| {
(
&f.account,
&f.key,
f.meta.compressed_size,
should_prioritize_file(&f.meta),
)
})
.map(|f| (&f.account, &f.key, f.meta.compressed_size, f.meta.max_ts))
.collect_vec(),
&mut scan_stats,
"parquet",
+29 -45
View File
@@ -48,22 +48,19 @@ use tantivy::Directory;
use tokio::sync::Semaphore;
use tracing::Instrument;
use crate::{
job::{self, should_prioritize_file},
service::{
db, file_list,
search::{
datafusion::exec,
generate_search_schema_diff,
index::IndexCondition,
inspector::{SearchInspectorFieldsBuilder, search_inspector_fields},
tantivy::puffin_directory::{
caching_directory::CachingDirectory,
convert_puffin_file_to_tantivy_dir,
footer_cache::FooterCache,
reader::{PuffinDirReader, warm_up_terms},
reader_cache,
},
use crate::service::{
db, file_list,
search::{
datafusion::exec,
generate_search_schema_diff,
index::IndexCondition,
inspector::{SearchInspectorFieldsBuilder, search_inspector_fields},
tantivy::puffin_directory::{
caching_directory::CachingDirectory,
convert_puffin_file_to_tantivy_dir,
footer_cache::FooterCache,
reader::{PuffinDirReader, warm_up_terms},
reader_cache,
},
},
};
@@ -260,14 +257,7 @@ pub async fn search(
&query.trace_id,
&files
.iter()
.map(|f| {
(
&f.account,
&f.key,
f.meta.compressed_size,
should_prioritize_file(&f.meta),
)
})
.map(|f| (&f.account, &f.key, f.meta.compressed_size, f.meta.max_ts))
.collect_vec(),
&mut scan_stats,
"parquet",
@@ -414,14 +404,14 @@ pub async fn search(
#[tracing::instrument(name = "service:search:grpc:storage:cache_files", skip_all)]
pub async fn cache_files(
trace_id: &str,
files: &[(&String, &String, i64, bool)],
files: &[(&String, &String, i64, i64)],
scan_stats: &mut ScanStats,
file_type: &str,
) -> Result<(file_data::CacheType, u64, u64), Error> {
// check how many files already cached
let mut cached_files = HashSet::with_capacity(files.len());
let (mut cache_hits, mut cache_misses) = (0, 0);
for (_account, file, _size, _to_priority_queue) in files.iter() {
for (_account, file, _size, _ts) in files.iter() {
if file_data::memory::exist(file).await {
scan_stats.querier_memory_cached_files += 1;
cached_files.insert(file);
@@ -461,26 +451,27 @@ pub async fn cache_files(
let trace_id = trace_id.to_string();
let files = files
.iter()
.filter_map(|(account, file, size, to_priority_queue)| {
.filter_map(|(account, file, size, ts)| {
if cached_files.contains(&file) {
None
} else {
Some((
account.to_string(),
file.to_string(),
*size,
*to_priority_queue,
))
Some((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 (account, file, size, to_priority_queue) in files {
if let Err(e) =
job::queue_download(account, file.clone(), size, cache_type, to_priority_queue)
.await
for (account, file, size, ts) in files {
if let Err(e) = crate::job::queue_download(
trace_id.clone(),
account,
file.clone(),
size,
ts,
cache_type,
)
.await
{
log::error!(
"[trace_id {trace_id}] error in queuing file {file} for background download: {e}"
@@ -543,14 +534,7 @@ pub async fn filter_file_list_by_tantivy_index(
&query.trace_id,
&index_file_names
.iter()
.map(|(ttv_file, f)| {
(
&f.account,
ttv_file,
f.meta.index_size,
should_prioritize_file(&f.meta),
)
})
.map(|(ttv_file, f)| (&f.account, ttv_file, f.meta.index_size, f.meta.max_ts))
.collect_vec(),
&mut scan_stats,
"index",