From d16b813ca7c9239e12fcfd818d1cd0b092cffdd3 Mon Sep 17 00:00:00 2001 From: Siddharth <111498619+siddharthteli12@users.noreply.github.com> Date: Thu, 29 May 2025 06:53:57 +0530 Subject: [PATCH] refactor: Writer new method & Reader test (#6917) --- src/ingester/src/writer.rs | 19 ++++----- src/wal/benches/write.rs | 8 ++-- src/wal/src/errors.rs | 4 ++ src/wal/src/writer.rs | 38 +++++------------ src/wal/tests/wal.rs | 83 +++++++++++++++++++++++--------------- 5 files changed, 77 insertions(+), 75 deletions(-) diff --git a/src/ingester/src/writer.rs b/src/ingester/src/writer.rs index 014306bdf4..57c16f6c44 100644 --- a/src/ingester/src/writer.rs +++ b/src/ingester/src/writer.rs @@ -31,7 +31,7 @@ use hashbrown::HashSet; use once_cell::sync::Lazy; use snafu::ResultExt; use tokio::sync::{RwLock, mpsc}; -use wal::Writer as WalWriter; +use wal::{Writer as WalWriter, build_file_path}; use crate::{ ReadRecordBatchEntry, WriterSignal, @@ -213,12 +213,10 @@ impl Writer { key: key.clone(), wal: Arc::new(RwLock::new( WalWriter::new( - wal_dir, - &key.org_id, - &key.stream_type, - wal_id.to_string(), + build_file_path(wal_dir, &key.org_id, &key.stream_type, wal_id.to_string()), cfg.limit.max_file_size_on_disk as u64, cfg.limit.wal_write_buffer_size, + None, ) .expect("wal file create error"), )), @@ -416,12 +414,15 @@ impl Writer { wal_id ); let new_wal = WalWriter::new( - wal_dir, - &self.key.org_id, - &self.key.stream_type, - wal_id.to_string(), + build_file_path( + wal_dir, + &self.key.org_id, + &self.key.stream_type, + wal_id.to_string(), + ), cfg.limit.max_file_size_on_disk as u64, cfg.limit.wal_write_buffer_size, + None, ) .context(WalSnafu)?; wal.sync().context(WalSnafu)?; // sync wal before rotation diff --git a/src/wal/benches/write.rs b/src/wal/benches/write.rs index 2b45309df4..2b52d1a741 100644 --- a/src/wal/benches/write.rs +++ b/src/wal/benches/write.rs @@ -15,7 +15,7 @@ use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; use tempfile::tempdir; -use wal::Writer; +use wal::{Writer, build_file_path}; pub fn write_benchmark(c: &mut Criterion) { let dir = tempdir().unwrap(); @@ -24,12 +24,10 @@ pub fn write_benchmark(c: &mut Criterion) { for buf_size in [4096, 8192, 16384, 32768] { for entry_size in [4096, 16384, 32768, 65536] { let mut writer = Writer::new( - dir, - "org", - "stream", - entry_size.to_string(), + build_file_path(dir, "org", "stream", entry_size.to_string()), 1024_1024, buf_size, + None, ) .unwrap(); let data = vec![42u8; entry_size]; diff --git a/src/wal/src/errors.rs b/src/wal/src/errors.rs index c7af0cc038..bb15210a07 100644 --- a/src/wal/src/errors.rs +++ b/src/wal/src/errors.rs @@ -86,4 +86,8 @@ pub enum Error { WriteQueueFull { idx: usize, }, + #[snafu(display("Provided path '{}' has no parent directory", path.display()))] + NoParentDir { + path: PathBuf, + }, } diff --git a/src/wal/src/writer.rs b/src/wal/src/writer.rs index ae73773104..cb2b788e7a 100644 --- a/src/wal/src/writer.rs +++ b/src/wal/src/writer.rs @@ -15,9 +15,9 @@ use std::{ fs::{File, Metadata, OpenOptions, create_dir_all, remove_file}, - io, - io::{BufWriter, Seek, SeekFrom, Write}, - path::PathBuf, + io::{self, BufWriter, Seek, SeekFrom, Write}, + ops::Deref, + path::{Path, PathBuf}, }; use byteorder::{BigEndian, WriteBytesExt}; @@ -37,35 +37,15 @@ pub struct Writer { impl Writer { pub fn new( - root_dir: impl Into, - org_id: &str, - stream_type: &str, - id: String, - init_size: u64, - buffer_size: usize, - ) -> Result { - Self::build( - root_dir, - org_id, - stream_type, - id, - init_size, - buffer_size, - None, - ) - } - - pub fn build( - root_dir: impl Into, - org_id: &str, - stream_type: &str, - id: String, + path: impl Into + Clone + Deref + AsRef, init_size: u64, buffer_size: usize, header: Option, ) -> Result { - let path = super::build_file_path(root_dir, org_id, stream_type, id); - create_dir_all(path.parent().unwrap()).context(FileOpenSnafu { path: path.clone() })?; + create_dir_all(path.parent().ok_or_else(|| Error::NoParentDir { + path: path.clone().into(), + })?) + .context(FileOpenSnafu { path: path.clone() })?; let mut f = OpenOptions::new() .write(true) .create(true) @@ -109,7 +89,7 @@ impl Writer { } Ok(Self { - path, + path: path.into(), f: BufWriter::with_capacity(buffer_size, f), bytes_written, uncompressed_bytes_written: bytes_written, diff --git a/src/wal/tests/wal.rs b/src/wal/tests/wal.rs index 03d14e5818..3af86734d5 100644 --- a/src/wal/tests/wal.rs +++ b/src/wal/tests/wal.rs @@ -19,22 +19,21 @@ use std::{ }; use tempfile::tempdir; -use wal::{Reader, Writer, build_file_path}; +use wal::{ReadFrom, Reader, Writer, build_file_path}; #[test] fn test_wal_new() { let entry_num = 100; let dir = tempdir().unwrap(); let dir = dir.path(); - let mut writer = - Writer::new(dir, "org", "stream", "1".to_string(), 1024_1024, 8 * 1024).unwrap(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path.clone(), 1024_1024, 8 * 1024, None).unwrap(); for i in 0..entry_num { let data = format!("hello world {}", i); writer.write(data.as_bytes()).unwrap(); } writer.close().unwrap(); - let path = build_file_path(dir, "org", "stream", "1".to_string()); let mut reader = Reader::from_path(path).unwrap(); for i in 0..entry_num { let data = format!("hello world {}", i); @@ -51,24 +50,14 @@ fn test_wal_build() { let mut header = wal::FileHeader::new(); header.insert("key1".into(), "value1".into()); header.insert("key2".into(), "value2".into()); - - let mut writer = Writer::build( - dir, - "org", - "stream", - "1".to_string(), - 1024_1024, - 8 * 1024, - Some(header), - ) - .unwrap(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path.clone(), 1024_1024, 8 * 1024, Some(header)).unwrap(); for i in 0..entry_num { let data = format!("hello world {}", i); writer.write(data.as_bytes()).unwrap(); } writer.close().unwrap(); - let path = build_file_path(dir, "org", "stream", "1".to_string()); let mut reader = Reader::from_path(path).unwrap(); let header = reader.header(); assert_eq!(header.get("key1"), Some(&"value1".to_string())); @@ -87,15 +76,14 @@ fn test_position() { let entry_num = 100; let dir = tempdir().unwrap(); let dir = dir.path(); - let mut writer = - Writer::new(dir, "org", "stream", "1".to_string(), 1024_1024, 8 * 1024).unwrap(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path.clone(), 1024_1024, 8 * 1024, None).unwrap(); for i in 0..entry_num { let data = format!("hello world {}", i); writer.write(data.as_bytes()).unwrap(); } writer.close().unwrap(); - let path = build_file_path(dir, "org", "stream", "1".to_string()); let mut reader = Reader::from_path(path).unwrap(); let mut pos = wal::FILE_TYPE_IDENTIFIER_LEN as u64 + wal::WAL_FILE_HEADER_LEN as u64; for _ in 0..entry_num { @@ -111,8 +99,8 @@ fn test_metadata() { let entry_num = 5; let dir = tempdir().unwrap(); let dir = dir.path(); - let mut writer = - Writer::new(dir, "org", "stream", "1".to_string(), 1024_1024, 8 * 1024).unwrap(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path, 1024_1024, 8 * 1024, None).unwrap(); for i in 0..entry_num { let data = format!("hello world {}", i); writer.write(data.as_bytes()).unwrap(); @@ -138,17 +126,8 @@ fn test_realtime_write_and_read() { let mut header = wal::FileHeader::new(); header.insert("key1".into(), "value1".into()); header.insert("key2".into(), "value2".into()); - - let mut writer = Writer::build( - dir, - "org", - "stream", - "1".to_string(), - 0, - 8 * 1024, - Some(header), - ) - .unwrap(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path, 0, 8 * 1024, Some(header)).unwrap(); std::thread::spawn(move || { for i in 0..entry_num { @@ -197,3 +176,43 @@ fn test_realtime_write_and_read() { } } } + +#[test] +fn test_reader_from_path_position_and_metadata() { + let entry_num = 5; + let dir = tempdir().unwrap(); + let dir = dir.path(); + let path = build_file_path(dir, "org", "stream", "1".to_string()); + let mut writer = Writer::new(path.clone(), 1024_1024, 8 * 1024, None).unwrap(); + + for i in 0..entry_num { + let data = format!("hello {}", i); + writer.write(data.as_bytes()).unwrap(); + } + writer.close().unwrap(); + + let fs_metadata = std::fs::metadata(&path).unwrap(); + + let mut reader = Reader::from_path_position(path.clone(), ReadFrom::Beginning).unwrap(); + let entry = reader.read_entry().unwrap().unwrap(); + assert_eq!(entry, b"hello 0"); + + let mut reader = Reader::from_path_position(path.clone(), ReadFrom::End).unwrap(); + + for _ in 0..entry_num { + let entry = reader.read_entry().unwrap(); + assert!(entry.is_none()); + } + + let start_pos = wal::FILE_TYPE_IDENTIFIER_LEN as u64 + wal::WAL_FILE_HEADER_LEN as u64; + let mut reader = Reader::from_path_position(path, ReadFrom::Checkpoint(start_pos)).unwrap(); + + for i in 0..entry_num { + let entry = reader.read_entry().unwrap().unwrap(); + assert_eq!(entry, format!("hello {}", i).as_bytes()); + } + assert!(reader.read_entry().unwrap().is_none()); + + let reader_metadata = reader.metadata().unwrap(); + assert_eq!(reader_metadata.len(), fs_metadata.len()); +}