diff --git a/BUILD.bazel b/BUILD.bazel index d2413b354..fe64d9a5b 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -745,6 +745,15 @@ cc_test( ], ) +cc_test( + name = "object_locator_test", + srcs = ["src/ray/gcs/gcs_server/test/object_locator_test.cc"], + copts = COPTS, + deps = [ + ":gcs_server_lib", + ], +) + cc_library( name = "object_manager", srcs = glob([ diff --git a/src/ray/gcs/gcs_server/object_locator.cc b/src/ray/gcs/gcs_server/object_locator.cc new file mode 100644 index 000000000..3c2e4f794 --- /dev/null +++ b/src/ray/gcs/gcs_server/object_locator.cc @@ -0,0 +1,129 @@ +#include "ray/gcs/gcs_server/object_locator.h" + +namespace ray { + +namespace gcs { + +ObjectLocator::ObjectLocator() {} + +ObjectLocator::~ObjectLocator() {} + +void ObjectLocator::AddObjectsLocation(const ClientID &node_id, + const std::unordered_set &object_ids) { + // TODO(micafan) Optimize the lock when necessary. + // Maybe use read/write lock. Or reduce the granularity of the lock. + absl::MutexLock lock(&mutex_); + + auto *node_hold_objects = GetNodeHoldObjectSet(node_id, /* create_if_not_exist */ true); + node_hold_objects->insert(object_ids.begin(), object_ids.end()); + + for (const auto &object_id : object_ids) { + auto *object_locations = + GetObjectLocationSet(object_id, /* create_if_not_exist */ true); + object_locations->emplace(node_id); + } +} + +void ObjectLocator::AddObjectLocation(const ObjectID &object_id, + const ClientID &node_id) { + absl::MutexLock lock(&mutex_); + + auto *node_hold_objects = GetNodeHoldObjectSet(node_id, /* create_if_not_exist */ true); + node_hold_objects->emplace(object_id); + + auto *object_locations = + GetObjectLocationSet(object_id, /* create_if_not_exist */ true); + object_locations->emplace(node_id); +} + +std::unordered_set ObjectLocator::GetObjectLocations( + const ObjectID &object_id) { + absl::MutexLock lock(&mutex_); + + auto *object_locations = GetObjectLocationSet(object_id); + if (object_locations) { + return *object_locations; + } + return std::unordered_set{}; +} + +void ObjectLocator::RemoveNode(const ClientID &node_id) { + absl::MutexLock lock(&mutex_); + + ObjectSet node_hold_objects; + auto it = node_to_objects_.find(node_id); + if (it != node_to_objects_.end()) { + node_hold_objects.swap(it->second); + node_to_objects_.erase(it); + } + + if (node_hold_objects.empty()) { + return; + } + + for (const auto &object_id : node_hold_objects) { + auto *object_locations = GetObjectLocationSet(object_id); + if (object_locations) { + object_locations->erase(node_id); + if (object_locations->empty()) { + object_to_locations_.erase(object_id); + } + } + } +} + +void ObjectLocator::RemoveObjectLocation(const ObjectID &object_id, + const ClientID &node_id) { + absl::MutexLock lock(&mutex_); + + auto *object_locations = GetObjectLocationSet(object_id); + if (object_locations) { + object_locations->erase(node_id); + if (object_locations->empty()) { + object_to_locations_.erase(object_id); + } + } + + auto *node_hold_objects = GetNodeHoldObjectSet(node_id); + if (node_hold_objects) { + node_hold_objects->erase(object_id); + if (node_hold_objects->empty()) { + node_to_objects_.erase(node_id); + } + } +} + +ObjectLocator::LocationSet *ObjectLocator::GetObjectLocationSet( + const ObjectID &object_id, bool create_if_not_exist) { + LocationSet *object_locations = nullptr; + + auto it = object_to_locations_.find(object_id); + if (it != object_to_locations_.end()) { + object_locations = &it->second; + } else if (create_if_not_exist) { + auto ret = object_to_locations_.emplace(std::make_pair(object_id, LocationSet{})); + RAY_CHECK(ret.second); + object_locations = &(ret.first->second); + } + + return object_locations; +} + +ObjectLocator::ObjectSet *ObjectLocator::GetNodeHoldObjectSet(const ClientID &node_id, + bool create_if_not_exist) { + ObjectSet *node_hold_objects = nullptr; + + auto it = node_to_objects_.find(node_id); + if (it != node_to_objects_.end()) { + node_hold_objects = &it->second; + } else if (create_if_not_exist) { + auto ret = node_to_objects_.emplace(std::make_pair(node_id, ObjectSet{})); + RAY_CHECK(ret.second); + node_hold_objects = &(ret.first->second); + } + return node_hold_objects; +} + +} // namespace gcs + +} // namespace ray diff --git a/src/ray/gcs/gcs_server/object_locator.h b/src/ray/gcs/gcs_server/object_locator.h new file mode 100644 index 000000000..b074908e8 --- /dev/null +++ b/src/ray/gcs/gcs_server/object_locator.h @@ -0,0 +1,93 @@ +#ifndef GCS_GCS_SERVER_OBJECT_LOCATOR_H +#define GCS_GCS_SERVER_OBJECT_LOCATOR_H + +#include +#include +#include +#include "absl/base/thread_annotations.h" +#include "absl/synchronization/mutex.h" +#include "ray/common/id.h" +#include "ray/util/logging.h" + +namespace ray { + +namespace gcs { + +class ObjectLocator { + public: + ObjectLocator(); + + ~ObjectLocator(); + + /// Add a location of objects. + /// + /// \param node_id The object location that will be added. + /// \param object_ids The ids of objects which location will be added. + void AddObjectsLocation(const ClientID &node_id, + const std::unordered_set &object_ids) + LOCKS_EXCLUDED(mutex_); + + /// Add a location of an object. + /// + /// \param object_id The id of object which location will be added. + /// \param node_id The object location that will be added. + void AddObjectLocation(const ObjectID &object_id, const ClientID &node_id) + LOCKS_EXCLUDED(mutex_); + + /// Get object's locations. + /// + /// \param object_id The id of object to lookup. + /// \return Object locations. + std::unordered_set GetObjectLocations(const ObjectID &object_id) + LOCKS_EXCLUDED(mutex_); + + /// Remove a node. + /// + /// \param node_id The node that will be removed. + void RemoveNode(const ClientID &node_id) LOCKS_EXCLUDED(mutex_); + + /// Remove object's location. + /// + /// \param object_id The id of the object which location will be removed. + /// \param node_id The location that will be removed. + void RemoveObjectLocation(const ObjectID &object_id, const ClientID &node_id) + LOCKS_EXCLUDED(mutex_); + + private: + typedef std::unordered_set LocationSet; + typedef std::unordered_set ObjectSet; + + /// Get object locations by object id from map. + /// Will create it if not exist and the flag create_if_not_exist is set to true. + /// + /// \param object_id The id of object to lookup. + /// \param create_if_not_exist Whether to create a new one if not exist. + /// \return LocationSet * + ObjectLocator::LocationSet *GetObjectLocationSet(const ObjectID &object_id, + bool create_if_not_exist = false) + EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + /// Get objects by node id from map. + /// Will create it if not exist and the flag create_if_not_exist is set to true. + /// + /// \param node_id The id of node to lookup. + /// \param create_if_not_exist Whether to create a new one if not exist. + /// \return ObjectSet * + ObjectLocator::ObjectSet *GetNodeHoldObjectSet(const ClientID &node_id, + bool create_if_not_exist = false) + EXCLUSIVE_LOCKS_REQUIRED(mutex_); + + mutable absl::Mutex mutex_; + + /// Mapping from object id to object locations. + std::unordered_map object_to_locations_ GUARDED_BY(mutex_); + + /// Mapping from node id to objects that held by the node. + std::unordered_map node_to_objects_ GUARDED_BY(mutex_); +}; + +} // namespace gcs + +} // namespace ray + +#endif // GCS_GCS_SERVER_OBJECT_LOCATOR_H diff --git a/src/ray/gcs/gcs_server/test/object_locator_test.cc b/src/ray/gcs/gcs_server/test/object_locator_test.cc new file mode 100644 index 000000000..50e4b5585 --- /dev/null +++ b/src/ray/gcs/gcs_server/test/object_locator_test.cc @@ -0,0 +1,99 @@ +#include "ray/gcs/gcs_server/object_locator.h" +#include "gtest/gtest.h" + +namespace ray { + +namespace gcs { + +class ObjectLocatorTest : public ::testing::Test { + public: + ObjectLocatorTest() {} + + void SetUp() override { GenTestData(); } + + void GenTestData() { + for (size_t i = 0; i < object_count_; ++i) { + ObjectID object_id = ObjectID::FromRandom(); + object_ids_.emplace(object_id); + } + for (size_t i = 0; i < node_count_; ++i) { + ClientID node_id = ClientID::FromRandom(); + node_ids_.emplace(node_id); + } + } + + void CheckLocations(const std::unordered_set &locations) { + ASSERT_EQ(locations.size(), node_ids_.size()); + for (const auto &location : locations) { + auto it = node_ids_.find(location); + ASSERT_TRUE(it != node_ids_.end()); + ASSERT_TRUE(location == *it); + } + } + + protected: + ObjectLocator object_locator_; + + size_t object_count_{5}; + size_t node_count_{10}; + std::unordered_set object_ids_; + std::unordered_set node_ids_; +}; + +TEST_F(ObjectLocatorTest, AddObjectsLocationAndGetLocationTest) { + for (const auto &node_id : node_ids_) { + object_locator_.AddObjectsLocation(node_id, object_ids_); + } + for (const auto &object_id : object_ids_) { + auto locations = object_locator_.GetObjectLocations(object_id); + CheckLocations(locations); + } +} + +TEST_F(ObjectLocatorTest, AddObjectLocationTest) { + for (const auto &object_id : object_ids_) { + for (const auto &node_id : node_ids_) { + object_locator_.AddObjectLocation(object_id, node_id); + } + } + + for (const auto &object_id : object_ids_) { + auto locations = object_locator_.GetObjectLocations(object_id); + CheckLocations(locations); + } +} + +TEST_F(ObjectLocatorTest, RemoveNodeTest) { + for (const auto &node_id : node_ids_) { + object_locator_.AddObjectsLocation(node_id, object_ids_); + } + + object_locator_.RemoveNode(*node_ids_.begin()); + auto locations = object_locator_.GetObjectLocations(*object_ids_.begin()); + ASSERT_EQ(locations.size() + 1, node_ids_.size()); + + locations.emplace(*node_ids_.begin()); + ASSERT_EQ(locations.size(), node_ids_.size()); +} + +TEST_F(ObjectLocatorTest, RemoveObjectLocationTest) { + for (const auto &node_id : node_ids_) { + object_locator_.AddObjectsLocation(node_id, object_ids_); + } + + object_locator_.RemoveObjectLocation(*object_ids_.begin(), *node_ids_.begin()); + auto locations = object_locator_.GetObjectLocations(*object_ids_.begin()); + ASSERT_EQ(locations.size() + 1, node_ids_.size()); + + locations.emplace(*node_ids_.begin()); + ASSERT_EQ(locations.size(), node_ids_.size()); +} + +} // namespace gcs + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}