Copybara import of the project:

--
c9acc0a6f677951db34f2b607bd60fc18a43f72b by Levi Burner <leviburner@gmail.com>:

Fix race condition in Python viewers set_X methods

--
b89ab8f7a0f628ee7a02ebfbac2a97f113112076 by Levi Burner <leviburner@gmail.com>:

use std::swap to replace some copying

--
02625bed0cffbfc9ea4537f74689f0bbcb8bf44c by Levi Burner <leviburner@gmail.com>:

fix whitespace

COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/2613 from aftersomemath:simulate-set-race 02625bed0cffbfc9ea4537f74689f0bbcb8bf44c
PiperOrigin-RevId: 766646168
Change-Id: I1e1dc16afbfb1e69fb958d54bd1075cc9aed3569

python/mujoco/simulate.cc

@@ -96,7 +96,6 @@ class SimulateWrapper {
 
   void Destroy() {
     if (simulate_) {
-      ClearImages();
       delete simulate_;
       simulate_ = nullptr;
       destroyed_.store(1);
@@ -140,40 +139,71 @@ class SimulateWrapper {
 
   void SetFigures(
       const std::vector<std::pair<mjrRect, py::object>>& viewports_figures) {
-    // Pairs of [viewport, figure], where viewport corresponds to the location
-    // of the figure on the viewer window.
-    std::vector<std::pair<mjrRect, mjvFigure>> user_figures;
-    for (const auto& [viewport, figure] : viewports_figures) {
-      mjvFigure casted_figure = *figure.cast<MjvFigureWrapper&>().get();
-      user_figures.push_back(std::make_pair(viewport, casted_figure));
+
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newfigurerequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
     }
 
-    // Set them all at once to prevent figure flickering.
-    simulate_->user_figures_ = user_figures;
+    // Pairs of [viewport, figure], where viewport corresponds to the location
+    // of the figure on the viewer window.
+    for (const auto& [viewport, figure] : viewports_figures) {
+      mjvFigure casted_figure = *figure.cast<MjvFigureWrapper&>().get();
+      simulate_->user_figures_new_.push_back(std::make_pair(viewport, casted_figure));
+    }
+
+    int value = 0;
+    simulate_->newfigurerequest.compare_exchange_strong(value, 1);
   }
 
-  void ClearFigures() { simulate_->user_figures_.clear(); }
+  void ClearFigures() {
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newfigurerequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+
+    simulate_->user_figures_new_.clear();
+
+    int value = 0;
+    simulate_->newfigurerequest.compare_exchange_strong(value, 1);
+  }
 
   void SetTexts(
       const std::vector<std::tuple<int, int, std::string, std::string>>&
           texts) {
-    // Collection of [font, gridpos, text1, text2] tuples for overlay text
-    std::vector<std::tuple<int, int, std::string, std::string>> user_texts;
-    for (const auto& [font, gridpos, text1, text2] : texts) {
-      user_texts.push_back(std::make_tuple(font, gridpos, text1, text2));
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newtextrequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
     }
 
-    // Set them all at once to prevent text flickering.
-    simulate_->user_texts_ = user_texts;
+    // Collection of [font, gridpos, text1, text2] tuples for overlay text
+    for (const auto& [font, gridpos, text1, text2] : texts) {
+      simulate_->user_texts_new_.push_back(std::make_tuple(font, gridpos, text1, text2));
+    }
+
+    int value = 0;
+    simulate_->newtextrequest.compare_exchange_strong(value, 1);
   }
 
-  void ClearTexts() { simulate_->user_texts_.clear(); }
+  void ClearTexts() {
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newtextrequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
+
+    simulate_->user_texts_new_.clear();
+
+    int value = 0;
+    simulate_->newtextrequest.compare_exchange_strong(value, 1);
+  }
 
   void SetImages(
     const std::vector<std::tuple<mjrRect, pybind11::array&>> viewports_images
   ) {
-    // Clear previous images to prevent memory leaks
-    ClearImages();
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newimagerequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
+    }
 
     for (const auto& [viewport, image] : viewports_images) {
       auto buf = image.request();
@@ -192,20 +222,28 @@ class SimulateWrapper {
       size_t width = buf.shape[1];
       size_t size = height * width * 3;
 
-      // Make a copy of the image data to prevent flickering
-      unsigned char* image_copy = new unsigned char[size];
-      std::memcpy(image_copy, buf.ptr, size);
+      // Make a copy of the image data since Python is
+      // not required to keep it
+      std::unique_ptr<unsigned char[]> image_copy(new unsigned char[size]());
+      std::memcpy(image_copy.get(), buf.ptr, size);
 
-      simulate_->user_images_.push_back(std::make_tuple(viewport, image_copy));
+      simulate_->user_images_new_.push_back(std::make_tuple(viewport, std::move(image_copy)));
     }
+
+    int value = 0;
+    simulate_->newimagerequest.compare_exchange_strong(value, 1);
   }
 
   void ClearImages() {
-    // Free memory for each image before clearing the vector
-    for (const auto& [viewport, image_ptr] : simulate_->user_images_) {
-      delete[] image_ptr;
+    // TODO: replace with atomic wait when we migrate to C++20
+    while (simulate_ && simulate_->newimagerequest.load() != 0) {
+      std::this_thread::sleep_for(std::chrono::milliseconds(1));
     }
-    simulate_->user_images_.clear();
+
+    simulate_->user_images_new_.clear();
+
+    int value = 0;
+    simulate_->newimagerequest.compare_exchange_strong(value, 1);
   }
 
  private:

simulate/simulate.cc

@@ -2632,18 +2632,36 @@ void Simulate::Render() {
   }
 
   // user figures
+  if (this->newfigurerequest.load() == 1) {
+    this->user_figures_.clear();
+    std::swap(this->user_figures_, this->user_figures_new_);
+    int value = 1;
+    this->newfigurerequest.compare_exchange_strong(value, 0);
+  }
   for (auto& [viewport, figure] : this->user_figures_) {
     ShowFigure(this, viewport, &figure);
   }
 
   // overlay text
+  if (this->newtextrequest.load() == 1) {
+    this->user_texts_.clear();
+    std::swap(this->user_texts_, this->user_texts_new_);
+    int value = 1;
+    this->newtextrequest.compare_exchange_strong(value, 0);
+  }
   for (auto& [font, gridpos, text1, text2] : this->user_texts_) {
     ShowOverlayText(this, rect, font, gridpos, text1, text2);
   }
 
   // user images
+  if (this->newimagerequest.load() == 1) {
+    this->user_images_.clear();
+    std::swap(this->user_images_, this->user_images_new_);
+    int value = 1;
+    this->newimagerequest.compare_exchange_strong(value, 0);
+  }
   for (auto& [viewport, image] : this->user_images_) {
-    ShowImage(this, viewport, image);
+    ShowImage(this, viewport, image.get());
   }
 
   // finalize

simulate/simulate.h

@@ -206,6 +206,9 @@ class Simulate {
   std::atomic_int droploadrequest = 0;
   std::atomic_int screenshotrequest = 0;
   std::atomic_int uiloadrequest = 0;
+  std::atomic_int newfigurerequest = 0;
+  std::atomic_int newtextrequest = 0;
+  std::atomic_int newimagerequest = 0;
 
   // loadrequest
   //   3: display a loading message
@@ -263,8 +266,11 @@ class Simulate {
   mjvScene* user_scn = nullptr;
   mjtByte user_scn_flags_prev_[mjNRNDFLAG];
   std::vector<std::pair<mjrRect, mjvFigure>> user_figures_;
+  std::vector<std::pair<mjrRect, mjvFigure>> user_figures_new_;
   std::vector<std::tuple<int, int, std::string, std::string>> user_texts_;
-  std::vector<std::tuple<mjrRect, unsigned char*>> user_images_;
+  std::vector<std::tuple<int, int, std::string, std::string>> user_texts_new_;
+  std::vector<std::tuple<mjrRect, std::unique_ptr<unsigned char[]>>> user_images_;
+  std::vector<std::tuple<mjrRect, std::unique_ptr<unsigned char[]>>> user_images_new_;
 
   // OpenGL rendering and UI
   int refresh_rate = 60;