diff --git a/src/tempdir.rs b/src/tempdir.rs index 9421310..85180d9 100644 --- a/src/tempdir.rs +++ b/src/tempdir.rs @@ -452,4 +452,19 @@ mod tests { assert!(tokio::fs::metadata(dir_path).await.is_err()); Ok(()) } + + #[tokio::test] + async fn test_manual_close() -> Result<(), Error> { + let dir = TempDir::new().await?; + + // The dir exists. + let dir_path = dir.dir_path().clone(); + assert!(tokio::fs::metadata(dir_path.clone()).await.is_ok()); + + // Deletes the directory. + dir.close().await?; + + assert!(tokio::fs::metadata(dir_path).await.is_err()); + Ok(()) + } } diff --git a/src/tempfile.rs b/src/tempfile.rs index 2a681c8..5dbbcd8 100644 --- a/src/tempfile.rs +++ b/src/tempfile.rs @@ -40,6 +40,7 @@ use std::mem::ManuallyDrop; use std::ops::{Deref, DerefMut}; use std::path::PathBuf; use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::task::{Context, Poll}; use tokio::fs::{File, OpenOptions}; @@ -79,6 +80,10 @@ struct TempFileCore { /// If set to `Ownership::Owned`, the file specified in `path` will be deleted /// when this instance is dropped. If set to `Ownership::Borrowed`, the file will be kept. ownership: Ownership, + + /// Indicates whether the file is opened by this handle. + /// This is used to skip closing operations in drop when the file was manually closed. + open: AtomicBool, } impl TempFile { @@ -322,13 +327,10 @@ impl TempFile { /// * `Ok(false)` if the file was not deleted because it is still used /// * `Err(_)` if deletion of the file failed (e.g. due to file locks on Windows) pub async fn close(mut self) -> Result { - // Ensure all file handles are closed before we attempt to delete the file itself via core. drop(unsafe { ManuallyDrop::take(&mut self.file) }); - - // Take core out of self and attempt to unwrap the Arc. This only succeeds if we - // are the last instance pointing at it. let core = unsafe { ManuallyDrop::take(&mut self.core) }; - if let Ok(core) = Arc::try_unwrap(core) { + + if let Ok(core) = Arc::into_inner(core) { core.close().await?; Ok(true) } else { @@ -395,6 +397,7 @@ impl TempFile { async fn new_internal(path: PathBuf, ownership: Ownership) -> Result { let core = TempFileCore { + open: AtomicBool::new(true), file: ManuallyDrop::new( OpenOptions::new() .create(ownership == Ownership::Owned) @@ -436,15 +439,24 @@ impl crate::AsyncClose for TempFile { impl TempFileCore { /// Attempt to close and remove this file. - pub async fn close(mut self) -> Result<(), Error> { + async fn close(&mut self) -> Result { // Ensure we don't drop borrowed directories. if self.ownership != Ownership::Owned { - return Ok(()); + return Ok(false); } - // Closing the file handle first, as otherwise the file might not be deleted. - let _file = unsafe { ManuallyDrop::take(&mut self.file) }; - Ok(tokio::fs::remove_file(&self.path).await?) + // Leave if the file was already closed (somehow). + if !self.open.swap(false, Ordering::SeqCst).is_err() { + return Ok(false); + } + + // Attempt to delete the file. If it fails, reset the marker to "open" in order to + // give the drop another chance. + if let Err(e) = tokio::fs::remove_file(&self.path).await { + self.open.store(true, Ordering::SeqCst).ok(); + return Err(e.into()); + } + Ok(true) } } @@ -456,7 +468,12 @@ impl Drop for TempFile { fn drop(&mut self) { // Ensure all file handles are closed before we attempt to delete the file itself via core. drop(unsafe { ManuallyDrop::take(&mut self.file) }); - drop(unsafe { ManuallyDrop::take(&mut self.core) }); + let core = unsafe { ManuallyDrop::take(&mut self.core) }; + + // Don't drop core manually if it was already handled by the manual drops. + if core.swap(false, Ordering::SeqCst).is_err() { + drop(core) + } } } @@ -473,6 +490,11 @@ impl Drop for TempFileCore { // Closing the file handle first, as otherwise the file might not be deleted. drop(unsafe { ManuallyDrop::take(&mut self.file) }); + // Skip the deletion if the file was already closed manually. + if !self.open.swap(false, Ordering::SeqCst) { + return; + } + // TODO: Use asynchronous variant if running in an async context. // Note that if TempFile is used from the executor's handle, // this may block the executor itself. @@ -591,4 +613,19 @@ mod tests { let name = RandomName::new(FILE_PREFIX); assert!(name.as_ref().starts_with(FILE_PREFIX)) } + + #[tokio::test] + async fn test_manual_close() -> Result<(), Error> { + let file = TempFile::new().await?; + + // The file exists. + let file_path = file.file_path().clone(); + assert!(tokio::fs::metadata(file_path.clone()).await.is_ok()); + + // Deletes the directory. + file.close().await?; + + assert!(tokio::fs::metadata(file_path).await.is_err()); + Ok(()) + } }