Saran Tunyasuvunakool
a07ae6f849
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
209 lines
6.4 KiB
Python
209 lines
6.4 KiB
Python
# Copyright 2026 DeepMind Technologies Limited
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Thread-safety tests for MuJoCo Python bindings.
|
|
|
|
Validates that the free-threading mutex guards in the bindings work correctly:
|
|
- Concurrent access to lazy-initialized indexers and struct lists.
|
|
- Concurrent callback setter/getter operations.
|
|
- Concurrent MjModel/MjData creation and destruction.
|
|
- Re-entrancy safety (no deadlock when __del__ re-enters callback setters).
|
|
"""
|
|
|
|
import threading
|
|
|
|
from absl.testing import absltest
|
|
import mujoco
|
|
|
|
|
|
# A model with contacts so data.contact is non-empty after stepping.
|
|
_CONTACT_XML = r"""
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom type="plane" size="1 1 0.1"/>
|
|
<body pos="0 0 0.05">
|
|
<freejoint/>
|
|
<geom type="sphere" size="0.05"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
_NUM_THREADS = 8
|
|
_ITERS_PER_THREAD = 50
|
|
|
|
|
|
class ConcurrentStressTest(absltest.TestCase):
|
|
"""Tests concurrent access to bindings under multiple threads."""
|
|
|
|
def test_concurrent_indexer_access(self):
|
|
"""Concurrent first-access of lazy indexers must not corrupt state."""
|
|
model = mujoco.MjModel.from_xml_string(_CONTACT_XML)
|
|
data = mujoco.MjData(model)
|
|
mujoco.mj_step(model, data)
|
|
errors = []
|
|
|
|
def access_indexers():
|
|
try:
|
|
for _ in range(_ITERS_PER_THREAD):
|
|
# Each of these triggers lazy init on first access.
|
|
_ = model.geom(0)
|
|
_ = data.qpos
|
|
_ = data.xpos
|
|
except Exception as e: # pylint: disable=broad-except
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=access_indexers)
|
|
for _ in range(_NUM_THREADS)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
self.assertEmpty(errors, f"Errors in concurrent indexer access: {errors}")
|
|
|
|
def test_concurrent_struct_list_access(self):
|
|
"""Concurrent access to contact struct lists must not crash."""
|
|
model = mujoco.MjModel.from_xml_string(_CONTACT_XML)
|
|
data = mujoco.MjData(model)
|
|
# Step to generate contacts.
|
|
mujoco.mj_step(model, data)
|
|
errors = []
|
|
|
|
def access_contacts():
|
|
try:
|
|
for _ in range(_ITERS_PER_THREAD):
|
|
ncon = data.ncon
|
|
if ncon > 0:
|
|
contacts = data.contact[:ncon]
|
|
_ = len(contacts)
|
|
except Exception as e: # pylint: disable=broad-except
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=access_contacts)
|
|
for _ in range(_NUM_THREADS)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
self.assertEmpty(
|
|
errors, f"Errors in concurrent struct list access: {errors}")
|
|
|
|
def test_concurrent_callback_set_get(self):
|
|
"""Concurrent callback setter/getter calls must not crash or deadlock."""
|
|
errors = []
|
|
|
|
def toggle_callback():
|
|
try:
|
|
for i in range(_ITERS_PER_THREAD):
|
|
if i % 2 == 0:
|
|
mujoco.set_mjcb_passive(lambda m, d: None)
|
|
else:
|
|
mujoco.set_mjcb_passive(None)
|
|
_ = mujoco.get_mjcb_passive()
|
|
except Exception as e: # pylint: disable=broad-except
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=toggle_callback)
|
|
for _ in range(_NUM_THREADS)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
# Clean up.
|
|
mujoco.set_mjcb_passive(None)
|
|
self.assertEmpty(
|
|
errors, f"Errors in concurrent callback set/get: {errors}")
|
|
|
|
def test_concurrent_model_data_lifecycle(self):
|
|
"""Concurrent MjModel/MjData creation and destruction must not crash."""
|
|
errors = []
|
|
|
|
def create_destroy():
|
|
try:
|
|
for _ in range(_ITERS_PER_THREAD):
|
|
m = mujoco.MjModel.from_xml_string(_CONTACT_XML)
|
|
d = mujoco.MjData(m)
|
|
mujoco.mj_step(m, d)
|
|
del d
|
|
del m
|
|
except Exception as e: # pylint: disable=broad-except
|
|
errors.append(e)
|
|
|
|
threads = [threading.Thread(target=create_destroy)
|
|
for _ in range(_NUM_THREADS)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
self.assertEmpty(
|
|
errors, f"Errors in concurrent model/data lifecycle: {errors}")
|
|
|
|
|
|
class ReentrancyTest(absltest.TestCase):
|
|
"""Tests that mutexes don't deadlock when __del__ re-enters setters."""
|
|
|
|
def test_callback_setter_from_del(self):
|
|
"""Setting a callback whose __del__ re-enters the setter must not hang."""
|
|
|
|
class ReentrantCallback:
|
|
"""A callable whose destructor re-enters the callback setter."""
|
|
|
|
def __call__(self, m, d):
|
|
pass
|
|
|
|
def __del__(self):
|
|
# When this object is destroyed by set_mjcb_passive(None) or by
|
|
# being replaced, __del__ will try to set the callback again.
|
|
# This must not deadlock.
|
|
try:
|
|
mujoco.set_mjcb_passive(None)
|
|
except Exception: # pylint: disable=broad-except
|
|
pass # Swallow — we just want to verify no deadlock.
|
|
|
|
mujoco.set_mjcb_passive(ReentrantCallback())
|
|
# This replaces the callback → drops last ref → triggers __del__ →
|
|
# re-enters set_mjcb_passive. Must complete without hanging.
|
|
mujoco.set_mjcb_passive(None)
|
|
|
|
# If we get here, no deadlock occurred.
|
|
self.assertIsNone(mujoco.get_mjcb_passive())
|
|
|
|
def test_callback_replacement_from_del(self):
|
|
"""Replacing a callback whose __del__ sets a new callback must not hang."""
|
|
|
|
class ChainedCallback:
|
|
"""A callable whose destructor sets a different callback."""
|
|
|
|
def __call__(self, m, d):
|
|
pass
|
|
|
|
def __del__(self):
|
|
try:
|
|
mujoco.set_mjcb_passive(lambda m, d: None)
|
|
except Exception: # pylint: disable=broad-except
|
|
pass
|
|
|
|
mujoco.set_mjcb_passive(ChainedCallback())
|
|
# Replace with a plain lambda — old ChainedCallback.__del__ fires.
|
|
mujoco.set_mjcb_passive(lambda m, d: None)
|
|
|
|
# Clean up.
|
|
cb = mujoco.get_mjcb_passive()
|
|
self.assertIsNotNone(cb)
|
|
mujoco.set_mjcb_passive(None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
absltest.main()
|