1use crate::{
2 domain::entities::Event,
3 error::{AllSourceError, Result},
4};
5use chrono::{DateTime, Utc};
6use parking_lot::RwLock;
7use serde::{Deserialize, Serialize};
8use std::{
9 collections::HashSet,
10 fs::{self, File, OpenOptions},
11 io::{BufRead, BufReader, BufWriter, Read, Seek, SeekFrom, Write},
12 path::{Path, PathBuf},
13 sync::Arc,
14};
15
16pub const QUARANTINE_DIR: &str = "quarantine";
19
20pub struct WriteAheadLog {
27 wal_dir: PathBuf,
29
30 current_file: Arc<RwLock<WALFile>>,
32
33 config: WALConfig,
35
36 stats: Arc<RwLock<WALStats>>,
38
39 sequence: Arc<RwLock<u64>>,
41
42 replication_tx: parking_lot::Mutex<Option<tokio::sync::broadcast::Sender<WALEntry>>>,
46
47 unreadable_segments: parking_lot::Mutex<HashSet<PathBuf>>,
49}
50
51#[derive(Debug, Clone)]
52pub struct WALConfig {
53 pub max_file_size: usize,
55
56 pub sync_on_write: bool,
58
59 pub max_wal_files: usize,
61
62 pub compress: bool,
64
65 pub fsync_interval_ms: Option<u64>,
70}
71
72impl Default for WALConfig {
73 fn default() -> Self {
74 Self {
75 max_file_size: 64 * 1024 * 1024, sync_on_write: true,
77 max_wal_files: 10,
78 compress: false,
79 fsync_interval_ms: None,
80 }
81 }
82}
83
84#[derive(Debug, Clone, Default, Serialize)]
85pub struct WALStats {
86 pub total_entries: u64,
87 pub total_bytes_written: u64,
88 pub current_file_size: usize,
89 pub files_rotated: u64,
90 pub files_cleaned: u64,
91 pub recovery_count: u64,
92}
93
94#[derive(Debug, Clone, Serialize, Deserialize)]
96pub struct WALEntry {
97 pub sequence: u64,
99
100 pub wal_timestamp: DateTime<Utc>,
102
103 pub event: Event,
105
106 pub checksum: u32,
108}
109
110impl WALEntry {
111 pub fn new(sequence: u64, event: Event) -> Self {
112 let mut entry = Self {
113 sequence,
114 wal_timestamp: Utc::now(),
115 event,
116 checksum: 0,
117 };
118 entry.checksum = entry.calculate_checksum();
119 entry
120 }
121
122 fn calculate_checksum(&self) -> u32 {
123 let data = format!("{}{}{}", self.sequence, self.wal_timestamp, self.event.id);
125 crc32fast::hash(data.as_bytes())
126 }
127
128 pub fn verify(&self) -> bool {
129 self.checksum == self.calculate_checksum()
130 }
131}
132
133fn ends_with_newline(file: &mut File) -> Result<bool> {
134 let io_err = |e: std::io::Error| {
135 AllSourceError::StorageError(format!("Failed to read WAL file tail: {e}"))
136 };
137 file.seek(SeekFrom::End(-1)).map_err(io_err)?;
138 let mut last = [0u8; 1];
139 file.read_exact(&mut last).map_err(io_err)?;
140 Ok(last[0] == b'\n')
141}
142
143struct WALFile {
145 path: PathBuf,
146 writer: BufWriter<File>,
147 size: usize,
148 created_at: DateTime<Utc>,
149}
150
151impl WALFile {
152 fn new(path: PathBuf) -> Result<Self> {
153 let mut file = OpenOptions::new()
154 .create(true)
155 .read(true)
156 .append(true)
157 .open(&path)
158 .map_err(|e| AllSourceError::StorageError(format!("Failed to open WAL file: {e}")))?;
159
160 let mut size = file.metadata().map_or(0, |m| m.len() as usize);
161
162 if size > 0 && !ends_with_newline(&mut file)? {
166 file.write_all(b"\n").map_err(|e| {
167 AllSourceError::StorageError(format!("Failed to terminate torn WAL line: {e}"))
168 })?;
169 size += 1;
170 }
171
172 Ok(Self {
173 path,
174 writer: BufWriter::new(file),
175 size,
176 created_at: Utc::now(),
177 })
178 }
179
180 fn write_entry(&mut self, entry: &WALEntry, sync: bool) -> Result<usize> {
181 let json = serde_json::to_string(entry)?;
183
184 let line = format!("{json}\n");
185 let bytes_written = line.len();
186
187 self.writer
188 .write_all(line.as_bytes())
189 .map_err(|e| AllSourceError::StorageError(format!("Failed to write to WAL: {e}")))?;
190
191 if sync {
192 self.writer
193 .flush()
194 .map_err(|e| AllSourceError::StorageError(format!("Failed to flush WAL: {e}")))?;
195
196 self.writer
197 .get_ref()
198 .sync_all()
199 .map_err(|e| AllSourceError::StorageError(format!("Failed to sync WAL: {e}")))?;
200 }
201
202 self.size += bytes_written;
203 Ok(bytes_written)
204 }
205
206 fn flush(&mut self) -> Result<()> {
207 self.writer
208 .flush()
209 .map_err(|e| AllSourceError::StorageError(format!("Failed to flush WAL: {e}")))?;
210 Ok(())
211 }
212}
213
214impl WriteAheadLog {
215 pub fn new(wal_dir: impl Into<PathBuf>, config: WALConfig) -> Result<Self> {
217 let wal_dir = wal_dir.into();
218
219 fs::create_dir_all(&wal_dir).map_err(|e| {
221 AllSourceError::StorageError(format!("Failed to create WAL directory: {e}"))
222 })?;
223
224 let initial_file_path = Self::generate_wal_filename(&wal_dir, 0);
226 let current_file = WALFile::new(initial_file_path)?;
227
228 tracing::info!("โ
WAL initialized at: {}", wal_dir.display());
229
230 Ok(Self {
231 wal_dir,
232 current_file: Arc::new(RwLock::new(current_file)),
233 config,
234 stats: Arc::new(RwLock::new(WALStats::default())),
235 sequence: Arc::new(RwLock::new(0)),
236 replication_tx: parking_lot::Mutex::new(None),
237 unreadable_segments: parking_lot::Mutex::new(HashSet::new()),
238 })
239 }
240
241 fn generate_wal_filename(dir: &Path, sequence: u64) -> PathBuf {
243 dir.join(format!("wal-{sequence:016x}.log"))
244 }
245
246 #[cfg_attr(feature = "hotpath", hotpath::measure)]
248 pub fn append(&self, event: Event) -> Result<u64> {
249 let mut seq = self.sequence.write();
251 *seq += 1;
252 let sequence = *seq;
253 drop(seq);
254
255 let entry = WALEntry::new(sequence, event);
257
258 let mut current = self.current_file.write();
260 let bytes_written = current.write_entry(&entry, self.config.sync_on_write)?;
261
262 let mut stats = self.stats.write();
264 stats.total_entries += 1;
265 stats.total_bytes_written += bytes_written as u64;
266 stats.current_file_size = current.size;
267 drop(stats);
268
269 if let Some(ref tx) = *self.replication_tx.lock() {
272 let _ = tx.send(entry);
273 }
274
275 let should_rotate = current.size >= self.config.max_file_size;
277 drop(current);
278
279 if should_rotate {
280 self.rotate()?;
281 }
282
283 tracing::trace!("WAL entry written: sequence={}", sequence);
284
285 Ok(sequence)
286 }
287
288 #[cfg_attr(feature = "hotpath", hotpath::measure)]
290 fn rotate(&self) -> Result<()> {
291 let seq = *self.sequence.read();
292 let new_file_path = Self::generate_wal_filename(&self.wal_dir, seq);
293
294 tracing::info!("๐ Rotating WAL to new file: {:?}", new_file_path);
295
296 let new_file = WALFile::new(new_file_path)?;
297
298 let mut current = self.current_file.write();
299 current.flush()?;
300 *current = new_file;
301
302 let mut stats = self.stats.write();
303 stats.files_rotated += 1;
304 stats.current_file_size = 0;
305 drop(stats);
306
307 self.warn_if_segments_accumulate()?;
308
309 Ok(())
310 }
311
312 fn warn_if_segments_accumulate(&self) -> Result<()> {
315 let count = self.list_wal_files()?.len();
316 if count > self.config.max_wal_files {
317 tracing::warn!(
318 "WAL holds {} segments (max_wal_files = {}); checkpoints are not keeping up, \
319 so the WAL keeps growing rather than dropping unflushed events",
320 count,
321 self.config.max_wal_files
322 );
323 }
324 Ok(())
325 }
326
327 fn retire_segment(&self, path: &Path) -> Result<()> {
330 if !self.unreadable_segments.lock().remove(path) {
331 fs::remove_file(path).map_err(|e| {
332 AllSourceError::StorageError(format!("Failed to remove WAL file: {e}"))
333 })?;
334 tracing::debug!("Removed WAL file: {:?}", path);
335 return Ok(());
336 }
337
338 let quarantine = self.wal_dir.join(QUARANTINE_DIR);
339 fs::create_dir_all(&quarantine).map_err(|e| {
340 AllSourceError::StorageError(format!("Failed to create WAL quarantine: {e}"))
341 })?;
342 let name = path.file_name().map_or_else(
343 || "wal-unknown.log".into(),
344 |n| n.to_string_lossy().into_owned(),
345 );
346 let target = quarantine.join(format!(
347 "{}-{name}",
348 Utc::now().format("%Y%m%dT%H%M%S%.6fZ")
349 ));
350 fs::rename(path, &target).map_err(|e| {
351 AllSourceError::StorageError(format!("Failed to quarantine WAL file: {e}"))
352 })?;
353 tracing::warn!(
354 "WAL segment {:?} held lines this binary could not read; kept at {:?} instead of \
355 deleting it",
356 path,
357 target
358 );
359 Ok(())
360 }
361
362 pub fn seal(&self) -> Result<PathBuf> {
369 let seq = *self.sequence.read();
370 let new_path = Self::generate_wal_filename(&self.wal_dir, seq);
371
372 let mut current = self.current_file.write();
373 current.flush()?;
374 if current.path != new_path {
375 *current = WALFile::new(new_path.clone())?;
376 self.stats.write().current_file_size = current.size;
377 }
378 Ok(new_path)
379 }
380
381 pub fn remove_sealed(&self, active: &Path) -> Result<()> {
384 let Some(active_name) = active.file_name() else {
385 return Ok(());
386 };
387 for path in self.list_wal_files()? {
388 if path.file_name().is_some_and(|name| name < active_name) {
389 self.retire_segment(&path)?;
390 }
391 }
392 Ok(())
393 }
394
395 fn list_wal_files(&self) -> Result<Vec<PathBuf>> {
397 let entries = fs::read_dir(&self.wal_dir).map_err(|e| {
398 AllSourceError::StorageError(format!("Failed to read WAL directory: {e}"))
399 })?;
400
401 let mut wal_files = Vec::new();
402 for entry in entries {
403 let entry = entry.map_err(|e| {
404 AllSourceError::StorageError(format!("Failed to read directory entry: {e}"))
405 })?;
406
407 let path = entry.path();
408 if let Some(name) = path.file_name()
409 && name.to_string_lossy().starts_with("wal-")
410 && name.to_string_lossy().ends_with(".log")
411 {
412 wal_files.push(path);
413 }
414 }
415
416 Ok(wal_files)
417 }
418
419 #[cfg_attr(feature = "hotpath", hotpath::measure)]
421 pub fn recover(&self) -> Result<Vec<Event>> {
422 tracing::info!("๐ Starting WAL recovery...");
423
424 let mut wal_files = self.list_wal_files()?;
425 wal_files.sort();
426
427 let mut recovered_events = Vec::new();
428 let mut max_sequence = 0u64;
429 let mut corrupted_entries = 0;
430
431 for wal_file_path in &wal_files {
432 tracing::debug!("Reading WAL file: {:?}", wal_file_path);
433
434 let file = File::open(wal_file_path).map_err(|e| {
435 AllSourceError::StorageError(format!("Failed to open WAL file for recovery: {e}"))
436 })?;
437
438 let reader = BufReader::new(file);
439 let corrupted_before = corrupted_entries;
440
441 for (line_num, line) in reader.lines().enumerate() {
442 let line = match line {
443 Ok(l) => l,
444 Err(e) => {
445 tracing::warn!(
446 "I/O error reading WAL line at {:?}:{}: {}",
447 wal_file_path,
448 line_num + 1,
449 e
450 );
451 corrupted_entries += 1;
452 continue;
453 }
454 };
455
456 if line.trim().is_empty() {
457 continue;
458 }
459
460 match serde_json::from_str::<WALEntry>(&line) {
461 Ok(entry) => {
462 if !entry.verify() {
464 tracing::warn!(
465 "Corrupted WAL entry at {:?}:{} (checksum mismatch)",
466 wal_file_path,
467 line_num + 1
468 );
469 corrupted_entries += 1;
470 continue;
471 }
472
473 max_sequence = max_sequence.max(entry.sequence);
474 recovered_events.push(entry.event);
475 }
476 Err(e) => {
477 tracing::warn!(
478 "Failed to parse WAL entry at {:?}:{}: {}",
479 wal_file_path,
480 line_num + 1,
481 e
482 );
483 corrupted_entries += 1;
484 }
485 }
486 }
487
488 if corrupted_entries > corrupted_before {
489 self.unreadable_segments
490 .lock()
491 .insert(wal_file_path.clone());
492 }
493 }
494
495 let mut seq = self.sequence.write();
497 *seq = max_sequence;
498 drop(seq);
499
500 let mut stats = self.stats.write();
502 stats.recovery_count += 1;
503 drop(stats);
504
505 tracing::info!(
506 "โ
WAL recovery complete: {} events recovered, {} corrupted entries",
507 recovered_events.len(),
508 corrupted_entries
509 );
510
511 Ok(recovered_events)
512 }
513
514 #[cfg_attr(feature = "hotpath", hotpath::measure)]
516 pub fn flush(&self) -> Result<()> {
517 let mut current = self.current_file.write();
518 current.flush()?;
519 Ok(())
520 }
521
522 #[cfg_attr(feature = "hotpath", hotpath::measure)]
528 pub fn sync(&self) -> Result<()> {
529 let mut current = self.current_file.write();
530 current
531 .writer
532 .flush()
533 .map_err(|e| AllSourceError::StorageError(format!("Failed to flush WAL: {e}")))?;
534 current
535 .writer
536 .get_ref()
537 .sync_all()
538 .map_err(|e| AllSourceError::StorageError(format!("Failed to sync WAL: {e}")))?;
539 Ok(())
540 }
541
542 pub fn fsync_interval_ms(&self) -> Option<u64> {
544 self.config.fsync_interval_ms
545 }
546
547 #[cfg_attr(feature = "hotpath", hotpath::measure)]
549 pub fn truncate(&self) -> Result<()> {
550 tracing::info!("๐งน Truncating WAL after checkpoint");
551
552 let mut current = self.current_file.write();
554 current.flush()?;
555
556 let wal_files = self.list_wal_files()?;
557 for file_path in wal_files {
558 self.retire_segment(&file_path)?;
559 }
560
561 let new_file_path = Self::generate_wal_filename(&self.wal_dir, 0);
563 *current = WALFile::new(new_file_path)?;
564
565 let mut seq = self.sequence.write();
567 *seq = 0;
568
569 tracing::info!("โ
WAL truncated successfully");
570
571 Ok(())
572 }
573
574 pub fn stats(&self) -> WALStats {
576 (*self.stats.read()).clone()
577 }
578
579 pub fn on_disk_stats(&self) -> Result<(u64, usize)> {
592 let wal_files = self.list_wal_files()?;
593 let mut total_bytes = 0u64;
594 for path in &wal_files {
595 if let Ok(metadata) = fs::metadata(path) {
596 total_bytes += metadata.len();
597 }
598 }
599 Ok((total_bytes, wal_files.len()))
600 }
601
602 pub fn current_sequence(&self) -> u64 {
604 *self.sequence.read()
605 }
606
607 pub fn oldest_sequence(&self) -> Option<u64> {
613 let Ok(mut wal_files) = self.list_wal_files() else {
614 return None;
615 };
616
617 if wal_files.is_empty() {
618 return None;
619 }
620
621 wal_files.sort();
622
623 for wal_file_path in &wal_files {
625 let Ok(file) = File::open(wal_file_path) else {
626 continue;
627 };
628 let reader = BufReader::new(file);
629 for line in reader.lines() {
630 let Ok(line) = line else {
631 continue;
632 };
633 if line.trim().is_empty() {
634 continue;
635 }
636 if let Ok(entry) = serde_json::from_str::<WALEntry>(&line) {
637 return Some(entry.sequence);
638 }
639 }
640 }
641
642 None
643 }
644
645 pub fn set_replication_tx(&self, tx: tokio::sync::broadcast::Sender<WALEntry>) {
650 *self.replication_tx.lock() = Some(tx);
651 }
652}
653
654#[cfg(test)]
655mod tests {
656 use super::*;
657 use serde_json::json;
658 use tempfile::TempDir;
659 use uuid::Uuid;
660
661 fn create_test_event() -> Event {
662 Event::reconstruct_from_strings(
663 Uuid::new_v4(),
664 "test.event".to_string(),
665 "test-entity".to_string(),
666 "default".to_string(),
667 json!({"test": "data"}),
668 Utc::now(),
669 None,
670 1,
671 )
672 }
673
674 #[test]
675 fn test_wal_creation() {
676 let temp_dir = TempDir::new().unwrap();
677 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default());
678 assert!(wal.is_ok());
679 }
680
681 #[test]
682 fn test_wal_append() {
683 let temp_dir = TempDir::new().unwrap();
684 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
685
686 let event = create_test_event();
687 let seq = wal.append(event);
688 assert!(seq.is_ok());
689 assert_eq!(seq.unwrap(), 1);
690
691 let stats = wal.stats();
692 assert_eq!(stats.total_entries, 1);
693 }
694
695 #[test]
696 fn test_wal_on_disk_stats() {
697 let temp_dir = TempDir::new().unwrap();
698 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
699
700 let (bytes0, segs0) = wal.on_disk_stats().unwrap();
703 assert_eq!(bytes0, 0);
704 assert!(segs0 <= 1);
705
706 for _ in 0..5 {
707 wal.append(create_test_event()).unwrap();
708 }
709 wal.flush().unwrap();
710
711 let (bytes, segs) = wal.on_disk_stats().unwrap();
712 assert!(segs >= 1, "expected at least one WAL segment, got {segs}");
713 assert!(
714 bytes > 0,
715 "expected non-zero WAL bytes after appends, got {bytes}"
716 );
717 }
718
719 #[test]
720 fn test_wal_recovery() {
721 let temp_dir = TempDir::new().unwrap();
722 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
723
724 for _ in 0..5 {
726 wal.append(create_test_event()).unwrap();
727 }
728
729 wal.flush().unwrap();
730
731 let wal2 = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
733 let recovered = wal2.recover().unwrap();
734
735 assert_eq!(recovered.len(), 5);
736 }
737
738 #[test]
739 fn test_wal_recovery_with_partial_write() {
740 let temp_dir = TempDir::new().unwrap();
741 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
742
743 for _ in 0..3 {
745 wal.append(create_test_event()).unwrap();
746 }
747 wal.flush().unwrap();
748
749 let wal_file_path = temp_dir.path().join("wal-0000000000000000.log");
751 use std::io::Write as _;
752 let mut f = std::fs::OpenOptions::new()
753 .append(true)
754 .open(&wal_file_path)
755 .unwrap();
756 f.write_all(b"{\"partial\": true, \"seq\"\n").unwrap(); drop(f);
758
759 let wal2 = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
761 let recovered = wal2.recover().unwrap();
762 assert_eq!(
763 recovered.len(),
764 3,
765 "Should recover only the 3 valid events, not the partial one"
766 );
767 }
768
769 #[test]
770 fn test_wal_rotation() {
771 let temp_dir = TempDir::new().unwrap();
772 let config = WALConfig {
773 max_file_size: 1024, ..Default::default()
775 };
776
777 let wal = WriteAheadLog::new(temp_dir.path(), config).unwrap();
778
779 for _ in 0..50 {
781 wal.append(create_test_event()).unwrap();
782 }
783
784 let stats = wal.stats();
785 assert!(stats.files_rotated > 0);
786 }
787
788 #[test]
789 fn test_wal_entry_checksum() {
790 let event = create_test_event();
791 let entry = WALEntry::new(1, event);
792
793 assert!(entry.verify());
794
795 let mut corrupted = entry.clone();
797 corrupted.checksum = 0;
798 assert!(!corrupted.verify());
799 }
800
801 #[test]
802 fn test_wal_fsync_interval_config() {
803 let config = WALConfig {
804 fsync_interval_ms: Some(100),
805 ..Default::default()
806 };
807 assert_eq!(config.fsync_interval_ms, Some(100));
808 assert_eq!(WALConfig::default().fsync_interval_ms, None);
810 }
811
812 #[test]
813 fn test_wal_sync_method() {
814 let temp_dir = TempDir::new().unwrap();
815 let config = WALConfig {
816 sync_on_write: false, ..Default::default()
818 };
819 let wal = WriteAheadLog::new(temp_dir.path(), config).unwrap();
820
821 for _ in 0..5 {
823 wal.append(create_test_event()).unwrap();
824 }
825
826 wal.sync().unwrap();
828
829 let wal2 = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
831 let recovered = wal2.recover().unwrap();
832 assert_eq!(recovered.len(), 5);
833 }
834
835 #[test]
836 fn test_wal_truncate() {
837 let temp_dir = TempDir::new().unwrap();
838 let wal = WriteAheadLog::new(temp_dir.path(), WALConfig::default()).unwrap();
839
840 for _ in 0..5 {
842 wal.append(create_test_event()).unwrap();
843 }
844
845 wal.truncate().unwrap();
847
848 assert_eq!(wal.current_sequence(), 0);
850
851 let recovered = wal.recover().unwrap();
853 assert_eq!(recovered.len(), 0);
854 }
855}