diff --git a/src/ray/raylet/scheduling/scheduling_ids.cc b/src/ray/raylet/scheduling/scheduling_ids.cc index 39ce1c944..4e098e141 100644 --- a/src/ray/raylet/scheduling/scheduling_ids.cc +++ b/src/ray/raylet/scheduling/scheduling_ids.cc @@ -17,6 +17,7 @@ namespace ray { int64_t StringIdMap::Get(const std::string &string_id) const { + absl::ReaderMutexLock lock(&mutex_); auto it = string_to_int_.find(string_id); if (it == string_to_int_.end()) { return -1; @@ -26,6 +27,7 @@ int64_t StringIdMap::Get(const std::string &string_id) const { }; std::string StringIdMap::Get(uint64_t id) const { + absl::ReaderMutexLock lock(&mutex_); std::string id_string; auto it = int_to_string_.find(id); if (it == int_to_string_.end()) { @@ -37,6 +39,7 @@ std::string StringIdMap::Get(uint64_t id) const { }; int64_t StringIdMap::Insert(const std::string &string_id, uint8_t max_id) { + absl::WriterMutexLock lock(&mutex_); auto sit = string_to_int_.find(string_id); if (sit == string_to_int_.end()) { int64_t id = hasher_(string_id); @@ -63,12 +66,17 @@ int64_t StringIdMap::Insert(const std::string &string_id, uint8_t max_id) { }; StringIdMap &StringIdMap::InsertOrDie(const std::string &string_id, int64_t value) { - RAY_CHECK(Get(string_id) == -1) << string_id << " or " << value << " already exist!"; - string_to_int_.emplace(string_id, value); - int_to_string_.emplace(value, string_id); + absl::WriterMutexLock lock(&mutex_); + RAY_CHECK(string_to_int_.emplace(string_id, value).second) + << string_id << " or " << value << " already exist!"; + RAY_CHECK(int_to_string_.emplace(value, string_id).second) + << string_id << " or " << value << " already exist!"; return *this; } -int64_t StringIdMap::Count() { return string_to_int_.size(); } +int64_t StringIdMap::Count() { + absl::ReaderMutexLock lock(&mutex_); + return string_to_int_.size(); +} } // namespace ray diff --git a/src/ray/raylet/scheduling/scheduling_ids.h b/src/ray/raylet/scheduling/scheduling_ids.h index e92f0ccc3..135bd35cf 100644 --- a/src/ray/raylet/scheduling/scheduling_ids.h +++ b/src/ray/raylet/scheduling/scheduling_ids.h @@ -18,6 +18,7 @@ #include #include "absl/container/flat_hash_map.h" +#include "absl/synchronization/mutex.h" #include "ray/util/logging.h" #include "ray/util/util.h" @@ -40,6 +41,7 @@ class StringIdMap { absl::flat_hash_map string_to_int_; absl::flat_hash_map int_to_string_; std::hash hasher_; + mutable absl::Mutex mutex_; public: StringIdMap(){}; @@ -86,13 +88,13 @@ enum class SchedulingIDTag { Node, Resource }; template class BaseSchedulingID { public: - explicit BaseSchedulingID(const std::string &name) : id_{GetMap()->Insert(name)} {} + explicit BaseSchedulingID(const std::string &name) : id_{GetMap().Insert(name)} {} explicit BaseSchedulingID(int64_t id) : id_{id} {} int64_t ToInt() const { return id_; } - std::string Binary() const { return GetMap()->Get(id_); } + std::string Binary() const { return GetMap().Get(id_); } bool operator==(const BaseSchedulingID &rhs) const { return id_ == rhs.id_; } @@ -106,8 +108,8 @@ class BaseSchedulingID { private: /// Meyer's singleton to store the StringIdMap. - static ThreadPrivate &GetMap() { - static ThreadPrivate map; + static StringIdMap &GetMap() { + static StringIdMap map; return map; } int64_t id_ = -1; @@ -129,15 +131,16 @@ inline std::ostream &operator<<( /// Specialization for SchedulingIDTag. Specifically, we populate /// the singleton map with PredefinedResources. template <> -inline ThreadPrivate &BaseSchedulingID::GetMap() { - static ThreadPrivate map{[]() { - StringIdMap map; - return map.InsertOrDie(kCPU_ResourceLabel, CPU) +inline StringIdMap &BaseSchedulingID::GetMap() { + static std::unique_ptr map{[]() { + std::unique_ptr map(new StringIdMap()); + map->InsertOrDie(kCPU_ResourceLabel, CPU) .InsertOrDie(kGPU_ResourceLabel, GPU) .InsertOrDie(kObjectStoreMemory_ResourceLabel, OBJECT_STORE_MEM) .InsertOrDie(kMemory_ResourceLabel, MEM); + return map; }()}; - return map; + return *map; } namespace scheduling {