Allow overriding tmpdir

Author: teohhanhui

src/fs_store.rs

@@ -66,19 +66,43 @@ use error::ErrorKind;
 /// half finished snapshot file. Deleting the inputs after the rename means that
 /// the worst case is that we have some leftover incremental files which will
 /// be deleted on the next compaction.
+#[derive(Debug)]
 pub struct FsStore {
     root: std::path::PathBuf,
+    tmpdir: std::path::PathBuf,
 }
 
 impl FsStore {
-    /// Create an [`FsStore`] from a [`std::path::PathBuf`]
+    /// Creates a new [`FsStore`] from a [`Path`].
     ///
     /// This will attempt to create the root directory and throw an error if
     /// it does not exist.
     pub fn open<P: AsRef<Path>>(root: P) -> Result<Self, std::io::Error> {
-        std::fs::create_dir_all(root.as_ref())?;
+        let root = root.as_ref();
+        std::fs::create_dir_all(root)?;
         Ok(Self {
-            root: root.as_ref().into(),
+            root: root.into(),
+            tmpdir: root.into(),
+        })
+    }
+
+    /// Overrides the tmpdir directory used for temporary files.
+    ///
+    /// The default is to use the root directory passed to [`FsStore::open`].
+    ///
+    /// The tmpdir used must be on the same mount point as the root directory,
+    /// otherwise the store will throw an error on writing data.
+    ///
+    /// # Errors
+    ///
+    /// This will attempt to create the tmpdir directory and throw an error if
+    /// it does not exist.
+    pub fn with_tmpdir<P: AsRef<Path>>(self, tmpdir: P) -> Result<Self, std::io::Error> {
+        let tmpdir = tmpdir.as_ref();
+        std::fs::create_dir_all(tmpdir)?;
+        Ok(Self {
+            tmpdir: tmpdir.into(),
+            ..self
         })
     }
 
@@ -151,7 +175,7 @@ impl FsStore {
         })?;
 
         let chunk_name = SavedChunkName::new_incremental(changes);
-        write_chunk(&self.root, &paths, changes, chunk_name)?;
+        write_chunk(&self.root, &paths, changes, chunk_name, &self.tmpdir)?;
 
         Ok(())
     }
@@ -171,7 +195,7 @@ impl FsStore {
         // Write the snapshot
         let output_chunk_name = SavedChunkName::new_snapshot(doc.get_heads());
         let chunk = doc.save();
-        write_chunk(&self.root, &paths, &chunk, output_chunk_name)?;
+        write_chunk(&self.root, &paths, &chunk, output_chunk_name, &self.tmpdir)?;
 
         // Remove all the old data
         for incremental in chunks.incrementals.keys() {
@@ -191,10 +215,11 @@ fn write_chunk(
     paths: &DocIdPaths,
     chunk: &[u8],
     name: SavedChunkName,
+    tmpdir: &Path,
 ) -> Result<(), Error> {
     // Write to a temp file and then rename to avoid partial writes
     let temp_dir =
-        tempfile::TempDir::new_in(root).map_err(|e| Error(ErrorKind::CreateTempFile(e)))?;
+        tempfile::TempDir::new_in(tmpdir).map_err(|e| Error(ErrorKind::CreateTempFile(e)))?;
     let temp_save_path = temp_dir.path().join(name.filename());
     let mut temp_save_file =
         File::create(&temp_save_path).map_err(|e| Error(ErrorKind::CreateTempFile(e)))?;

src/tokio/fs_storage.rs

@@ -1,17 +1,21 @@
-use std::sync::{Arc, Mutex};
+use std::{
+    path::Path,
+    sync::{Arc, Mutex},
+};
 
 use futures::{future::BoxFuture, FutureExt, TryFutureExt};
 
 use crate::{DocumentId, Storage, StorageError};
 
 /// A wrapper around [`crate::fs_store::FsStore`] that implements [`crate::Storage`]
+#[derive(Debug)]
 pub struct FsStorage {
     inner: Arc<Mutex<crate::fs_store::FsStore>>,
     handle: tokio::runtime::Handle,
 }
 
 impl FsStorage {
-    /// Create a new [`FsStorage`] from a [`std::path::PathBuf`]
+    /// Creates a new [`FsStorage`] from a [`Path`].
     ///
     /// # Errors
     ///
@@ -21,11 +25,34 @@ impl FsStorage {
     /// # Panics
     ///
     /// If there is not a tokio runtime available
-    pub fn open(root: std::path::PathBuf) -> Result<Self, std::io::Error> {
+    pub fn open<P: AsRef<Path>>(root: P) -> Result<Self, std::io::Error> {
         let handle = tokio::runtime::Handle::current();
         let inner = Arc::new(Mutex::new(crate::fs_store::FsStore::open(root)?));
         Ok(Self { inner, handle })
     }
+
+    /// Overrides the tmpdir directory used for temporary files.
+    ///
+    /// The default is to use the root directory passed to [`FsStorage::open`].
+    ///
+    /// The tmpdir used must be on the same mount point as the root directory,
+    /// otherwise the storage will throw an error on writing data.
+    ///
+    /// # Errors
+    ///
+    /// This will attempt to create the tmpdir directory and throw an error if
+    /// it does not exist.
+    pub fn with_tmpdir<P: AsRef<Path>>(self, tmpdir: P) -> Option<Result<Self, std::io::Error>> {
+        let Self { inner, handle } = self;
+        let inner = Arc::into_inner(inner)?.into_inner().ok()?;
+        let inner = inner.with_tmpdir(tmpdir);
+        let Ok(inner) = inner else {
+            let e = inner.unwrap_err();
+            return Some(Err(e));
+        };
+        let inner = Arc::new(Mutex::new(inner));
+        Some(Ok(Self { inner, handle }))
+    }
 }
 
 impl Storage for FsStorage {