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

Summary:
ObjectLibrary is shared between multiple DB instances, the
Register() could have race condition.

Pull Request resolved: https://github.com/facebook/rocksdb/pull/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
index 5a454d7755f45455e1e65670f2622672e3984405..946a26705a5539a3af02049a9b08f4887831443c 100644 (file)
--- 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
index d74d37d171710e0a1580cfbc52ae23e316582ade..79e4b5683b869ab913f43aa7a3afc8c74d85a990 100644 (file)
--- 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());