[GCS] Add ObjectLocator to gcs server (#7557)

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([

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<ObjectID> &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<ClientID> 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<ClientID>{};
+}
+
+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

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 <memory>
+#include <unordered_map>
+#include <unordered_set>
+#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<ObjectID> &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<ClientID> 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<ClientID> LocationSet;
+  typedef std::unordered_set<ObjectID> 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<ObjectID, LocationSet> object_to_locations_ GUARDED_BY(mutex_);
+
+  /// Mapping from node id to objects that held by the node.
+  std::unordered_map<ClientID, ObjectSet> node_to_objects_ GUARDED_BY(mutex_);
+};
+
+}  // namespace gcs
+
+}  // namespace ray
+
+#endif  // GCS_GCS_SERVER_OBJECT_LOCATOR_H

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<ClientID> &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<ObjectID> object_ids_;
+  std::unordered_set<ClientID> 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();
+}