Fix false positives from mjData stack leakage detection.

When built and run under address sanitizer (asan), `mj_markStack` and `mj_freeStack` are instrumented to detect leakage of mjData stack frames. When the compiler inlines several callees that call into mark/free in the same function, this instrumentation requires that the compiler retains separate mark/free calls for each original callee.

This change adds necessary optimization barriers to ensure that this is the case.

PiperOrigin-RevId: 609373993
Change-Id: I30dfd998fb66530735348e9715d92ee713f28210

include/mujoco/mujoco.h

@@ -19,7 +19,7 @@
 
 
 // this is a C-API
-#if defined(__cplusplus)
+#ifdef __cplusplus
 extern "C" {
 #endif
 
@@ -187,6 +187,8 @@ MJAPI void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_va
 // Reset data. If 0 <= key < nkey, set fields from specified keyframe.
 MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 
+#ifndef ADDRESS_SANITIZER
+
 // Mark a new frame on the mjData stack.
 MJAPI void mj_markStack(mjData* d);
 
@@ -194,6 +196,8 @@ MJAPI void mj_markStack(mjData* d);
 // to mj_markStack must no longer be used afterwards.
 MJAPI void mj_freeStack(mjData* d);
 
+#endif  // ADDRESS_SANITIZER
+
 // Allocate a number of bytes on mjData stack at a specific alignment.
 // Call mju_error on stack overflow.
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);
@@ -1340,8 +1344,36 @@ MJAPI void mju_defaultTask(mjTask* task);
 // Wait for a task to complete.
 MJAPI void mju_taskJoin(mjTask* task);
 
+//---------------------- Sanitizer instrumentation helpers -----------------------------------------
+//
+// Most MuJoCo users can ignore these functions, the following comments are aimed primarily at
+// MuJoCo developers.
+//
+// When built and run under address sanitizer (asan), mj_markStack and mj_freeStack are instrumented
+// to detect leakage of mjData stack frames. When the compiler inlines several callees that call
+// into mark/free into the same function, this instrumentation requires that the compiler retains
+// separate mark/free calls for each original callee. The memory-clobbered asm blocks act as a
+// barrier to prevent mark/free calls from being combined under optimization.
 
-#if defined(__cplusplus)
+#ifdef ADDRESS_SANITIZER
+
+void mj__markStack(mjData*) __attribute__((noinline));
+static inline void mj_markStack(mjData* d) __attribute__((always_inline)) {
+  asm volatile("" ::: "memory");
+  mj__markStack(d);
+  asm volatile("" ::: "memory");
+}
+
+void mj__freeStack(mjData*) __attribute__((noinline));
+static inline void mj_freeStack(mjData* d) __attribute__((always_inline)) {
+  asm volatile("" ::: "memory");
+  mj__freeStack(d);
+  asm volatile("" ::: "memory");
+}
+
+#endif  // ADDRESS_SANITIZER
+
+#ifdef __cplusplus
 }
 #endif
 

src/engine/engine_crossplatform.cc

@@ -53,12 +53,23 @@ void CheckRosetta() {
 #include <string>
 #include <string_view>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 
 namespace {
-std::string_view SymbolizeCached(void* pc) {
+const std::pair<std::string, std::string>&
+FuncNameAndDebugInfoCached(void* pc) {
+  static const std::unordered_set<std::string>* const kIgnoredInlinedFunctions =
+      []() {
+        return new std::unordered_set<std::string>{
+            "mj_freeStack",
+            "mj_markStack",
+        };
+      }();
+
   static auto* mu = new std::shared_mutex;
-  static auto* pc_to_func_name_map = new std::unordered_map<void*, std::string>;
+  static auto* pc_to_func_name_map =
+      new std::unordered_map<void*, std::pair<std::string, std::string>>;
 
   {
     std::shared_lock lock(*mu);
@@ -70,14 +81,43 @@ std::string_view SymbolizeCached(void* pc) {
 
   std::array<char, 256> buf;
   __sanitizer_symbolize_pc(pc, "%f", buf.data(), buf.size());
+
+  // buf contains sequence of null-terminated strings of inlined function names
+  // so we walk through the sequence until we find the first unignored function
+  std::string_view func_name(buf.data());
+  int idx = 0;
+  while (kIgnoredInlinedFunctions->find(func_name.data()) !=
+         kIgnoredInlinedFunctions->end()) {
+    func_name = func_name.data() + func_name.size() + 1;
+    ++idx;
+  }
+
+  std::array<char, 1024> buf2;
+  __sanitizer_symbolize_pc(pc, "%F at %S", buf2.data(), buf2.size());
+  std::string_view debug_info(buf2.data());
+  for (int i = 0; i < idx; ++i) {
+    debug_info = debug_info.data() + debug_info.size() + 1;
+  }
+
   {
     std::unique_lock lock(*mu);
-    return pc_to_func_name_map->emplace(pc, buf.data()).first->second;
+    return pc_to_func_name_map
+        ->emplace(
+            pc, std::make_pair(std::string(func_name), std::string(debug_info)))
+        .first->second;
   }
 }
+
+const std::string& SymbolizeCached(void* pc) {
+  return FuncNameAndDebugInfoCached(pc).first;
+}
+
+const std::string& DebugInfoCached(void* pc) {
+  return FuncNameAndDebugInfoCached(pc).second;
+}
 }  // namespace
 
-int _mj_comparePcFuncName(void* pc1, void* pc2) {
+int mj__comparePcFuncName(void* pc1, void* pc2) {
   static auto* mu = new std::shared_mutex;
   static auto* same_func_map = new std::map<std::pair<void*, void*>, bool>;
 
@@ -96,4 +136,8 @@ int _mj_comparePcFuncName(void* pc1, void* pc2) {
     return same_func_map->emplace(pc_pair, is_same).first->second;
   }
 }
+
+const char* mj__getPcDebugInfo(void* pc) {
+  return DebugInfoCached(pc).c_str();
+}
 #endif  // ADDRESS_SANITIZER

src/engine/engine_crossplatform.h

@@ -80,7 +80,8 @@ extern "C" {
 #endif
 
 #ifdef ADDRESS_SANITIZER
-int _mj_comparePcFuncName(void* pc1, void* pc2);
+int mj__comparePcFuncName(void* pc1, void* pc2);
+const char* mj__getPcDebugInfo(void* pc);
 #endif
 
 #ifdef __cplusplus

src/engine/engine_io.c

@@ -1439,10 +1439,12 @@ static inline void markstackinternal(mjData* d, mjStackInfo* stack_info) {
 
 
 // mjData mark stack frame
-#ifdef ADDRESS_SANITIZER
-__attribute__((noinline))
+#ifndef ADDRESS_SANITIZER
+void mj_markStack(mjData* d)
+#else
+void mj__markStack(mjData* d)
 #endif
-void mj_markStack(mjData* d) {
+{
   if (!d->threadpool) {
     mjStackInfo stack_info = get_stack_info_from_data(d);
     markstackinternal(d, &stack_info);
@@ -1469,15 +1471,10 @@ static inline void freestackinternal(mjStackInfo* stack_info) {
   mjStackFrame* s = (mjStackFrame*) stack_info->stack_base;
 #ifdef ADDRESS_SANITIZER
   // raise an error if caller function name doesn't match the most recent caller of mj_markStack
-  if (!_mj_comparePcFuncName(s->pc, __sanitizer_return_address())) {
-    #define mjSYMBOLIZELEN 256
-    char dbginfo[mjSYMBOLIZELEN];
-    __sanitizer_symbolize_pc(
-      s->pc, "mj_markStack %F at %S has no corresponding mj_freeStack",
-      dbginfo, sizeof(dbginfo));
-    dbginfo[mjSYMBOLIZELEN - 1] = '\0';
-    mjERROR("%s", dbginfo);
-    #undef mjSYMBOLIZELEN
+  if (!mj__comparePcFuncName(s->pc, __sanitizer_return_address())) {
+    mjERROR("mj_markStack %s has no corresponding mj_freeStack (detected %s)",
+            mj__getPcDebugInfo(s->pc),
+            mj__getPcDebugInfo(__sanitizer_return_address()));
   }
 #endif
 
@@ -1494,10 +1491,12 @@ static inline void freestackinternal(mjStackInfo* stack_info) {
 
 
 // mjData free stack frame
-#ifdef ADDRESS_SANITIZER
-__attribute__((noinline))
+#ifndef ADDRESS_SANITIZER
+void mj_freeStack(mjData* d)
+#else
+void mj__freeStack(mjData* d)
 #endif
-void mj_freeStack(mjData* d) {
+{
   if (!d->threadpool) {
     mjStackInfo stack_info = get_stack_info_from_data(d);
     freestackinternal(&stack_info);

src/engine/engine_io.h

@@ -110,12 +110,21 @@ MJAPI void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
 // mjData arena allocate
 MJAPI void* mj_arenaAllocByte(mjData* d, size_t bytes, size_t alignment);
 
+#ifndef ADDRESS_SANITIZER
+
 // mjData mark stack frame
 MJAPI void mj_markStack(mjData* d);
 
 // mjData free stack frame
 MJAPI void mj_freeStack(mjData* d);
 
+#else
+
+void mj__markStack(mjData* d) __attribute__((noinline));
+void mj__freeStack(mjData* d) __attribute__((noinline));
+
+#endif  // ADDRESS_SANITIZER
+
 // mjData stack allocate
 MJAPI void* mj_stackAllocByte(mjData* d, size_t bytes, size_t alignment);