Make MuJoCo Python bindings compatible with free-threading.

Introduce a new header `gil.h` defining `MutexLockIfGilDisabled` to support thread-safety in both standard and free-threaded CPython builds.

Protect critical shared states and registries:
- Guard global Python callback pointers in `callbacks.cc` using a mutex. Move `gil_scoped_acquire` into local blocks around refcount modifications to prevent `longjmp` from bypassing destructors.
- Protect raw pointer maps in `structs_wrappers.cc` with static mutexes.
- Replace TOCTOU race in `mjcb_time` initialization with thread-safe `std::call_once`.
- Add synchronization to lazy indexer array cache initialization in `indexers.cc` and `indexer_xmacro.h`.
- Protect vector mutations in `StructListBase::PopulateUpTo` in `structs.h` with a mutex.
- Revert unnecessary atomic changes to threadpool counters.
- Declare free-threading compatibility by passing `pybind11::mod_gil_not_used()` to all extension modules.

Fixes #3259
Fixes #3256
Fixes #2978

PiperOrigin-RevId: 941101502
Change-Id: Iec4ce58afcbc75d4b0be6a9a21fc8a47854242e3
This commit is contained in:
Saran Tunyasuvunakool
2026-07-01 08:08:10 -07:00
committed by Copybara-Service
parent cab191755a
commit a07ae6f849
21 changed files with 489 additions and 66 deletions
+11 -1
View File
@@ -16,6 +16,7 @@
#define MUJOCO_PYTHON_STRUCTS_H_
#include <algorithm>
#include <mutex>
#include <array>
#include <cctype>
#include <cstddef>
@@ -35,6 +36,7 @@
#include <absl/types/span.h>
#include <mujoco/mujoco.h>
#include <mujoco/mjxmacro.h>
#include "gil.h"
#include "indexers.h"
#include "raw.h"
#include <pybind11/numpy.h>
@@ -154,7 +156,12 @@ class StructListBase {
}
StructListBase(const StructListBase& other) = delete;
StructListBase(StructListBase&& other) = default;
StructListBase(StructListBase&& other)
: ptr_(other.ptr_),
num_(other.num_),
owner_(std::move(other.owner_)),
wrappers_(std::move(other.wrappers_)) {}
// populate_mutex_ is default-constructed (std::mutex is not movable)
virtual ~StructListBase() = default;
@@ -175,6 +182,8 @@ class StructListBase {
protected:
void PopulateUpTo(int n) {
MutexLockIfGilDisabled lock(populate_mutex_);
wrappers_.reserve(n + 1);
while (wrappers_.size() <= n) {
wrappers_.push_back(
std::make_shared<MjWrapper<T>>(&ptr_[wrappers_.size()], owner_));
@@ -202,6 +211,7 @@ class StructListBase {
// Using shared_ptr here so that we get identical Python objects when slicing.
std::vector<std::shared_ptr<MjWrapper<T>>> wrappers_;
mutable std::mutex populate_mutex_;
};
template <typename T>