Move plugin loading and scanning logic from Simulate to MuJoCo library.

PiperOrigin-RevId: 489020177
Change-Id: I1c29d931b47d5252e57e1d797cc7c17645859ffc

doc/changelog.rst

@@ -50,6 +50,10 @@ General
     <https://github.com/deepmind/mujoco/tree/main/test/benchmark/engine_core_smooth_benchmark_test.cc>`_.
   - 50% faster ``mju_dotSparse`` using manual loop unroll. See `engine_util_sparse_benchmark_test
     <https://github.com/deepmind/mujoco/tree/main/test/benchmark/engine_util_sparse_benchmark_test.cc>`_.
+- Add API functions ``mj_loadPluginLibrary`` and  ``mj_loadAllPluginLibraries``. The first function is identical to
+  ``dlopen`` on a POSIX system, and to ``LoadLibraryA`` on Windows. The second function scans a specified directory for
+  all dynamic libraries file and loads each library found. Dynamic libraries opened by these functions are assumed to
+  register one or more MuJoCo plugins on load.
 
 Simulate
 ^^^^^^^^

include/mujoco/mjplugin.h

@@ -94,4 +94,7 @@ typedef struct mjpPlugin_ mjpPlugin;
 
 #endif  // defined(_MSC_VER)
 
+// function pointer type for mj_loadAllPluginLibraries callback
+typedef void (*mjfPluginLibraryLoadCallback)(const char* filename, int first, int count);
+
 #endif  // MUJOCO_INCLUDE_MJPLUGIN_H_

include/mujoco/mujoco.h

@@ -458,6 +458,14 @@ MJAPI void mj_setTotalmass(mjModel* m, mjtNum newmass);
 // NULL: invalid plugin instance ID or attribute name
 MJAPI const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attrib);
 
+// Load a dynamic library. The dynamic library is assumed to register one or more plugins.
+MJAPI void mj_loadPluginLibrary(const char* path);
+
+// Scan a directory and load all dynamic libraries. Dynamic libraries in the specified directory
+// are assumed to register one or more plugins. Optionally, if a callback is specified, it is called
+// for each dynamic library encountered that registers plugins.
+MJAPI void mj_loadAllPluginLibraries(const char* directory, mjfPluginLibraryLoadCallback callback);
+
 // Return version number: 1.0.2 is encoded as 102.
 MJAPI int mj_version(void);
 

introspect/functions.py

@@ -2698,6 +2698,38 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Return a config attribute value of a plugin instance; NULL: invalid plugin instance ID or attribute name',  # pylint: disable=line-too-long
      )),
+    ('mj_loadPluginLibrary',
+     FunctionDecl(
+         name='mj_loadPluginLibrary',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='path',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Load a dynamic library. The dynamic library is assumed to register one or more plugins.',  # pylint: disable=line-too-long
+     )),
+    ('mj_loadAllPluginLibraries',
+     FunctionDecl(
+         name='mj_loadAllPluginLibraries',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='directory',
+                 type=PointerType(
+                     inner_type=ValueType(name='char', is_const=True),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='callback',
+                 type=ValueType(name='mjfPluginLibraryLoadCallback'),
+             ),
+         ),
+         doc='Scan a directory and load all dynamic libraries. Dynamic libraries in the specified directory are assumed to register one or more plugins. Optionally, if a callback is specified, it is called for each dynamic library encountered that registers plugins.',  # pylint: disable=line-too-long
+     )),
     ('mj_version',
      FunctionDecl(
          name='mj_version',

python/mujoco/functions.cc

@@ -535,6 +535,13 @@ PYBIND11_MODULE(_functions, pymodule) {
   Def<traits::mj_local2Global>(pymodule);
   Def<traits::mj_getTotalmass>(pymodule);
   Def<traits::mj_setTotalmass>(pymodule);
+  Def<traits::mj_loadPluginLibrary>(pymodule);
+  DEF_WITH_OMITTED_PY_ARGS(traits::mj_loadAllPluginLibraries, "callback")(
+      pymodule,
+      [](const std::string& directory) {
+        InterceptMjErrors(::mj_loadAllPluginLibraries)(
+            directory.c_str(), nullptr);
+      });
   Def<traits::mj_version>(pymodule);
   Def<traits::mj_versionString>(pymodule);
 

simulate/main.cc

@@ -39,8 +39,6 @@ extern "C" {
   #if defined(__APPLE__)
     #include <mach-o/dyld.h>
   #endif
-  #include <dirent.h>
-  #include <dlfcn.h>
   #include <sys/errno.h>
   #include <unistd.h>
 #endif
@@ -164,15 +162,9 @@ std::string getExecutableDir() {
 }
 
 
-#if defined(_WIN32) || defined(__CYGWIN__)
-using unique_dlhandle = std::unique_ptr<std::remove_pointer_t<HMODULE>, decltype(&FreeLibrary)>;
-#else
-using unique_dlhandle = std::unique_ptr<void, decltype(&dlclose)>;
-#endif
-
 
 // scan for libraries in the plugin directory to load additional plugins
-std::vector<unique_dlhandle> scanPluginLibraries() {
+void scanPluginLibraries() {
   // check and print plugins that are linked directly into the executable
   int nplugin = mjp_pluginCount();
   if (nplugin) {
@@ -185,91 +177,26 @@ std::vector<unique_dlhandle> scanPluginLibraries() {
   // define platform-specific strings
 #if defined(_WIN32) || defined(__CYGWIN__)
   const std::string sep = "\\";
-  const std::string dso_suffix = ".dll";
 #else
   const std::string sep = "/";
-  #if defined(__APPLE__)
-    const std::string dso_suffix = ".dylib";
-  #else
-    const std::string dso_suffix = ".so";
-  #endif
 #endif
 
-  // output vectors containing DSO handles
-  std::vector<unique_dlhandle> dso_handles;
-
-  // platform-independent routine for checking and printing plugins registered by a dynamic library
-  const auto check_and_print_plugins = [&](const std::string& name, unique_dlhandle&& dlhandle) {
-    if (!dlhandle) {
-      return;
-    }
-
-    const int nplugin_new = mjp_pluginCount();
-    if (nplugin_new > nplugin) {
-      dso_handles.push_back(std::move(dlhandle));
-
-      // print all newly registered plugins
-      std::printf("Plugins registered by library '%s':\n", name.c_str());
-      for (int i = nplugin; i < nplugin_new; ++i) {
-        std::printf("    %s\n", mjp_getPluginAtSlot(i)->name);
-      }
-
-      // update counter for plugins registered so far
-      nplugin = nplugin_new;
-    }
-  };
 
   // try to open the ${EXECDIR}/plugin directory
   // ${EXECDIR} is the directory containing the simulate binary itself
   const std::string executable_dir = getExecutableDir();
   if (executable_dir.empty()) {
-    return dso_handles;
+    return;
   }
 
   const std::string plugin_dir = getExecutableDir() + sep + MUJOCO_PLUGIN_DIR;
-
+  mj_loadAllPluginLibraries(
-#if defined(_WIN32) || defined(__CYGWIN__)
+      plugin_dir.c_str(), +[](const char* filename, int first, int count) {
-  WIN32_FIND_DATAA find_data;
+        std::printf("Plugins registered by library '%s':\n", filename);
-  HANDLE hfile = FindFirstFileA((plugin_dir + sep + "*.dll").c_str(), &find_data);
+        for (int i = first; i < first + count; ++i) {
-  if (!hfile) {
+          std::printf("    %s\n", mjp_getPluginAtSlot(i)->name);
-    return dso_handles;
+        }
-  }
+      });
-
-  // go through each file in the directory
-  bool keep_going = true;
-  while (keep_going) {
-    const std::string name(find_data.cFileName);
-    // load the library and check for plugins
-    const std::string dso_path = plugin_dir + sep + name;
-    check_and_print_plugins(
-        name, unique_dlhandle(LoadLibraryA(dso_path.c_str()), &FreeLibrary));
-    keep_going = FindNextFileA(hfile, &find_data);
-  }
-  FindClose(hfile);
-#else
-  DIR* dirp = opendir(plugin_dir.c_str());
-  if (!dirp) {
-    return dso_handles;
-  }
-
-  // go through each entry in the directory
-  for (struct dirent* dp; (dp = readdir(dirp));) {
-    // only look at regular files (i.e. skip symlinks, pipes, directories, etc.)
-    if (dp->d_type == DT_REG) {
-      const std::string name(dp->d_name);
-      if (name.size() > dso_suffix.size() &&
-          name.substr(name.size() - dso_suffix.size()) == dso_suffix) {
-        // load the library and check for plugins
-        const std::string dso_path = plugin_dir + sep + name;
-        check_and_print_plugins(
-            name, unique_dlhandle(dlopen(dso_path.c_str(), RTLD_NOW | RTLD_LOCAL), &dlclose));
-      }
-    }
-  }
-  closedir(dirp);
-#endif
-
-  return dso_handles;
 }
 
 
@@ -519,7 +446,7 @@ int main(int argc, const char** argv) {
   }
 
   // scan for libraries in the plugin directory to load additional plugins
-  std::vector<unique_dlhandle> dso_handles = scanPluginLibraries();
+  scanPluginLibraries();
 
   // simulate object encapsulates the UI
   auto sim = std::make_unique<mj::Simulate>();

src/engine/engine_plugin.cc

@@ -32,6 +32,15 @@
 #include <utility>
 #include <vector>
 
+extern "C" {
+#if defined(_WIN32) || defined(__CYGWIN__)
+  #include <windows.h>
+#else
+  #include <dirent.h>
+  #include <dlfcn.h>
+#endif
+}
+
 #ifdef __APPLE__
 #include <Availability.h>
 #if !defined(MAC_OS_X_VERSION_MIN_REQUIRED) && defined(__MAC_OS_X_VERSION_MIN_REQUIRED)
@@ -79,6 +88,30 @@ static_assert(
 
 using Mutex = std::shared_mutex;
 
+class ReentrantWriteLock {
+ public:
+  ReentrantWriteLock(Mutex& mutex) : mutex_(mutex) {
+    if (LockCountOnCurrentThread() == 0) {
+      mutex_.lock();
+    }
+    ++LockCountOnCurrentThread();
+  }
+
+  ~ReentrantWriteLock() {
+    if (--LockCountOnCurrentThread() == 0) {
+      mutex_.unlock();
+    }
+  }
+
+ private:
+  Mutex& mutex_;
+
+  static int& LockCountOnCurrentThread() noexcept {
+    thread_local int counter = 0;
+    return counter;
+  }
+};
+
 class Global {
  public:
   Global() {
@@ -94,6 +127,10 @@ class Global {
     return *std::launder(reinterpret_cast<Mutex*>(&mutex_));
   }
 
+  ReentrantWriteLock lock_mutex_exclusively() {
+    return ReentrantWriteLock(mutex());
+  }
+
  private:
   PluginTable table_;
   std::atomic_int count_;
@@ -229,9 +266,8 @@ int mjp_registerPlugin(const mjpPlugin* plugin) {
       }
     }
 
-    // exclusively lock the global plugin table
     Global& global = GetGlobal();
-    std::unique_lock lock(global.mutex());
+    auto lock = global.lock_mutex_exclusively();
 
     int count = global.count().load(std::memory_order_acquire);
     int local_idx = 0;
@@ -429,3 +465,85 @@ const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attr
 
   return nullptr;
 }
+
+// load plugins from a dynamic library
+void mj_loadPluginLibrary(const char* path) {
+#if defined(_WIN32) || defined(__CYGWIN__)
+  LoadLibraryA(path);
+#else
+  dlopen(path, RTLD_NOW | RTLD_LOCAL);
+#endif
+}
+
+// scan a directory and load all dynamic libraries
+void mj_loadAllPluginLibraries(const char* directory,
+                               mjfPluginLibraryLoadCallback callback) {
+  auto load_dso_and_call_callback = [&](const std::string& filename,
+                                        const std::string& dso_path) {
+    int nplugin_before;
+    int nplugin_after;
+
+    Global& global = GetGlobal();
+    {
+      auto lock = global.lock_mutex_exclusively();
+      nplugin_before = mjp_pluginCount();
+      mj_loadPluginLibrary(dso_path.c_str());
+      nplugin_after = mjp_pluginCount();
+    }
+
+    if (callback) {
+      int count = nplugin_after - nplugin_before;
+      int first = count ? nplugin_before : -1;
+      callback(filename.c_str(), first, count);
+    }
+  };
+
+  // define platform-specific strings
+#if defined(_WIN32) || defined(__CYGWIN__)
+  const std::string sep = "\\";
+  WIN32_FIND_DATAA find_data;
+  HANDLE hfile = FindFirstFileA(
+      (directory + sep + "*.dll").c_str(), &find_data);
+  if (!hfile) {
+    return;
+  }
+
+  // go through each file in the directory
+  bool keep_going = true;
+  while (keep_going) {
+    const std::string name(find_data.cFileName);
+    // load the library and check for plugins
+    const std::string dso_path = directory + sep + name;
+    load_dso_and_call_callback(name.c_str(), dso_path.c_str());
+    keep_going = FindNextFileA(hfile, &find_data);
+  }
+  FindClose(hfile);
+#else
+  const std::string sep = "/";
+  #if defined(__APPLE__)
+    const std::string dso_suffix = ".dylib";
+  #else
+    const std::string dso_suffix = ".so";
+  #endif
+
+  DIR* dirp = opendir(directory);
+  if (!dirp) {
+    return;
+  }
+
+  // go through each entry in the directory
+  for (struct dirent* dp; (dp = readdir(dirp));) {
+    // only look at regular files (skip symlinks, pipes, directories, etc.)
+    if (dp->d_type == DT_REG) {
+      const std::string name(dp->d_name);
+      if (name.size() > dso_suffix.size() &&
+          name.substr(name.size() - dso_suffix.size()) == dso_suffix) {
+        // load the library
+        const std::string dso_path = directory + sep + name;
+        load_dso_and_call_callback(name.c_str(), dso_path.c_str());
+      }
+    }
+  }
+  closedir(dirp);
+#endif
+}

src/engine/engine_plugin.h

@@ -41,6 +41,12 @@ MJAPI const mjpPlugin* mjp_getPluginAtSlot(int slot);
 // NULL: invalid plugin instance ID or attribute name
 MJAPI const char* mj_getPluginConfig(const mjModel* m, int plugin_id, const char* attrib);
 
+// load plugins from a dynamic library
+MJAPI void mj_loadPluginLibrary(const char* path);
+
+// scan a directory and load all dynamic libraries
+MJAPI void mj_loadAllPluginLibraries(const char* directory, mjfPluginLibraryLoadCallback callback);
+
 // =================================================================================================
 // MuJoCo-internal functions beyond this point.
 // "Unsafe" suffix indicates that improper use of these functions may result in data races.

unity/Runtime/Bindings/MujocoBindings.cs

@@ -3035,6 +3035,9 @@ public static unsafe extern void mj_setTotalmass(mjModel_* m, double newmass);
 [return: MarshalAs(UnmanagedType.LPStr)]
 public static unsafe extern string mj_getPluginConfig(mjModel_* m, int plugin_id, [MarshalAs(UnmanagedType.LPStr)]string attrib);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_loadPluginLibrary([MarshalAs(UnmanagedType.LPStr)]string path);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern int mj_version();