Skip to main content

fakecloud_persistence/
atomic.rs

1use std::fs::{File, OpenOptions};
2use std::io::{self, Write};
3use std::path::{Path, PathBuf};
4
5use serde::Serialize;
6
7fn tmp_path(path: &Path) -> PathBuf {
8    let mut os = path.as_os_str().to_owned();
9    os.push(".tmp");
10    PathBuf::from(os)
11}
12
13fn fsync_parent(path: &Path) -> io::Result<()> {
14    if let Some(parent) = path.parent() {
15        if !parent.as_os_str().is_empty() {
16            let dir = File::open(parent)?;
17            dir.sync_all()?;
18        }
19    }
20    Ok(())
21}
22
23fn write_atomic_bytes_inner(tmp: &Path, path: &Path, bytes: &[u8]) -> io::Result<()> {
24    {
25        let mut f = OpenOptions::new()
26            .write(true)
27            .create(true)
28            .truncate(true)
29            .open(tmp)?;
30        f.write_all(bytes)?;
31        f.sync_all()?;
32    }
33    std::fs::rename(tmp, path)?;
34    fsync_parent(path)?;
35    Ok(())
36}
37
38pub fn write_atomic_bytes(path: &Path, bytes: &[u8]) -> io::Result<()> {
39    let tmp = tmp_path(path);
40    match write_atomic_bytes_inner(&tmp, path, bytes) {
41        Ok(()) => Ok(()),
42        Err(e) => {
43            let _ = std::fs::remove_file(&tmp);
44            Err(e)
45        }
46    }
47}
48
49pub fn write_atomic_toml<T: Serialize>(path: &Path, value: &T) -> io::Result<()> {
50    let text = toml::to_string_pretty(value)
51        .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?;
52    write_atomic_bytes(path, text.as_bytes())
53}
54
55fn write_atomic_from_file_inner(src: &Path, dst: &Path) -> io::Result<()> {
56    {
57        let f = File::open(src)?;
58        f.sync_all()?;
59    }
60    std::fs::rename(src, dst)?;
61    fsync_parent(dst)?;
62    Ok(())
63}
64
65pub fn write_atomic_from_file(src: &Path, dst: &Path) -> io::Result<()> {
66    match write_atomic_from_file_inner(src, dst) {
67        Ok(()) => Ok(()),
68        Err(e) => {
69            // Best-effort cleanup: remove any stray tmp the caller might see.
70            let tmp = tmp_path(dst);
71            let _ = std::fs::remove_file(&tmp);
72            Err(e)
73        }
74    }
75}
76
77fn write_atomic_copy_from_file_inner(tmp: &Path, src: &Path, dst: &Path) -> io::Result<()> {
78    {
79        let mut input = File::open(src)?;
80        let mut out = OpenOptions::new()
81            .write(true)
82            .create(true)
83            .truncate(true)
84            .open(tmp)?;
85        io::copy(&mut input, &mut out)?;
86        out.sync_all()?;
87    }
88    std::fs::rename(tmp, dst)?;
89    fsync_parent(dst)?;
90    Ok(())
91}
92
93/// Copy `src` into `dst` atomically, leaving `src` untouched. Used by the
94/// S3 store to replicate disk-backed object bodies without round-tripping
95/// through RAM.
96pub fn write_atomic_copy_from_file(src: &Path, dst: &Path) -> io::Result<()> {
97    let tmp = tmp_path(dst);
98    match write_atomic_copy_from_file_inner(&tmp, src, dst) {
99        Ok(()) => Ok(()),
100        Err(e) => {
101            let _ = std::fs::remove_file(&tmp);
102            Err(e)
103        }
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn failed_write_leaves_no_tmp() {
113        // Writing into a non-existent parent directory should fail without
114        // leaving a lingering `.tmp` sibling. Use a tempdir so the test is
115        // hermetic.
116        let tmp = tempfile::tempdir().unwrap();
117        let bogus = tmp.path().join("does/not/exist/target.bin");
118        let err = write_atomic_bytes(&bogus, b"hello").unwrap_err();
119        let tmp_sibling = tmp_path(&bogus);
120        assert!(!tmp_sibling.exists(), "stray tmp: {:?}", tmp_sibling);
121        let _ = err;
122    }
123
124    #[test]
125    fn write_atomic_bytes_round_trip() {
126        let tmp = tempfile::tempdir().unwrap();
127        let path = tmp.path().join("out.bin");
128        write_atomic_bytes(&path, b"hello world").unwrap();
129        assert_eq!(std::fs::read(&path).unwrap(), b"hello world");
130    }
131
132    #[test]
133    fn write_atomic_bytes_overwrites() {
134        let tmp = tempfile::tempdir().unwrap();
135        let path = tmp.path().join("out.bin");
136        write_atomic_bytes(&path, b"v1").unwrap();
137        write_atomic_bytes(&path, b"v2").unwrap();
138        assert_eq!(std::fs::read(&path).unwrap(), b"v2");
139    }
140
141    #[test]
142    fn write_atomic_toml_round_trip() {
143        #[derive(serde::Serialize, serde::Deserialize, PartialEq, Debug)]
144        struct Config {
145            name: String,
146            count: i64,
147        }
148        let tmp = tempfile::tempdir().unwrap();
149        let path = tmp.path().join("cfg.toml");
150        let cfg = Config {
151            name: "test".to_string(),
152            count: 42,
153        };
154        write_atomic_toml(&path, &cfg).unwrap();
155        let content = std::fs::read_to_string(&path).unwrap();
156        assert!(content.contains("name"));
157        assert!(content.contains("test"));
158    }
159}