Fix an race condition during multiple DB opening (facebook#8574)

jay-zhuang · facebook-github-bot · commit c4a503f3dfc2 · 2021-07-22T13:43:06.000-07:00
Summary: ObjectLibrary is shared between multiple DB instances, the Register() could have race condition. Pull Request resolved: facebook#8574 Test Plan: pass the failed test Reviewed By: ajkr Differential Revision: D29855096 Pulled By: jay-zhuang fbshipit-source-id: 541eed0bd495d2c963d858d81e7eabf1ba16153c
diff --git a/include/rocksdb/utilities/object_registry.h b/include/rocksdb/utilities/object_registry.h
@@ -9,10 +9,12 @@
 
 #include <functional>
 #include <memory>
+#include <mutex>
 #include <regex>
 #include <string>
 #include <unordered_map>
 #include <vector>
+
 #include "rocksdb/status.h"
 
 namespace ROCKSDB_NAMESPACE {
@@ -109,6 +111,8 @@ class ObjectLibrary {
   // Adds the input entry to the list for the given type
   void AddEntry(const std::string& type, std::unique_ptr<Entry>& entry);
 
+  // Protects the entry map
+  mutable std::mutex mu_;
   // ** FactoryFunctions for this loader, organized by type
   std::unordered_map<std::string, std::vector<std::unique_ptr<Entry>>> entries_;
 
diff --git a/utilities/object_registry.cc b/utilities/object_registry.cc
@@ -15,6 +15,7 @@ namespace ROCKSDB_NAMESPACE {
 // Otherwise, nullptr is returned
 const ObjectLibrary::Entry *ObjectLibrary::FindEntry(
     const std::string &type, const std::string &name) const {
+  std::unique_lock<std::mutex> lock(mu_);
   auto entries = entries_.find(type);
   if (entries != entries_.end()) {
     for (const auto &entry : entries->second) {
@@ -28,11 +29,13 @@ const ObjectLibrary::Entry *ObjectLibrary::FindEntry(
 
 void ObjectLibrary::AddEntry(const std::string &type,
                              std::unique_ptr<Entry> &entry) {
+  std::unique_lock<std::mutex> lock(mu_);
   auto &entries = entries_[type];
   entries.emplace_back(std::move(entry));
 }
 
 size_t ObjectLibrary::GetFactoryCount(size_t *types) const {
+  std::unique_lock<std::mutex> lock(mu_);
   *types = entries_.size();
   size_t factories = 0;
   for (const auto &e : entries_) {
@@ -42,6 +45,7 @@ size_t ObjectLibrary::GetFactoryCount(size_t *types) const {
 }
 
 void ObjectLibrary::Dump(Logger *logger) const {
+  std::unique_lock<std::mutex> lock(mu_);
   for (const auto &iter : entries_) {
     ROCKS_LOG_HEADER(logger, "    Registered factories for type[%s] ",
                      iter.first.c_str());
