From b72e85e0fe080fa5cce5ecca8ed8a74923ce7aff Mon Sep 17 00:00:00 2001 From: Hengfei Yang Date: Fri, 12 Sep 2025 11:56:19 +0800 Subject: [PATCH] feat: support nats ttl marker (#8377) --- Cargo.lock | 4 +- Cargo.toml | 2 +- src/common/infra/wal.rs | 418 --------------------------------------- src/config/src/config.rs | 8 +- src/infra/src/db/nats.rs | 6 +- src/main.rs | 1 - 6 files changed, 14 insertions(+), 425 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 98b28e7f41..afe3cc977e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -939,9 +939,9 @@ dependencies = [ [[package]] name = "async-nats" -version = "0.39.0" +version = "0.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a798aab0c0203b31d67d501e5ed1f3ac6c36a329899ce47fc93c3bea53f3ae89" +checksum = "08f6da6d49a956424ca4e28fe93656f790d748b469eaccbc7488fec545315180" dependencies = [ "base64 0.22.1", "bytes", diff --git a/Cargo.toml b/Cargo.toml index f3da1521b0..e752e35011 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -235,7 +235,7 @@ actix-tls = { version = "3.4", features = [ anyhow = "1.0" arc-swap = "1.7.1" argon2 = { version = "0.5", features = ["alloc", "password-hash"] } -async-nats = "0.39" +async-nats = "0.42" async-stream = "0.3.6" async-trait = "0.1" async-recursion = "1.0" diff --git a/src/common/infra/wal.rs b/src/common/infra/wal.rs index 15b4da6d1f..caa576aa28 100644 --- a/src/common/infra/wal.rs +++ b/src/common/infra/wal.rs @@ -13,25 +13,8 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::{path::Path, sync::Arc}; - -use chrono::{DateTime, Datelike, TimeZone, Utc}; -use config::{ - FILE_EXT_JSON, get_config, ider, - meta::stream::{PartitionTimeLevel, StreamParams, StreamType}, - metrics, - utils::async_file::get_file_contents, -}; use hashbrown::HashMap; use once_cell::sync::Lazy; -use tokio::{ - fs::{File, OpenOptions, create_dir_all}, - io::AsyncWriteExt, - sync::RwLock, -}; - -// MANAGER for manage using WAL files, in use, should not move to s3 -static MANAGER: Lazy = Lazy::new(Manager::new); // SEARCHING_FILES for searching files, in use, should not move to s3 static SEARCHING_FILES: Lazy> = @@ -41,8 +24,6 @@ static SEARCHING_FILES: Lazy> = static SEARCHING_REQUESTS: Lazy>>> = Lazy::new(Default::default); -type RwData = RwLock>>; - struct SearchingFileLocker { inner: HashMap, } @@ -86,295 +67,11 @@ impl SearchingFileLocker { } } -struct Manager { - data: Arc>, -} - -pub struct RwFile { - file: Option>, - org_id: String, - stream_name: String, - stream_type: StreamType, - dir: String, - name: String, - expired: i64, -} - pub fn init() -> Result<(), anyhow::Error> { - _ = MANAGER.data.len(); _ = SEARCHING_FILES.read().len(); Ok(()) } -pub async fn get_or_create( - thread_id: usize, - stream: StreamParams, - partition_time_level: Option, - key: &str, -) -> Arc { - MANAGER - .get_or_create(thread_id, stream, partition_time_level, key) - .await -} - -pub async fn check_in_use(stream: StreamParams, file_name: &str) -> bool { - MANAGER.check_in_use(stream, file_name).await -} - -pub async fn flush_all_to_disk() { - for data in MANAGER.data.iter() { - for (_, file) in data.read().await.iter() { - file.sync().await; - } - } -} - -impl Default for Manager { - fn default() -> Self { - Self::new() - } -} - -impl Manager { - pub fn new() -> Manager { - let size = get_config().limit.cpu_num; - let mut data = Vec::with_capacity(size); - for _ in 0..size { - data.push(RwLock::new(HashMap::>::default())); - } - Self { - data: Arc::new(data), - } - } - - pub async fn get( - &self, - thread_id: usize, - stream: StreamParams, - key: &str, - ) -> Option> { - let full_key = format!( - "{}/{}/{}/{key}", - stream.org_id, stream.stream_type, stream.stream_name - ); - let locker = self.data.get(thread_id)?; - let manager = locker.read().await; - let file = manager.get(&full_key)?.clone(); - - drop(manager); - - // check size & ttl - if file.size().await >= (get_config().limit.max_file_size_on_disk as i64) - || file.expired() <= Utc::now().timestamp() - { - let mut manager = locker.write().await; - manager.remove(&full_key); - manager.shrink_to_fit(); - file.sync().await; - return None; - } - - Some(file) - } - - pub async fn create( - &self, - thread_id: usize, - stream: StreamParams, - partition_time_level: Option, - key: &str, - ) -> Arc { - let stream_type = stream.stream_type; - let full_key = format!( - "{}/{}/{}/{key}", - stream.org_id, stream.stream_type, stream.stream_name - ); - let mut data = self.data.get(thread_id).unwrap().write().await; - if let Some(f) = data.get(&full_key) { - return f.clone(); - } - let file = Arc::new(RwFile::new(thread_id, stream, partition_time_level, key).await); - if !stream_type.eq(&StreamType::EnrichmentTables) { - data.insert(full_key, file.clone()); - }; - file - } - - pub async fn get_or_create( - &self, - thread_id: usize, - stream: StreamParams, - partition_time_level: Option, - key: &str, - ) -> Arc { - if let Some(file) = self.get(thread_id, stream.clone(), key).await { - file - } else { - self.create(thread_id, stream, partition_time_level, key) - .await - } - } - - pub async fn check_in_use(&self, stream: StreamParams, file_name: &str) -> bool { - let mut columns = file_name.split('/'); - let thread_id: usize = columns - .next() - .unwrap() - .parse() - .unwrap_or_else(|_| panic!("need a thread id, but the file is: {file_name}")); - // Remove the last element which is the file name - // and join the rest to form the key - columns.next_back(); - let key = itertools::Itertools::join(&mut columns, "/"); - self.get(thread_id, stream, &key) - .await - .is_some_and(|file| file.name() == file_name) - } -} - -impl RwFile { - async fn new( - thread_id: usize, - stream: StreamParams, - partition_time_level: Option, - key: &str, - ) -> RwFile { - let cfg = get_config(); - let mut dir_path = format!( - "{}files/{}/{}/{}/", - &cfg.common.data_wal_dir, stream.org_id, stream.stream_type, stream.stream_name - ); - // Hack for file_list - let file_list_prefix = "/files//file_list//"; - if dir_path.contains(file_list_prefix) { - dir_path = dir_path.replace(file_list_prefix, "/file_list/"); - } - let id = ider::generate(); - let file_name = format!("{thread_id}/{key}/{id}{FILE_EXT_JSON}"); - let file_path = format!("{dir_path}{file_name}"); - create_dir_all(Path::new(&file_path).parent().unwrap()) - .await - .unwrap(); - - let f = OpenOptions::new() - .create(true) - .append(true) - .open(&file_path) - .await - .unwrap_or_else(|e| panic!("open wal file [{file_path}] error: {e}")); - let file = Some(RwLock::new(f)); - - let time_now: DateTime = Utc::now(); - let level_duration = partition_time_level.unwrap_or_default().duration(); - let ttl = if !cfg.limit.ignore_file_retention_by_stream && level_duration > 0 { - let time_end_day = Utc - .with_ymd_and_hms( - time_now.year(), - time_now.month(), - time_now.day(), - 23, - 59, - 59, - ) - .unwrap() - .timestamp(); - let expired = time_now.timestamp() + level_duration; - if expired > time_end_day { - // if the file expired time is tomorrow, it should be deleted at 23:59:59 + - // 10min - time_end_day + cfg.limit.max_file_retention_time as i64 - } else { - expired - } - } else { - time_now.timestamp() + cfg.limit.max_file_retention_time as i64 - }; - - RwFile { - file, - org_id: stream.org_id.to_string(), - stream_name: stream.stream_name.to_string(), - stream_type: stream.stream_type, - dir: dir_path, - name: file_name, - expired: ttl, - } - } - - #[inline] - pub async fn write(&self, data: &[u8]) { - // metrics - metrics::INGEST_WAL_USED_BYTES - .with_label_values(&[&self.org_id, self.stream_type.as_str()]) - .add(data.len() as i64); - metrics::INGEST_WAL_WRITE_BYTES - .with_label_values(&[&self.org_id, self.stream_type.as_str()]) - .inc_by(data.len() as u64); - - self.file - .as_ref() - .unwrap() - .write() - .await - .write_all(data) - .await - .unwrap(); - } - - #[inline] - pub async fn read(&self) -> Result, std::io::Error> { - get_file_contents(&self.full_name(), None).await - } - - #[inline] - pub async fn sync(&self) { - self.file - .as_ref() - .unwrap() - .write() - .await - .sync_all() - .await - .unwrap() - } - - #[inline] - pub async fn size(&self) -> i64 { - self.file - .as_ref() - .unwrap() - .read() - .await - .metadata() - .await - .unwrap() - .len() as i64 - } - - #[inline] - pub fn name(&self) -> &str { - &self.name - } - - #[inline] - pub fn wal_name(&self) -> String { - format!( - "files/{}/{}/{}/{}", - self.org_id, self.stream_type, self.stream_name, self.name - ) - } - - #[inline] - pub fn full_name(&self) -> String { - format!("{}{}", self.dir, self.name) - } - - #[inline] - pub fn expired(&self) -> i64 { - self.expired - } -} - pub fn lock_files(files: &[String]) { let mut locker = SEARCHING_FILES.write(); for file in files.iter() { @@ -424,38 +121,6 @@ mod tests { use super::*; - #[tokio::test] - async fn test_wal_manager() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key1"; - let file = get_or_create(thread_id, stream, None, key).await; - let data = "test_data".to_string().into_bytes(); - file.write(&data).await; - assert_eq!(file.read().await.unwrap(), data); - assert_eq!(file.size().await, data.len() as i64); - assert!(file.name().contains(&format!("{thread_id}/{key}"))); - } - - #[tokio::test] - async fn test_wal_rw_file() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key2"; - let file = RwFile::new(thread_id, stream, None, key).await; - let data = "test_data".to_string().into_bytes(); - file.write(&data).await; - assert_eq!(file.read().await.unwrap(), data); - assert_eq!(file.size().await, data.len() as i64); - assert!(file.name().contains(&format!("{thread_id}/{key}"))); - } - #[tokio::test] async fn test_wal_file_locking() { let files = vec![ @@ -493,87 +158,4 @@ mod tests { assert!(!lock_files_exists(&files[0])); assert!(!lock_files_exists(&files[1])); } - - #[tokio::test] - async fn test_wal_file_expiration() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key5"; - - // Create a file with short retention time - let file = get_or_create(thread_id, stream.clone(), None, key).await; - let initial_expired = file.expired(); - - // Wait a bit and check if file is still accessible - tokio::time::sleep(std::time::Duration::from_secs(1)).await; - let file_after = get_or_create(thread_id, stream.clone(), None, key).await; - assert_eq!(file_after.expired(), initial_expired); - } - - #[tokio::test] - async fn test_wal_manager_operations() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key6"; - - // Test get_or_create - let file1 = get_or_create(thread_id, stream.clone(), None, key).await; - let file2 = get_or_create(thread_id, stream.clone(), None, key).await; - assert_eq!(file1.name(), file2.name()); - - // Test check_in_use - assert!(check_in_use(stream.clone(), file1.name()).await); - assert!(!check_in_use(stream.clone(), "1/2025/06/06/01/md5/test_key.json").await); - } - - #[tokio::test] - async fn test_wal_file_size_limits() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key7"; - - let file = get_or_create(thread_id, stream.clone(), None, key).await; - - // Write some data - let data = "test_data".to_string().into_bytes(); - file.write(&data).await; - file.sync().await; - - // Check size - assert_eq!(file.size().await, data.len() as i64); - - // Sync to disk - file.sync().await; - } - - #[tokio::test] - async fn test_wal_file_operations() { - let thread_id = 1; - let org_id = "test_org"; - let stream_name = "test_stream"; - let stream_type = StreamType::Logs; - let stream = StreamParams::new(org_id, stream_name, stream_type); - let key = "test_key8"; - - let file = RwFile::new(thread_id, stream, None, key).await; - - // Test write and read - let data = "test_data".to_string().into_bytes(); - file.write(&data).await; - assert_eq!(file.read().await.unwrap(), data); - - // Test file names - assert!(file.name().contains(&format!("{thread_id}/{key}"))); - assert!(file.wal_name().contains(&format!("{thread_id}/{key}"))); - assert!(file.full_name().contains(&format!("{thread_id}/{key}"))); - } } diff --git a/src/config/src/config.rs b/src/config/src/config.rs index 9361df85a2..59211967b1 100644 --- a/src/config/src/config.rs +++ b/src/config/src/config.rs @@ -1226,8 +1226,6 @@ pub struct Limit { pub ingest_allowed_in_future: i64, #[env_config(name = "ZO_INGEST_FLATTEN_LEVEL", default = 3)] // default flatten level pub ingest_flatten_level: u32, - #[env_config(name = "ZO_IGNORE_FILE_RETENTION_BY_STREAM", default = false)] - pub ignore_file_retention_by_stream: bool, #[env_config(name = "ZO_LOGS_FILE_RETENTION", default = "hourly")] pub logs_file_retention: String, #[env_config(name = "ZO_TRACES_FILE_RETENTION", default = "hourly")] @@ -1733,6 +1731,12 @@ pub struct Nats { default = 2048 )] pub queue_max_size: i64, + #[env_config( + name = "ZO_NATS_V211_SUPPORT", + help = "Support NATS v2.11.x", + default = false + )] + pub v211_support: bool, } #[derive(Debug, Default, EnvConfig)] diff --git a/src/infra/src/db/nats.rs b/src/infra/src/db/nats.rs index 6e6daea489..69a4432921 100644 --- a/src/infra/src/db/nats.rs +++ b/src/infra/src/db/nats.rs @@ -86,7 +86,11 @@ async fn get_bucket_by_key<'a>( if bucket_name == "nodes" || bucket_name == "clusters" { // if changed ttl need recreate the bucket // CMD: nats kv del -f o2_nodes - bucket.max_age = Duration::from_secs(cfg.limit.node_heartbeat_ttl as u64); + let ttl = Duration::from_secs(cfg.limit.node_heartbeat_ttl as u64); + bucket.max_age = ttl; + if cfg.nats.v211_support { + bucket.limit_markers = Some(ttl); + } } let kv = jetstream.create_key_value(bucket).await.map_err(|e| { Error::Message(format!( diff --git a/src/main.rs b/src/main.rs index b03ee5b7de..c5e2a22d1b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -299,7 +299,6 @@ async fn main() -> Result<(), anyhow::Error> { _ = metadata::close().await; // flush WAL cache to disk _ = ingester::flush_all().await; - common_infra::wal::flush_all_to_disk().await; // flush compact offset cache to disk disk _ = db::compact::files::sync_cache_to_db().await; // flush db