diff --git a/src/cli/basic/cli.rs b/src/cli/basic/cli.rs index d169fcf8e5..42e2e5122b 100644 --- a/src/cli/basic/cli.rs +++ b/src/cli/basic/cli.rs @@ -231,6 +231,14 @@ pub async fn cli() -> Result { .required(true) .help("snowflake id"), ]), + clap::Command::new("consistent-hash").about("consistent hash").args([ + clap::Arg::new("file") + .short('f') + .long("file") + .required(true) + .num_args(1..) + .help("file"), + ]), clap::Command::new("upgrade-db") .about("upgrade db table schemas").args(dataArgs()), ]) @@ -495,6 +503,14 @@ pub async fn cli() -> Result { let td = t.format("%Y-%m-%dT%H:%M:%SZ").to_string(); println!("datetimes: {}", td); } + "consistent-hash" => { + let files = command + .get_many::("file") + .unwrap() + .collect::>(); + let files = files.iter().map(|f| f.to_string()).collect::>(); + super::http::consistent_hash(files).await?; + } "upgrade-db" => { crate::migration::init_db().await?; } diff --git a/src/cli/basic/http.rs b/src/cli/basic/http.rs index f89305df22..fb0a91efe9 100644 --- a/src/cli/basic/http.rs +++ b/src/cli/basic/http.rs @@ -96,6 +96,7 @@ pub async fn node_offline() -> Result<(), anyhow::Error> { } Ok(()) } + pub async fn node_online() -> Result<(), anyhow::Error> { let url = "/node/enable?value=true"; let response = request(url, None, reqwest::Method::PUT).await?; @@ -262,6 +263,35 @@ pub async fn local_node_metrics() -> Result<(), anyhow::Error> { Ok(()) } +pub async fn consistent_hash(files: Vec) -> Result<(), anyhow::Error> { + let url = "/node/consistent_hash"; + let req = config::meta::search::HashFileRequest { files }; + let body = serde_json::to_vec(&req)?; + let response = request(url, Some(body), reqwest::Method::POST).await?; + let Some(body) = response else { + return Err(anyhow::anyhow!("consistent hash failed")); + }; + let response: config::meta::search::HashFileResponse = serde_json::from_str(&body)?; + + // Create header row with all column names + let mut table = Table::new(); + table.add_row(Row::new(vec![ + Cell::new("File"), + Cell::new("Querier"), + Cell::new("Alert Querier"), + ])); + for (file, nodes) in response.files.iter() { + table.add_row(Row::new(vec![ + Cell::new(file), + Cell::new(nodes.get("querier_interactive").unwrap_or(&"".to_string())), + Cell::new(nodes.get("querier_background").unwrap_or(&"".to_string())), + ])); + } + + table.printstd(); + Ok(()) +} + async fn request( url: &str, body: Option>, @@ -273,11 +303,12 @@ async fn request( let url = format!("{}{}", local.http_addr, url); let user = cfg.auth.root_user_email.clone(); let password = cfg.auth.root_user_password.clone(); - let mut response = client.request(method, url).basic_auth(user, Some(password)); + let mut req = client.request(method, url).basic_auth(user, Some(password)); if let Some(body) = body { - response = response.body(body); + req = req.header("Content-Type", "application/json"); + req = req.body(body); } - let response = response.send().await?; + let response = req.send().await?; let status = response.status(); let body = response.text().await?; if status.is_success() { diff --git a/src/common/infra/cluster/etcd.rs b/src/common/infra/cluster/etcd.rs index 92317839bd..a6bb3f150a 100644 --- a/src/common/infra/cluster/etcd.rs +++ b/src/common/infra/cluster/etcd.rs @@ -250,7 +250,7 @@ pub(crate) async fn set_status(status: NodeStatus, new_lease_id: bool) -> Result }; // update node status metrics - node.metrics = super::update_node_status_metrics(); + node.metrics = super::update_node_status_metrics().await; let val = json::to_string(&node).unwrap(); diff --git a/src/common/infra/cluster/mod.rs b/src/common/infra/cluster/mod.rs index 14960fbb96..c954f08385 100644 --- a/src/common/infra/cluster/mod.rs +++ b/src/common/infra/cluster/mod.rs @@ -163,6 +163,22 @@ pub async fn print_consistent_hash() -> HashMap map } +pub async fn count_consistent_hash() -> HashMap { + let mut map = HashMap::new(); + let r = QUERIER_INTERACTIVE_CONSISTENT_HASH.read().await; + map.insert("querier_interactive".to_string(), r.len()); + drop(r); + let r = QUERIER_BACKGROUND_CONSISTENT_HASH.read().await; + map.insert("querier_background".to_string(), r.len()); + drop(r); + let r = COMPACTOR_CONSISTENT_HASH.read().await; + map.insert("compactor".to_string(), r.len()); + drop(r); + let r = FLATTEN_COMPACTOR_CONSISTENT_HASH.read().await; + map.insert("flatten_compactor".to_string(), r.len()); + map +} + /// Register and keep alive the node to cluster pub async fn register_and_keep_alive() -> Result<()> { let cfg = get_config(); @@ -616,7 +632,7 @@ async fn set_node_status_metrics(node: &Node) { } } -fn update_node_status_metrics() -> NodeMetrics { +async fn update_node_status_metrics() -> NodeMetrics { let node_status = get_node_metrics(); config::metrics::NODE_UP @@ -650,6 +666,14 @@ fn update_node_status_metrics() -> NodeMetrics { .with_label_values(&["resets"]) .set(node_status.tcp_conns_resets as i64); + // update node consistent hash metrics + let consistent_hash = count_consistent_hash().await; + for (k, v) in consistent_hash { + config::metrics::NODE_CONSISTENT_HASH + .with_label_values(&[k.as_str()]) + .set(v as i64); + } + node_status } diff --git a/src/common/infra/cluster/nats.rs b/src/common/infra/cluster/nats.rs index fcb0db7b6f..0b3b2ec465 100644 --- a/src/common/infra/cluster/nats.rs +++ b/src/common/infra/cluster/nats.rs @@ -260,7 +260,7 @@ pub(crate) async fn set_status(status: NodeStatus) -> Result<()> { }; // update node status metrics - node.metrics = super::update_node_status_metrics(); + node.metrics = super::update_node_status_metrics().await; let val = json::to_string(&node).unwrap(); diff --git a/src/config/src/meta/search.rs b/src/config/src/meta/search.rs index 76c274fffc..5d0ecd1864 100644 --- a/src/config/src/meta/search.rs +++ b/src/config/src/meta/search.rs @@ -16,6 +16,7 @@ use std::collections::VecDeque; use bytes::Bytes as BytesImpl; +use hashbrown::HashMap; use proto::cluster_rpc; use serde::{Deserialize, Deserializer, Serialize}; use utoipa::ToSchema; @@ -1203,6 +1204,16 @@ pub struct ValuesRequest { pub sql: String, } +#[derive(Debug, Deserialize, Clone, Serialize)] +pub struct HashFileRequest { + pub files: Vec, +} + +#[derive(Debug, Default, Deserialize, Clone, Serialize)] +pub struct HashFileResponse { + pub files: HashMap>, +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/config/src/meta/stream.rs b/src/config/src/meta/stream.rs index db2d87f375..12bebebcfa 100644 --- a/src/config/src/meta/stream.rs +++ b/src/config/src/meta/stream.rs @@ -168,6 +168,7 @@ pub struct ListStreamParams { #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct FileKey { + pub id: i64, pub account: String, pub key: String, pub meta: FileMeta, @@ -176,8 +177,9 @@ pub struct FileKey { } impl FileKey { - pub fn new(account: String, key: String, meta: FileMeta, deleted: bool) -> Self { + pub fn new(id: i64, account: String, key: String, meta: FileMeta, deleted: bool) -> Self { Self { + id, account, key, meta, @@ -188,6 +190,7 @@ impl FileKey { pub fn from_file_name(file: &str) -> Self { Self { + id: 0, account: String::default(), key: file.to_string(), meta: FileMeta::default(), @@ -474,6 +477,7 @@ impl From<&cluster_rpc::FileMeta> for FileMeta { impl From<&FileKey> for cluster_rpc::FileKey { fn from(req: &FileKey) -> Self { cluster_rpc::FileKey { + id: req.id, account: req.account.clone(), key: req.key.clone(), meta: Some(cluster_rpc::FileMeta::from(&req.meta)), @@ -486,6 +490,7 @@ impl From<&FileKey> for cluster_rpc::FileKey { impl From<&cluster_rpc::FileKey> for FileKey { fn from(req: &cluster_rpc::FileKey) -> Self { FileKey { + id: req.id, account: req.account.clone(), key: req.key.clone(), meta: FileMeta::from(req.meta.as_ref().unwrap()), diff --git a/src/config/src/metrics.rs b/src/config/src/metrics.rs index 2248914969..dcd562e481 100644 --- a/src/config/src/metrics.rs +++ b/src/config/src/metrics.rs @@ -789,6 +789,16 @@ pub static NODE_TCP_CONNECTIONS: Lazy = Lazy::new(|| { .expect("Metric created") }); +pub static NODE_CONSISTENT_HASH: Lazy = Lazy::new(|| { + IntGaugeVec::new( + Opts::new("node_consistent_hash", "Consistent hash") + .namespace(NAMESPACE) + .const_labels(create_const_labels()), + &["type"], + ) + .expect("Metric created") +}); + // query disk cache metrics pub static QUERY_DISK_CACHE_HIT_COUNT: Lazy = Lazy::new(|| { IntCounterVec::new( @@ -1052,6 +1062,9 @@ fn register_metrics(registry: &Registry) { registry .register(Box::new(NODE_TCP_CONNECTIONS.clone())) .expect("Metric registered"); + registry + .register(Box::new(NODE_CONSISTENT_HASH.clone())) + .expect("Metric registered"); // query disk cache metrics registry diff --git a/src/handler/http/request/status/mod.rs b/src/handler/http/request/status/mod.rs index 70f853bc1c..ddce59d49e 100644 --- a/src/handler/http/request/status/mod.rs +++ b/src/handler/http/request/status/mod.rs @@ -20,20 +20,23 @@ use actix_web::{ cookie::{Cookie, SameSite}, get, head, http::header, - put, web, + post, put, web, }; use arrow_schema::Schema; use config::{ Config, META_ORG_ID, QUICK_MODEL_FIELDS, SQL_FULL_TEXT_SEARCH_FIELDS, TIMESTAMP_COL_NAME, cluster::LOCAL_NODE, get_config, get_instance_id, - meta::{cluster::NodeStatus, function::ZoFunction}, + meta::{ + cluster::{NodeStatus, Role, RoleGroup}, + function::ZoFunction, + search::{HashFileRequest, HashFileResponse}, + }, utils::{base64, json, schema_ext::SchemaExt}, }; use hashbrown::HashMap; use infra::{ - cache::{self, file_data::disk::FileType}, - file_list, + cache, file_list, schema::{STREAM_SCHEMAS, STREAM_SCHEMAS_LATEST}, }; use serde::Serialize; @@ -349,11 +352,13 @@ pub async fn cache_status() -> Result { let mem_file_num = cache::file_data::memory::len().await; let (mem_max_size, mem_cur_size) = cache::file_data::memory::stats().await; - let disk_file_num = cache::file_data::disk::len(FileType::DATA).await; - let (disk_max_size, disk_cur_size) = cache::file_data::disk::stats(FileType::DATA).await; - let disk_result_file_num = cache::file_data::disk::len(FileType::RESULT).await; + let disk_file_num = cache::file_data::disk::len(cache::file_data::disk::FileType::DATA).await; + let (disk_max_size, disk_cur_size) = + cache::file_data::disk::stats(cache::file_data::disk::FileType::DATA).await; + let disk_result_file_num = + cache::file_data::disk::len(cache::file_data::disk::FileType::RESULT).await; let (disk_result_max_size, disk_result_cur_size) = - cache::file_data::disk::stats(FileType::RESULT).await; + cache::file_data::disk::stats(cache::file_data::disk::FileType::RESULT).await; stats.insert( "FILE_DATA", json::json!({ @@ -820,3 +825,33 @@ async fn node_metrics() -> Result { let metrics = config::utils::sysinfo::get_node_metrics(); Ok(MetaHttpResponse::json(metrics)) } + +#[post("/consistent_hash")] +async fn consistent_hash(body: web::Json) -> Result { + let mut ret = HashFileResponse::default(); + for file in body.files.iter() { + let mut nodes = HashMap::new(); + nodes.insert( + "querier_interactive".to_string(), + cluster::get_node_from_consistent_hash( + file, + &Role::Querier, + Some(RoleGroup::Interactive), + ) + .await + .unwrap_or_default(), + ); + nodes.insert( + "querier_background".to_string(), + cluster::get_node_from_consistent_hash( + file, + &Role::Querier, + Some(RoleGroup::Background), + ) + .await + .unwrap_or_default(), + ); + ret.files.insert(file.clone(), nodes); + } + Ok(MetaHttpResponse::json(ret)) +} diff --git a/src/handler/http/router/mod.rs b/src/handler/http/router/mod.rs index 75542da9cf..ffc09dee44 100644 --- a/src/handler/http/router/mod.rs +++ b/src/handler/http/router/mod.rs @@ -261,7 +261,8 @@ pub fn get_basic_routes(svc: &mut web::ServiceConfig) { .service(status::enable_node) .service(status::flush_node) .service(status::list_node) - .service(status::node_metrics), + .service(status::node_metrics) + .service(status::consistent_hash), ); if get_config().common.swagger_enabled { diff --git a/src/infra/src/file_list/mod.rs b/src/infra/src/file_list/mod.rs index 005d89ccbf..8cad5bd458 100644 --- a/src/infra/src/file_list/mod.rs +++ b/src/infra/src/file_list/mod.rs @@ -52,11 +52,11 @@ pub fn connect_local_cache() -> Box { pub trait FileList: Sync + Send + 'static { async fn create_table(&self) -> Result<()>; async fn create_table_index(&self) -> Result<()>; - async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()>; - async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()>; + async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result; + async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result; async fn remove(&self, file: &str) -> Result<()>; async fn batch_add(&self, files: &[FileKey]) -> Result<()>; - async fn batch_add_with_id(&self, files: &[(i64, &FileKey)]) -> Result<()>; + async fn batch_add_with_id(&self, files: &[FileKey]) -> Result<()>; async fn batch_add_history(&self, files: &[FileKey]) -> Result<()>; async fn update_dump_records(&self, dump_file: &FileKey, dumped_ids: &[i64]) -> Result<()>; async fn batch_process(&self, files: &[FileKey]) -> Result<()>; @@ -71,7 +71,7 @@ pub trait FileList: Sync + Send + 'static { async fn contains(&self, file: &str) -> Result; async fn update_flattened(&self, file: &str, flattened: bool) -> Result<()>; async fn update_compressed_size(&self, file: &str, size: i64) -> Result<()>; - async fn list(&self) -> Result>; + async fn list(&self) -> Result>; async fn query( &self, org_id: &str, @@ -80,15 +80,15 @@ pub trait FileList: Sync + Send + 'static { time_level: PartitionTimeLevel, time_range: Option<(i64, i64)>, flattened: Option, - ) -> Result>; + ) -> Result>; async fn query_for_merge( &self, org_id: &str, stream_type: StreamType, stream_name: &str, date_range: Option<(String, String)>, - ) -> Result>; - async fn query_by_ids(&self, ids: &[i64]) -> Result>; + ) -> Result>; + async fn query_by_ids(&self, ids: &[i64]) -> Result>; async fn query_ids( &self, org_id: &str, @@ -96,6 +96,7 @@ pub trait FileList: Sync + Send + 'static { stream_name: &str, time_range: Option<(i64, i64)>, ) -> Result>; + async fn query_ids_by_files(&self, files: &[FileKey]) -> Result>; async fn query_old_data_hours( &self, org_id: &str, @@ -192,12 +193,12 @@ pub async fn create_table_index() -> Result<()> { } #[inline] -pub async fn add(account: &str, file: &str, meta: &FileMeta) -> Result<()> { +pub async fn add(account: &str, file: &str, meta: &FileMeta) -> Result { CLIENT.add(account, file, meta).await } #[inline] -pub async fn add_history(account: &str, file: &str, meta: &FileMeta) -> Result<()> { +pub async fn add_history(account: &str, file: &str, meta: &FileMeta) -> Result { CLIENT.add_history(account, file, meta).await } @@ -261,7 +262,7 @@ pub async fn update_compressed_size(file: &str, size: i64) -> Result<()> { } #[inline] -pub async fn list() -> Result> { +pub async fn list() -> Result> { CLIENT.list().await } @@ -274,7 +275,7 @@ pub async fn query( time_level: PartitionTimeLevel, time_range: Option<(i64, i64)>, flattened: Option, -) -> Result> { +) -> Result> { validate_time_range(time_range)?; CLIENT .query( @@ -295,7 +296,7 @@ pub async fn query_for_merge( stream_type: StreamType, stream_name: &str, date_range: Option<(String, String)>, -) -> Result> { +) -> Result> { CLIENT .query_for_merge(org_id, stream_type, stream_name, date_range) .await @@ -303,7 +304,7 @@ pub async fn query_for_merge( #[inline] #[tracing::instrument(name = "infra:file_list:query_db_by_ids", skip_all)] -pub async fn query_by_ids(ids: &[i64]) -> Result> { +pub async fn query_by_ids(ids: &[i64]) -> Result> { CLIENT.query_by_ids(ids).await } @@ -321,6 +322,12 @@ pub async fn query_ids( .await } +#[inline] +#[tracing::instrument(name = "infra:file_list:db:query_ids_by_files")] +pub async fn query_ids_by_files(files: &[FileKey]) -> Result> { + CLIENT.query_ids_by_files(files).await +} + #[inline] #[tracing::instrument(name = "infra:file_list:db:query_old_data_hours")] pub async fn query_old_data_hours( @@ -572,16 +579,29 @@ pub struct FileRecord { pub flattened: bool, } -impl From<&FileRecord> for FileMeta { - fn from(record: &FileRecord) -> Self { +impl From<&FileRecord> for FileKey { + fn from(r: &FileRecord) -> Self { Self { - min_ts: record.min_ts, - max_ts: record.max_ts, - records: record.records, - original_size: record.original_size, - compressed_size: record.compressed_size, - index_size: record.index_size, - flattened: record.flattened, + id: r.id, + account: r.account.to_string(), + key: "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, + meta: r.into(), + deleted: r.deleted, + segment_ids: None, + } + } +} + +impl From<&FileRecord> for FileMeta { + fn from(r: &FileRecord) -> Self { + Self { + min_ts: r.min_ts, + max_ts: r.max_ts, + records: r.records, + original_size: r.original_size, + compressed_size: r.compressed_size, + index_size: r.index_size, + flattened: r.flattened, } } } @@ -655,3 +675,9 @@ pub struct FileId { #[sqlx(default)] pub deleted: bool, } + +#[derive(Clone, Debug, Default, sqlx::FromRow)] +pub struct FileIdWithFile { + pub id: i64, + pub file: String, +} diff --git a/src/infra/src/file_list/mysql.rs b/src/infra/src/file_list/mysql.rs index 2ec333f10e..eee43cbe60 100644 --- a/src/infra/src/file_list/mysql.rs +++ b/src/infra/src/file_list/mysql.rs @@ -63,11 +63,11 @@ impl super::FileList for MysqlFileList { create_table_index().await } - async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list", account, file, meta).await } - async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list_history", account, file, meta) .await } @@ -94,7 +94,7 @@ impl super::FileList for MysqlFileList { self.inner_batch_process("file_list", files).await } - async fn batch_add_with_id(&self, _files: &[(i64, &FileKey)]) -> Result<()> { + async fn batch_add_with_id(&self, _files: &[FileKey]) -> Result<()> { unimplemented!("Unsupported") } @@ -361,7 +361,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat Ok(()) } - async fn list(&self) -> Result> { + async fn list(&self) -> Result> { return Ok(vec![]); // disallow list all data } @@ -373,7 +373,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat _time_level: PartitionTimeLevel, time_range: Option<(i64, i64)>, flattened: Option, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = time_range { if start == 0 && end == 0 { return Ok(Vec::new()); @@ -390,7 +390,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat let ret = if flattened.is_some() { sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = ? AND flattened = ? LIMIT 1000; "#, @@ -404,7 +404,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let max_ts_upper_bound = super::calculate_max_ts_upper_bound(time_end, stream_type); sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = ? AND max_ts >= ? AND max_ts <= ? AND min_ts <= ?; "#, @@ -422,17 +422,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .observe(time); Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } @@ -442,7 +432,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s stream_type: StreamType, stream_name: &str, date_range: Option<(String, String)>, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = date_range.as_ref() { if start.is_empty() && end.is_empty() { return Ok(Vec::new()); @@ -459,7 +449,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let (date_start, date_end) = date_range.unwrap_or(("".to_string(), "".to_string())); let ret = sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = ? AND date >= ? AND date <= ?; "#, @@ -475,21 +465,11 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .observe(time); Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } - async fn query_by_ids(&self, ids: &[i64]) -> Result> { + async fn query_by_ids(&self, ids: &[i64]) -> Result> { if ids.is_empty() { return Ok(Vec::default()); } @@ -522,17 +502,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s ret.extend_from_slice(&res?); } - Ok(ret - .into_iter() - .map(|r| { - ( - r.id, - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - FileMeta::from(&r), - ) - }) - .collect()) + Ok(ret.iter().map(|r| r.into()).collect()) } async fn query_ids( @@ -587,7 +557,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .bind(time_end) .fetch_all(&pool) .await - })); + })); } let mut rets = Vec::new(); @@ -609,6 +579,45 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s Ok(rets) } + async fn query_ids_by_files(&self, files: &[FileKey]) -> Result> { + let mut ret = stdHashMap::with_capacity(files.len()); + // group by date + let mut stream_files = HashMap::new(); + let mut files_map = HashMap::with_capacity(files.len()); + for file in files { + if file.id > 0 { + ret.insert(file.key.clone(), file.id); + continue; + } + let (stream_key, date_key, file_name) = + parse_file_key_columns(&file.key).map_err(|e| Error::Message(e.to_string()))?; + let stream_entry = stream_files.entry(stream_key).or_insert(HashMap::new()); + let date_entry = stream_entry.entry(date_key).or_insert(Vec::new()); + date_entry.push(file_name.clone()); + files_map.insert(file_name, &file.key); + } + for (stream_key, stream_files) in stream_files { + let pool = CLIENT_RO.clone(); + for (date_key, files) in stream_files { + let sql = format!( + "SELECT id, file FROM file_list WHERE stream = ? AND date = ? AND file IN ({});", + files.join(",") + ); + let query_res = sqlx::query_as::<_, super::FileIdWithFile>(&sql) + .bind(&stream_key) + .bind(&date_key) + .fetch_all(&pool) + .await?; + for file in query_res { + if let Some(file_name) = files_map.get(&file.file) { + ret.insert(file_name.to_string(), file.id); + } + } + } + } + Ok(ret) + } + async fn query_old_data_hours( &self, org_id: &str, @@ -1521,7 +1530,7 @@ impl MysqlFileList { account: &str, file: &str, meta: &FileMeta, - ) -> Result<()> { + ) -> Result { let pool = CLIENT.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; @@ -1551,12 +1560,12 @@ INSERT IGNORE INTO {table} (account, org, stream, date, file, deleted, min_ts, m .execute(&pool) .await { Err(sqlx::Error::Database(e)) => if e.is_unique_violation() { - Ok(()) + Ok(0) } else { Err(Error::Message(e.to_string())) }, Err(e) => Err(e.into()), - Ok(_) => Ok(()), + Ok(v) => Ok(v.last_insert_id() as i64), } } @@ -1615,6 +1624,10 @@ INSERT IGNORE INTO {table} (account, org, stream, date, file, deleted, min_ts, m // get ids of the files let mut ids = Vec::with_capacity(files.len()); for file in files { + if file.id > 0 { + ids.push(file.id.to_string()); + continue; + } let (stream_key, date_key, file_name) = parse_file_key_columns(&file.key) .map_err(|e| Error::Message(e.to_string()))?; DB_QUERY_NUMS diff --git a/src/infra/src/file_list/postgres.rs b/src/infra/src/file_list/postgres.rs index e34d907806..a10cd71970 100644 --- a/src/infra/src/file_list/postgres.rs +++ b/src/infra/src/file_list/postgres.rs @@ -63,11 +63,11 @@ impl super::FileList for PostgresFileList { create_table_index().await } - async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list", account, file, meta).await } - async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list_history", account, file, meta) .await } @@ -95,7 +95,7 @@ impl super::FileList for PostgresFileList { self.inner_batch_process("file_list", files).await } - async fn batch_add_with_id(&self, _files: &[(i64, &FileKey)]) -> Result<()> { + async fn batch_add_with_id(&self, _files: &[FileKey]) -> Result<()> { unimplemented!("Unsupported") } @@ -374,7 +374,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat Ok(()) } - async fn list(&self) -> Result> { + async fn list(&self) -> Result> { return Ok(vec![]); // disallow list all data } @@ -386,7 +386,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat _time_level: PartitionTimeLevel, time_range: Option<(i64, i64)>, flattened: Option, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = time_range { if start == 0 && end == 0 { return Ok(Vec::new()); @@ -403,7 +403,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat let ret = if flattened.is_some() { sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND flattened = $2 LIMIT 1000; "# @@ -415,7 +415,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let (time_start, time_end) = time_range.unwrap_or((0, 0)); let max_ts_upper_bound = super::calculate_max_ts_upper_bound(time_end, stream_type); let sql = r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND max_ts >= $2 AND max_ts <= $3 AND min_ts <= $4; "#; @@ -434,17 +434,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .observe(time); Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } @@ -454,7 +444,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s stream_type: StreamType, stream_name: &str, date_range: Option<(String, String)>, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = date_range.as_ref() { if start.is_empty() && end.is_empty() { return Ok(Vec::new()); @@ -471,7 +461,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let (date_start, date_end) = date_range.unwrap_or(("".to_string(), "".to_string())); let sql = r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND date >= $2 AND date <= $3; "#; @@ -488,21 +478,11 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .observe(time); Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } - async fn query_by_ids(&self, ids: &[i64]) -> Result> { + async fn query_by_ids(&self, ids: &[i64]) -> Result> { if ids.is_empty() { return Ok(Vec::default()); } @@ -535,17 +515,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s ret.extend_from_slice(&res?); } - Ok(ret - .into_iter() - .map(|r| { - ( - r.id, - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - FileMeta::from(&r), - ) - }) - .collect()) + Ok(ret.iter().map(|r| r.into()).collect()) } async fn query_ids( @@ -622,6 +592,45 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s Ok(rets) } + async fn query_ids_by_files(&self, files: &[FileKey]) -> Result> { + let mut ret = stdHashMap::with_capacity(files.len()); + // group by date + let mut stream_files = HashMap::new(); + let mut files_map = HashMap::with_capacity(files.len()); + for file in files { + if file.id > 0 { + ret.insert(file.key.clone(), file.id); + continue; + } + let (stream_key, date_key, file_name) = + parse_file_key_columns(&file.key).map_err(|e| Error::Message(e.to_string()))?; + let stream_entry = stream_files.entry(stream_key).or_insert(HashMap::new()); + let date_entry = stream_entry.entry(date_key).or_insert(Vec::new()); + date_entry.push(file_name.clone()); + files_map.insert(file_name, &file.key); + } + for (stream_key, stream_files) in stream_files { + let pool = CLIENT_RO.clone(); + for (date_key, files) in stream_files { + let sql = format!( + "SELECT id, file FROM file_list WHERE stream = $1 AND date = $2 AND file IN ({});", + files.join(",") + ); + let query_res = sqlx::query_as::<_, super::FileIdWithFile>(&sql) + .bind(&stream_key) + .bind(&date_key) + .fetch_all(&pool) + .await?; + for file in query_res { + if let Some(file_name) = files_map.get(&file.file) { + ret.insert(file_name.to_string(), file.id); + } + } + } + } + Ok(ret) + } + async fn query_old_data_hours( &self, org_id: &str, @@ -1469,7 +1478,7 @@ impl PostgresFileList { account: &str, file: &str, meta: &FileMeta, - ) -> Result<()> { + ) -> Result { let pool = CLIENT.clone(); let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; @@ -1477,39 +1486,40 @@ impl PostgresFileList { DB_QUERY_NUMS .with_label_values(&["insert", table, ""]) .inc(); - match - sqlx - ::query( - format!( - r#" + let ret: std::result::Result, sea_orm::SqlxError> = sqlx::query_scalar( + format!( + r#" INSERT INTO {table} (account, org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13) - ON CONFLICT DO NOTHING; - "# - ).as_str() - ) - .bind(account) - .bind(org_id) - .bind(stream_key) - .bind(date_key) - .bind(file_name) - .bind(false) - .bind(meta.min_ts) - .bind(meta.max_ts) - .bind(meta.records) - .bind(meta.original_size) - .bind(meta.compressed_size) - .bind(meta.index_size) - .bind(meta.flattened) - .execute(&pool).await - { - Err(sqlx::Error::Database(e)) => if e.is_unique_violation() { - Ok(()) - } else { - Err(Error::Message(e.to_string())) + ON CONFLICT DO NOTHING + RETURNING id; + "# + ).as_str() + ) + .bind(account) + .bind(org_id) + .bind(stream_key) + .bind(date_key) + .bind(file_name) + .bind(false) + .bind(meta.min_ts) + .bind(meta.max_ts) + .bind(meta.records) + .bind(meta.original_size) + .bind(meta.compressed_size) + .bind(meta.index_size) + .bind(meta.flattened) + .fetch_one(&pool).await; + match ret { + Err(sqlx::Error::Database(e)) => { + if e.is_unique_violation() { + Ok(0) + } else { + Err(Error::Message(e.to_string())) + } } Err(e) => Err(e.into()), - Ok(_) => Ok(()), + Ok(v) => Ok(v.unwrap_or(0)), } } @@ -1568,6 +1578,10 @@ INSERT INTO {table} (account, org, stream, date, file, deleted, min_ts, max_ts, // get ids of the files let mut ids = Vec::with_capacity(files.len()); for file in files { + if file.id > 0 { + ids.push(file.id.to_string()); + continue; + } let (stream_key, date_key, file_name) = parse_file_key_columns(&file.key) .map_err(|e| Error::Message(e.to_string()))?; DB_QUERY_NUMS diff --git a/src/infra/src/file_list/sqlite.rs b/src/infra/src/file_list/sqlite.rs index 9e4b1da6c5..39a3fb9692 100644 --- a/src/infra/src/file_list/sqlite.rs +++ b/src/infra/src/file_list/sqlite.rs @@ -61,11 +61,11 @@ impl super::FileList for SqliteFileList { create_table_index().await } - async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list", account, file, meta).await } - async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result<()> { + async fn add_history(&self, account: &str, file: &str, meta: &FileMeta) -> Result { self.inner_add("file_list_history", account, file, meta) .await } @@ -92,12 +92,8 @@ impl super::FileList for SqliteFileList { self.inner_batch_process("file_list", files).await } - async fn batch_add_with_id(&self, files: &[(i64, &FileKey)]) -> Result<()> { - let files = files - .iter() - .map(|(id, f)| (Some(*id), *f)) - .collect::>(); - self.inner_batch_process_with_id("file_list", &files).await + async fn batch_add_with_id(&self, files: &[FileKey]) -> Result<()> { + self.inner_batch_process("file_list", files).await } async fn batch_add_history(&self, files: &[FileKey]) -> Result<()> { @@ -329,23 +325,14 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat Ok(()) } - async fn list(&self) -> Result> { + async fn list(&self) -> Result> { let pool = CLIENT_RO.clone(); let ret = sqlx::query_as::<_, super::FileRecord>( - r#"SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list;"#, + r#"SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list;"#, ) .fetch_all(&pool) .await?; - Ok(ret - .iter() - .map(|r| { - ( - r.account.to_string(), - format!("files/{}/{}/{}", r.stream, r.date, r.file), - r.into(), - ) - }) - .collect()) + Ok(ret.iter().map(|r| r.into()).collect()) } async fn query( @@ -356,7 +343,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat _time_level: PartitionTimeLevel, time_range: Option<(i64, i64)>, flattened: Option, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = time_range { if start == 0 && end == 0 { return Ok(Vec::new()); @@ -369,7 +356,7 @@ SELECT min_ts, max_ts, records, original_size, compressed_size, index_size, flat let ret = if flattened.is_some() { sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND flattened = $2 LIMIT 1000; "#, @@ -383,7 +370,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let max_ts_upper_bound = super::calculate_max_ts_upper_bound(time_end, stream_type); sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND max_ts >= $2 AND max_ts <= $3 AND min_ts <= $4; "#, @@ -397,17 +384,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s }; Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } @@ -417,7 +394,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s stream_type: StreamType, stream_name: &str, date_range: Option<(String, String)>, - ) -> Result> { + ) -> Result> { if let Some((start, end)) = date_range.as_ref() { if start.is_empty() && end.is_empty() { return Ok(Vec::new()); @@ -430,7 +407,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s let (date_start, date_end) = date_range.unwrap_or(("".to_string(), "".to_string())); let ret = sqlx::query_as::<_, super::FileRecord>( r#" -SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened +SELECT id, account, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened FROM file_list WHERE stream = $1 AND date >= $2 AND date <= $3; "#, @@ -442,21 +419,11 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s .await; Ok(ret? .iter() - .filter_map(|r| { - if r.deleted { - None - } else { - Some(( - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - )) - } - }) + .filter_map(|r| if r.deleted { None } else { Some(r.into()) }) .collect()) } - async fn query_by_ids(&self, ids: &[i64]) -> Result> { + async fn query_by_ids(&self, ids: &[i64]) -> Result> { if ids.is_empty() { return Ok(Vec::default()); } @@ -482,17 +449,7 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s ret.extend_from_slice(&res); } - Ok(ret - .iter() - .map(|r| { - ( - r.id, - r.account.to_string(), - "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - r.into(), - ) - }) - .collect()) + Ok(ret.iter().map(|r| r.into()).collect()) } async fn query_ids( @@ -561,6 +518,45 @@ SELECT account, stream, date, file, deleted, min_ts, max_ts, records, original_s Ok(rets) } + async fn query_ids_by_files(&self, files: &[FileKey]) -> Result> { + let mut ret = stdHashMap::with_capacity(files.len()); + // group by date + let mut stream_files = HashMap::new(); + let mut files_map = HashMap::with_capacity(files.len()); + for file in files { + if file.id > 0 { + ret.insert(file.key.clone(), file.id); + continue; + } + let (stream_key, date_key, file_name) = + parse_file_key_columns(&file.key).map_err(|e| Error::Message(e.to_string()))?; + let stream_entry = stream_files.entry(stream_key).or_insert(HashMap::new()); + let date_entry = stream_entry.entry(date_key).or_insert(Vec::new()); + date_entry.push(file_name.clone()); + files_map.insert(file_name, &file.key); + } + for (stream_key, stream_files) in stream_files { + let pool = CLIENT_RO.clone(); + for (date_key, files) in stream_files { + let sql = format!( + "SELECT id, file FROM file_list WHERE stream = $1 AND date = $2 AND file IN ({});", + files.join(",") + ); + let query_res = sqlx::query_as::<_, super::FileIdWithFile>(&sql) + .bind(&stream_key) + .bind(&date_key) + .fetch_all(&pool) + .await?; + for file in query_res { + if let Some(file_name) = files_map.get(&file.file) { + ret.insert(file_name.to_string(), file.id); + } + } + } + } + Ok(ret) + } + async fn query_old_data_hours( &self, org_id: &str, @@ -1299,7 +1295,7 @@ impl SqliteFileList { account: &str, file: &str, meta: &FileMeta, - ) -> Result<()> { + ) -> Result { self.inner_add_with_id(table, None, account, file, meta) .await } @@ -1311,7 +1307,7 @@ impl SqliteFileList { account: &str, file: &str, meta: &FileMeta, - ) -> Result<()> { + ) -> Result { let (stream_key, date_key, file_name) = parse_file_key_columns(file).map_err(|e| Error::Message(e.to_string()))?; let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); @@ -1340,25 +1336,16 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ .execute(&*client) .await { Err(sqlx::Error::Database(e)) => if e.is_unique_violation() { - Ok(()) + Ok(0) } else { Err(Error::Message(e.to_string())) }, Err(e) => Err(e.into()), - Ok(_) => Ok(()), + Ok(v) => Ok(v.last_insert_rowid()), } } async fn inner_batch_process(&self, table: &str, files: &[FileKey]) -> Result<()> { - let files: Vec<(Option, _)> = files.iter().map(|f| (None, f)).collect::>(); - self.inner_batch_process_with_id(table, &files).await - } - - async fn inner_batch_process_with_id( - &self, - table: &str, - files: &[(Option, &FileKey)], - ) -> Result<()> { if files.is_empty() { return Ok(()); } @@ -1367,14 +1354,15 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ let client = client.lock().await; let mut tx = client.begin().await?; - let add_items = files.iter().filter(|v| !v.1.deleted).collect::>(); + let add_items = files.iter().filter(|f| !f.deleted).collect::>(); if !add_items.is_empty() { let chunks = add_items.chunks(100); for files in chunks { let mut query_builder: QueryBuilder = QueryBuilder::new( format!("INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ts, records, original_size, compressed_size, index_size, flattened)").as_str(), ); - query_builder.push_values(files, |mut b, (id, item)| { + query_builder.push_values(files, |mut b, item| { + let id = if item.id > 0 { Some(item.id) } else { None }; let (stream_key, date_key, file_name) = parse_file_key_columns(&item.key).expect("parse file key failed"); let org_id = stream_key[..stream_key.find('/').unwrap()].to_string(); @@ -1405,14 +1393,18 @@ INSERT INTO {table} (id, account, org, stream, date, file, deleted, min_ts, max_ } } - let del_items = files.iter().filter(|v| v.1.deleted).collect::>(); + let del_items = files.iter().filter(|f| f.deleted).collect::>(); if !del_items.is_empty() { let chunks = del_items.chunks(1000); for files in chunks { // get ids of the files let mut ids = Vec::with_capacity(files.len()); for file in files { - let (stream_key, date_key, file_name) = parse_file_key_columns(&file.1.key) + if file.id > 0 { + ids.push(file.id.to_string()); + continue; + } + let (stream_key, date_key, file_name) = parse_file_key_columns(&file.key) .map_err(|e| Error::Message(e.to_string()))?; let query_res: std::result::Result, sea_orm::SqlxError> = sqlx::query_scalar( r#"SELECT id FROM file_list WHERE stream = $1 AND date = $2 AND file = $3;"#, diff --git a/src/job/file_list_dump.rs b/src/job/file_list_dump.rs index 63b3872d91..ad9d777eac 100644 --- a/src/job/file_list_dump.rs +++ b/src/job/file_list_dump.rs @@ -436,6 +436,7 @@ async fn generate_dump( infra::storage::put(&account, &file_key, buf.into()).await?; let dump_file = FileKey { + id: 0, account: account.clone(), key: file_key.clone(), meta: meta.clone(), diff --git a/src/job/files/broadcast.rs b/src/job/files/broadcast.rs index e0822b6814..48356172a8 100644 --- a/src/job/files/broadcast.rs +++ b/src/job/files/broadcast.rs @@ -30,7 +30,7 @@ pub async fn run() -> Result<(), anyhow::Error> { } q.drain(..).collect::>() }; - if let Err(e) = broadcast::send(&files, None).await { + if let Err(e) = broadcast::send(&files).await { log::error!("[broadcast] local queue to nodes error: {}", e); } } diff --git a/src/job/files/parquet.rs b/src/job/files/parquet.rs index 59416bfde3..cd84ff966f 100644 --- a/src/job/files/parquet.rs +++ b/src/job/files/parquet.rs @@ -328,6 +328,7 @@ async fn prepare_files( let prefix = columns.join("/"); let partition = partition_files_with_size.entry(prefix).or_default(); partition.push(FileKey::new( + 0, "".to_string(), // here we don't need it file_key.clone(), parquet_meta, @@ -536,8 +537,9 @@ async fn move_files( } // write file list to storage - let ret = db::file_list::set(&account, &new_file_name, Some(new_file_meta), false).await; - if let Err(e) = ret { + if let Err(e) = + db::file_list::set(&account, &new_file_name, Some(new_file_meta), false).await + { log::error!( "[INGESTER:JOB] Failed write parquet file meta: {}, error: {}", new_file_name, diff --git a/src/migration/file_list.rs b/src/migration/file_list.rs index 3746d7f931..48fb560898 100644 --- a/src/migration/file_list.rs +++ b/src/migration/file_list.rs @@ -15,7 +15,7 @@ use config::{ get_config, - meta::stream::{ALL_STREAM_TYPES, FileKey, PartitionTimeLevel}, + meta::stream::{ALL_STREAM_TYPES, PartitionTimeLevel}, utils::time::{BASE_TIME, now_micros}, }; use infra::file_list as infra_file_list; @@ -63,13 +63,7 @@ pub async fn run(from: &str, to: &str) -> Result<(), anyhow::Error> { ) .await .expect("load file_list failed"); - let put_items = files - .into_iter() - .map(|(account, file_key, file_meta)| { - FileKey::new(account, file_key, file_meta, false) - }) - .collect::>(); - dest.batch_add(&put_items) + dest.batch_add(&files) .await .expect("load list_list into db failed"); } diff --git a/src/proto/proto/cluster/common.proto b/src/proto/proto/cluster/common.proto index d388dea26d..400e8b0beb 100644 --- a/src/proto/proto/cluster/common.proto +++ b/src/proto/proto/cluster/common.proto @@ -45,11 +45,12 @@ message FileList { } message FileKey { - string key = 1; - FileMeta meta = 2; - bool deleted = 3; - optional bytes segment_ids = 4; - string account = 5; + int64 id = 1; + string account = 2; + string key = 3; + FileMeta meta = 4; + bool deleted = 5; + optional bytes segment_ids = 6; } message IdxFileName { diff --git a/src/proto/src/generated/cluster.rs b/src/proto/src/generated/cluster.rs index 10bf3bb898..3ffaf5d6b7 100644 --- a/src/proto/src/generated/cluster.rs +++ b/src/proto/src/generated/cluster.rs @@ -86,16 +86,18 @@ pub struct FileList { #[allow(clippy::derive_partial_eq_without_eq)] #[derive(Clone, PartialEq, ::prost::Message)] pub struct FileKey { - #[prost(string, tag = "1")] - pub key: ::prost::alloc::string::String, - #[prost(message, optional, tag = "2")] - pub meta: ::core::option::Option, - #[prost(bool, tag = "3")] - pub deleted: bool, - #[prost(bytes = "vec", optional, tag = "4")] - pub segment_ids: ::core::option::Option<::prost::alloc::vec::Vec>, - #[prost(string, tag = "5")] + #[prost(int64, tag = "1")] + pub id: i64, + #[prost(string, tag = "2")] pub account: ::prost::alloc::string::String, + #[prost(string, tag = "3")] + pub key: ::prost::alloc::string::String, + #[prost(message, optional, tag = "4")] + pub meta: ::core::option::Option, + #[prost(bool, tag = "5")] + pub deleted: bool, + #[prost(bytes = "vec", optional, tag = "6")] + pub segment_ids: ::core::option::Option<::prost::alloc::vec::Vec>, } #[derive(Eq)] #[allow(clippy::derive_partial_eq_without_eq)] diff --git a/src/service/compact/flatten.rs b/src/service/compact/flatten.rs index 0d3bb42874..05d46e31e5 100644 --- a/src/service/compact/flatten.rs +++ b/src/service/compact/flatten.rs @@ -139,17 +139,10 @@ pub async fn generate_by_stream( start.elapsed().as_millis() ); - for (account, key, meta) in files { - if PROCESSING_FILES.read().contains(&key) { + for file in files { + if PROCESSING_FILES.read().contains(&file.key) { continue; } - let file = FileKey { - account, - key, - meta, - deleted: false, - segment_ids: None, - }; // add into queue PROCESSING_FILES.write().insert(file.key.clone()); worker_tx.send(file).await?; diff --git a/src/service/compact/merge.rs b/src/service/compact/merge.rs index a2afc558ba..740913b737 100644 --- a/src/service/compact/merge.rs +++ b/src/service/compact/merge.rs @@ -644,11 +644,9 @@ pub async fn merge_by_stream( for file in delete_file_list { events.push(FileKey { - account: file.account.clone(), - key: file.key.clone(), - meta: file.meta.clone(), deleted: true, segment_ids: None, + ..file.clone() }); } events.sort_by(|a, b| a.key.cmp(&b.key)); @@ -999,7 +997,7 @@ pub async fn merge_files( ) .await?; } - new_files.push(FileKey::new(account, new_file_key, new_file_meta, false)); + new_files.push(FileKey::new(0, account, new_file_key, new_file_meta, false)); } MergeParquetResult::Multiple { bufs, file_metas } => { for (buf, file_meta) in bufs.into_iter().zip(file_metas.into_iter()) { @@ -1041,7 +1039,7 @@ pub async fn merge_files( .await?; } - new_files.push(FileKey::new(account, new_file_key, new_file_meta, false)); + new_files.push(FileKey::new(0, account, new_file_key, new_file_meta, false)); } log::info!( "[COMPACTOR:WORKER:{thread_id}] merged {} files into a new file: {:?}, original_size: {}, compressed_size: {}, took: {} ms", @@ -1114,6 +1112,7 @@ async fn generate_inverted_index( if let Err(e) = write_file_list( org_id, &[FileKey { + id: 0, account, key: file_name.clone(), meta: filemeta, @@ -1203,7 +1202,15 @@ async fn write_file_list(org_id: &str, events: &[FileKey]) -> Result<(), anyhow: if success { // send broadcast to other nodes if get_config().cache_latest_files.enabled { - if let Err(e) = db::file_list::broadcast::send(events, None).await { + // get id for all the new files + let file_ids = infra_file_list::query_ids_by_files(events).await?; + let mut events = events.to_vec(); + for event in events.iter_mut() { + if let Some(id) = file_ids.get(&event.key) { + event.id = *id; + } + } + if let Err(e) = db::file_list::broadcast::send(&events).await { log::error!("[COMPACTOR] send broadcast for file_list failed: {}", e); } } diff --git a/src/service/db/file_list/broadcast.rs b/src/service/db/file_list/broadcast.rs index c02176771c..8480716c0d 100644 --- a/src/service/db/file_list/broadcast.rs +++ b/src/service/db/file_list/broadcast.rs @@ -41,25 +41,16 @@ static EVENTS: Lazy>> = type EventChannel = Arc>>; /// send an event to broadcast, will create a new channel for each nodes -pub async fn send(items: &[FileKey], node_uuid: Option) -> Result<(), anyhow::Error> { +pub async fn send(items: &[FileKey]) -> Result<(), anyhow::Error> { let cfg = get_config(); if cfg.common.local_mode || items.is_empty() { return Ok(()); } - let nodes = if let Some(node_uuid) = node_uuid { - cluster::get_node_by_uuid(&node_uuid) - .await - .map(|node| vec![node]) - .unwrap_or_default() - } else { - cluster::get_cached_nodes(|node| { - node.scheduled - && (node.is_querier()) - && (node.status == NodeStatus::Prepare || node.status == NodeStatus::Online) - }) - .await - .unwrap_or_default() - }; + let nodes = cluster::get_cached_nodes(|node| { + node.scheduled && node.is_querier() && node.status == NodeStatus::Online + }) + .await + .unwrap_or_default(); let mut events = EVENTS.write().await; for node in nodes { if node.uuid.eq(&LOCAL_NODE.uuid) { @@ -74,7 +65,7 @@ pub async fn send(items: &[FileKey], node_uuid: Option) -> Result<(), an } // check if the item is for interactive node if let Some(node_name) = cluster::get_node_from_consistent_hash( - &item.key, + &item.id.to_string(), &Role::Querier, Some(RoleGroup::Interactive), ) @@ -87,7 +78,7 @@ pub async fn send(items: &[FileKey], node_uuid: Option) -> Result<(), an } // check if the item is for background node if let Some(node_name) = cluster::get_node_from_consistent_hash( - &item.key, + &item.id.to_string(), &Role::Querier, Some(RoleGroup::Background), ) diff --git a/src/service/db/file_list/mod.rs b/src/service/db/file_list/mod.rs index 008c5975c4..88777068b9 100644 --- a/src/service/db/file_list/mod.rs +++ b/src/service/db/file_list/mod.rs @@ -46,7 +46,8 @@ pub static BLOCKED_ORGS: Lazy> = Lazy::new(|| { }); pub async fn set(account: &str, key: &str, meta: Option, deleted: bool) -> Result<()> { - let file_data = FileKey::new( + let mut file_data = FileKey::new( + 0, account.to_string(), key.to_string(), meta.clone().unwrap_or_default(), @@ -56,11 +57,15 @@ pub async fn set(account: &str, key: &str, meta: Option, deleted: bool // write into file_list storage // retry 5 times for _ in 0..5 { - if let Err(e) = progress(account, key, meta.as_ref(), deleted).await { - log::error!("[FILE_LIST] Error saving file to storage, retrying: {}", e); - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - } else { - break; + match progress(account, key, meta.as_ref(), deleted).await { + Ok(id) => { + file_data.id = id; + break; + } + Err(e) => { + log::error!("[FILE_LIST] Error saving file to storage, retrying: {}", e); + tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; + } } } @@ -75,26 +80,24 @@ pub async fn set(account: &str, key: &str, meta: Option, deleted: bool Ok(()) } -async fn progress(account: &str, key: &str, data: Option<&FileMeta>, delete: bool) -> Result<()> { +async fn progress(account: &str, key: &str, data: Option<&FileMeta>, delete: bool) -> Result { + let mut id = 0; if delete { if let Err(e) = infra::file_list::remove(key).await { - log::error!( - "service:db:file_list: delete {}, set_file_to_cache error: {}", - key, - e - ); + log::error!("service:db:file_list: delete {}, remove error: {}", key, e); } - } else { - if let Err(e) = infra::file_list::add(account, key, data.unwrap()).await { - log::error!( - "service:db:file_list: add {}, set_file_to_cache error: {}", - key, - e - ); + } else if let Some(data) = data { + match infra::file_list::add(account, key, data).await { + Ok(v) => { + id = v; + } + Err(e) => { + log::error!("service:db:file_list: add {}, add error: {}", key, e); + } } // update stream stats realtime if config::get_config().common.local_mode { - if let Err(e) = infra::cache::stats::incr_stream_stats(key, data.unwrap()) { + if let Err(e) = infra::cache::stats::incr_stream_stats(key, data) { log::error!( "service:db:file_list: add {}, incr_stream_stats error: {}", key, @@ -104,7 +107,7 @@ async fn progress(account: &str, key: &str, data: Option<&FileMeta>, delete: boo } } - Ok(()) + Ok(id) } pub async fn cache_stats() -> Result<()> { diff --git a/src/service/file_list.rs b/src/service/file_list.rs index 291286689d..49cbf1401f 100644 --- a/src/service/file_list.rs +++ b/src/service/file_list.rs @@ -19,7 +19,7 @@ use config::{ get_config, meta::{ search::ScanStats, - stream::{FileKey, FileMeta, PartitionTimeLevel, StreamType}, + stream::{FileKey, PartitionTimeLevel, StreamType}, }, metrics::{FILE_LIST_CACHE_HIT_COUNT, FILE_LIST_ID_SELECT_COUNT}, utils::file::get_file_meta as util_get_file_meta, @@ -48,7 +48,7 @@ pub async fn query( time_max: i64, ) -> Result> { let cfg = get_config(); - let files = file_list::query( + let mut files = file_list::query( org_id, stream_type, stream_name, @@ -73,8 +73,8 @@ pub async fn query( .collect::>(); let missing_files: usize = files .iter() - .map(|(_account, name, _)| { - if dumped_file_names.contains(name) { + .map(|f| { + if dumped_file_names.contains(&f.key) { 0 } else { 1 @@ -87,42 +87,19 @@ pub async fn query( ); } } - let mut file_keys = Vec::with_capacity(files.len()); - for (account, key, meta) in files { - file_keys.push(FileKey { - account, - key, - meta, - deleted: false, - segment_ids: None, - }); - } + if !cfg.common.file_list_dump_dual_write { // we only consider these files in case of dual write disabled, // because with dual write there are some edge cases which cannot be sovled even with id // de-dup so the data gets counted twice. - for file in dumped_files { - file_keys.push(FileKey { - account: file.account, - key: "files/".to_string() + &file.stream + "/" + &file.date + "/" + &file.file, - meta: FileMeta { - min_ts: file.min_ts, - max_ts: file.max_ts, - records: file.records, - original_size: file.original_size, - compressed_size: file.compressed_size, - index_size: file.index_size, - flattened: file.flattened, - }, - deleted: file.deleted, - segment_ids: None, - }) + for file in dumped_files.iter() { + files.push(file.into()) } } - file_keys.par_sort_unstable_by(|a, b| a.key.cmp(&b.key)); - file_keys.dedup_by(|a, b| a.key == b.key); - Ok(file_keys) + files.par_sort_unstable_by(|a, b| a.key.cmp(&b.key)); + files.dedup_by(|a, b| a.key == b.key); + Ok(files) } /// NOTE: This will not query the files from file_dump. If you also want files from the dump, use @@ -148,20 +125,10 @@ pub async fn query_for_merge( Some((date_start.to_string(), date_end.to_string())), ) .await?; - let mut file_keys = Vec::with_capacity(files.len()); - for (account, key, meta) in files { - file_keys.push(FileKey { - account, - key, - meta, - deleted: false, - segment_ids: None, - }); - } // we don't need to query from dump here, because // we expected all dumped files to be already compacted, // and the compaction marks old files as deleted, which is not possible with dump - Ok(file_keys) + Ok(files) } #[tracing::instrument(name = "service::file_list::query_by_ids", skip_all)] @@ -191,10 +158,7 @@ pub async fn query_by_ids( Vec::new() } }; - let cached_ids = cached_files - .iter() - .map(|(id, ..)| *id) - .collect::>(); + let cached_ids = cached_files.iter().map(|f| f.id).collect::>(); log::info!( "{}", search_inspector_fields( @@ -221,18 +185,8 @@ pub async fn query_by_ids( .with_label_values(&[]) .set(cached_ids.len() as i64); - let mut file_keys = Vec::with_capacity(ids.len()); - for (_, account, key, meta) in cached_files { - file_keys.push(FileKey { - account, - key, - meta, - deleted: false, - segment_ids: None, - }); - } ( - file_keys, + cached_files, ids_set.difference(&cached_ids).cloned().collect::>(), ) } else { @@ -242,21 +196,6 @@ pub async fn query_by_ids( // 2. query from remote db let start = std::time::Instant::now(); let db_files = file_list::query_by_ids(&ids).await?; - let db_files = db_files - .into_iter() - .map(|(id, account, key, meta)| { - ( - id, - FileKey { - account, - key, - meta, - deleted: false, - segment_ids: None, - }, - ) - }) - .collect::>(); log::info!( "{}", search_inspector_fields( @@ -288,37 +227,15 @@ pub async fn query_by_ids( ) .await?; - let dumped_files: Vec<_> = dumped_files - .into_iter() + let dumped_files: Vec = dumped_files + .iter() .filter(|r| ids.contains(&r.id)) - .map(|r| { - ( - r.id, - FileKey { - account: r.account, - key: "files/".to_string() + &r.stream + "/" + &r.date + "/" + &r.file, - meta: FileMeta { - min_ts: r.min_ts, - max_ts: r.max_ts, - records: r.records, - original_size: r.original_size, - compressed_size: r.compressed_size, - index_size: r.index_size, - flattened: r.flattened, - }, - deleted: r.deleted, - segment_ids: None, - }, - ) - }) + .map(|r| r.into()) .collect(); if cfg.common.file_list_dump_enabled && cfg.common.file_list_dump_debug_check { - let dump_ids = dumped_files - .iter() - .map(|entry| entry.0) - .collect::>(); - let db_ids = db_files.iter().map(|entry| entry.0).collect::>(); + let dump_ids = dumped_files.iter().map(|f| f.id).collect::>(); + let db_ids = db_files.iter().map(|f| f.id).collect::>(); let missing_files: usize = db_ids .iter() .map(|id| if dump_ids.contains(id) { 0 } else { 1 }) @@ -333,7 +250,6 @@ pub async fn query_by_ids( // 3. set the local cache if !cfg.common.local_mode { let start = std::time::Instant::now(); - let db_files: Vec<_> = db_files.iter().map(|(id, f)| (*id, f)).collect(); if let Err(e) = file_list::LOCAL_CACHE.batch_add_with_id(&db_files).await { log::error!( "[trace_id {trace_id}] file_list set cache failed for db files: {:?}", @@ -342,7 +258,6 @@ pub async fn query_by_ids( } if !cfg.common.file_list_dump_dual_write { - let dumped_files: Vec<_> = dumped_files.iter().map(|(id, f)| (*id, f)).collect(); if let Err(e) = file_list::LOCAL_CACHE .batch_add_with_id(&dumped_files) .await @@ -374,9 +289,9 @@ pub async fn query_by_ids( } // 4. merge the results - files.extend(db_files.into_iter().map(|(_, f)| f)); + files.extend(db_files); if !cfg.common.file_list_dump_dual_write { - files.extend(dumped_files.into_iter().map(|(_, f)| f)); + files.extend(dumped_files); } files.par_sort_unstable_by(|a, b| a.key.cmp(&b.key)); files.dedup_by(|a, b| a.key == b.key); @@ -421,7 +336,7 @@ pub async fn query_ids( } } if !cfg.common.file_list_dump_dual_write { - files.extend(dumped_files.into_iter()); + files.extend(dumped_files); } files.par_sort_unstable_by(|a, b| a.id.cmp(&b.id)); files.dedup_by(|a, b| a.id == b.id); @@ -458,6 +373,7 @@ pub fn calculate_local_files_size(files: &[String]) -> Result { pub async fn delete_parquet_file(account: &str, key: &str, file_list_only: bool) -> Result<()> { // delete from file list in metastore file_list::batch_process(&[FileKey::new( + 0, account.to_string(), key.to_string(), Default::default(), diff --git a/src/service/file_list_dump.rs b/src/service/file_list_dump.rs index b6d9b3d22a..3a4956294d 100644 --- a/src/service/file_list_dump.rs +++ b/src/service/file_list_dump.rs @@ -230,16 +230,7 @@ pub async fn query( let db_time = db_start.elapsed().as_millis(); let process_start = std::time::Instant::now(); - let dump_files: Vec<_> = dump_files - .into_iter() - .map(|f| FileKey { - account: f.account.to_string(), - key: format!("files/{}/{}/{}", stream_key, f.date, f.file), - meta: (&f).into(), - deleted: false, - segment_ids: None, - }) - .collect(); + let dump_files: Vec<_> = dump_files.iter().map(|f| f.into()).collect(); let max_ts_upper_bound = calculate_max_ts_upper_bound(range.1, stream_type); let stream_key = format!("{org}/{stream_type}/{stream}"); @@ -283,16 +274,7 @@ pub async fn get_ids_in_range( let db_time = db_start.elapsed().as_millis(); let process_start = std::time::Instant::now(); - let dump_files: Vec<_> = dump_files - .into_iter() - .map(|f| FileKey { - account: f.account.to_string(), - key: format!("files/{}/{}/{}", stream_key, f.date, f.file), - meta: (&f).into(), - deleted: false, - segment_ids: None, - }) - .collect(); + let dump_files: Vec<_> = dump_files.iter().map(|f| f.into()).collect(); let max_ts_upper_bound = calculate_max_ts_upper_bound(range.1, stream_type); let stream_key = format!("{org}/{stream_type}/{stream}"); @@ -367,12 +349,11 @@ async fn move_and_delete( let items: Vec<_> = list .iter() .chain(dump_files.iter()) - .map(|f| FileKey { - account: f.account.to_string(), - key: format!("files/{}/{}/{}", f.stream, f.date, f.file), - meta: f.into(), - deleted: true, - segment_ids: None, + .map(|f| { + let mut f = FileKey::from(f); + f.deleted = true; + f.segment_ids = None; + f }) .collect(); if let Err(e) = infra::file_list::batch_process(&items).await { @@ -461,16 +442,7 @@ pub async fn stats( return Ok(vec![]); } - let dump_files: Vec<_> = dump_files - .into_iter() - .map(|f| FileKey { - account: f.account.to_string(), - key: format!("files/{}/{}/{}", f.stream, f.date, f.file), - meta: (&f).into(), - deleted: false, - segment_ids: None, - }) - .collect(); + let dump_files: Vec<_> = dump_files.iter().map(|f| f.into()).collect(); let sql = format!( r#"