Skip to main content

sh_layer2/checkpoint_system/
atomic.rs

1//! # Atomic File Writer
2//!
3//! 原子文件写入实现,确保数据持久化的完整性。
4
5use std::path::{Path, PathBuf};
6
7use crate::types::{Layer2Error, Layer2Result};
8
9/// 临时文件前缀
10const TEMP_FILE_PREFIX: &str = ".tmp_checkpoint_";
11
12/// 原子文件写入器
13///
14/// 使用临时文件 + 重命名模式实现原子写入。
15/// 在 Windows 上使用 os.replace(),在 NTFS 上是原子的。
16/// 在 Unix 上使用 rename(),在相同文件系统上是原子的。
17pub struct AtomicFileWriter {
18    sync_on_write: bool,
19    verify_write: bool,
20}
21
22impl AtomicFileWriter {
23    /// 创建新的原子写入器
24    pub fn new() -> Self {
25        Self {
26            sync_on_write: true,
27            verify_write: true,
28        }
29    }
30
31    /// 配置是否在写入后同步
32    pub fn with_sync(mut self, sync: bool) -> Self {
33        self.sync_on_write = sync;
34        self
35    }
36
37    /// 配置是否验证写入
38    pub fn with_verify(mut self, verify: bool) -> Self {
39        self.verify_write = verify;
40        self
41    }
42
43    /// 原子写入内容到文件
44    ///
45    /// # Arguments
46    /// * `filepath` - 目标文件路径
47    /// * `content` - 要写入的内容
48    ///
49    /// # Returns
50    /// 成功返回 Ok(()),失败返回错误信息
51    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        // 确保父目录存在
61        std::fs::create_dir_all(parent_dir)?;
62
63        // 生成临时文件路径(在同一目录下,保证同一文件系统)
64        let temp_filename = format!("{}{}", TEMP_FILE_PREFIX, uuid::Uuid::new_v4());
65        let temp_path = parent_dir.join(temp_filename);
66
67        // 执行写入
68        let result = self.do_write(&filepath, &temp_path, content);
69
70        // 清理临时文件(如果失败)
71        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        // 1. 写入临时文件
83        {
84            let mut file = File::create(temp_path)?;
85            file.write_all(content.as_bytes())?;
86
87            // 同步到磁盘
88            if self.sync_on_write {
89                file.sync_all()?;
90            }
91        }
92
93        // 2. 验证写入(可选)
94        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        // 3. 原子重命名
105        std::fs::rename(temp_path, filepath)?;
106
107        // 4. 同步目录(Unix only)
108        #[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    /// 安全删除文件
122    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}