wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-06 13:31:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Resubscribe object table info when gcs service restart (#8639)

Browse Source
This commit is contained in:
Tao Wang
2020-06-01 10:42:26 +08:00
committed by GitHub
parent d8a081a185
commit 1df408d6ed
6 changed files with 104 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/ray/gcs/accessor.h
+6
View File
@@ -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;
};
src/ray/gcs/gcs_client/service_based_accessor.cc
+47 -33
View File
@@ -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<rpc::ObjectTableData> 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<rpc::ObjectTableData> 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<rpc::ObjectTableData> &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<rpc::ObjectTableData> &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;
}
src/ray/gcs/gcs_client/service_based_accessor.h
+6
View File
@@ -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<ObjectID, SubscribeOperation> subscribe_object_operations_;
ServiceBasedGcsClient *client_impl_;
Sequencer<ObjectID> sequencer_;
src/ray/gcs/gcs_client/service_based_gcs_client.cc
+1
View File
@@ -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());
};
src/ray/gcs/gcs_client/test/service_based_gcs_client_test.cc
+42
View File
@@ -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<int> object1_change_count(0);
std::atomic<int> 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.
src/ray/gcs/redis_accessor.h
+2
View File
@@ -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};