From 0b4843d4514cd8b7e653990025d0ecd5e80d56ba Mon Sep 17 00:00:00 2001
From: Yuhang Wei <weiyuhang3@huawei.com>
Date: Tue, 20 Feb 2024 10:18:27 +0800
Subject: [PATCH 1/2] fix: mutex locking in agent_impl.rs

Signed-off-by: Yuhang Wei <weiyuhang3@huawei.com>
---
 KubeOS-Rust/agent/src/rpc/agent_impl.rs | 32 +++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/KubeOS-Rust/agent/src/rpc/agent_impl.rs b/KubeOS-Rust/agent/src/rpc/agent_impl.rs
index 5f3a3259..ab826413 100644
--- a/KubeOS-Rust/agent/src/rpc/agent_impl.rs
+++ b/KubeOS-Rust/agent/src/rpc/agent_impl.rs
@@ -57,7 +57,10 @@ impl Default for AgentImpl {
 
 impl AgentImpl {
     fn prepare_upgrade_impl(&self, req: UpgradeRequest) -> Result<Response> {
-        let _lock = self.mutex.lock().unwrap();
+        let lock = self.mutex.try_lock();
+        if lock.is_err() {
+            bail!("os-agent is processing another request");
+        }
         debug!("Received an 'prepare upgrade' request: {:?}", req);
         info!("Start preparing for upgrading to version: {}", req.version);
 
@@ -76,7 +79,10 @@ impl AgentImpl {
     }
 
     fn upgrade_impl(&self) -> Result<Response> {
-        let _lock = self.mutex.lock().unwrap();
+        let lock = self.mutex.try_lock();
+        if lock.is_err() {
+            bail!("os-agent is processing another request");
+        }
         info!("Start to upgrade");
         let command_executor = RealCommandExecutor {};
         let (_, next_partition_info) = get_partition_info(&command_executor)?;
@@ -91,7 +97,10 @@ impl AgentImpl {
     }
 
     fn configure_impl(&self, mut req: ConfigureRequest) -> Result<Response> {
-        let _lock = self.mutex.lock().unwrap();
+        let lock = self.mutex.try_lock();
+        if lock.is_err() {
+            bail!("os-agent is processing another request");
+        }
         debug!("Received a 'configure' request: {:?}", req);
         info!("Start to configure");
         let config_map = &*CONFIG_TEMPLATE;
@@ -108,7 +117,10 @@ impl AgentImpl {
     }
 
     fn rollback_impl(&self) -> Result<Response> {
-        let _lock = self.mutex.lock().unwrap();
+        let lock = self.mutex.try_lock();
+        if lock.is_err() {
+            bail!("os-agent is processing another request");
+        }
         info!("Start to rollback");
         let command_executor = RealCommandExecutor {};
         let (_, next_partition_info) = get_partition_info(&command_executor)?;
@@ -172,6 +184,18 @@ mod test {
         };
         let res = agent.configure(req);
         assert!(res.is_err());
+
+        // test lock
+        let _lock = agent.mutex.lock().unwrap();
+        let req = ConfigureRequest {
+            configs: vec![Sysconfig {
+                model: "kernel.sysctl".to_string(),
+                config_path: "".to_string(),
+                contents: HashMap::new(),
+            }],
+        };
+        let res = agent.configure(req);
+        assert!(res.is_err());
     }
 
     #[test]
-- 
2.34.1