Skip to main content

swh_graph/utils/
atomic_file.rs

1/*
2 * Copyright (C) 2026  The Software Heritage developers
3 * See the AUTHORS file at the top-level directory of this distribution
4 * License: GNU General Public License version 3, or any later version
5 * See top-level LICENSE file for more information
6 */
7
8use std::fs::File;
9use std::io::IoSlice;
10use std::io::Result;
11use std::path::{Path, PathBuf};
12
13/// Wrapper around [`File`] that writes to a temporary file and
14/// moves it to the path when closed
15pub struct AtomicFile {
16    tmp_path: PathBuf,
17    path: PathBuf,
18    /// `None` if already closed
19    file: Option<File>,
20}
21
22impl AtomicFile {
23    pub fn create_new<P: AsRef<Path>>(path: P) -> Result<Self> {
24        use rand::distr::{Alphanumeric, SampleString};
25
26        let path = path.as_ref().to_path_buf();
27        let file_name = path.file_name().unwrap_or_default().to_string_lossy();
28        let random_string = Alphanumeric.sample_string(&mut rand::rng(), 16);
29        let tmp_path = path.with_file_name(format!(".tmp.{file_name}.{random_string}"));
30        Ok(Self {
31            file: Some(File::create_new(&tmp_path)?),
32            tmp_path,
33            path,
34        })
35    }
36
37    pub fn commit(mut self) -> Result<()> {
38        if self.file.is_some() {
39            std::fs::rename(&self.tmp_path, &self.path)?;
40            self.file = None;
41        }
42        Ok(())
43    }
44
45    pub fn rollback(mut self) -> Result<()> {
46        if self.file.is_some() {
47            std::fs::remove_file(&self.tmp_path)?;
48            self.file = None;
49        }
50        Ok(())
51    }
52
53    fn file(&mut self) -> Result<&mut File> {
54        self.file.as_mut().ok_or_else(|| {
55            std::io::Error::other(anyhow::anyhow!(
56                "File {} was already committed to {} and closed",
57                self.tmp_path.display(),
58                self.path.display()
59            ))
60        })
61    }
62}
63
64impl Drop for AtomicFile {
65    fn drop(&mut self) {
66        if self.file.is_some() {
67            let mut other = AtomicFile {
68                tmp_path: PathBuf::from("/"),
69                path: PathBuf::from("/"),
70                file: None,
71            };
72            std::mem::swap(&mut other, self);
73            other
74                .rollback()
75                .expect("Error while removing dropped AtomicFile");
76        }
77    }
78}
79
80impl std::io::Write for AtomicFile {
81    fn write(&mut self, buf: &[u8]) -> Result<usize> {
82        self.file()?.write(buf)
83    }
84    fn flush(&mut self) -> Result<()> {
85        self.file()?.flush()
86    }
87
88    fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> Result<usize> {
89        self.file()?.write_vectored(bufs)
90    }
91    fn write_all(&mut self, buf: &[u8]) -> Result<()> {
92        self.file()?.write_all(buf)
93    }
94    fn write_fmt(&mut self, args: std::fmt::Arguments<'_>) -> Result<()> {
95        self.file()?.write_fmt(args)
96    }
97    fn by_ref(&mut self) -> &mut Self
98    where
99        Self: Sized,
100    {
101        self
102    }
103}