// SPDX-License-Identifier: MIT // // Copyright IBM Corp. 2024 use std::{ ffi::OsStr, fs::{rename, File, OpenOptions}, io::{self, Seek, SeekFrom, Write}, os::unix::fs::OpenOptionsExt, path::Path, }; use pv::{Error, FileAccessErrorType, PvCoreError, Result}; /// This type helps to perform atomic operations by writing to a temporary file /// and renaming it to the actual filename when the [`AtomicFile::finish`] /// function is called. If the [`AtomicFile::finish`] function is never called, /// the temporary file is automatically removed when it goes out of scope. It /// utilizes the `renameat2` libc function and its semantics. #[derive(Debug)] pub struct AtomicFile { /// Do not change the order! See comment in [`TempPath::drop`]. file: F, path: TempPath, } impl AsRef for AtomicFile { fn as_ref(&self) -> &F { &self.file } } impl AsMut for AtomicFile { fn as_mut(&mut self) -> &mut F { &mut self.file } } /// Enum used for more verbosity. #[derive(Debug, PartialEq, Eq)] pub enum AtomicFileOperation { /// Replace existing file Replace, /// Do not replace existing file NoReplace, } impl AtomicFile {} impl AtomicFile { /// Creates a new [`AtomicFile`] at the given `output` path using /// `open_options`. /// /// # Errors /// /// This function will return an error if the temporary file could not be /// created. /// /// # Example /// /// ``` /// # use utils::AtomicFile; /// /// let file = AtomicFile::new("test", &mut std::fs::OpenOptions::new()).unwrap(); /// ``` pub fn new>(output: P, open_options: &mut OpenOptions) -> Result { Self::with_extension(output, "part", open_options) } /// Creates a new [`AtomicFile`] at the given `output` path using /// `open_options` and `suffix` as the suffix for the temporary file. /// /// # Errors /// /// This function will return an error if the temporary file could not be /// created. /// /// # Example /// /// ``` /// # use utils::AtomicFile; /// /// let file = AtomicFile::with_extension("test", ".incomplete", &mut std::fs::OpenOptions::new()).unwrap(); /// ``` pub fn with_extension, S: AsRef>( output: P, tmp_suffix: S, open_options: &mut OpenOptions, ) -> Result { let path = TempPath::new(output, tmp_suffix); open_options .read(true) .write(true) .create_new(true) .mode(0o600); match open_options.open(&path.temp_path) { Ok(file) => Ok(Self { path, file }), Err(err) => { let path_copy = path.temp_path.to_path_buf(); path.forget(); Err(Error::PvCore(PvCoreError::FileAccess { ty: FileAccessErrorType::Create, path: path_copy, source: err, })) } } } /// Renames the file to the actual file name. If `overwrite` is set to /// [`AtomicFileOperation::Replace`] the file is overwritten. /// /// # Errors /// /// This function will return an error if the rename operation fails. pub fn finish(mut self, overwrite: AtomicFileOperation) -> Result<()> { self.file.flush()?; self.file.sync_all()?; drop(self.file); self.path.persist(overwrite) } /// Discard all changes and delete the temporary file. /// /// # Errors /// /// This function will return an error if the temporary file could not be /// deleted. pub fn discard(mut self) -> io::Result<()> { self.file.flush()?; self.file.sync_all()?; drop(self.file); self.path.remove_file() } } impl Seek for AtomicFile { /// Seek to the offset (in bytes) in the underlying file. fn seek(&mut self, pos: SeekFrom) -> io::Result { self.as_mut().seek(pos) } } impl Write for AtomicFile { fn write(&mut self, buf: &[u8]) -> io::Result { self.as_mut().write(buf) } fn flush(&mut self) -> io::Result<()> { self.as_mut().flush() } } #[derive(Debug)] struct TempPath { temp_path: Box, path: Box, } impl TempPath { fn new, S: AsRef>(path: P, extension: S) -> Self { let mut temp_path = path.as_ref().to_path_buf(); match temp_path.extension() { Some(ext) => { let mut ext = ext.to_os_string(); ext.push("."); ext.push(extension.as_ref()); temp_path.set_extension(ext) } None => temp_path.set_extension(extension.as_ref()), }; Self { temp_path: temp_path.into(), path: path.as_ref().into(), } } fn forget(self) { // Make sure that [`Self::drop`] is not called. Self { .. } = self; } /// Removes the temporary file pub fn remove_file(self) -> io::Result<()> { let result = std::fs::remove_file(&self.temp_path); self.forget(); result } fn persist(self, operation: AtomicFileOperation) -> Result<()> { if let Ok(true) = self.path.try_exists() { if operation == AtomicFileOperation::NoReplace { return Err(Error::Io(io::Error::from(io::ErrorKind::AlreadyExists))); } } rename(&self.temp_path, &self.path).map_err(|e| PvCoreError::FileAccessRename { src: self.temp_path.as_ref().to_str().unwrap().to_string(), dst: self.path.as_ref().to_str().unwrap().to_string(), source: e, })?; self.forget(); Ok(()) } } impl Drop for TempPath { // When this is called while dropping an [`AtomicFile`], // the field [`AtomicFile::file`] is already dropped. fn drop(&mut self) { let _ = std::fs::remove_file(&self.temp_path); } } #[cfg(test)] mod tests { use std::io::Write; use super::{AtomicFile, AtomicFileOperation}; use crate::TemporaryDirectory; #[test] fn atomicfile_basic_functionality() { let tmp_dir = TemporaryDirectory::new().expect("should work"); let tmp_file = tmp_dir.path().join("atomic_basic"); let data = b"test_data"; assert!(!tmp_file.exists()); let mut writer = AtomicFile::with_extension(&tmp_file, "basic", &mut std::fs::OpenOptions::new()) .unwrap(); writer.write_all(data).unwrap(); writer.finish(AtomicFileOperation::Replace).unwrap(); assert!(tmp_file.exists()); let writer = AtomicFile::with_extension(&tmp_file, "basic", &mut std::fs::OpenOptions::new()) .unwrap(); writer.finish(AtomicFileOperation::NoReplace).expect_err( "Creating the file is expected to fail", ); std::fs::remove_file(tmp_file).expect("should not fail"); } #[test] fn atomicfile_discard() { let tmp_dir = TemporaryDirectory::new().expect("should work"); let tmp_file = tmp_dir.path().join("atomic_close"); let data = b"test_data"; assert!(!tmp_file.exists()); let mut writer = AtomicFile::with_extension(&tmp_file, "close", &mut std::fs::OpenOptions::new()) .unwrap(); writer.write_all(data).unwrap(); writer.discard().expect("bla"); assert!(!tmp_file.exists()); } #[test] fn atomicfile_drop() { let tmp_dir = TemporaryDirectory::new().expect("should work"); let tmp_file = tmp_dir.path().join("atomic_close"); let data = b"test_data"; assert!(!tmp_file.exists()); let mut writer = AtomicFile::with_extension(&tmp_file, "close", &mut std::fs::OpenOptions::new()) .unwrap(); writer.write_all(data).unwrap(); drop(writer); assert!(!tmp_file.exists()); } }