feat: implement locked temp file write version

Author: achsanalfitra

Cargo.lock

@@ -87,6 +87,7 @@ email_address = "0.2.9"
 fake = "4"
 fantoccini = "0.22"
 form_urlencoded = "1"
+fs4 = { version = "0.13.1", features = ["tokio", "sync"] }
 futures = { version = "0.3", default-features = false }
 futures-core = { version = "0.3", default-features = false }
 futures-util = { version = "0.3", default-features = false }

Cargo.toml

@@ -34,6 +34,7 @@ digest.workspace = true
 email_address.workspace = true
 fake = { workspace = true, optional = true, features = ["derive", "chrono"] }
 form_urlencoded.workspace = true
+fs4.workspace = true
 futures-core.workspace = true
 futures-util.workspace = true
 hex.workspace = true

cot/Cargo.toml

@@ -47,6 +47,8 @@ use std::path::Path;
 
 use blake3::hash;
 use chrono::{DateTime, Utc};
+use fs4::fs_std::FileExt;
+use fs4::tokio::AsyncFileExt;
 use serde_json::Value;
 use thiserror::Error;
 use tokio::fs::OpenOptions;
@@ -158,11 +160,14 @@ impl FileStore {
         std::fs::create_dir_all(&self.dir_path)
             .map_err(|e| FileCacheStoreError::DirCreation(Box::new(e)))?;
 
-        // try to cleanup the leftovers .tmp files due to crash
         if let Ok(entries) = std::fs::read_dir(&self.dir_path) {
             for entry in entries.flatten() {
                 let path = entry.path();
-                if path.extension().is_some_and(|ext| ext == TEMPFILE_SUFFIX) {
+
+                if path.extension().is_some_and(|ext| ext == TEMPFILE_SUFFIX)
+                    && let Ok(file) = std::fs::File::open(&path)
+                    && file.try_lock_exclusive().is_ok()
+                {
                     let _ = std::fs::remove_file(path);
                 }
             }
@@ -190,11 +195,22 @@ impl FileStore {
         file.sync_data()
             .await
             .map_err(|e| FileCacheStoreError::Io(Box::new(e)))?;
-        tokio::fs::rename(file_path, self.dir_path.join(&key_hash))
+
+        file.unlock_async()
             .await
             .map_err(|e| FileCacheStoreError::Io(Box::new(e)))?;
 
-        Ok(())
+        match tokio::fs::rename(&file_path, self.dir_path.join(&key_hash)).await {
+            Ok(()) => Ok(()),
+            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
+                // if the source temp file is gone, it means another thread
+                // likely renamed it already or the directory was cleared.
+                // In a thundering herd, this is actually a "success" state
+                // because the data is already there.
+                Ok(())
+            }
+            Err(e) => Err(FileCacheStoreError::Io(Box::new(e)))?,
+        }
     }
 
     async fn read(&self, key: &str) -> CacheStoreResult<Option<Value>> {
@@ -324,36 +340,33 @@ impl FileStore {
         &self,
         key_hash: &str,
     ) -> CacheStoreResult<(tokio::fs::File, std::path::PathBuf)> {
-        let mut temp_path = self.dir_path.join(format!("{key_hash}.{TEMPFILE_SUFFIX}"));
+        let temp_path = self.dir_path.join(format!("{key_hash}.{TEMPFILE_SUFFIX}"));
 
         let temp_file = loop {
             match OpenOptions::new()
                 .write(true)
-                .create_new(true)
-                .truncate(true)
+                .read(true)
+                .create(true)
+                .truncate(false)
                 .open(&temp_path)
                 .await
             {
                 Ok(handle) => {
-                    break handle;
+                    if handle.try_lock_exclusive().is_ok() {
+                        handle
+                            .set_len(0)
+                            .await
+                            .map_err(|e| FileCacheStoreError::TempFileCreation(Box::new(e)))?;
+
+                        break handle;
+                    }
                 }
                 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                     // if this fails then we should bail regardless
                     self.create_dir_root()
                         .await
                         .map_err(|e| FileCacheStoreError::DirCreation(Box::new(e)))?;
                 }
-                Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
-                    // we continuously append a .tmp to a the filename
-                    // until we find an empty spot
-                    let file_name = temp_path
-                        .file_name()
-                        .ok_or_else(|| FileCacheStoreError::TempFileCreation(Box::new(e)))?;
-                    let mut os_name = file_name.to_os_string();
-                    os_name.push(format!(".{TEMPFILE_SUFFIX}"));
-
-                    temp_path.set_file_name(os_name);
-                }
                 Err(e) => {
                     return Err(FileCacheStoreError::TempFileCreation(Box::new(e)))?;
                 }
@@ -660,6 +673,52 @@ mod tests {
         let _ = tokio::fs::remove_dir_all(&path).await;
     }
 
+    #[cot::test]
+    async fn test_interference_during_write() {
+        let path = make_store_path();
+        let store = FileStore::new(path.clone()).expect("failed to init store");
+
+        let key = "test_key".to_string();
+        let value = serde_json::json!({ "id": 1, "message": "hello world" });
+
+        let num_task = 10;
+        let barrier = Arc::new(Barrier::new(num_task + 1));
+        let mut handles = Vec::with_capacity(num_task);
+
+        for _ in 0..num_task {
+            let b = barrier.clone();
+            let k = key.clone();
+            let s = store.clone();
+            let v = value.clone();
+
+            handles.push(tokio::spawn(async move {
+                b.wait().await;
+                s.insert(k, v, Timeout::Never)
+                    .await
+                    .expect("failed to insert data to store");
+
+                sleep(Duration::from_millis(10)).await;
+            }));
+        }
+
+        barrier.wait().await;
+
+        tokio::task::yield_now().await;
+
+        let _store_2 = FileStore::new(path.clone()).expect("failed to init store");
+
+        for handle in handles {
+            handle.await.expect("task panicked");
+        }
+
+        let retrieved = store.read(&key).await.expect("failed to read from store");
+        if let Some(found) = retrieved {
+            assert_eq!(found, value);
+        }
+
+        let _ = tokio::fs::remove_dir_all(&path).await;
+    }
+
     #[cot::test]
     async fn test_thundering_write() {
         let path = make_store_path();