make REGISTER thread safe and idempotent (#520)

dmlc · Mar 24, 2019 · 3ffea86 · 3ffea86
1 parent 865865d
commit 3ffea86
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/include/dmlc/registry.h b/include/dmlc/registry.h
@@ -76,8 +76,10 @@ class Registry {
    * \return ref to the registered entry, used to set properties
    */
   inline EntryType &__REGISTER__(const std::string& name) {
-    CHECK_EQ(fmap_.count(name), 0U)
-        << name << " already registered";
+    std::lock_guard<std::mutex> guard(registering_mutex);
+    if (fmap_.count(name) > 0) {
+      return *fmap_[name];
+    }
     EntryType *e = new EntryType();
     e->name = name;
     fmap_[name] = e;
@@ -111,6 +113,8 @@ class Registry {
   std::vector<const EntryType*> const_list_;
   /*! \brief map of name->function */
   std::map<std::string, EntryType*> fmap_;
+  /*! \brief lock guarding the registering*/
+  std::mutex registering_mutex;
   /*! \brief constructor */
   Registry() {}
   /*! \brief destructor */