Fix undefined behavior on simultaneous write in daemon state file

cw-orch-daemon/Cargo.toml

@@ -67,6 +67,10 @@ async-recursion = "1.0.5"
 # Gzip
 flate2 = { version = "1.0.26" }
 
+# Lock daemon
+file-lock = { version = "2.1.10" }
+once_cell = { version = "1.19.0" }
+
 [dev-dependencies]
 cw-orch-daemon = { path = "." }
 uid = "0.1.7"

cw-orch-daemon/src/json_file.rs

@@ -1,8 +1,88 @@
 use crate::DaemonError;
+use file_lock::{FileLock, FileOptions};
 use serde_json::{from_reader, json, Value};
-use std::fs::{File, OpenOptions};
+use std::{
+    fs::{File, OpenOptions},
+    io::Seek,
+};
+
+pub struct JsonFileState {
+    lock: FileLock,
+    json: Value,
+}
+
+impl JsonFileState {
+    /// Lock a new file
+    pub fn new(filename: &str) -> Self {
+        // open file pointer set read/write permissions to true
+        // create it if it does not exists
+        // dont truncate it
+
+        let options = FileOptions::new()
+            .create(true)
+            .read(true)
+            .write(true)
+            .truncate(false);
+
+        let lock: FileLock = FileLock::lock(filename, true, options).unwrap();
+
+        // return empty json object if file is empty
+        // return file content if not
+        let json: Value = if lock.file.metadata().unwrap().len().eq(&0) {
+            json!({})
+        } else {
+            from_reader(&lock.file).unwrap()
+        };
+
+        JsonFileState { lock, json }
+    }
+
+    /// Prepare json for further writes
+    pub fn prepare(&mut self, chain_id: &str, network_id: &str, deploy_id: &str) {
+        let json = &mut self.json;
+        // check and add network_id path if it's missing
+        if json.get(network_id).is_none() {
+            json[network_id] = json!({});
+        }
+
+        // add deployment_id to chain_id path
+        if json[network_id].get(chain_id).is_none() {
+            json[network_id][chain_id] = json!({
+                deploy_id: {},
+                "code_ids": {}
+            });
+        }
+    }
+
+    pub fn state(&self) -> Value {
+        self.json.clone()
+    }
+
+    /// Get a value for read
+    // pub fn get(&self, chain_id: &str, network_id: &str) -> &Value {
+    //     self.json[network_id].get(chain_id).unwrap()
+    // }
+
+    /// Give a value to write
+    pub fn get_mut(&mut self, chain_id: &str, network_id: &str) -> &mut Value {
+        self.json[network_id].get_mut(chain_id).unwrap()
+    }
+
+    /// Force write to a file
+    pub fn force_write(&mut self) {
+        self.lock.file.rewind().unwrap();
+        serde_json::to_writer_pretty(&self.lock.file, &self.json).unwrap();
+    }
+}
+
+// Write json when dropping
+impl Drop for JsonFileState {
+    fn drop(&mut self) {
+        self.force_write()
+    }
+}
 
 pub fn write(filename: &String, chain_id: &String, network_id: &String, deploy_id: &String) {
     // open file pointer set read/write permissions to true
     // create it if it does not exists
     // dont truncate it

cw-orch-daemon/src/state.rs

@@ -1,5 +1,5 @@
 use super::error::DaemonError;
-use crate::{channel::GrpcChannel, networks::ChainKind};
+use crate::{channel::GrpcChannel, json_file::JsonFileState, networks::ChainKind};
 
 use cosmwasm_std::Addr;
 use cw_orch_core::{
@@ -9,17 +9,21 @@
     CwEnvError, CwOrchEnvVars,
 };
 use ibc_chain_registry::chain::ChainData;
+use once_cell::sync::Lazy;
 use serde::Serialize;
 use serde_json::{json, Value};
-use std::{collections::HashMap, fs::File, path::Path};
+use std::{collections::HashMap, path::Path, sync::Mutex};
 use tonic::transport::Channel;
 
+pub(crate) static GLOBAL_WRITE_STATE: Lazy<Mutex<HashMap<String, (u64, JsonFileState)>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
 /// Stores the chain information and deployment state.
 /// Uses a simple JSON file to store the deployment information locally.
-#[derive(Clone, Debug)]
+#[derive(Debug)]
 pub struct DaemonState {
     /// this is passed via env var STATE_FILE
-    pub json_file_path: String,
+    json_file_path: String,
     /// Deployment identifier
     pub deployment_id: String,
     /// gRPC channel
@@ -30,6 +34,24 @@
     pub read_only: bool,
 }
 
+impl Clone for DaemonState {
+    fn clone(&self) -> Self {
+        let new_self = Self {
+            json_file_path: self.json_file_path.clone(),
+            deployment_id: self.deployment_id.clone(),
+            grpc_channel: self.grpc_channel.clone(),
+            chain_data: self.chain_data.clone(),
+            read_only: self.read_only.clone(),
+        };
+        // Increase DaemonStates count for this file
+        let mut lock = GLOBAL_WRITE_STATE.lock().unwrap();
+        let (count, _) = lock.get_mut(&self.json_file_path).unwrap();
+        *count += 1;
+
+        new_self
+    }
+}
+
 impl DaemonState {
     /// Creates a new state from the given chain data and deployment id.
     /// Attempts to connect to any of the provided gRPC endpoints.
@@ -85,7 +107,7 @@
 
         // build daemon state
         let state = DaemonState {
-            json_file_path,
+            json_file_path: json_file_path.clone(),
             deployment_id,
             grpc_channel,
             chain_data,
@@ -99,10 +121,23 @@
                 state.json_file_path
             );
 
-            // write json state file
-            crate::json_file::write(
-                &state.json_file_path,
-                &state.chain_data.chain_id.to_string(),
+            let mut lock = GLOBAL_WRITE_STATE.lock().unwrap();
+            // Lock file if first time write
+            let file_state = match lock.entry(json_file_path.clone()) {
+                // Increase count if already locked this file
+                std::collections::hash_map::Entry::Occupied(o) => {
+                    let (count, lock) = o.into_mut();
+                    *count += 1;
+                    lock
+                }
+                // Insert as 1 count of DaemonStates pointing to this file if it's first open
+                std::collections::hash_map::Entry::Vacant(v) => {
+                    let (_, lock) = v.insert((1, JsonFileState::new(&json_file_path)));
+                    lock
+                }
+            };
+            file_state.prepare(
+                state.chain_data.chain_id.as_str(),
                 &state.chain_data.chain_name,
                 &state.deployment_id,
             );
@@ -148,7 +183,14 @@
     }
     /// Get the state filepath and read it as json
     fn read_state(&self) -> Result<serde_json::Value, DaemonError> {
-        crate::json_file::read(&self.json_file_path)
+        // Check if already open in write mode {
+        let lock = GLOBAL_WRITE_STATE.lock().unwrap();
+        if let Some((_, j)) = lock.get(&self.json_file_path) {
+            Ok(j.state())
+        } else {
+            // Or just read it from a file
+            crate::json_file::read(&self.json_file_path)
+        }
     }
 
     /// Retrieve a stateful value using the chainId and networkId
@@ -168,14 +210,40 @@
             return Err(DaemonError::StateReadOnly);
         }
 
-        let mut json = self.read_state()?;
+        let mut lock = GLOBAL_WRITE_STATE.lock().unwrap();
 
-        json[&self.chain_data.chain_name][&self.chain_data.chain_id.to_string()][key]
-            [contract_id] = json!(value);
+        let (_, file_state) = lock.get_mut(&self.json_file_path).unwrap();
+        let val = file_state.get_mut(
+            self.chain_data.chain_id.as_str(),
+            &self.chain_data.chain_name,
+        );
+        val[key][contract_id] = json!(value);
 
-        serde_json::to_writer_pretty(File::create(&self.json_file_path).unwrap(), &json)?;
         Ok(())
     }
+
+    /// Forcefully write current json to a file
+    pub fn force_write(&self) -> Result<(), DaemonError> {
+        let mut lock = GLOBAL_WRITE_STATE.lock().unwrap();
+        let (_, file_state) = lock.get_mut(&self.json_file_path).unwrap();
+        file_state.force_write();
+        Ok(())
+    }
+}
+
+impl Drop for DaemonState {
+    fn drop(&mut self) {
+        let mut lock = GLOBAL_WRITE_STATE.lock().unwrap();
+
+        // Decrease open count
+        let (count, _) = lock.get_mut(&self.json_file_path).unwrap();
+        *count -= 1;
+
+        // If we get to zero count - write to a file
+        if *count == 0 {
+            lock.remove(&self.json_file_path);
+        }
+    }
 }
 
 impl StateInterface for DaemonState {

cw-orch-daemon/tests/daemon_state.rs

@@ -0,0 +1,61 @@
+use cw_orch_core::env::STATE_FILE_ENV_NAME;
+use cw_orch_daemon::{ChainRegistryData, DaemonState};
+use cw_orch_networks::networks::JUNO_1;
+use tokio::runtime::Runtime;
+
+#[test]
+fn simultaneous_read() {
+    let runtime = Runtime::new().unwrap();
+
+    let chain_data: ChainRegistryData = JUNO_1.into();
+    std::env::set_var(STATE_FILE_ENV_NAME, "./tests/test.json");
+
+    let daemon_state = runtime
+        .block_on(DaemonState::new(chain_data, "test".to_owned(), false))
+        .unwrap();
+    daemon_state.set("test", "test", "test").unwrap();
+
+    let mut handles = vec![];
+    for _ in 0..25 {
+        let daemon_state = daemon_state.clone();
+        let handle = std::thread::spawn(move || daemon_state.get("test").unwrap());
+        handles.push(handle);
+    }
+    for handle in handles {
+        handle.join().unwrap();
+    }
+}
+
+#[test]
+fn simultaneous_write() {
+    let runtime = Runtime::new().unwrap();
+
+    let chain_data: ChainRegistryData = JUNO_1.into();
+    std::env::set_var(STATE_FILE_ENV_NAME, "./tests/test.json");
+
+    let daemon_state = runtime
+        .block_on(DaemonState::new(chain_data, "test".to_owned(), false))
+        .unwrap();
+
+    let mut handles = vec![];
+    for i in 0..25 {
+        let daemon_state = daemon_state.clone();
+        let handle = std::thread::spawn(move || {
+            daemon_state
+                .set("test", &format!("test{i}"), format!("test-{i}"))
+                .unwrap();
+        });
+        handles.push(handle);
+    }
+
+    let mut maybe_err = Ok(());
+    // Finish all handles
+    for handle in handles {
+        let result = handle.join();
+        if result.is_err() {
+            maybe_err = result;
+        }
+    }
+    // Error if at least one failed
+    maybe_err.unwrap()
+}