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

Author: havess

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>`__

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);

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);
 

python/mujoco/introspect/codegen/generate_functions.py

@@ -99,8 +99,13 @@ def _make_comment(self, node: ClangJsonNode) -> str:
       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
 

src/user/user_api.cc

@@ -860,16 +860,18 @@ mjsFrame* mjs_bodyToFrame(mjsBody** body) {
   return &frameC->spec;
 }
 
-
+void mjs_setUserValue(mjsElement* element, const char* key, const void* data) {
+  mjs_setUserValueWithCleanup(element, key, data, nullptr);
+}
 
 // set user payload
-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*)) {
   mjCBase* baseC = static_cast<mjCBase*>(element);
-  baseC->SetUserValue(key, data);
+  baseC->SetUserValue(key, data, cleanup);
 }
 
-
-
 // return user payload or NULL if none found
 const void* mjs_getUserValue(mjsElement* element, const char* key) {
   mjCBase* baseC = static_cast<mjCBase*>(element);

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);
 

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;
 }
 
 

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<std::string, const void*> user_payload_;
+  std::unordered_map<std::string, UserValue> user_payload_;
 };
 
 

test/user/user_api_test.cc

@@ -2845,6 +2845,14 @@ TEST_F(MujocoTest, UserValue) {
   EXPECT_STREQ(static_cast<const char*>(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<const std::string*>(data); });
+  payload = mjs_getUserValue(body->element, "key");
+  EXPECT_STREQ(static_cast<const std::string*>(payload)->c_str(),
+               heap_data->c_str());
   mj_deleteSpec(spec);
 }