sh_layer2/checkpoint_system/
atomic.rs1use std::path::{Path, PathBuf};
6
7use crate::types::{Layer2Error, Layer2Result};
8
9const TEMP_FILE_PREFIX: &str = ".tmp_checkpoint_";
11
12pub struct AtomicFileWriter {
18 sync_on_write: bool,
19 verify_write: bool,
20}
21
22impl AtomicFileWriter {
23 pub fn new() -> Self {
25 Self {
26 sync_on_write: true,
27 verify_write: true,
28 }
29 }
30
31 pub fn with_sync(mut self, sync: bool) -> Self {
33 self.sync_on_write = sync;
34 self
35 }
36
37 pub fn with_verify(mut self, verify: bool) -> Self {
39 self.verify_write = verify;
40 self
41 }
42
43 pub fn write_atomic(&self, filepath: &Path, content: &str) -> Layer2Result<()> {
52 let filepath = PathBuf::from(filepath);
53 let parent_dir = filepath.parent().ok_or_else(|| {
54 Layer2Error::Io(std::io::Error::new(
55 std::io::ErrorKind::InvalidInput,
56 "Invalid file path",
57 ))
58 })?;
59
60 std::fs::create_dir_all(parent_dir)?;
62
63 let temp_filename = format!("{}{}", TEMP_FILE_PREFIX, uuid::Uuid::new_v4());
65 let temp_path = parent_dir.join(temp_filename);
66
67 let result = self.do_write(&filepath, &temp_path, content);
69
70 if result.is_err() {
72 let _ = std::fs::remove_file(&temp_path);
73 }
74
75 result
76 }
77
78 fn do_write(&self, filepath: &Path, temp_path: &Path, content: &str) -> Layer2Result<()> {
79 use std::fs::File;
80 use std::io::Write;
81
82 {
84 let mut file = File::create(temp_path)?;
85 file.write_all(content.as_bytes())?;
86
87 if self.sync_on_write {
89 file.sync_all()?;
90 }
91 }
92
93 if self.verify_write {
95 let written = std::fs::read_to_string(temp_path)?;
96 if written != content {
97 return Err(Layer2Error::CheckpointCorrupted(
98 "Write verification failed: content mismatch".to_string(),
99 )
100 .into());
101 }
102 }
103
104 std::fs::rename(temp_path, filepath)?;
106
107 #[cfg(unix)]
109 if self.sync_on_write {
110 use std::os::unix::fs::OpenOptionsExt;
111 let dir_fd = std::fs::OpenOptions::new()
112 .read(true)
113 .custom_flags(libc::O_DIRECTORY)
114 .open(filepath.parent().unwrap())?;
115 dir_fd.sync_all()?;
116 }
117
118 Ok(())
119 }
120
121 pub fn safe_remove(&self, path: &Path) -> Layer2Result<()> {
123 match std::fs::remove_file(path) {
124 Ok(_) => Ok(()),
125 Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
126 Err(e) => Err(Layer2Error::Io(e).into()),
127 }
128 }
129}
130
131impl Default for AtomicFileWriter {
132 fn default() -> Self {
133 Self::new()
134 }
135}
136
137#[cfg(test)]
138mod tests {
139 use super::*;
140 use tempfile::TempDir;
141
142 #[test]
143 fn test_atomic_write() {
144 let temp_dir = TempDir::new().unwrap();
145 let file_path = temp_dir.path().join("test.json");
146
147 let writer = AtomicFileWriter::new();
148 writer.write_atomic(&file_path, "test content").unwrap();
149
150 let content = std::fs::read_to_string(&file_path).unwrap();
151 assert_eq!(content, "test content");
152 }
153
154 #[test]
155 fn test_atomic_write_creates_parent() {
156 let temp_dir = TempDir::new().unwrap();
157 let file_path = temp_dir.path().join("nested/dir/test.json");
158
159 let writer = AtomicFileWriter::new();
160 writer.write_atomic(&file_path, "test").unwrap();
161
162 assert!(file_path.exists());
163 }
164}