diff --git a/src/ray/gcs/accessor.h b/src/ray/gcs/accessor.h index bfe1bf655..46bdbf5e8 100644 --- a/src/ray/gcs/accessor.h +++ b/src/ray/gcs/accessor.h @@ -378,6 +378,12 @@ class ObjectInfoAccessor { /// \return Status virtual Status AsyncUnsubscribeToLocations(const ObjectID &object_id) = 0; + /// Reestablish subscription. + /// This should be called when GCS server restarts from a failure. + /// + /// \return Status + virtual Status AsyncReSubscribe() = 0; + protected: ObjectInfoAccessor() = default; }; diff --git a/src/ray/gcs/gcs_client/service_based_accessor.cc b/src/ray/gcs/gcs_client/service_based_accessor.cc index d16766bee..f1efdf55a 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.cc +++ b/src/ray/gcs/gcs_client/service_based_accessor.cc @@ -1095,46 +1095,60 @@ Status ServiceBasedObjectInfoAccessor::AsyncSubscribeToLocations( RAY_LOG(DEBUG) << "Subscribing object location, object id = " << object_id; RAY_CHECK(subscribe != nullptr) << "Failed to subscribe object location, object id = " << object_id; - auto on_subscribe = [object_id, subscribe](const std::string &id, - const std::string &data) { - rpc::ObjectLocationChange object_location_change; - object_location_change.ParseFromString(data); - std::vector object_data_vector; - object_data_vector.emplace_back(object_location_change.data()); - auto change_mode = object_location_change.is_add() ? rpc::GcsChangeMode::APPEND_OR_ADD - : rpc::GcsChangeMode::REMOVE; - gcs::ObjectChangeNotification notification(change_mode, object_data_vector); - subscribe(object_id, notification); + auto subscribe_operation = [this, object_id, subscribe](const StatusCallback &done) { + auto on_subscribe = [object_id, subscribe](const std::string &id, + const std::string &data) { + rpc::ObjectLocationChange object_location_change; + object_location_change.ParseFromString(data); + std::vector object_data_vector; + object_data_vector.emplace_back(object_location_change.data()); + auto change_mode = object_location_change.is_add() + ? rpc::GcsChangeMode::APPEND_OR_ADD + : rpc::GcsChangeMode::REMOVE; + gcs::ObjectChangeNotification notification(change_mode, object_data_vector); + subscribe(object_id, notification); + }; + auto on_done = [this, object_id, subscribe, done](const Status &status) { + if (status.ok()) { + auto callback = [object_id, subscribe, done]( + const Status &status, + const std::vector &result) { + if (status.ok()) { + gcs::ObjectChangeNotification notification(rpc::GcsChangeMode::APPEND_OR_ADD, + result); + subscribe(object_id, notification); + } + if (done) { + done(status); + } + }; + RAY_CHECK_OK(AsyncGetLocations(object_id, callback)); + } else if (done) { + done(status); + } + }; + auto status = client_impl_->GetGcsPubSub().Subscribe(OBJECT_CHANNEL, object_id.Hex(), + on_subscribe, on_done); + RAY_LOG(DEBUG) << "Finished subscribing object location, object id = " << object_id; + return status; }; - auto on_done = [this, object_id, subscribe, done](const Status &status) { - if (status.ok()) { - auto callback = [object_id, subscribe, done]( - const Status &status, - const std::vector &result) { - if (status.ok()) { - gcs::ObjectChangeNotification notification(rpc::GcsChangeMode::APPEND_OR_ADD, - result); - subscribe(object_id, notification); - } - if (done) { - done(status); - } - }; - RAY_CHECK_OK(AsyncGetLocations(object_id, callback)); - } else if (done) { - done(status); - } - }; - auto status = client_impl_->GetGcsPubSub().Subscribe(OBJECT_CHANNEL, object_id.Hex(), - on_subscribe, on_done); - RAY_LOG(DEBUG) << "Finished subscribing object location, object id = " << object_id; - return status; + subscribe_object_operations_[object_id] = subscribe_operation; + return subscribe_operation(done); +} + +Status ServiceBasedObjectInfoAccessor::AsyncReSubscribe() { + RAY_LOG(INFO) << "Reestablishing subscription for object locations."; + for (auto &item : subscribe_object_operations_) { + RAY_CHECK_OK(item.second(nullptr)); + } + return Status::OK(); } Status ServiceBasedObjectInfoAccessor::AsyncUnsubscribeToLocations( const ObjectID &object_id) { RAY_LOG(DEBUG) << "Unsubscribing object location, object id = " << object_id; auto status = client_impl_->GetGcsPubSub().Unsubscribe(OBJECT_CHANNEL, object_id.Hex()); + subscribe_object_operations_.erase(object_id); RAY_LOG(DEBUG) << "Finished unsubscribing object location, object id = " << object_id; return status; } diff --git a/src/ray/gcs/gcs_client/service_based_accessor.h b/src/ray/gcs/gcs_client/service_based_accessor.h index 2c1077f23..db912d4ee 100644 --- a/src/ray/gcs/gcs_client/service_based_accessor.h +++ b/src/ray/gcs/gcs_client/service_based_accessor.h @@ -296,7 +296,13 @@ class ServiceBasedObjectInfoAccessor : public ObjectInfoAccessor { Status AsyncUnsubscribeToLocations(const ObjectID &object_id) override; + Status AsyncReSubscribe() override; + private: + /// Save the subscribe operation in this function, so we can call it again when GCS + /// restarts from a failure. + std::unordered_map subscribe_object_operations_; + ServiceBasedGcsClient *client_impl_; Sequencer sequencer_; diff --git a/src/ray/gcs/gcs_client/service_based_gcs_client.cc b/src/ray/gcs/gcs_client/service_based_gcs_client.cc index 005b8c376..7a72c0d0c 100644 --- a/src/ray/gcs/gcs_client/service_based_gcs_client.cc +++ b/src/ray/gcs/gcs_client/service_based_gcs_client.cc @@ -52,6 +52,7 @@ Status ServiceBasedGcsClient::Connect(boost::asio::io_service &io_service) { RAY_CHECK_OK(actor_accessor_->AsyncReSubscribe()); RAY_CHECK_OK(node_accessor_->AsyncReSubscribe()); RAY_CHECK_OK(task_accessor_->AsyncReSubscribe()); + RAY_CHECK_OK(object_accessor_->AsyncReSubscribe()); RAY_CHECK_OK(worker_accessor_->AsyncReSubscribe()); }; diff --git a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc index 7f0a3d212..e32f7f247 100644 --- a/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc +++ b/src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc @@ -880,6 +880,48 @@ TEST_F(ServiceBasedGcsClientTest, TestActorTableReSubscribe) { WaitPendingDone(actor2_update_count, 1); } +TEST_F(ServiceBasedGcsClientTest, TestObjectTableReSubscribe) { + ObjectID object1_id = ObjectID::FromRandom(); + ObjectID object2_id = ObjectID::FromRandom(); + ClientID node_id = ClientID::FromRandom(); + + // Subscribe to any update of an object's location. + std::atomic object1_change_count(0); + std::atomic object2_change_count(0); + ASSERT_TRUE(SubscribeToLocations( + object1_id, [&object1_change_count](const ObjectID &object_id, + const gcs::ObjectChangeNotification &result) { + if (!result.GetData().empty()) { + ++object1_change_count; + } + })); + ASSERT_TRUE(SubscribeToLocations( + object2_id, [&object2_change_count](const ObjectID &object_id, + const gcs::ObjectChangeNotification &result) { + if (!result.GetData().empty()) { + ++object2_change_count; + } + })); + + ASSERT_TRUE(AddLocation(object1_id, node_id)); + WaitPendingDone(object1_change_count, 1); + ASSERT_TRUE(AddLocation(object2_id, node_id)); + WaitPendingDone(object2_change_count, 1); + + // Cancel subscription to any update of an object's location. + UnsubscribeToLocations(object1_id); + usleep(100 * 1000); + + // Restart GCS. + RestartGcsServer(); + + // Add location of object to GCS again and check if resubscribe works. + ASSERT_TRUE(AddLocation(object1_id, node_id)); + WaitPendingDone(object1_change_count, 1); + ASSERT_TRUE(AddLocation(object2_id, node_id)); + WaitPendingDone(object2_change_count, 2); +} + TEST_F(ServiceBasedGcsClientTest, TestNodeTableReSubscribe) { // Test that subscription of the node table can still work when GCS server restarts. // Subscribe to node addition and removal events from GCS and cache those information. diff --git a/src/ray/gcs/redis_accessor.h b/src/ray/gcs/redis_accessor.h index 42ea0f5f7..82ecb5cc0 100644 --- a/src/ray/gcs/redis_accessor.h +++ b/src/ray/gcs/redis_accessor.h @@ -296,6 +296,8 @@ class RedisObjectInfoAccessor : public ObjectInfoAccessor { Status AsyncUnsubscribeToLocations(const ObjectID &object_id) override; + Status AsyncReSubscribe() override { return Status::NotImplemented(""); } + private: RedisGcsClient *client_impl_{nullptr};