hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 10:31:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

[Core] make StringIdMap thread safe (#22893)

Browse source 
* make StringIdMap thread safe

* fix comment

Co-authored-by: 黑驰 <senlin.zsl@antgroup.com>
This commit is contained in:
ZhuSenlin 2022-03-09 01:23:41 +08:00 • committed by GitHub
parent f5995dccdf
commit 1e4d7bc1f4
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
src/ray/raylet/scheduling/scheduling_ids.cc
View file

@ -17,6 +17,7 @@
namespace ray { namespace ray {
int64_t StringIdMap::Get(const std::string &string_id) const { int64_t StringIdMap::Get(const std::string &string_id) const {
absl::ReaderMutexLock lock(&mutex_);
auto it = string_to_int_.find(string_id); auto it = string_to_int_.find(string_id);
if (it == string_to_int_.end()) { if (it == string_to_int_.end()) {
return -1; return -1;
@ -26,6 +27,7 @@ int64_t StringIdMap::Get(const std::string &string_id) const {
}; };
std::string StringIdMap::Get(uint64_t id) const { std::string StringIdMap::Get(uint64_t id) const {
absl::ReaderMutexLock lock(&mutex_);
std::string id_string; std::string id_string;
auto it = int_to_string_.find(id); auto it = int_to_string_.find(id);
if (it == int_to_string_.end()) { 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) { 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); auto sit = string_to_int_.find(string_id);
if (sit == string_to_int_.end()) { if (sit == string_to_int_.end()) {
int64_t id = hasher_(string_id); 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) { StringIdMap &StringIdMap::InsertOrDie(const std::string &string_id, int64_t value) {
RAY_CHECK(Get(string_id) == -1) << string_id << " or " << value << " already exist!"; absl::WriterMutexLock lock(&mutex_);
string_to_int_.emplace(string_id, value); RAY_CHECK(string_to_int_.emplace(string_id, value).second)
int_to_string_.emplace(value, string_id); << string_id << " or " << value << " already exist!";
RAY_CHECK(int_to_string_.emplace(value, string_id).second)
<< string_id << " or " << value << " already exist!";
return *this; 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 } // namespace ray

21
src/ray/raylet/scheduling/scheduling_ids.h
View file

@ -18,6 +18,7 @@
#include <string> #include <string>
#include "absl/container/flat_hash_map.h" #include "absl/container/flat_hash_map.h"
#include "absl/synchronization/mutex.h"
#include "ray/util/logging.h" #include "ray/util/logging.h"
#include "ray/util/util.h" #include "ray/util/util.h"
@ -40,6 +41,7 @@ class StringIdMap {
absl::flat_hash_map<std::string, int64_t> string_to_int_; absl::flat_hash_map<std::string, int64_t> string_to_int_;
absl::flat_hash_map<int64_t, std::string> int_to_string_; absl::flat_hash_map<int64_t, std::string> int_to_string_;
std::hash<std::string> hasher_; std::hash<std::string> hasher_;
mutable absl::Mutex mutex_;
public: public:
StringIdMap(){}; StringIdMap(){};
@ -86,13 +88,13 @@ enum class SchedulingIDTag { Node, Resource };
template <SchedulingIDTag T> template <SchedulingIDTag T>
class BaseSchedulingID { class BaseSchedulingID {
public: 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} {} explicit BaseSchedulingID(int64_t id) : id_{id} {}
int64_t ToInt() const { return 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_; } bool operator==(const BaseSchedulingID &rhs) const { return id_ == rhs.id_; }
@ -106,8 +108,8 @@ class BaseSchedulingID {
private: private:
/// Meyer's singleton to store the StringIdMap. /// Meyer's singleton to store the StringIdMap.
static ThreadPrivate<StringIdMap> &GetMap() { static StringIdMap &GetMap() {
static ThreadPrivate<StringIdMap> map; static StringIdMap map;
return map; return map;
} }
int64_t id_ = -1; int64_t id_ = -1;
@ -129,15 +131,16 @@ inline std::ostream &operator<<(
/// Specialization for SchedulingIDTag. Specifically, we populate /// Specialization for SchedulingIDTag. Specifically, we populate
/// the singleton map with PredefinedResources. /// the singleton map with PredefinedResources.
template <> template <>
inline ThreadPrivate<StringIdMap> &BaseSchedulingID<SchedulingIDTag::Resource>::GetMap() { inline StringIdMap &BaseSchedulingID<SchedulingIDTag::Resource>::GetMap() {
static ThreadPrivate<StringIdMap> map{[]() { static std::unique_ptr<StringIdMap> map{[]() {
StringIdMap map; std::unique_ptr<StringIdMap> map(new StringIdMap());
return map.InsertOrDie(kCPU_ResourceLabel, CPU) map->InsertOrDie(kCPU_ResourceLabel, CPU)
.InsertOrDie(kGPU_ResourceLabel, GPU) .InsertOrDie(kGPU_ResourceLabel, GPU)
.InsertOrDie(kObjectStoreMemory_ResourceLabel, OBJECT_STORE_MEM) .InsertOrDie(kObjectStoreMemory_ResourceLabel, OBJECT_STORE_MEM)
.InsertOrDie(kMemory_ResourceLabel, MEM); .InsertOrDie(kMemory_ResourceLabel, MEM);
return map;
}()}; }()};
return map; return *map;
} }
namespace scheduling { namespace scheduling {