refactor: Writer new method & Reader test (#6917)

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

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];

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,
+    },
 }

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<PathBuf>,
-        org_id: &str,
-        stream_type: &str,
-        id: String,
-        init_size: u64,
-        buffer_size: usize,
-    ) -> Result<Self> {
-        Self::build(
-            root_dir,
-            org_id,
-            stream_type,
-            id,
-            init_size,
-            buffer_size,
-            None,
-        )
-    }
-
-    pub fn build(
-        root_dir: impl Into<PathBuf>,
-        org_id: &str,
-        stream_type: &str,
-        id: String,
+        path: impl Into<PathBuf> + Clone + Deref<Target = Path> + AsRef<Path>,
         init_size: u64,
         buffer_size: usize,
         header: Option<FileHeader>,
     ) -> Result<Self> {
-        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,

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());
+}