use std::fs::{self, File, OpenOptions}; use std::io::Write; use std::path::{Path, PathBuf}; use crate::replication::ShardId; use super::error::WalError; /// Format a segment file name from the shard ID and first sequence number. /// /// For `ShardId::SINGLE` (shard 0), produces the v1 format: /// `wal-00000000000000000001.seg` /// /// For any other shard, produces the v2 format: /// `wal-s00003-00000000000000000001.seg` /// /// Zero-padded fields ensure lexicographic ordering matches numeric ordering. #[must_use] pub fn segment_filename(shard_id: ShardId, first_seq: u64) -> String { if shard_id == ShardId::SINGLE { format!("wal-{first_seq:020}.seg") } else { format!("wal-s{:05}-{:020}.seg", shard_id.0, first_seq) } } /// Parse the shard ID and first sequence number from a WAL segment filename. /// /// Accepts both formats: /// - `wal-{first_seq:020}.seg` (single-shard, v1) -> `(ShardId::SINGLE, first_seq)` /// - `wal-s{shard_id:05}-{first_seq:020}.seg` (multi-shard, v2) -> `(ShardId(n), first_seq)` /// /// Returns `None` if the filename does not match either format. #[must_use] pub fn parse_segment_filename(filename: &str) -> Option<(ShardId, u64)> { let name = filename.strip_suffix(".seg")?; // Multi-shard format: wal-s{shard_id}-{first_seq} if let Some(rest) = name.strip_prefix("wal-s") { let dash = rest.find('-')?; let shard_id: u16 = rest[..dash].parse().ok()?; let first_seq: u64 = rest[dash + 1..].parse().ok()?; return Some((ShardId(shard_id), first_seq)); } // Single-shard format: wal-{first_seq} if let Some(seq_str) = name.strip_prefix("wal-") { let first_seq: u64 = seq_str.parse().ok()?; return Some((ShardId::SINGLE, first_seq)); } None } /// Parse the first sequence number from a segment file name. /// /// Backward-compatible wrapper around [`parse_segment_filename`] that discards /// the shard ID. Used by [`list_segments`] and legacy callers. /// /// Returns `None` if the file name does not match the expected pattern. #[must_use] pub fn parse_segment_seq(filename: &str) -> Option { parse_segment_filename(filename).map(|(_, seq)| seq) } /// List all WAL segment files in the directory, sorted by first sequence number. /// /// # Errors /// /// Returns `WalError::Io` on filesystem failure. pub fn list_segments(dir: &Path) -> Result, WalError> { let mut segments = Vec::new(); let entries = match fs::read_dir(dir) { Ok(e) => e, Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(segments), Err(e) => return Err(WalError::Io(e)), }; for entry in entries { let entry = entry?; let filename = entry.file_name(); let Some(name) = filename.to_str() else { continue; }; if let Some(seq) = parse_segment_seq(name) { segments.push((seq, entry.path())); } } segments.sort_by_key(|(seq, _)| *seq); Ok(segments) } /// List WAL segment files for a specific shard, sorted by first sequence number. /// /// When `shard_id` is `ShardId::SINGLE`, includes all single-shard format segments. /// When `shard_id > 0`, includes only segments with that shard prefix. /// /// # Errors /// /// Returns `WalError::Io` on filesystem failure. pub fn list_segments_for_shard( dir: &Path, shard_id: ShardId, ) -> Result, WalError> { let mut segments = Vec::new(); let entries = match fs::read_dir(dir) { Ok(e) => e, Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(segments), Err(e) => return Err(WalError::Io(e)), }; for entry in entries { let entry = entry?; let file_name = entry.file_name(); let Some(name) = file_name.to_str() else { continue; }; if let Some((seg_shard, seq)) = parse_segment_filename(name) && seg_shard == shard_id { segments.push((seq, entry.path())); } } segments.sort_by_key(|(seq, _)| *seq); Ok(segments) } /// Manages the current writable WAL segment file. /// /// Handles creation of new segment files, tracking file size for rotation, /// and data sync. Rotation is triggered externally by the writer when /// the segment exceeds `max_size`. pub struct SegmentWriter { dir: PathBuf, file: File, current_size: u64, max_size: u64, first_seq: u64, /// The last sequence number written to this segment. last_seq: u64, } impl SegmentWriter { /// Open or create a segment file for writing. /// /// If `first_seq` identifies an existing segment, it is opened for append. /// Otherwise, a new file is created. /// /// # Errors /// /// Returns `WalError::Io` on filesystem failure. pub fn open(dir: &Path, first_seq: u64, max_size: u64) -> Result { // Single-node: all segments use the v1 (unsharded) naming format. let filename = segment_filename(ShardId::SINGLE, first_seq); let path = dir.join(&filename); let is_new = !path.exists(); let file = OpenOptions::new().create(true).append(true).open(&path)?; // Fsync the parent directory so the new directory entry is durable. // Without this, a crash after file creation but before the directory // metadata is flushed could lose the segment file entirely. if is_new { let dir_fd = File::open(dir)?; dir_fd.sync_all()?; } let metadata = file.metadata()?; let current_size = metadata.len(); Ok(Self { dir: dir.to_path_buf(), file, current_size, max_size, first_seq, last_seq: first_seq, }) } /// Write a raw batch of bytes to the current segment. /// /// Returns the file offset where the batch was written. /// /// # Errors /// /// Returns `WalError::Io` on write failure. pub fn write_batch_bytes(&mut self, bytes: &[u8]) -> Result { let offset = self.current_size; self.file.write_all(bytes)?; self.current_size += bytes.len() as u64; Ok(offset) } /// Sync all written data to stable storage. /// /// Uses `File::sync_data()` which maps to `fdatasync` on Linux and /// `fsync` on macOS. This is the safe Rust equivalent. /// /// # Errors /// /// Returns `WalError::Io` on sync failure. pub fn sync(&self) -> Result<(), WalError> { self.file.sync_data()?; Ok(()) } /// Whether the segment has reached its size threshold and should be rotated. #[must_use] pub const fn needs_rotation(&self) -> bool { self.current_size >= self.max_size } /// The first sequence number in this segment. #[must_use] pub const fn first_seq(&self) -> u64 { self.first_seq } /// The last sequence number written to this segment. #[must_use] pub const fn last_seq(&self) -> u64 { self.last_seq } /// Update the last sequence number written to this segment. pub const fn set_last_seq(&mut self, seq: u64) { self.last_seq = seq; } /// The current file size in bytes. #[must_use] pub const fn current_size(&self) -> u64 { self.current_size } /// Create a new segment file and return a writer for it. /// /// Finalizes the current segment (syncs it) and opens a new one. /// /// # Errors /// /// Returns `WalError::Io` on filesystem failure. pub fn rotate(&mut self, new_first_seq: u64) -> Result<(), WalError> { // Sync current segment before rotation self.sync()?; // Single-node: all segments use the v1 (unsharded) naming format. let filename = segment_filename(ShardId::SINGLE, new_first_seq); let path = self.dir.join(&filename); let file = OpenOptions::new().create(true).append(true).open(&path)?; // Fsync the parent directory so the new segment's directory entry // is durable. Without this, a crash after file creation but before // the directory metadata is flushed could lose the new segment. let dir_fd = File::open(&self.dir)?; dir_fd.sync_all()?; self.file = file; self.current_size = 0; self.first_seq = new_first_seq; self.last_seq = new_first_seq; Ok(()) } } /// Delete all segment files whose first sequence number is less than `before_seq`. /// /// # Errors /// /// Returns `WalError::Io` on filesystem failure. Partial deletion may occur /// if an error is encountered mid-way. pub fn delete_segments_before(dir: &Path, before_seq: u64) -> Result { let segments = list_segments(dir)?; let mut deleted = 0; for (seq, path) in segments { if seq < before_seq { fs::remove_file(&path)?; deleted += 1; } } Ok(deleted) } #[cfg(test)] mod tests { use super::*; #[test] fn segment_filename_format() { assert_eq!( segment_filename(ShardId::SINGLE, 1), "wal-00000000000000000001.seg" ); assert_eq!( segment_filename(ShardId::SINGLE, 0), "wal-00000000000000000000.seg" ); assert_eq!( segment_filename(ShardId::SINGLE, u64::MAX), "wal-18446744073709551615.seg" ); } #[test] fn parse_segment_seq_valid() { assert_eq!(parse_segment_seq("wal-00000000000000000001.seg"), Some(1)); assert_eq!(parse_segment_seq("wal-00000000000000000000.seg"), Some(0)); } #[test] fn parse_segment_seq_invalid() { assert_eq!(parse_segment_seq("not-a-segment.txt"), None); assert_eq!(parse_segment_seq("wal-.seg"), None); assert_eq!(parse_segment_seq("wal-abc.seg"), None); assert_eq!(parse_segment_seq(""), None); } #[test] fn write_and_check_size() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let mut writer = SegmentWriter::open(dir.path(), 1, 1024).expect("open should succeed"); assert_eq!(writer.current_size(), 0); let data = [0xABu8; 100]; writer .write_batch_bytes(&data) .expect("write should succeed"); assert_eq!(writer.current_size(), 100); } #[test] fn rotation_creates_new_file() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let mut writer = SegmentWriter::open(dir.path(), 1, 100).expect("open should succeed"); writer .write_batch_bytes(&[0u8; 50]) .expect("write should succeed"); writer.rotate(100).expect("rotate should succeed"); assert_eq!(writer.current_size(), 0); assert_eq!(writer.first_seq(), 100); let segments = list_segments(dir.path()).expect("list should succeed"); assert_eq!(segments.len(), 2); assert_eq!(segments[0].0, 1); assert_eq!(segments[1].0, 100); } #[test] fn needs_rotation_threshold() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let mut writer = SegmentWriter::open(dir.path(), 1, 100).expect("open should succeed"); assert!(!writer.needs_rotation()); writer .write_batch_bytes(&[0u8; 100]) .expect("write should succeed"); assert!(writer.needs_rotation()); } #[test] fn list_segments_sorted() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); // Create segments out of order let _ = SegmentWriter::open(dir.path(), 300, 1024); let _ = SegmentWriter::open(dir.path(), 100, 1024); let _ = SegmentWriter::open(dir.path(), 200, 1024); let segments = list_segments(dir.path()).expect("list should succeed"); assert_eq!(segments.len(), 3); assert_eq!(segments[0].0, 100); assert_eq!(segments[1].0, 200); assert_eq!(segments[2].0, 300); } #[test] fn list_segments_empty_dir() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let segments = list_segments(dir.path()).expect("list should succeed"); assert!(segments.is_empty()); } #[test] fn list_segments_ignores_non_segment_files() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); fs::write(dir.path().join("checkpoint.meta"), "seq=1\nts=1\n") .expect("write should succeed"); fs::write(dir.path().join("random.txt"), "hello").expect("write should succeed"); let _ = SegmentWriter::open(dir.path(), 1, 1024); let segments = list_segments(dir.path()).expect("list should succeed"); assert_eq!(segments.len(), 1); } #[test] fn delete_segments_before_removes_older() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let _ = SegmentWriter::open(dir.path(), 1, 1024); let _ = SegmentWriter::open(dir.path(), 100, 1024); let _ = SegmentWriter::open(dir.path(), 200, 1024); let deleted = delete_segments_before(dir.path(), 200).expect("delete should succeed"); assert_eq!(deleted, 2); let remaining = list_segments(dir.path()).expect("list should succeed"); assert_eq!(remaining.len(), 1); assert_eq!(remaining[0].0, 200); } #[test] fn delete_segments_before_none_to_delete() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let _ = SegmentWriter::open(dir.path(), 100, 1024); let deleted = delete_segments_before(dir.path(), 50).expect("delete should succeed"); assert_eq!(deleted, 0); } #[test] fn sync_does_not_error() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let writer = SegmentWriter::open(dir.path(), 1, 1024).expect("open should succeed"); writer.sync().expect("sync should succeed"); } // ── Multi-shard naming tests ────────────────────────────────────────── #[test] fn segment_filename_multi_shard() { let name = segment_filename(ShardId(3), 42); assert_eq!(name, "wal-s00003-00000000000000000042.seg"); } #[test] fn segment_filename_single_shard_backward_compat() { // ShardId::SINGLE retains old format exactly assert_eq!( segment_filename(ShardId::SINGLE, 1), "wal-00000000000000000001.seg" ); assert_eq!( segment_filename(ShardId(0), 42), "wal-00000000000000000042.seg" ); } #[test] fn parse_segment_filename_both_formats() { assert_eq!( parse_segment_filename("wal-00000000000000000001.seg"), Some((ShardId::SINGLE, 1)) ); assert_eq!( parse_segment_filename("wal-s00003-00000000000000000042.seg"), Some((ShardId(3), 42)) ); assert_eq!(parse_segment_filename("not-a-segment.txt"), None); } #[test] fn parse_segment_filename_edge_cases() { // Empty filename assert_eq!(parse_segment_filename(""), None); // Missing .seg suffix assert_eq!(parse_segment_filename("wal-00000000000000000001"), None); // Shard format but missing seq assert_eq!(parse_segment_filename("wal-s00003.seg"), None); // Non-numeric shard assert_eq!( parse_segment_filename("wal-sXYZ-00000000000000000001.seg"), None ); } #[test] fn list_segments_for_shard_filters_correctly() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); // Create single-shard segment files (v1 format) let _ = SegmentWriter::open(dir.path(), 1, 1024); let _ = SegmentWriter::open(dir.path(), 100, 1024); // Create multi-shard segment files (v2 format) manually let s3_name = segment_filename(ShardId(3), 50); fs::write(dir.path().join(s3_name), []).expect("write should succeed"); let s3_name2 = segment_filename(ShardId(3), 200); fs::write(dir.path().join(s3_name2), []).expect("write should succeed"); let s5_name = segment_filename(ShardId(5), 75); fs::write(dir.path().join(s5_name), []).expect("write should succeed"); // list_segments returns ALL segment files (backward compat) let all = list_segments(dir.path()).expect("list should succeed"); assert_eq!(all.len(), 5); // list_segments_for_shard returns only matching shard let single = list_segments_for_shard(dir.path(), ShardId::SINGLE).expect("list should succeed"); assert_eq!(single.len(), 2); assert_eq!(single[0].0, 1); assert_eq!(single[1].0, 100); let shard3 = list_segments_for_shard(dir.path(), ShardId(3)).expect("list should succeed"); assert_eq!(shard3.len(), 2); assert_eq!(shard3[0].0, 50); assert_eq!(shard3[1].0, 200); let shard5 = list_segments_for_shard(dir.path(), ShardId(5)).expect("list should succeed"); assert_eq!(shard5.len(), 1); assert_eq!(shard5[0].0, 75); // Non-existent shard returns empty let empty = list_segments_for_shard(dir.path(), ShardId(99)).expect("list should succeed"); assert!(empty.is_empty()); } #[test] fn list_segments_for_shard_missing_dir() { let dir = tempfile::tempdir().expect("tempdir creation should succeed"); let missing = dir.path().join("does-not-exist"); let result = list_segments_for_shard(&missing, ShardId::SINGLE) .expect("should handle missing dir gracefully"); assert!(result.is_empty()); } mod proptests { use super::*; use proptest::prelude::*; proptest! { #[test] fn filename_roundtrip(seq: u64) { let name = segment_filename(ShardId::SINGLE, seq); let parsed = parse_segment_seq(&name); prop_assert_eq!(parsed, Some(seq)); } #[test] fn shard_filename_roundtrip(shard_id in 0u16..100u16, seq in proptest::num::u64::ANY) { let shard = ShardId(shard_id); let name = segment_filename(shard, seq); let parsed = parse_segment_filename(&name); prop_assert_eq!(parsed, Some((shard, seq))); } } } }