From 26166fb485686f881c6c937cbef23b5fd03cbd95 Mon Sep 17 00:00:00 2001
From: fan-ziqi <fanziqi614@gmail.com>
Date: Wed, 16 Oct 2024 14:46:37 +0800
Subject: [PATCH] fix: add mutex protection for robot_date to ensure thread
 safety in multi-threaded environments

---
 src/rl_sar/library/rl_sdk/rl_sdk.hpp | 2 ++
 src/rl_sar/src/rl_real_a1.cpp        | 4 ++++
 src/rl_sar/src/rl_real_go2.cpp       | 4 ++++
 src/rl_sar/src/rl_sim.cpp            | 4 ++++
 4 files changed, 14 insertions(+)

diff --git a/src/rl_sar/library/rl_sdk/rl_sdk.hpp b/src/rl_sar/library/rl_sdk/rl_sdk.hpp
index 47ac2d2..5f86f76 100644
--- a/src/rl_sar/library/rl_sdk/rl_sdk.hpp
+++ b/src/rl_sar/library/rl_sdk/rl_sdk.hpp
@@ -5,6 +5,7 @@
 #include <iostream>
 #include <string>
 #include <unistd.h>
+#include <mutex>
 
 #include <yaml-cpp/yaml.h>
 
@@ -123,6 +124,7 @@ public:
 
     RobotState<double> robot_state;
     RobotCommand<double> robot_command;
+    std::mutex robot_state_mutex;
 
     // init
     void InitObservations();
diff --git a/src/rl_sar/src/rl_real_a1.cpp b/src/rl_sar/src/rl_real_a1.cpp
index f9dd0a8..1052954 100644
--- a/src/rl_sar/src/rl_real_a1.cpp
+++ b/src/rl_sar/src/rl_real_a1.cpp
@@ -139,6 +139,8 @@ void RL_Real::SetCommand(const RobotCommand<double> *command)
 
 void RL_Real::RobotControl()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     this->motiontime++;
 
     this->GetState(&this->robot_state);
@@ -148,6 +150,8 @@ void RL_Real::RobotControl()
 
 void RL_Real::RunModel()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     if (this->running_state == STATE_RL_RUNNING)
     {
         this->obs.ang_vel = torch::tensor(this->robot_state.imu.gyroscope).unsqueeze(0);
diff --git a/src/rl_sar/src/rl_real_go2.cpp b/src/rl_sar/src/rl_real_go2.cpp
index b4171b4..f414b1d 100644
--- a/src/rl_sar/src/rl_real_go2.cpp
+++ b/src/rl_sar/src/rl_real_go2.cpp
@@ -144,6 +144,8 @@ void RL_Real::SetCommand(const RobotCommand<double> *command)
 
 void RL_Real::RobotControl()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     this->motiontime++;
 
     this->GetState(&this->robot_state);
@@ -153,6 +155,8 @@ void RL_Real::RobotControl()
 
 void RL_Real::RunModel()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     if (this->running_state == STATE_RL_RUNNING)
     {
         this->obs.ang_vel = torch::tensor(this->robot_state.imu.gyroscope).unsqueeze(0);
diff --git a/src/rl_sar/src/rl_sim.cpp b/src/rl_sar/src/rl_sim.cpp
index 887c6ba..69ebcc2 100644
--- a/src/rl_sar/src/rl_sim.cpp
+++ b/src/rl_sar/src/rl_sim.cpp
@@ -153,6 +153,8 @@ void RL_Sim::SetCommand(const RobotCommand<double> *command)
 
 void RL_Sim::RobotControl()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     if (this->control.control_state == STATE_RESET_SIMULATION)
     {
         gazebo_msgs::SetModelState set_model_state;
@@ -216,6 +218,8 @@ void RL_Sim::JointStatesCallback(const sensor_msgs::JointState::ConstPtr &msg)
 
 void RL_Sim::RunModel()
 {
+    std::lock_guard<std::mutex> lock(robot_state_mutex);
+
     if (this->running_state == STATE_RL_RUNNING && simulation_running)
     {
         this->obs.lin_vel = torch::tensor({{this->vel.linear.x, this->vel.linear.y, this->vel.linear.z}});