From 924ee3070ae8addb16a3a00e037113601be2111d Mon Sep 17 00:00:00 2001 From: Sam Haves Date: Fri, 9 May 2025 11:02:09 -0700 Subject: [PATCH] Modify user values in mujoco to allow users to provide cleanup functions to avoid memory leaks. For the MJCF -> USD plugin, and I suspect other usecases for user values, it's necessary to give ownership of the object to Mujoco. However since they get type erased, we can't clean up the data automatically for the user. Instead this allows c++ clients to provide a cleanup function with their data. PiperOrigin-RevId: 756832010 Change-Id: I80b8e7822e1a0e399a0d19dcaa57ee69b3ccc16a --- doc/APIreference/functions.rst | 11 +++++ doc/includes/references.h | 3 ++ include/mujoco/mujoco.h | 7 ++++ .../introspect/codegen/generate_functions.py | 9 +++- src/user/user_api.cc | 14 ++++--- src/user/user_api.h | 5 +++ src/user/user_objects.cc | 9 ++-- src/user/user_objects.h | 41 ++++++++++++++++++- test/user/user_api_test.cc | 8 ++++ 9 files changed, 92 insertions(+), 15 deletions(-) diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 9606e94d..e13ad19d 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -4459,6 +4459,17 @@ Transform body into a frame. Set user payload, overriding the existing value for the specified key if present. +.. _mjs_setUserValueWithCleanup: + +`mjs_setUserValueWithCleanup <#mjs_setUserValueWithCleanup>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_setUserValueWithCleanup + +Set user payload, overriding the existing value for the specified key if +present. This version differs from mjs_setUserValue in that it takes a +cleanup function that will be called when the user payload is deleted. + .. _mjs_getUserValue: `mjs_getUserValue <#mjs_getUserValue>`__ diff --git a/doc/includes/references.h b/doc/includes/references.h index 6ebc759a..0689e150 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3732,6 +3732,9 @@ const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* s const mjsOrientation* orientation); mjsFrame* mjs_bodyToFrame(mjsBody** body); void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); const void* mjs_getUserValue(mjsElement* element, const char* key); void mjs_deleteUserValue(mjsElement* element, const char* key); void mjs_defaultSpec(mjSpec* spec); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 63bfafcc..5bca71ea 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1644,6 +1644,13 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload, overriding the existing value for the specified key if present. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload, overriding the existing value for the specified key if +// present. This version differs from mjs_setUserValue in that it takes a +// cleanup function that will be called when the user payload is deleted. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/python/mujoco/introspect/codegen/generate_functions.py b/python/mujoco/introspect/codegen/generate_functions.py index 1966e328..06beeaf2 100644 --- a/python/mujoco/introspect/codegen/generate_functions.py +++ b/python/mujoco/introspect/codegen/generate_functions.py @@ -99,8 +99,13 @@ class MjFunctionVisitor: return ''.join(strings) def visit(self, node: ClangJsonNode) -> None: - if (node.get('kind') == 'FunctionDecl' and - node.get('name', '').startswith('mj')): + # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated + # objects and doesn't need a python wrapper. + if ( + node.get('kind') == 'FunctionDecl' + and node.get('name', '').startswith('mj') + and node.get('name', '') != 'mjs_setUserValueWithCleanup' + ): func_decl = self._make_function(node) self._functions[func_decl.name] = func_decl diff --git a/src/user/user_api.cc b/src/user/user_api.cc index e103aa8e..41cba2d2 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -860,15 +860,17 @@ mjsFrame* mjs_bodyToFrame(mjsBody** body) { return &frameC->spec; } - - -// set user payload void mjs_setUserValue(mjsElement* element, const char* key, const void* data) { - mjCBase* baseC = static_cast(element); - baseC->SetUserValue(key, data); + mjs_setUserValueWithCleanup(element, key, data, nullptr); } - +// set user payload +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)) { + mjCBase* baseC = static_cast(element); + baseC->SetUserValue(key, data, cleanup); +} // return user payload or NULL if none found const void* mjs_getUserValue(mjsElement* element, const char* key) { diff --git a/src/user/user_api.h b/src/user/user_api.h index e1b73d86..cd228147 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -374,6 +374,11 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 82853307..51c0f63a 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -784,15 +784,14 @@ void mjCBase::SetFrame(mjCFrame* _frame) { frame = _frame; } - -void mjCBase::SetUserValue(std::string_view key, const void* data) { - user_payload_[std::string(key)] = data; +void mjCBase::SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)) { + user_payload_[std::string(key)] = UserValue(data, cleanup); } - const void* mjCBase::GetUserValue(std::string_view key) { auto found = user_payload_.find(std::string(key)); - return found != user_payload_.end() ? found->second : nullptr; + return found != user_payload_.end() ? found->second.value : nullptr; } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index bfec0797..9cdae417 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -281,7 +281,8 @@ class mjCBase : public mjCBase_ { } // Set and get user payload - void SetUserValue(std::string_view key, const void* data); + void SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)); const void* GetUserValue(std::string_view key); void DeleteUserValue(std::string_view key); @@ -292,8 +293,44 @@ class mjCBase : public mjCBase_ { // reference count for allowing deleting an attached object int refcount = 1; + // Arbitrary user value that cleans up the data when destroyed. + struct UserValue { + const void* value = nullptr; + void (*cleanup)(const void*) = nullptr; + + UserValue() {} + UserValue(const void* value, void (*cleanup)(const void*)) + : value(value), cleanup(cleanup) {} + UserValue(const UserValue& other) = delete; + UserValue& operator=(const UserValue& other) = delete; + + UserValue(UserValue&& other) : value(other.value), cleanup(other.cleanup) { + other.value = nullptr; + other.cleanup = nullptr; + } + + UserValue& operator=(UserValue&& other) { + if (this != &other) { + if (cleanup && value) { + cleanup(value); + } + value = other.value; + cleanup = other.cleanup; + other.value = nullptr; + other.cleanup = nullptr; + } + return *this; + } + + ~UserValue() { + if (cleanup && value) { + cleanup(value); + } + } + }; + // user payload - std::unordered_map user_payload_; + std::unordered_map user_payload_; }; diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 4db9f25b..1669ea20 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -2845,6 +2845,14 @@ TEST_F(MujocoTest, UserValue) { EXPECT_STREQ(static_cast(payload), data.c_str()); mjs_deleteUserValue(body->element, "key"); EXPECT_THAT(mjs_getUserValue(body->element, "key"), IsNull()); + + std::string* heap_data = new std::string("heap_data"); + mjs_setUserValueWithCleanup( + body->element, "key", heap_data, + [](const void* data) { delete static_cast(data); }); + payload = mjs_getUserValue(body->element, "key"); + EXPECT_STREQ(static_cast(payload)->c_str(), + heap_data->c_str()); mj_deleteSpec(spec); }