[core] Kill GcsUnregisterSubscriber + publisher retry tests + doc (#56038)

Signed-off-by: dayshah <dhyey2019@gmail.com>

Author: dayshah

src/fakes/ray/pubsub/publisher.h

@@ -21,24 +21,20 @@ namespace pubsub {
 
 class FakePublisher : public Publisher {
  public:
-  bool RegisterSubscription(const rpc::ChannelType channel_type,
+  void RegisterSubscription(const rpc::ChannelType channel_type,
                             const UniqueID &subscriber_id,
-                            const std::optional<std::string> &key_id) override {
-    return true;
-  }
+                            const std::optional<std::string> &key_id) override {}
 
   void Publish(rpc::PubMessage pub_message) override {}
 
   void PublishFailure(const rpc::ChannelType channel_type,
                       const std::string &key_id) override {}
 
-  bool UnregisterSubscription(const rpc::ChannelType channel_type,
+  void UnregisterSubscription(const rpc::ChannelType channel_type,
                               const UniqueID &subscriber_id,
-                              const std::optional<std::string> &key_id) override {
-    return true;
-  }
+                              const std::optional<std::string> &key_id) override {}
 
-  void UnregisterSubscriber(const UniqueID &subscriber_id) override { return; }
+  void UnregisterSubscriber(const UniqueID &subscriber_id) override {}
 
   std::string DebugString() const override { return "FakePublisher"; }
 };

src/mock/ray/pubsub/publisher.h

@@ -29,7 +29,7 @@ class MockPublisher : public PublisherInterface {
                google::protobuf::RepeatedPtrField<rpc::PubMessage> *pub_messages,
                rpc::SendReplyCallback send_reply_callback),
               (override));
-  MOCK_METHOD(bool,
+  MOCK_METHOD(void,
               RegisterSubscription,
               (const rpc::ChannelType channel_type,
                const UniqueID &subscriber_id,
@@ -40,7 +40,7 @@ class MockPublisher : public PublisherInterface {
               PublishFailure,
               (const rpc::ChannelType channel_type, const std::string &key_id),
               (override));
-  MOCK_METHOD(bool,
+  MOCK_METHOD(void,
               UnregisterSubscription,
               (const rpc::ChannelType channel_type,
                const UniqueID &subscriber_id,

src/ray/core_worker/tests/reference_count_test.cc

@@ -107,7 +107,7 @@ using SubscriptionFailureCallbackMap =
 // static maps are used to simulate distirubted environment.
 static SubscriptionCallbackMap subscription_callback_map;
 static SubscriptionFailureCallbackMap subscription_failure_callback_map;
-static pubsub::pub_internal::SubscriptionIndex directory(
+static pubsub::SubscriptionIndex directory(
     rpc::ChannelType::WORKER_OBJECT_LOCATIONS_CHANNEL);
 
 static std::string GenerateID(UniqueID publisher_id, UniqueID subscriber_id) {
@@ -127,7 +127,7 @@ using PublisherFactoryFn =
 
 class MockDistributedSubscriber : public pubsub::SubscriberInterface {
  public:
-  MockDistributedSubscriber(pubsub::pub_internal::SubscriptionIndex *dict,
+  MockDistributedSubscriber(pubsub::SubscriptionIndex *dict,
                             SubscriptionCallbackMap *sub_callback_map,
                             SubscriptionFailureCallbackMap *sub_failure_callback_map,
                             UniqueID subscriber_id,
@@ -136,7 +136,7 @@ class MockDistributedSubscriber : public pubsub::SubscriberInterface {
         subscription_callback_map_(sub_callback_map),
         subscription_failure_callback_map_(sub_failure_callback_map),
         subscriber_id_(subscriber_id),
-        subscriber_(std::make_unique<pubsub::pub_internal::SubscriberState>(
+        subscriber_(std::make_unique<pubsub::SubscriberState>(
             subscriber_id,
             /*get_time_ms=*/[]() { return 1.0; },
             /*subscriber_timeout_ms=*/1000,
@@ -207,17 +207,17 @@ class MockDistributedSubscriber : public pubsub::SubscriberInterface {
     return "";
   }
 
-  pubsub::pub_internal::SubscriptionIndex *directory_;
+  pubsub::SubscriptionIndex *directory_;
   SubscriptionCallbackMap *subscription_callback_map_;
   SubscriptionFailureCallbackMap *subscription_failure_callback_map_;
   UniqueID subscriber_id_;
-  std::unique_ptr<pubsub::pub_internal::SubscriberState> subscriber_;
+  std::unique_ptr<pubsub::SubscriberState> subscriber_;
   PublisherFactoryFn client_factory_;
 };
 
 class MockDistributedPublisher : public pubsub::PublisherInterface {
  public:
-  MockDistributedPublisher(pubsub::pub_internal::SubscriptionIndex *dict,
+  MockDistributedPublisher(pubsub::SubscriptionIndex *dict,
                            SubscriptionCallbackMap *sub_callback_map,
                            SubscriptionFailureCallbackMap *sub_failure_callback_map,
                            WorkerID publisher_id)
@@ -227,11 +227,10 @@ class MockDistributedPublisher : public pubsub::PublisherInterface {
         publisher_id_(publisher_id) {}
   ~MockDistributedPublisher() = default;
 
-  bool RegisterSubscription(const rpc::ChannelType channel_type,
+  void RegisterSubscription(const rpc::ChannelType channel_type,
                             const UniqueID &subscriber_id,
                             const std::optional<std::string> &key_id_binary) override {
     RAY_CHECK(false) << "No need to implement it for testing.";
-    return false;
   }
 
   void PublishFailure(const rpc::ChannelType channel_type,
@@ -256,13 +255,11 @@ class MockDistributedPublisher : public pubsub::PublisherInterface {
     }
   }
 
-  bool UnregisterSubscription(const rpc::ChannelType channel_type,
+  void UnregisterSubscription(const rpc::ChannelType channel_type,
                               const UniqueID &subscriber_id,
-                              const std::optional<std::string> &key_id_binary) override {
-    return true;
-  }
+                              const std::optional<std::string> &key_id_binary) override {}
 
-  void UnregisterSubscriber(const UniqueID &subscriber_id) override { return; }
+  void UnregisterSubscriber(const UniqueID &subscriber_id) override {}
 
   void ConnectToSubscriber(
       const rpc::PubsubLongPollingRequest &request,
@@ -272,7 +269,7 @@ class MockDistributedPublisher : public pubsub::PublisherInterface {
 
   std::string DebugString() const override { return ""; }
 
-  pubsub::pub_internal::SubscriptionIndex *directory_;
+  pubsub::SubscriptionIndex *directory_;
   SubscriptionCallbackMap *subscription_callback_map_;
   SubscriptionFailureCallbackMap *subscription_failure_callback_map_;
   WorkerID publisher_id_;

src/ray/gcs/gcs_server/grpc_service_interfaces.h

@@ -143,10 +143,6 @@ class InternalPubSubGcsServiceHandler {
   virtual void HandleGcsSubscriberCommandBatch(GcsSubscriberCommandBatchRequest request,
                                                GcsSubscriberCommandBatchReply *reply,
                                                SendReplyCallback send_reply_callback) = 0;
-
-  virtual void HandleGcsUnregisterSubscriber(GcsUnregisterSubscriberRequest request,
-                                             GcsUnregisterSubscriberReply *reply,
-                                             SendReplyCallback send_reply_callback) = 0;
 };
 
 class JobInfoGcsServiceHandler {

src/ray/gcs/gcs_server/grpc_services.cc

@@ -79,8 +79,6 @@ void InternalPubSubGrpcService::InitServerCallFactories(
       InternalPubSubGcsService, GcsSubscriberPoll, max_active_rpcs_per_handler_);
   RPC_SERVICE_HANDLER(
       InternalPubSubGcsService, GcsSubscriberCommandBatch, max_active_rpcs_per_handler_);
-  RPC_SERVICE_HANDLER(
-      InternalPubSubGcsService, GcsUnregisterSubscriber, max_active_rpcs_per_handler_);
 }
 
 void JobInfoGrpcService::InitServerCallFactories(

src/ray/gcs/gcs_server/pubsub_handler.cc

@@ -93,15 +93,6 @@ void InternalPubSubHandler::HandleGcsSubscriberCommandBatch(
   send_reply_callback(Status::OK(), nullptr, nullptr);
 }
 
-void InternalPubSubHandler::HandleGcsUnregisterSubscriber(
-    rpc::GcsUnregisterSubscriberRequest request,
-    rpc::GcsUnregisterSubscriberReply *reply,
-    rpc::SendReplyCallback send_reply_callback) {
-  const auto subscriber_id = UniqueID::FromBinary(request.subscriber_id());
-  gcs_publisher_.GetPublisher().UnregisterSubscriber(subscriber_id);
-  send_reply_callback(Status::OK(), nullptr, nullptr);
-}
-
 void InternalPubSubHandler::AsyncRemoveSubscriberFrom(const std::string &sender_id) {
   io_service_.post(
       [this, sender_id]() {

src/ray/gcs/gcs_server/pubsub_handler.h

@@ -44,10 +44,6 @@ class InternalPubSubHandler : public rpc::InternalPubSubGcsServiceHandler {
                                        rpc::GcsSubscriberCommandBatchReply *reply,
                                        rpc::SendReplyCallback send_reply_callback) final;
 
-  void HandleGcsUnregisterSubscriber(rpc::GcsUnregisterSubscriberRequest request,
-                                     rpc::GcsUnregisterSubscriberReply *reply,
-                                     rpc::SendReplyCallback send_reply_callback) final;
-
   /// This function is only for external callers. Internally, can just erase from
   /// sender_to_subscribers_ and everything should be on the Publisher's io_service_.
   void AsyncRemoveSubscriberFrom(const std::string &sender_id);

src/ray/gcs/pubsub/gcs_pub_sub.cc

@@ -335,10 +335,14 @@ Status PythonGcsSubscriber::Close() {
 
   grpc::ClientContext context;
 
-  rpc::GcsUnregisterSubscriberRequest request;
+  rpc::GcsSubscriberCommandBatchRequest request;
   request.set_subscriber_id(subscriber_id_);
-  rpc::GcsUnregisterSubscriberReply reply;
-  grpc::Status status = pubsub_stub_->GcsUnregisterSubscriber(&context, request, &reply);
+  auto *command = request.add_commands();
+  command->set_channel_type(channel_type_);
+  command->mutable_unsubscribe_message();
+  rpc::GcsSubscriberCommandBatchReply reply;
+  grpc::Status status =
+      pubsub_stub_->GcsSubscriberCommandBatch(&context, request, &reply);
 
   if (!status.ok()) {
     RAY_LOG(WARNING) << "Error while unregistering the subscriber: "

src/ray/object_manager/ownership_object_directory.cc

@@ -346,7 +346,6 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations(
     auto failure_callback = [this, owner_address](const std::string &object_id_binary,
                                                   const Status &status) {
       const auto obj_id = ObjectID::FromBinary(object_id_binary);
-      rpc::WorkerObjectLocationsPubMessage location_info;
       if (!status.ok()) {
         RAY_LOG(INFO).WithField(obj_id)
             << "Failed to get the location: " << status.ToString();
@@ -362,9 +361,10 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations(
       // Location lookup can fail if the owner is reachable but no longer has a
       // record of this ObjectRef, most likely due to an issue with the
       // distributed reference counting protocol.
-      ObjectLocationSubscriptionCallback(location_info,
-                                         obj_id,
-                                         /*location_lookup_failed*/ true);
+      ObjectLocationSubscriptionCallback(
+          /*location_info=*/rpc::WorkerObjectLocationsPubMessage{},
+          obj_id,
+          /*location_lookup_failed*/ true);
     };
 
     auto sub_message = std::make_unique<rpc::SubMessage>();

src/ray/protobuf/gcs_service.proto

@@ -680,12 +680,6 @@ message GcsSubscriberCommandBatchReply {
   GcsStatus status = 100;
 }
 
-message GcsUnregisterSubscriberRequest {
-  bytes subscriber_id = 1;
-}
-
-message GcsUnregisterSubscriberReply {}
-
 /// This supports subscribing updates from GCS with long poll, and registering /
 /// de-registering subscribers.
 service InternalPubSubGcsService {
@@ -699,10 +693,6 @@ service InternalPubSubGcsService {
   /// A batch of subscribe / unsubscribe requests sent by the subscriber.
   rpc GcsSubscriberCommandBatch(GcsSubscriberCommandBatchRequest)
       returns (GcsSubscriberCommandBatchReply);
-  /// Unregister a subscriber from GCS, removing all subscriptions as well as the
-  /// subscriber itself.
-  rpc GcsUnregisterSubscriber(GcsUnregisterSubscriberRequest)
-      returns (GcsUnregisterSubscriberReply);
 }
 
 message GetAllResourceUsageRequest {}