diff --git a/BUILD.bazel b/BUILD.bazel index ef5217901..021bfaccd 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -843,6 +843,22 @@ cc_test( ], ) +cc_test( + name = "redis_object_info_accessor_test", + srcs = ["src/ray/gcs/test/redis_object_info_accessor_test.cc"], + args = ["$(location redis-server) $(location redis-cli) $(location libray_redis_module.so)"], + copts = COPTS, + data = [ + "//:libray_redis_module.so", + "//:redis-cli", + "//:redis-server", + ], + deps = [ + ":gcs", + "@com_google_googletest//:gtest_main", + ], +) + cc_test( name = "subscription_executor_test", srcs = ["src/ray/gcs/test/subscription_executor_test.cc"], diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index e21770b71..0b88e7042 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -3,6 +3,7 @@ #include "ray/common/id.h" #include "ray/gcs/callback.h" +#include "ray/gcs/entry_change_notification.h" #include "ray/protobuf/gcs.pb.h" namespace ray { @@ -176,6 +177,64 @@ class TaskInfoAccessor { TaskInfoAccessor() = default; }; +/// `ObjectInfoAccessor` is a sub-interface of `GcsClient`. +/// This class includes all the methods that are related to accessing +/// object information in the GCS. +class ObjectInfoAccessor { + public: + virtual ~ObjectInfoAccessor() {} + + /// Get object's locations from GCS asynchronously. + /// + /// \param object_id The ID of object to lookup in GCS. + /// \param callback Callback that will be called after lookup finishes. + /// \return Status + virtual Status AsyncGetLocations( + const ObjectID &object_id, + const MultiItemCallback &callback) = 0; + + /// Add location of object to GCS asynchronously. + /// + /// \param object_id The ID of object which location will be added to GCS. + /// \param node_id The location that will be added to GCS. + /// \param callback Callback that will be called after object has been added to GCS. + /// \return Status + virtual Status AsyncAddLocation(const ObjectID &object_id, const ClientID &node_id, + const StatusCallback &callback) = 0; + + /// Remove location of object from GCS asynchronously. + /// + /// \param object_id The ID of object which location will be removed from GCS. + /// \param node_id The location that will be removed from GCS. + /// \param callback Callback that will be called after the delete finished. + /// \return Status + virtual Status AsyncRemoveLocation(const ObjectID &object_id, const ClientID &node_id, + const StatusCallback &callback) = 0; + + /// Subscribe to any update of an object's location. + /// + /// \param object_id The ID of the object to be subscribed to. + /// \param subscribe Callback that will be called each time when the object's + /// location is updated. + /// \param done Callback that will be called when subscription is complete. + /// \return Status + virtual Status AsyncSubscribeToLocations( + const ObjectID &object_id, + const SubscribeCallback &subscribe, + const StatusCallback &done) = 0; + + /// Cancel subscription to any update of an object's location. + /// + /// \param object_id The ID of the object to be unsubscribed to. + /// \param done Callback that will be called when unsubscription is complete. + /// \return Status + virtual Status AsyncUnsubscribeToLocations(const ObjectID &object_id, + const StatusCallback &done) = 0; + + protected: + ObjectInfoAccessor() = default; +}; + /// \class NodeInfoAccessor /// `NodeInfoAccessor` is a sub-interface of `GcsClient`. /// This class includes all the methods that are related to accessing diff --git a/src/ray/gcs/entry_change_notification.h b/src/ray/gcs/entry_change_notification.h new file mode 100644 index 000000000..8e1bb4336 --- /dev/null +++ b/src/ray/gcs/entry_change_notification.h @@ -0,0 +1,57 @@ +#ifndef RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H +#define RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H + +#include +#include + +namespace ray { + +namespace gcs { + +/// \class EntryChangeNotification +/// EntryChangeNotification class is a template class which represent +/// notification of entry change from GCS. +template +class EntryChangeNotification { + public: + EntryChangeNotification(rpc::GcsChangeMode change_mode, std::vector data) + : change_mode_(change_mode), data_(std::move(data)) {} + + EntryChangeNotification(EntryChangeNotification &&other) { + change_mode_ = other.change_mode_; + data_ = std::move(other.data_); + } + + EntryChangeNotification &operator=(EntryChangeNotification &&other) { + change_mode_ = other.change_mode_; + data_ = std::move(other.data_); + } + + /// Whether the entry data is removed from GCS. + bool IsRemoved() const { return change_mode_ == rpc::GcsChangeMode::REMOVE; } + + /// Whether the entry data is added to GCS. + bool IsAdded() const { return change_mode_ == rpc::GcsChangeMode::APPEND_OR_ADD; } + + /// Get change mode of this notification. For test only. + /// + /// \return rpc::GcsChangeMode + rpc::GcsChangeMode GetGcsChangeMode() const { return change_mode_; } + + /// Get data of this notification. + /// + /// \return Data + const std::vector &GetData() const { return data_; } + + private: + rpc::GcsChangeMode change_mode_; + std::vector data_; +}; + +typedef EntryChangeNotification ObjectChangeNotification; + +} // namespace gcs + +} // namespace ray + +#endif // RAY_GCS_ENTRY_CHANGE_NOTIFICATION_H diff --git a/src/ray/gcs/gcs_client.h b/src/ray/gcs/gcs_client.h index 687ae3792..5f47ce1fc 100644 --- a/src/ray/gcs/gcs_client.h +++ b/src/ray/gcs/gcs_client.h @@ -74,6 +74,13 @@ class GcsClient : public std::enable_shared_from_this { return *job_accessor_; } + /// Get the sub-interface for accessing object information in GCS. + /// This function is thread safe. + ObjectInfoAccessor &Objects() { + RAY_CHECK(object_accessor_ != nullptr); + return *object_accessor_; + } + /// Get the sub-interface for accessing node information in GCS. /// This function is thread safe. NodeInfoAccessor &Nodes() { @@ -101,6 +108,7 @@ class GcsClient : public std::enable_shared_from_this { std::unique_ptr actor_accessor_; std::unique_ptr job_accessor_; + std::unique_ptr object_accessor_; std::unique_ptr node_accessor_; std::unique_ptr task_accessor_; }; diff --git a/src/ray/gcs/redis_accessor.cc b/src/ray/gcs/redis_accessor.cc index 40ac2b0d5..d99b2a087 100644 --- a/src/ray/gcs/redis_accessor.cc +++ b/src/ray/gcs/redis_accessor.cc @@ -126,12 +126,12 @@ Status RedisActorInfoAccessor::AsyncSubscribe( const ActorID &actor_id, const SubscribeCallback &subscribe, const StatusCallback &done) { RAY_CHECK(subscribe != nullptr); - return actor_sub_executor_.AsyncSubscribe(node_id_, actor_id, subscribe, done); + return actor_sub_executor_.AsyncSubscribe(subscribe_id_, actor_id, subscribe, done); } Status RedisActorInfoAccessor::AsyncUnsubscribe(const ActorID &actor_id, const StatusCallback &done) { - return actor_sub_executor_.AsyncUnsubscribe(node_id_, actor_id, done); + return actor_sub_executor_.AsyncUnsubscribe(subscribe_id_, actor_id, done); } RedisJobInfoAccessor::RedisJobInfoAccessor(RedisGcsClient *client_impl) @@ -228,6 +228,70 @@ Status RedisTaskInfoAccessor::AsyncUnsubscribe(const TaskID &task_id, return task_sub_executor_.AsyncUnsubscribe(subscribe_id_, task_id, done); } +RedisObjectInfoAccessor::RedisObjectInfoAccessor(RedisGcsClient *client_impl) + : client_impl_(client_impl), object_sub_executor_(client_impl->object_table()) {} + +Status RedisObjectInfoAccessor::AsyncGetLocations( + const ObjectID &object_id, const MultiItemCallback &callback) { + RAY_CHECK(callback != nullptr); + auto on_done = [callback](RedisGcsClient *client, const ObjectID &object_id, + const std::vector &data) { + callback(Status::OK(), data); + }; + + ObjectTable &object_table = client_impl_->object_table(); + return object_table.Lookup(JobID::Nil(), object_id, on_done); +} + +Status RedisObjectInfoAccessor::AsyncAddLocation(const ObjectID &object_id, + const ClientID &node_id, + const StatusCallback &callback) { + std::function + on_done = nullptr; + if (callback != nullptr) { + on_done = [callback](RedisGcsClient *client, const ObjectID &object_id, + const ObjectTableData &data) { callback(Status::OK()); }; + } + + std::shared_ptr data_ptr = std::make_shared(); + data_ptr->set_manager(node_id.Binary()); + + ObjectTable &object_table = client_impl_->object_table(); + return object_table.Add(JobID::Nil(), object_id, data_ptr, on_done); +} + +Status RedisObjectInfoAccessor::AsyncRemoveLocation(const ObjectID &object_id, + const ClientID &node_id, + const StatusCallback &callback) { + std::function + on_done = nullptr; + if (callback != nullptr) { + on_done = [callback](RedisGcsClient *client, const ObjectID &object_id, + const ObjectTableData &data) { callback(Status::OK()); }; + } + + std::shared_ptr data_ptr = std::make_shared(); + data_ptr->set_manager(node_id.Binary()); + + ObjectTable &object_table = client_impl_->object_table(); + return object_table.Remove(JobID::Nil(), object_id, data_ptr, on_done); +} + +Status RedisObjectInfoAccessor::AsyncSubscribeToLocations( + const ObjectID &object_id, + const SubscribeCallback &subscribe, + const StatusCallback &done) { + RAY_CHECK(subscribe != nullptr); + return object_sub_executor_.AsyncSubscribe(subscribe_id_, object_id, subscribe, done); +} + +Status RedisObjectInfoAccessor::AsyncUnsubscribeToLocations(const ObjectID &object_id, + const StatusCallback &done) { + return object_sub_executor_.AsyncUnsubscribe(subscribe_id_, object_id, done); +} + RedisNodeInfoAccessor::RedisNodeInfoAccessor(RedisGcsClient *client_impl) : client_impl_(client_impl) {} diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index 071add72a..55e55667a 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -54,7 +54,7 @@ class RedisActorInfoAccessor : public ActorInfoAccessor { // And because the new GCS Client will no longer hold the local ClientID, so we use // random ClientID instead. // TODO(micafan): Remove this random id, once GCS becomes a service. - ClientID node_id_{ClientID::FromRandom()}; + ClientID subscribe_id_{ClientID::FromRandom()}; typedef SubscriptionExecutor ActorSubscriptionExecutor; @@ -132,6 +132,48 @@ class RedisTaskInfoAccessor : public TaskInfoAccessor { TaskSubscriptionExecutor task_sub_executor_; }; +/// \class RedisObjectInfoAccessor +/// RedisObjectInfoAccessor is an implementation of `ObjectInfoAccessor` +/// that uses Redis as the backend storage. +class RedisObjectInfoAccessor : public ObjectInfoAccessor { + public: + explicit RedisObjectInfoAccessor(RedisGcsClient *client_impl); + + virtual ~RedisObjectInfoAccessor() {} + + Status AsyncGetLocations(const ObjectID &object_id, + const MultiItemCallback &callback) override; + + Status AsyncAddLocation(const ObjectID &object_id, const ClientID &node_id, + const StatusCallback &callback) override; + + Status AsyncRemoveLocation(const ObjectID &object_id, const ClientID &node_id, + const StatusCallback &callback) override; + + Status AsyncSubscribeToLocations( + const ObjectID &object_id, + const SubscribeCallback &subscribe, + const StatusCallback &done) override; + + Status AsyncUnsubscribeToLocations(const ObjectID &object_id, + const StatusCallback &done) override; + + private: + RedisGcsClient *client_impl_{nullptr}; + + // Use a random ClientID for object subscription. Because: + // If we use ClientID::Nil, GCS will still send all objects' updates to this GCS Client. + // Even we can filter out irrelevant updates, but there will be extra overhead. + // And because the new GCS Client will no longer hold the local ClientID, so we use + // random ClientID instead. + // TODO(micafan): Remove this random id, once GCS becomes a service. + ClientID subscribe_id_{ClientID::FromRandom()}; + + typedef SubscriptionExecutor + ObjectSubscriptionExecutor; + ObjectSubscriptionExecutor object_sub_executor_; +}; + /// \class RedisNodeInfoAccessor /// RedisNodeInfoAccessor is an implementation of `NodeInfoAccessor` /// that uses Redis as the backend storage. diff --git a/src/ray/gcs/redis_gcs_client.cc b/src/ray/gcs/redis_gcs_client.cc index 9858d5e2e..0fcf9dad7 100644 --- a/src/ray/gcs/redis_gcs_client.cc +++ b/src/ray/gcs/redis_gcs_client.cc @@ -144,6 +144,7 @@ Status RedisGcsClient::Connect(boost::asio::io_service &io_service) { actor_accessor_.reset(new RedisActorInfoAccessor(this)); job_accessor_.reset(new RedisJobInfoAccessor(this)); + object_accessor_.reset(new RedisObjectInfoAccessor(this)); node_accessor_.reset(new RedisNodeInfoAccessor(this)); task_accessor_.reset(new RedisTaskInfoAccessor(this)); diff --git a/src/ray/gcs/redis_gcs_client.h b/src/ray/gcs/redis_gcs_client.h index d4317af4f..e9fd78d1c 100644 --- a/src/ray/gcs/redis_gcs_client.h +++ b/src/ray/gcs/redis_gcs_client.h @@ -24,10 +24,12 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { friend class RedisJobInfoAccessor; friend class RedisTaskInfoAccessor; friend class RedisNodeInfoAccessor; + friend class RedisObjectInfoAccessor; friend class SubscriptionExecutorTest; friend class LogSubscribeTestHelper; friend class TaskTableTestHelper; friend class ClientTableTestHelper; + friend class SetTestHelper; public: /// Constructor of RedisGcsClient. @@ -58,7 +60,6 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { void Disconnect(); // TODO: Some API for getting the error on the driver - ObjectTable &object_table(); TaskReconstructionLog &task_reconstruction_log(); TaskLeaseTable &task_lease_table(); HeartbeatTable &heartbeat_table(); @@ -98,6 +99,8 @@ class RAY_EXPORT RedisGcsClient : public GcsClient { ActorTable &actor_table(); /// This method will be deprecated, use method Jobs() instead. JobTable &job_table(); + /// This method will be deprecated, use method Objects() instead + ObjectTable &object_table(); /// This method will be deprecated, use method Nodes() instead. ClientTable &client_table(); /// This method will be deprecated, use method Tasks() instead. diff --git a/src/ray/gcs/subscription_executor.cc b/src/ray/gcs/subscription_executor.cc index 8d5b60856..a5a46a071 100644 --- a/src/ray/gcs/subscription_executor.cc +++ b/src/ray/gcs/subscription_executor.cc @@ -52,8 +52,6 @@ Status SubscriptionExecutor::AsyncSubscribeAll( return; } - RAY_LOG(DEBUG) << "Subscribe received update of id " << id; - SubscribeCallback sub_one_callback = nullptr; SubscribeCallback sub_all_callback = nullptr; { @@ -190,6 +188,7 @@ template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; template class SubscriptionExecutor; +template class SubscriptionExecutor; } // namespace gcs diff --git a/src/ray/gcs/tables.cc b/src/ray/gcs/tables.cc index d3ef086bd..7aaa7fb51 100644 --- a/src/ray/gcs/tables.cc +++ b/src/ray/gcs/tables.cc @@ -361,6 +361,21 @@ Status Set::Remove(const JobID &job_id, const ID &id, prefix_, pubsub_channel_, std::move(callback)); } +template +Status Set::Subscribe(const JobID &job_id, const ClientID &client_id, + const NotificationCallback &subscribe, + const SubscriptionCallback &done) { + auto on_subscribe = [subscribe](RedisGcsClient *client, const ID &id, + const GcsChangeMode change_mode, + const std::vector &data) { + EntryChangeNotification change_notification(change_mode, data); + std::vector> notification_vec; + notification_vec.emplace_back(std::move(change_notification)); + subscribe(client, id, notification_vec); + }; + return Log::Subscribe(job_id, client_id, on_subscribe, done); +} + template std::string Set::DebugString() const { std::stringstream result; diff --git a/src/ray/gcs/tables.h b/src/ray/gcs/tables.h index f1feb11b1..87b106bea 100644 --- a/src/ray/gcs/tables.h +++ b/src/ray/gcs/tables.h @@ -12,6 +12,7 @@ #include "ray/util/logging.h" #include "ray/gcs/callback.h" +#include "ray/gcs/entry_change_notification.h" #include "ray/gcs/redis_context.h" #include "ray/protobuf/gcs.pb.h" @@ -93,9 +94,11 @@ class Log : public LogInterface, virtual public PubsubInterface { public: using Callback = std::function &data)>; + using NotificationCallback = std::function &data)>; + /// The callback to call when a write to a key succeeds. using WriteCallback = typename LogInterface::WriteCallback; /// The callback to call when a SUBSCRIBE call completes and we are ready to @@ -164,6 +167,7 @@ class Log : public LogInterface, virtual public PubsubInterface { /// called with an empty vector, then there was no data at the key. /// \return Status Status Lookup(const JobID &job_id, const ID &id, const Callback &lookup); + /// Subscribe to any Append operations to this table. The caller may choose /// requests notifications for. This may only be called once per Log /// @@ -210,6 +214,28 @@ class Log : public LogInterface, virtual public PubsubInterface { Status CancelNotifications(const JobID &job_id, const ID &id, const ClientID &client_id, const StatusCallback &done); + /// Subscribe to any modifications to the key. The caller may choose + /// to subscribe to all modifications, or to subscribe only to keys that it + /// requests notifications for. This may only be called once per Log + /// instance. This function is different from public version due to + /// an additional parameter change_mode in NotificationCallback. Therefore this + /// function supports notifications of remove operations. + /// + /// \param job_id The ID of the job. + /// \param client_id The type of update to listen to. If this is nil, then a + /// message for each Add to the table will be received. Else, only + /// messages for the given client will be received. In the latter + /// case, the client may request notifications on specific keys in the + /// table via `RequestNotifications`. + /// \param subscribe Callback that is called on each received message. If the + /// callback is called with an empty vector, then there was no data at the key. + /// \param done Callback that is called when subscription is complete and we + /// are ready to receive messages. + /// \return Status + Status Subscribe(const JobID &job_id, const ClientID &client_id, + const NotificationCallback &subscribe, + const SubscriptionCallback &done); + /// Delete an entire key from redis. /// /// \param job_id The ID of the job. @@ -235,28 +261,6 @@ class Log : public LogInterface, virtual public PubsubInterface { return shard_contexts_[index(id) % shard_contexts_.size()]; } - /// Subscribe to any modifications to the key. The caller may choose - /// to subscribe to all modifications, or to subscribe only to keys that it - /// requests notifications for. This may only be called once per Log - /// instance. This function is different from public version due to - /// an additional parameter change_mode in NotificationCallback. Therefore this - /// function supports notifications of remove operations. - /// - /// \param job_id The ID of the job. - /// \param client_id The type of update to listen to. If this is nil, then a - /// message for each Add to the table will be received. Else, only - /// messages for the given client will be received. In the latter - /// case, the client may request notifications on specific keys in the - /// table via `RequestNotifications`. - /// \param subscribe Callback that is called on each received message. If the - /// callback is called with an empty vector, then there was no data at the key. - /// \param done Callback that is called when subscription is complete and we - /// are ready to receive messages. - /// \return Status - Status Subscribe(const JobID &job_id, const ClientID &client_id, - const NotificationCallback &subscribe, - const SubscriptionCallback &done); - /// The connection to the GCS. std::vector> shard_contexts_; /// The GCS client. @@ -435,7 +439,6 @@ class Set : private Log, public: using Callback = typename Log::Callback; using WriteCallback = typename Log::WriteCallback; - using NotificationCallback = typename Log::NotificationCallback; using SubscriptionCallback = typename Log::SubscriptionCallback; Set(const std::vector> &contexts, RedisGcsClient *client) @@ -468,11 +471,24 @@ class Set : private Log, Status Remove(const JobID &job_id, const ID &id, const std::shared_ptr &data, const WriteCallback &done); + using NotificationCallback = + std::function &data)>; + /// Subscribe to any add or remove operations to this table. + /// + /// \param job_id The ID of the job. + /// \param client_id The type of update to listen to. If this is nil, then a + /// message for each add or remove to the table will be received. Else, only + /// messages for the given client will be received. In the latter + /// case, the client may request notifications on specific keys in the + /// table via `RequestNotifications`. + /// \param subscribe Callback that is called on each received message. + /// \param done Callback that is called when subscription is complete and we + /// are ready to receive messages. + /// \return Status Status Subscribe(const JobID &job_id, const ClientID &client_id, const NotificationCallback &subscribe, - const SubscriptionCallback &done) { - return Log::Subscribe(job_id, client_id, subscribe, done); - } + const SubscriptionCallback &done); /// Returns debug string for class. /// diff --git a/src/ray/gcs/test/redis_gcs_client_test.cc b/src/ray/gcs/test/redis_gcs_client_test.cc index f2cb8244b..1f5d16033 100644 --- a/src/ray/gcs/test/redis_gcs_client_test.cc +++ b/src/ray/gcs/test/redis_gcs_client_test.cc @@ -470,71 +470,341 @@ TEST_F(TestGcsWithAsio, TestLogAppendAt) { TestLogAppendAt(job_id_, client_); } -void TestSet(const JobID &job_id, std::shared_ptr client) { - // Add some entries to the set at an object ID. - ObjectID object_id = ObjectID::FromRandom(); - std::vector managers = {"abc", "def", "ghi"}; - for (auto &manager : managers) { - auto data = std::make_shared(); - data->set_manager(manager); - // Check that we added the correct object entries. - auto add_callback = [object_id, data](gcs::RedisGcsClient *client, const ObjectID &id, - const ObjectTableData &d) { +class SetTestHelper { + public: + static void TestSet(const JobID &job_id, std::shared_ptr client) { + // Add some entries to the set at an object ID. + ObjectID object_id = ObjectID::FromRandom(); + std::vector managers = {"abc", "def", "ghi"}; + for (auto &manager : managers) { + auto data = std::make_shared(); + data->set_manager(manager); + // Check that we added the correct object entries. + auto add_callback = [object_id, data](gcs::RedisGcsClient *client, + const ObjectID &id, + const ObjectTableData &d) { + ASSERT_EQ(id, object_id); + ASSERT_EQ(data->manager(), d.manager()); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback)); + } + + // Check that lookup returns the added object entries. + auto lookup_callback = [object_id, managers]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector &data) { ASSERT_EQ(id, object_id); - ASSERT_EQ(data->manager(), d.manager()); + ASSERT_EQ(data.size(), managers.size()); test->IncrementNumCallbacks(); }; - RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback)); - } - // Check that lookup returns the added object entries. - auto lookup_callback = [object_id, managers](gcs::RedisGcsClient *client, - const ObjectID &id, - const std::vector &data) { - ASSERT_EQ(id, object_id); - ASSERT_EQ(data.size(), managers.size()); - test->IncrementNumCallbacks(); - }; + // Do a lookup at the object ID. + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); - // Do a lookup at the object ID. - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); + for (auto &manager : managers) { + auto data = std::make_shared(); + data->set_manager(manager); + // Check that we added the correct object entries. + auto remove_entry_callback = [object_id, data](gcs::RedisGcsClient *client, + const ObjectID &id, + const ObjectTableData &d) { + ASSERT_EQ(id, object_id); + ASSERT_EQ(data->manager(), d.manager()); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK( + client->object_table().Remove(job_id, object_id, data, remove_entry_callback)); + } - for (auto &manager : managers) { - auto data = std::make_shared(); - data->set_manager(manager); - // Check that we added the correct object entries. - auto remove_entry_callback = [object_id, data](gcs::RedisGcsClient *client, - const ObjectID &id, - const ObjectTableData &d) { + // Check that the entries are removed. + auto lookup_callback2 = [object_id, managers]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector &data) { ASSERT_EQ(id, object_id); - ASSERT_EQ(data->manager(), d.manager()); + ASSERT_EQ(data.size(), 0); test->IncrementNumCallbacks(); + test->Stop(); }; - RAY_CHECK_OK( - client->object_table().Remove(job_id, object_id, data, remove_entry_callback)); + + // Do a lookup at the object ID. + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback2)); + // Run the event loop. The loop will only stop if the Lookup callback is + // called (or an assertion failure). + test->Start(); + ASSERT_EQ(test->NumCallbacks(), managers.size() * 2 + 2); } - // Check that the entries are removed. - auto lookup_callback2 = [object_id, managers]( - gcs::RedisGcsClient *client, const ObjectID &id, - const std::vector &data) { - ASSERT_EQ(id, object_id); - ASSERT_EQ(data.size(), 0); - test->IncrementNumCallbacks(); - test->Stop(); - }; + static void TestDeleteKeysFromSet( + const JobID &job_id, std::shared_ptr client, + std::vector> &data_vector) { + std::vector ids; + ObjectID object_id; + for (auto &data : data_vector) { + object_id = ObjectID::FromRandom(); + ids.push_back(object_id); + // Check that we added the correct object entries. + auto add_callback = [object_id, data](gcs::RedisGcsClient *client, + const ObjectID &id, + const ObjectTableData &d) { + ASSERT_EQ(id, object_id); + ASSERT_EQ(data->manager(), d.manager()); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback)); + } + for (const auto &object_id : ids) { + // Check that lookup returns the added object entries. + auto lookup_callback = [object_id, data_vector]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector &data) { + ASSERT_EQ(id, object_id); + ASSERT_EQ(data.size(), 1); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); + } + if (ids.size() == 1) { + client->object_table().Delete(job_id, ids[0]); + } else { + client->object_table().Delete(job_id, ids); + } + for (const auto &object_id : ids) { + auto lookup_callback = [object_id](gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector &data) { + ASSERT_EQ(id, object_id); + ASSERT_TRUE(data.size() == 0); + test->IncrementNumCallbacks(); + }; + RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); + } + } - // Do a lookup at the object ID. - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback2)); - // Run the event loop. The loop will only stop if the Lookup callback is - // called (or an assertion failure). - test->Start(); - ASSERT_EQ(test->NumCallbacks(), managers.size() * 2 + 2); -} + static void TestSetSubscribeAll(const JobID &job_id, + std::shared_ptr client) { + std::vector object_ids; + for (int i = 0; i < 3; i++) { + object_ids.emplace_back(ObjectID::FromRandom()); + } + std::vector managers = {"abc", "def", "ghi"}; + + // Callback for a notification. + auto notification_callback = + [object_ids, managers]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector ¬ifications) { + if (test->NumCallbacks() < 3 * 3) { + ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD); + } else { + ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::REMOVE); + } + ASSERT_EQ(id, object_ids[test->NumCallbacks() / 3 % 3]); + // Check that we get notifications in the same order as the writes. + for (const auto &entry : notifications[0].GetData()) { + ASSERT_EQ(entry.manager(), managers[test->NumCallbacks() % 3]); + test->IncrementNumCallbacks(); + } + if (test->NumCallbacks() == object_ids.size() * 3 * 2) { + test->Stop(); + } + }; + + // Callback for subscription success. We are guaranteed to receive + // notifications after this is called. + auto subscribe_callback = [job_id, object_ids, + managers](gcs::RedisGcsClient *client) { + // We have subscribed. Do the writes to the table. + for (size_t i = 0; i < object_ids.size(); i++) { + for (size_t j = 0; j < managers.size(); j++) { + auto data = std::make_shared(); + data->set_manager(managers[j]); + for (int k = 0; k < 3; k++) { + // Add the same entry several times. + // Expect no notification if the entry already exists. + RAY_CHECK_OK( + client->object_table().Add(job_id, object_ids[i], data, nullptr)); + } + } + } + for (size_t i = 0; i < object_ids.size(); i++) { + for (size_t j = 0; j < managers.size(); j++) { + auto data = std::make_shared(); + data->set_manager(managers[j]); + for (int k = 0; k < 3; k++) { + // Remove the same entry several times. + // Expect no notification if the entry doesn't exist. + RAY_CHECK_OK( + client->object_table().Remove(job_id, object_ids[i], data, nullptr)); + } + } + } + }; + + // Subscribe to all driver table notifications. Once we have successfully + // subscribed, we will append to the key several times and check that we get + // notified for each. + RAY_CHECK_OK(client->object_table().Subscribe( + job_id, ClientID::Nil(), notification_callback, subscribe_callback)); + + // Run the event loop. The loop will only stop if the registered subscription + // callback is called (or an assertion failure). + test->Start(); + // Check that we received one notification callback for each write. + ASSERT_EQ(test->NumCallbacks(), object_ids.size() * 3 * 2); + } + + static void TestSetSubscribeId(const JobID &job_id, + std::shared_ptr client) { + // Add a set entry. + ObjectID object_id1 = ObjectID::FromRandom(); + std::vector managers1 = {"abc", "def", "ghi"}; + auto data1 = std::make_shared(); + data1->set_manager(managers1[0]); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data1, nullptr)); + + // Add a set entry at a second key. + ObjectID object_id2 = ObjectID::FromRandom(); + std::vector managers2 = {"jkl", "mno", "pqr"}; + auto data2 = std::make_shared(); + data2->set_manager(managers2[0]); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data2, nullptr)); + + // The callback for a notification from the table. This should only be + // received for keys that we requested notifications for. + auto notification_callback = + [object_id2, managers2]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector ¬ifications) { + ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD); + // Check that we only get notifications for the requested key. + ASSERT_EQ(id, object_id2); + // Check that we get notifications in the same order as the writes. + for (const auto &entry : notifications[0].GetData()) { + ASSERT_EQ(entry.manager(), managers2[test->NumCallbacks()]); + test->IncrementNumCallbacks(); + } + if (test->NumCallbacks() == managers2.size()) { + test->Stop(); + } + }; + + // The callback for subscription success. Once we've subscribed, request + // notifications for only one of the keys, then write to both keys. + auto subscribe_callback = [job_id, object_id1, object_id2, managers1, + managers2](gcs::RedisGcsClient *client) { + // Request notifications for one of the keys. + RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id2, + local_client_id, nullptr)); + // Write both keys. We should only receive notifications for the key that + // we requested them for. + auto remaining = std::vector(++managers1.begin(), managers1.end()); + for (const auto &manager : remaining) { + auto data = std::make_shared(); + data->set_manager(manager); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data, nullptr)); + } + remaining = std::vector(++managers2.begin(), managers2.end()); + for (const auto &manager : remaining) { + auto data = std::make_shared(); + data->set_manager(manager); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data, nullptr)); + } + }; + + // Subscribe to notifications for this client. This allows us to request and + // receive notifications for specific keys. + RAY_CHECK_OK(client->object_table().Subscribe( + job_id, local_client_id, notification_callback, subscribe_callback)); + // Run the event loop. The loop will only stop if the registered subscription + // callback is called for the requested key. + test->Start(); + // Check that we received one notification callback for each write to the + // requested key. + ASSERT_EQ(test->NumCallbacks(), managers2.size()); + } + + static void TestSetSubscribeCancel(const JobID &job_id, + std::shared_ptr client) { + // Add a set entry. + ObjectID object_id = ObjectID::FromRandom(); + std::vector managers = {"jkl", "mno", "pqr"}; + auto data = std::make_shared(); + data->set_manager(managers[0]); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr)); + + // The callback for a notification from the object table. This should only be + // received for the object that we requested notifications for. + auto notification_callback = + [object_id, managers]( + gcs::RedisGcsClient *client, const ObjectID &id, + const std::vector ¬ifications) { + ASSERT_EQ(notifications[0].GetGcsChangeMode(), GcsChangeMode::APPEND_OR_ADD); + ASSERT_EQ(id, object_id); + // Check that we get a duplicate notification for the first write. We get a + // duplicate notification because notifications + // are canceled after the first write, then requested again. + const std::vector &data = notifications[0].GetData(); + if (data.size() == 1) { + // first notification + ASSERT_EQ(data[0].manager(), managers[0]); + test->IncrementNumCallbacks(); + } else { + // second notification + ASSERT_EQ(data.size(), managers.size()); + std::unordered_set managers_set(managers.begin(), + managers.end()); + std::unordered_set data_managers_set; + for (const auto &entry : data) { + data_managers_set.insert(entry.manager()); + test->IncrementNumCallbacks(); + } + ASSERT_EQ(managers_set, data_managers_set); + } + if (test->NumCallbacks() == managers.size() + 1) { + test->Stop(); + } + }; + + // The callback for a notification from the table. This should only be + // received for keys that we requested notifications for. + auto subscribe_callback = [job_id, object_id, managers](gcs::RedisGcsClient *client) { + // Request notifications, then cancel immediately. We should receive a + // notification for the current value at the key. + RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id, + local_client_id, nullptr)); + RAY_CHECK_OK(client->object_table().CancelNotifications(job_id, object_id, + local_client_id, nullptr)); + // Add to the key. Since we canceled notifications, we should not + // receive a notification for these writes. + auto remaining = std::vector(++managers.begin(), managers.end()); + for (const auto &manager : remaining) { + auto data = std::make_shared(); + data->set_manager(manager); + RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr)); + } + // Request notifications again. We should receive a notification for the + // current values at the key. + RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id, + local_client_id, nullptr)); + }; + + // Subscribe to notifications for this client. This allows us to request and + // receive notifications for specific keys. + RAY_CHECK_OK(client->object_table().Subscribe( + job_id, local_client_id, notification_callback, subscribe_callback)); + // Run the event loop. The loop will only stop if the registered subscription + // callback is called for the requested key. + test->Start(); + // Check that we received a notification callback for the first append to the + // key, then a notification for all of the appends, because we cancel + // notifications in between. + ASSERT_EQ(test->NumCallbacks(), managers.size() + 1); + } +}; TEST_F(TestGcsWithAsio, TestSet) { test = this; - TestSet(job_id_, client_); + SetTestHelper::TestSet(job_id_, client_); } void TestDeleteKeysFromLog( @@ -584,50 +854,6 @@ void TestDeleteKeysFromLog( } } -void TestDeleteKeysFromSet(const JobID &job_id, - std::shared_ptr client, - std::vector> &data_vector) { - std::vector ids; - ObjectID object_id; - for (auto &data : data_vector) { - object_id = ObjectID::FromRandom(); - ids.push_back(object_id); - // Check that we added the correct object entries. - auto add_callback = [object_id, data](gcs::RedisGcsClient *client, const ObjectID &id, - const ObjectTableData &d) { - ASSERT_EQ(id, object_id); - ASSERT_EQ(data->manager(), d.manager()); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, add_callback)); - } - for (const auto &object_id : ids) { - // Check that lookup returns the added object entries. - auto lookup_callback = [object_id, data_vector]( - gcs::RedisGcsClient *client, const ObjectID &id, - const std::vector &data) { - ASSERT_EQ(id, object_id); - ASSERT_EQ(data.size(), 1); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); - } - if (ids.size() == 1) { - client->object_table().Delete(job_id, ids[0]); - } else { - client->object_table().Delete(job_id, ids); - } - for (const auto &object_id : ids) { - auto lookup_callback = [object_id](gcs::RedisGcsClient *client, const ObjectID &id, - const std::vector &data) { - ASSERT_EQ(id, object_id); - ASSERT_TRUE(data.size() == 0); - test->IncrementNumCallbacks(); - }; - RAY_CHECK_OK(client->object_table().Lookup(job_id, object_id, lookup_callback)); - } -} - // Test delete function for keys of Log or Table. void TestDeleteKeys(const JobID &job_id, std::shared_ptr client) { // Test delete function for keys of Log. @@ -695,20 +921,20 @@ void TestDeleteKeys(const JobID &job_id, std::shared_ptr cl // Test one element case. AppendObjectData(1); ASSERT_EQ(object_vector.size(), 1); - TestDeleteKeysFromSet(job_id, client, object_vector); + SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector); // Test the case for more than one elements and less than // maximum_gcs_deletion_batch_size. AppendObjectData(RayConfig::instance().maximum_gcs_deletion_batch_size() / 2); ASSERT_GT(object_vector.size(), 1); ASSERT_LT(object_vector.size(), RayConfig::instance().maximum_gcs_deletion_batch_size()); - TestDeleteKeysFromSet(job_id, client, object_vector); + SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector); // Test the case for more than maximum_gcs_deletion_batch_size. // The Delete function will split the data into two commands. AppendObjectData(RayConfig::instance().maximum_gcs_deletion_batch_size() / 2); ASSERT_GT(object_vector.size(), RayConfig::instance().maximum_gcs_deletion_batch_size()); - TestDeleteKeysFromSet(job_id, client, object_vector); + SetTestHelper::TestDeleteKeysFromSet(job_id, client, object_vector); } TEST_F(TestGcsWithAsio, TestDeleteKey) { @@ -903,80 +1129,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeAll) { LogSubscribeTestHelper::TestLogSubscribeAll(job_id_, client_); } -void TestSetSubscribeAll(const JobID &job_id, - std::shared_ptr client) { - std::vector object_ids; - for (int i = 0; i < 3; i++) { - object_ids.emplace_back(ObjectID::FromRandom()); - } - std::vector managers = {"abc", "def", "ghi"}; - - // Callback for a notification. - auto notification_callback = [object_ids, managers]( - gcs::RedisGcsClient *client, const ObjectID &id, - const GcsChangeMode change_mode, - const std::vector data) { - if (test->NumCallbacks() < 3 * 3) { - ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD); - } else { - ASSERT_EQ(change_mode, GcsChangeMode::REMOVE); - } - ASSERT_EQ(id, object_ids[test->NumCallbacks() / 3 % 3]); - // Check that we get notifications in the same order as the writes. - for (const auto &entry : data) { - ASSERT_EQ(entry.manager(), managers[test->NumCallbacks() % 3]); - test->IncrementNumCallbacks(); - } - if (test->NumCallbacks() == object_ids.size() * 3 * 2) { - test->Stop(); - } - }; - - // Callback for subscription success. We are guaranteed to receive - // notifications after this is called. - auto subscribe_callback = [job_id, object_ids, managers](gcs::RedisGcsClient *client) { - // We have subscribed. Do the writes to the table. - for (size_t i = 0; i < object_ids.size(); i++) { - for (size_t j = 0; j < managers.size(); j++) { - auto data = std::make_shared(); - data->set_manager(managers[j]); - for (int k = 0; k < 3; k++) { - // Add the same entry several times. - // Expect no notification if the entry already exists. - RAY_CHECK_OK(client->object_table().Add(job_id, object_ids[i], data, nullptr)); - } - } - } - for (size_t i = 0; i < object_ids.size(); i++) { - for (size_t j = 0; j < managers.size(); j++) { - auto data = std::make_shared(); - data->set_manager(managers[j]); - for (int k = 0; k < 3; k++) { - // Remove the same entry several times. - // Expect no notification if the entry doesn't exist. - RAY_CHECK_OK( - client->object_table().Remove(job_id, object_ids[i], data, nullptr)); - } - } - } - }; - - // Subscribe to all driver table notifications. Once we have successfully - // subscribed, we will append to the key several times and check that we get - // notified for each. - RAY_CHECK_OK(client->object_table().Subscribe( - job_id, ClientID::Nil(), notification_callback, subscribe_callback)); - - // Run the event loop. The loop will only stop if the registered subscription - // callback is called (or an assertion failure). - test->Start(); - // Check that we received one notification callback for each write. - ASSERT_EQ(test->NumCallbacks(), object_ids.size() * 3 * 2); -} - TEST_F(TestGcsWithAsio, TestSetSubscribeAll) { test = this; - TestSetSubscribeAll(job_id_, client_); + SetTestHelper::TestSetSubscribeAll(job_id_, client_); } TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableSubscribeId); @@ -986,79 +1141,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeId) { LogSubscribeTestHelper::TestLogSubscribeId(job_id_, client_); } -void TestSetSubscribeId(const JobID &job_id, - std::shared_ptr client) { - // Add a set entry. - ObjectID object_id1 = ObjectID::FromRandom(); - std::vector managers1 = {"abc", "def", "ghi"}; - auto data1 = std::make_shared(); - data1->set_manager(managers1[0]); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data1, nullptr)); - - // Add a set entry at a second key. - ObjectID object_id2 = ObjectID::FromRandom(); - std::vector managers2 = {"jkl", "mno", "pqr"}; - auto data2 = std::make_shared(); - data2->set_manager(managers2[0]); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data2, nullptr)); - - // The callback for a notification from the table. This should only be - // received for keys that we requested notifications for. - auto notification_callback = [object_id2, managers2]( - gcs::RedisGcsClient *client, const ObjectID &id, - const GcsChangeMode change_mode, - const std::vector &data) { - ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD); - // Check that we only get notifications for the requested key. - ASSERT_EQ(id, object_id2); - // Check that we get notifications in the same order as the writes. - for (const auto &entry : data) { - ASSERT_EQ(entry.manager(), managers2[test->NumCallbacks()]); - test->IncrementNumCallbacks(); - } - if (test->NumCallbacks() == managers2.size()) { - test->Stop(); - } - }; - - // The callback for subscription success. Once we've subscribed, request - // notifications for only one of the keys, then write to both keys. - auto subscribe_callback = [job_id, object_id1, object_id2, managers1, - managers2](gcs::RedisGcsClient *client) { - // Request notifications for one of the keys. - RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id2, - local_client_id, nullptr)); - // Write both keys. We should only receive notifications for the key that - // we requested them for. - auto remaining = std::vector(++managers1.begin(), managers1.end()); - for (const auto &manager : remaining) { - auto data = std::make_shared(); - data->set_manager(manager); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id1, data, nullptr)); - } - remaining = std::vector(++managers2.begin(), managers2.end()); - for (const auto &manager : remaining) { - auto data = std::make_shared(); - data->set_manager(manager); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id2, data, nullptr)); - } - }; - - // Subscribe to notifications for this client. This allows us to request and - // receive notifications for specific keys. - RAY_CHECK_OK(client->object_table().Subscribe( - job_id, local_client_id, notification_callback, subscribe_callback)); - // Run the event loop. The loop will only stop if the registered subscription - // callback is called for the requested key. - test->Start(); - // Check that we received one notification callback for each write to the - // requested key. - ASSERT_EQ(test->NumCallbacks(), managers2.size()); -} - TEST_F(TestGcsWithAsio, TestSetSubscribeId) { test = this; - TestSetSubscribeId(job_id_, client_); + SetTestHelper::TestSetSubscribeId(job_id_, client_); } TEST_TASK_TABLE_MACRO(TestGcsWithAsio, TestTableSubscribeCancel); @@ -1068,85 +1153,9 @@ TEST_F(TestGcsWithAsio, TestLogSubscribeCancel) { LogSubscribeTestHelper::TestLogSubscribeCancel(job_id_, client_); } -void TestSetSubscribeCancel(const JobID &job_id, - std::shared_ptr client) { - // Add a set entry. - ObjectID object_id = ObjectID::FromRandom(); - std::vector managers = {"jkl", "mno", "pqr"}; - auto data = std::make_shared(); - data->set_manager(managers[0]); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr)); - - // The callback for a notification from the object table. This should only be - // received for the object that we requested notifications for. - auto notification_callback = [object_id, managers]( - gcs::RedisGcsClient *client, const ObjectID &id, - const GcsChangeMode change_mode, - const std::vector &data) { - ASSERT_EQ(change_mode, GcsChangeMode::APPEND_OR_ADD); - ASSERT_EQ(id, object_id); - // Check that we get a duplicate notification for the first write. We get a - // duplicate notification because notifications - // are canceled after the first write, then requested again. - if (data.size() == 1) { - // first notification - ASSERT_EQ(data[0].manager(), managers[0]); - test->IncrementNumCallbacks(); - } else { - // second notification - ASSERT_EQ(data.size(), managers.size()); - std::unordered_set managers_set(managers.begin(), managers.end()); - std::unordered_set data_managers_set; - for (const auto &entry : data) { - data_managers_set.insert(entry.manager()); - test->IncrementNumCallbacks(); - } - ASSERT_EQ(managers_set, data_managers_set); - } - if (test->NumCallbacks() == managers.size() + 1) { - test->Stop(); - } - }; - - // The callback for a notification from the table. This should only be - // received for keys that we requested notifications for. - auto subscribe_callback = [job_id, object_id, managers](gcs::RedisGcsClient *client) { - // Request notifications, then cancel immediately. We should receive a - // notification for the current value at the key. - RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id, - local_client_id, nullptr)); - RAY_CHECK_OK(client->object_table().CancelNotifications(job_id, object_id, - local_client_id, nullptr)); - // Add to the key. Since we canceled notifications, we should not - // receive a notification for these writes. - auto remaining = std::vector(++managers.begin(), managers.end()); - for (const auto &manager : remaining) { - auto data = std::make_shared(); - data->set_manager(manager); - RAY_CHECK_OK(client->object_table().Add(job_id, object_id, data, nullptr)); - } - // Request notifications again. We should receive a notification for the - // current values at the key. - RAY_CHECK_OK(client->object_table().RequestNotifications(job_id, object_id, - local_client_id, nullptr)); - }; - - // Subscribe to notifications for this client. This allows us to request and - // receive notifications for specific keys. - RAY_CHECK_OK(client->object_table().Subscribe( - job_id, local_client_id, notification_callback, subscribe_callback)); - // Run the event loop. The loop will only stop if the registered subscription - // callback is called for the requested key. - test->Start(); - // Check that we received a notification callback for the first append to the - // key, then a notification for all of the appends, because we cancel - // notifications in between. - ASSERT_EQ(test->NumCallbacks(), managers.size() + 1); -} - TEST_F(TestGcsWithAsio, TestSetSubscribeCancel) { test = this; - TestSetSubscribeCancel(job_id_, client_); + SetTestHelper::TestSetSubscribeCancel(job_id_, client_); } /// A helper class for ClientTable testing. diff --git a/src/ray/gcs/test/redis_object_info_accessor_test.cc b/src/ray/gcs/test/redis_object_info_accessor_test.cc new file mode 100644 index 000000000..3dc7247f6 --- /dev/null +++ b/src/ray/gcs/test/redis_object_info_accessor_test.cc @@ -0,0 +1,138 @@ +#include +#include +#include "gtest/gtest.h" +#include "ray/gcs/redis_accessor.h" +#include "ray/gcs/redis_gcs_client.h" +#include "ray/gcs/test/accessor_test_base.h" +#include "ray/util/test_util.h" + +namespace ray { + +namespace gcs { + +class RedisObjectInfoAccessorTest : public AccessorTestBase { + protected: + void GenTestData() { + for (size_t i = 0; i < object_count_; ++i) { + ObjectVector object_vec; + for (size_t j = 0; j < copy_count_; ++j) { + auto object = std::make_shared(); + ClientID node_id = ClientID::FromRandom(); + object->set_manager(node_id.Binary()); + object_vec.emplace_back(std::move(object)); + } + ObjectID id = ObjectID::FromRandom(); + object_id_to_data_[id] = object_vec; + } + } + + typedef std::vector> ObjectVector; + std::unordered_map object_id_to_data_; + + size_t object_count_{100}; + size_t copy_count_{5}; +}; + +TEST_F(RedisObjectInfoAccessorTest, TestGetAddRemove) { + ObjectInfoAccessor &object_accessor = gcs_client_->Objects(); + // add && get + // add + for (const auto &elem : object_id_to_data_) { + for (const auto &item : elem.second) { + ++pending_count_; + ClientID node_id = ClientID::FromBinary(item->manager()); + RAY_CHECK_OK( + object_accessor.AsyncAddLocation(elem.first, node_id, [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + })); + } + } + WaitPendingDone(wait_pending_timeout_); + // get + for (const auto &elem : object_id_to_data_) { + ++pending_count_; + size_t total_size = elem.second.size(); + RAY_CHECK_OK(object_accessor.AsyncGetLocations( + elem.first, + [this, total_size](Status status, const std::vector &result) { + RAY_CHECK_OK(status); + RAY_CHECK(total_size == result.size()); + --pending_count_; + })); + } + WaitPendingDone(wait_pending_timeout_); + + RAY_LOG(INFO) << "Case Add && Get done."; + + // subscribe && delete + // subscribe + std::atomic sub_pending_count(0); + auto subscribe = [this, &sub_pending_count](const ObjectID &object_id, + const ObjectChangeNotification &result) { + const auto it = object_id_to_data_.find(object_id); + ASSERT_TRUE(it != object_id_to_data_.end()); + static size_t response_count = 1; + size_t cur_count = response_count <= object_count_ ? copy_count_ : 1; + ASSERT_EQ(result.GetData().size(), cur_count); + rpc::GcsChangeMode change_mode = response_count <= object_count_ + ? rpc::GcsChangeMode::APPEND_OR_ADD + : rpc::GcsChangeMode::REMOVE; + ASSERT_EQ(change_mode, result.GetGcsChangeMode()); + ++response_count; + --sub_pending_count; + }; + for (const auto &elem : object_id_to_data_) { + ++pending_count_; + ++sub_pending_count; + RAY_CHECK_OK(object_accessor.AsyncSubscribeToLocations(elem.first, subscribe, + [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + })); + } + WaitPendingDone(wait_pending_timeout_); + WaitPendingDone(sub_pending_count, wait_pending_timeout_); + // delete + for (const auto &elem : object_id_to_data_) { + ++pending_count_; + ++sub_pending_count; + const ObjectVector &object_vec = elem.second; + ClientID node_id = ClientID::FromBinary(object_vec[0]->manager()); + RAY_CHECK_OK( + object_accessor.AsyncRemoveLocation(elem.first, node_id, [this](Status status) { + RAY_CHECK_OK(status); + --pending_count_; + })); + } + WaitPendingDone(wait_pending_timeout_); + WaitPendingDone(sub_pending_count, wait_pending_timeout_); + // get + for (const auto &elem : object_id_to_data_) { + ++pending_count_; + size_t total_size = elem.second.size(); + RAY_CHECK_OK(object_accessor.AsyncGetLocations( + elem.first, + [this, total_size](Status status, const std::vector &result) { + RAY_CHECK_OK(status); + ASSERT_EQ(total_size - 1, result.size()); + --pending_count_; + })); + } + WaitPendingDone(wait_pending_timeout_); + + RAY_LOG(INFO) << "Case Subscribe && Delete done."; +} + +} // namespace gcs + +} // namespace ray + +int main(int argc, char **argv) { + ::testing::InitGoogleTest(&argc, argv); + RAY_CHECK(argc == 4); + ray::REDIS_SERVER_EXEC_PATH = argv[1]; + ray::REDIS_CLIENT_EXEC_PATH = argv[2]; + ray::REDIS_MODULE_LIBRARY_PATH = argv[3]; + return RUN_ALL_TESTS(); +} diff --git a/src/ray/gcs/test/subscription_executor_test.cc b/src/ray/gcs/test/subscription_executor_test.cc index 40b19a70e..4e5be1a73 100644 --- a/src/ray/gcs/test/subscription_executor_test.cc +++ b/src/ray/gcs/test/subscription_executor_test.cc @@ -1,6 +1,7 @@ #include "ray/gcs/subscription_executor.h" #include "gtest/gtest.h" #include "ray/gcs/callback.h" +#include "ray/gcs/entry_change_notification.h" #include "ray/gcs/redis_gcs_client.h" #include "ray/gcs/test/accessor_test_base.h" @@ -190,6 +191,7 @@ TEST_F(SubscriptionExecutorTest, UnsubscribeTest) { WaitPendingDone(do_sub_pending_count_, wait_pending_timeout_); sub_pending_count_ = id_to_data_.size(); AsyncRegisterActorToGcs(); + WaitPendingDone(pending_count_, wait_pending_timeout_); WaitPendingDone(sub_pending_count_, wait_pending_timeout_); } diff --git a/src/ray/object_manager/object_directory.cc b/src/ray/object_manager/object_directory.cc index f4aab2832..7d510e6af 100644 --- a/src/ray/object_manager/object_directory.cc +++ b/src/ray/object_manager/object_directory.cc @@ -15,7 +15,7 @@ using ray::rpc::ObjectTableData; /// Process a notification of the object table entries and store the result in /// node_ids. This assumes that node_ids already contains the result of the /// object table entries up to but not including this notification. -void UpdateObjectLocations(const GcsChangeMode change_mode, +void UpdateObjectLocations(bool is_added, const std::vector &location_updates, std::shared_ptr gcs_client, std::unordered_set *node_ids) { @@ -24,7 +24,7 @@ void UpdateObjectLocations(const GcsChangeMode change_mode, // addition or deletion. for (const auto &object_table_data : location_updates) { ClientID node_id = ClientID::FromBinary(object_table_data.manager()); - if (change_mode != GcsChangeMode::REMOVE) { + if (is_added) { node_ids->insert(node_id); } else { node_ids->erase(node_id); @@ -42,52 +42,12 @@ void UpdateObjectLocations(const GcsChangeMode change_mode, } // namespace -void ObjectDirectory::RegisterBackend() { - auto object_notification_callback = - [this](gcs::RedisGcsClient *client, const ObjectID &object_id, - const GcsChangeMode change_mode, - const std::vector &location_updates) { - // Objects are added to this map in SubscribeObjectLocations. - auto it = listeners_.find(object_id); - // Do nothing for objects we are not listening for. - if (it == listeners_.end()) { - return; - } - - // Once this flag is set to true, it should never go back to false. - it->second.subscribed = true; - - // Update entries for this object. - UpdateObjectLocations(change_mode, location_updates, gcs_client_, - &it->second.current_object_locations); - // Copy the callbacks so that the callbacks can unsubscribe without interrupting - // looping over the callbacks. - auto callbacks = it->second.callbacks; - // Call all callbacks associated with the object id locations we have - // received. This notifies the client even if the list of locations is - // empty, since this may indicate that the objects have been evicted from - // all nodes. - for (const auto &callback_pair : callbacks) { - // It is safe to call the callback directly since this is already running - // in the subscription callback stack. - callback_pair.second(object_id, it->second.current_object_locations); - } - }; - RAY_CHECK_OK(gcs_client_->object_table().Subscribe( - JobID::Nil(), gcs_client_->Nodes().GetSelfId(), object_notification_callback, - nullptr)); -} - ray::Status ObjectDirectory::ReportObjectAdded( const ObjectID &object_id, const ClientID &client_id, const object_manager::protocol::ObjectInfoT &object_info) { RAY_LOG(DEBUG) << "Reporting object added to GCS " << object_id; - // Append the addition entry to the object table. - auto data = std::make_shared(); - data->set_manager(client_id.Binary()); - data->set_object_size(object_info.data_size); ray::Status status = - gcs_client_->object_table().Add(JobID::Nil(), object_id, data, nullptr); + gcs_client_->Objects().AsyncAddLocation(object_id, client_id, nullptr); return status; } @@ -95,12 +55,8 @@ ray::Status ObjectDirectory::ReportObjectRemoved( const ObjectID &object_id, const ClientID &client_id, const object_manager::protocol::ObjectInfoT &object_info) { RAY_LOG(DEBUG) << "Reporting object removed to GCS " << object_id; - // Append the eviction entry to the object table. - auto data = std::make_shared(); - data->set_manager(client_id.Binary()); - data->set_object_size(object_info.data_size); ray::Status status = - gcs_client_->object_table().Remove(JobID::Nil(), object_id, data, nullptr); + gcs_client_->Objects().AsyncRemoveLocation(object_id, client_id, nullptr); return status; }; @@ -136,7 +92,7 @@ void ObjectDirectory::HandleClientRemoved(const ClientID &client_id) { if (listener.second.current_object_locations.count(client_id) > 0) { // If the subscribed object has the removed client as a location, update // its locations with an empty update so that the location will be removed. - UpdateObjectLocations(GcsChangeMode::APPEND_OR_ADD, {}, gcs_client_, + UpdateObjectLocations(/*is_added*/ true, {}, gcs_client_, &listener.second.current_object_locations); // Re-call all the subscribed callbacks for the object, since its // locations have changed. @@ -156,10 +112,41 @@ ray::Status ObjectDirectory::SubscribeObjectLocations(const UniqueID &callback_i auto it = listeners_.find(object_id); if (it == listeners_.end()) { it = listeners_.emplace(object_id, LocationListenerState()).first; - status = gcs_client_->object_table().RequestNotifications( - JobID::Nil(), object_id, gcs_client_->Nodes().GetSelfId(), - /*done*/ nullptr); + + auto object_notification_callback = + [this](const ObjectID &object_id, + const gcs::ObjectChangeNotification &object_notification) { + // Objects are added to this map in SubscribeObjectLocations. + auto it = listeners_.find(object_id); + // Do nothing for objects we are not listening for. + if (it == listeners_.end()) { + return; + } + + // Once this flag is set to true, it should never go back to false. + it->second.subscribed = true; + + // Update entries for this object. + UpdateObjectLocations(object_notification.IsAdded(), + object_notification.GetData(), gcs_client_, + &it->second.current_object_locations); + // Copy the callbacks so that the callbacks can unsubscribe without interrupting + // looping over the callbacks. + auto callbacks = it->second.callbacks; + // Call all callbacks associated with the object id locations we have + // received. This notifies the client even if the list of locations is + // empty, since this may indicate that the objects have been evicted from + // all nodes. + for (const auto &callback_pair : callbacks) { + // It is safe to call the callback directly since this is already running + // in the subscription callback stack. + callback_pair.second(object_id, it->second.current_object_locations); + } + }; + status = gcs_client_->Objects().AsyncSubscribeToLocations( + object_id, object_notification_callback, /*done*/ nullptr); } + auto &listener_state = it->second; // TODO(hme): Make this fatal after implementing Pull suppression. if (listener_state.callbacks.count(callback_id) > 0) { @@ -185,8 +172,8 @@ ray::Status ObjectDirectory::UnsubscribeObjectLocations(const UniqueID &callback } entry->second.callbacks.erase(callback_id); if (entry->second.callbacks.empty()) { - status = gcs_client_->object_table().CancelNotifications( - JobID::Nil(), object_id, gcs_client_->Nodes().GetSelfId(), /*done*/ nullptr); + status = + gcs_client_->Objects().AsyncUnsubscribeToLocations(object_id, /*done*/ nullptr); listeners_.erase(entry); } return status; @@ -208,14 +195,16 @@ ray::Status ObjectDirectory::LookupLocations(const ObjectID &object_id, // We do not have any locations cached due to a concurrent // SubscribeObjectLocations call, so look up the object's locations // directly from the GCS. - status = gcs_client_->object_table().Lookup( - JobID::Nil(), object_id, - [this, callback](gcs::RedisGcsClient *client, const ObjectID &object_id, - const std::vector &location_updates) { + status = gcs_client_->Objects().AsyncGetLocations( + object_id, + [this, object_id, callback]( + Status status, const std::vector &location_updates) { + RAY_CHECK(status.ok()) + << "Failed to get object location from GCS: " << status.message(); // Build the set of current locations based on the entries in the log. std::unordered_set node_ids; - UpdateObjectLocations(GcsChangeMode::APPEND_OR_ADD, location_updates, - gcs_client_, &node_ids); + UpdateObjectLocations(/*is_added*/ true, location_updates, gcs_client_, + &node_ids); // It is safe to call the callback directly since this is already running // in the GCS client's lookup callback stack. callback(object_id, node_ids); diff --git a/src/ray/object_manager/object_directory.h b/src/ray/object_manager/object_directory.h index d6df226a2..669d9b7fa 100644 --- a/src/ray/object_manager/object_directory.h +++ b/src/ray/object_manager/object_directory.h @@ -33,8 +33,6 @@ class ObjectDirectoryInterface { public: virtual ~ObjectDirectoryInterface() {} - virtual void RegisterBackend() = 0; - /// Lookup how to connect to a remote object manager. /// /// \param connection_info The connection information to fill out. This @@ -135,8 +133,6 @@ class ObjectDirectory : public ObjectDirectoryInterface { virtual ~ObjectDirectory() {} - void RegisterBackend() override; - void LookupRemoteConnectionInfo(RemoteConnectionInfo &connection_info) const override; std::vector LookupAllRemoteConnections() const override; diff --git a/src/ray/object_manager/object_manager.cc b/src/ray/object_manager/object_manager.cc index 11eeeed29..b5a279c68 100644 --- a/src/ray/object_manager/object_manager.cc +++ b/src/ray/object_manager/object_manager.cc @@ -38,8 +38,6 @@ ObjectManager::ObjectManager(asio::io_service &main_service, const ClientID &sel ObjectManager::~ObjectManager() { StopRpcService(); } -void ObjectManager::RegisterGcs() { object_directory_->RegisterBackend(); } - void ObjectManager::RunRpcService() { rpc_service_.run(); } void ObjectManager::StartRpcService() { diff --git a/src/ray/object_manager/object_manager.h b/src/ray/object_manager/object_manager.h index f3c082c1f..cf8ddd7f1 100644 --- a/src/ray/object_manager/object_manager.h +++ b/src/ray/object_manager/object_manager.h @@ -157,9 +157,6 @@ class ObjectManager : public ObjectManagerInterface, ~ObjectManager(); - /// Register GCS-related functionality. - void RegisterGcs(); - /// Subscribe to notifications of objects added to local store. /// Upon subscribing, the callback will be invoked for all objects that /// diff --git a/src/ray/object_manager/test/object_manager_stress_test.cc b/src/ray/object_manager/test/object_manager_stress_test.cc index ba993176f..4efd7f45f 100644 --- a/src/ray/object_manager/test/object_manager_stress_test.cc +++ b/src/ray/object_manager/test/object_manager_stress_test.cc @@ -55,7 +55,6 @@ class MockServer { node_info.set_object_manager_port(object_manager_port); ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info); - object_manager_.RegisterGcs(); return status; } diff --git a/src/ray/object_manager/test/object_manager_test.cc b/src/ray/object_manager/test/object_manager_test.cc index ee3bcb0a7..f552a2d1c 100644 --- a/src/ray/object_manager/test/object_manager_test.cc +++ b/src/ray/object_manager/test/object_manager_test.cc @@ -49,7 +49,6 @@ class MockServer { node_info.set_object_manager_port(object_manager_port); ray::Status status = gcs_client_->Nodes().RegisterSelf(node_info); - object_manager_.RegisterGcs(); return status; } diff --git a/src/ray/raylet/mock_gcs_client.cc b/src/ray/raylet/mock_gcs_client.cc index 1a75b6593..4155447fe 100644 --- a/src/ray/raylet/mock_gcs_client.cc +++ b/src/ray/raylet/mock_gcs_client.cc @@ -114,8 +114,6 @@ ClientID GcsClient::Register(const std::string &ip, uint16_t port) { return client_id; } -ObjectTable &GcsClient::object_table() { return *object_table_; } - ClientTable &GcsClient::client_table() { return *client_table_; } } // namespace ray diff --git a/src/ray/raylet/mock_gcs_client.h b/src/ray/raylet/mock_gcs_client.h index a08e74f66..42256658a 100644 --- a/src/ray/raylet/mock_gcs_client.h +++ b/src/ray/raylet/mock_gcs_client.h @@ -80,7 +80,6 @@ class GcsClient { } // Register the ip and port of the connecting client. ClientID Register(const std::string &ip, uint16_t port); - ObjectTable &object_table(); ClientTable &client_table(); private: diff --git a/src/ray/raylet/node_manager.cc b/src/ray/raylet/node_manager.cc index d55b061a8..1f9f072f1 100644 --- a/src/ray/raylet/node_manager.cc +++ b/src/ray/raylet/node_manager.cc @@ -136,8 +136,6 @@ NodeManager::NodeManager(boost::asio::io_service &io_service, } ray::Status NodeManager::RegisterGcs() { - object_manager_.RegisterGcs(); - const auto task_lease_notification_callback = [this](gcs::RedisGcsClient *client, const TaskID &task_id, const TaskLeaseData &task_lease) { diff --git a/src/ray/raylet/reconstruction_policy_test.cc b/src/ray/raylet/reconstruction_policy_test.cc index 9a8e141b1..8556d51b4 100644 --- a/src/ray/raylet/reconstruction_policy_test.cc +++ b/src/ray/raylet/reconstruction_policy_test.cc @@ -64,7 +64,7 @@ class MockObjectDirectory : public ObjectDirectoryInterface { std::string DebugString() const override { return ""; } - MOCK_METHOD0(RegisterBackend, void(void)); + MOCK_METHOD0(GetLocalClientID, ray::ClientID()); MOCK_CONST_METHOD1(LookupRemoteConnectionInfo, void(RemoteConnectionInfo &)); MOCK_CONST_METHOD0(LookupAllRemoteConnections, std::vector()); MOCK_METHOD3(SubscribeObjectLocations,