[Core] Ownership-based Object Directory - Changed infinite short-poll location subscription to long-poll. (#13841)

src/ray/core_worker/core_worker.cc

@@ -2219,19 +2219,25 @@ void CoreWorker::HandleGetObjectLocationsOwner(
     return;
   }
   auto object_id = ObjectID::FromBinary(request.object_id());
-  absl::optional<absl::flat_hash_set<NodeID>> node_ids =
-      reference_counter_->GetObjectLocations(object_id);
-  Status status;
-  if (node_ids.has_value()) {
-    for (const auto &node_id : node_ids.value()) {
+  const auto &callback = [object_id, reply, send_reply_callback](
+                             const absl::flat_hash_set<NodeID> &locations,
+                             int64_t object_size, int64_t current_version) {
+    RAY_LOG(DEBUG) << "Replying to HandleGetObjectLocationsOwner for " << object_id
+                   << " with location update version " << current_version << ", "
+                   << locations.size() << " locations, and " << object_size
+                   << " object size.";
+    for (const auto &node_id : locations) {
       reply->add_node_ids(node_id.Binary());
     }
-    status = Status::OK();
-  } else {
-    status = Status::ObjectNotFound("Object " + object_id.Hex() + " not found");
+    reply->set_object_size(object_size);
+    reply->set_current_version(current_version);
+    send_reply_callback(Status::OK(), nullptr, nullptr);
+  };
+  auto status = reference_counter_->SubscribeObjectLocations(
+      object_id, request.last_version(), callback);
+  if (!status.ok()) {
+    send_reply_callback(status, nullptr, nullptr);
   }
-  reply->set_object_size(reference_counter_->GetObjectSize(object_id));
-  send_reply_callback(status, nullptr, nullptr);
 }
 
 void CoreWorker::HandleWaitForRefRemoved(const rpc::WaitForRefRemovedRequest &request,

src/ray/core_worker/reference_count.cc

@@ -185,6 +185,7 @@ void ReferenceCounter::UpdateObjectSize(const ObjectID &object_id, int64_t objec
   auto it = object_id_refs_.find(object_id);
   if (it != object_id_refs_.end()) {
     it->second.object_size = object_size;
+    PushToLocationSubscribers(it);
   }
 }
 
@@ -915,11 +916,12 @@ bool ReferenceCounter::AddObjectLocation(const ObjectID &object_id,
   absl::MutexLock lock(&mutex_);
   auto it = object_id_refs_.find(object_id);
   if (it == object_id_refs_.end()) {
-    RAY_LOG(WARNING) << "Tried to add an object location for an object " << object_id
-                     << " that doesn't exist in the reference table";
+    RAY_LOG(INFO) << "Tried to add an object location for an object " << object_id
+                  << " that doesn't exist in the reference table";
     return false;
   }
   it->second.locations.insert(node_id);
+  PushToLocationSubscribers(it);
   return true;
 }
 
@@ -928,11 +930,12 @@ bool ReferenceCounter::RemoveObjectLocation(const ObjectID &object_id,
   absl::MutexLock lock(&mutex_);
   auto it = object_id_refs_.find(object_id);
   if (it == object_id_refs_.end()) {
-    RAY_LOG(WARNING) << "Tried to remove an object location for an object " << object_id
-                     << " that doesn't exist in the reference table";
+    RAY_LOG(INFO) << "Tried to remove an object location for an object " << object_id
+                  << " that doesn't exist in the reference table";
     return false;
   }
   it->second.locations.erase(node_id);
+  PushToLocationSubscribers(it);
   return true;
 }
 
@@ -1003,6 +1006,39 @@ absl::optional<LocalityData> ReferenceCounter::GetLocalityData(
   return locality_data;
 }
 
+void ReferenceCounter::PushToLocationSubscribers(ReferenceTable::iterator it) {
+  const auto callbacks = it->second.location_subscription_callbacks;
+  it->second.location_subscription_callbacks.clear();
+  it->second.location_version++;
+  for (const auto callback : callbacks) {
+    callback(it->second.locations, it->second.object_size, it->second.location_version);
+  }
+}
+
+Status ReferenceCounter::SubscribeObjectLocations(
+    const ObjectID &object_id, int64_t last_location_version,
+    const LocationSubscriptionCallback &callback) {
+  absl::MutexLock lock(&mutex_);
+  auto it = object_id_refs_.find(object_id);
+  if (it == object_id_refs_.end()) {
+    RAY_LOG(INFO) << "Tried to register a location subscriber for an object " << object_id
+                  << " that doesn't exist in the reference table."
+                  << " The object has probably already been freed.";
+    return Status::ObjectNotFound("Object " + object_id.Hex() + " not found");
+  }
+
+  if (last_location_version < it->second.location_version) {
+    // If the last location version is less than the current location version, we
+    // already have location data that the subscriber hasn't seen yet, so we immediately
+    // invoke the callback.
+    callback(it->second.locations, it->second.object_size, it->second.location_version);
+  } else {
+    // Otherwise, save the callback for later invocation.
+    it->second.location_subscription_callbacks.push_back(callback);
+  }
+  return Status::OK();
+}
+
 ReferenceCounter::Reference ReferenceCounter::Reference::FromProto(
     const rpc::ObjectReferenceCount &ref_count) {
   Reference ref;

src/ray/core_worker/reference_count.h

@@ -49,6 +49,10 @@ class ReferenceCounterInterface {
   virtual ~ReferenceCounterInterface() {}
 };
 
+// Callback for location subscriptions.
+using LocationSubscriptionCallback =
+    std::function<void(const absl::flat_hash_set<NodeID> &, int64_t, int64_t)>;
+
 /// Class used by the core worker to keep track of ObjectID reference counts for garbage
 /// collection. This class is thread safe.
 class ReferenceCounter : public ReferenceCounterInterface,
@@ -397,6 +401,19 @@ class ReferenceCounter : public ReferenceCounterInterface,
   absl::optional<absl::flat_hash_set<NodeID>> GetObjectLocations(
       const ObjectID &object_id) LOCKS_EXCLUDED(mutex_);
 
+  /// Subscribe to object location changes that are more recent than the given version.
+  /// The provided callback will be invoked when new locations become available.
+  ///
+  /// \param[in] object_id The object whose locations we want.
+  /// \param[in] last_location_version The version of the last location update the
+  /// caller received. Only more recent location updates will be returned.
+  /// \param[in] callback The callback to invoke with the location update.
+  /// \return The status of the location get.
+  Status SubscribeObjectLocations(const ObjectID &object_id,
+                                  int64_t last_location_version,
+                                  const LocationSubscriptionCallback &callback)
+      LOCKS_EXCLUDED(mutex_);
+
   /// Get an object's size. This will return 0 if the object is out of scope.
   ///
   /// \param[in] object_id The object whose size to get.
@@ -492,13 +509,17 @@ class ReferenceCounter : public ReferenceCounterInterface,
     /// process is a borrower, the borrower must add the owner's address before
     /// using the ObjectID.
     absl::optional<rpc::Address> owner_address;
-    // If this object is owned by us and stored in plasma, and reference
-    // counting is enabled, then some raylet must be pinning the object value.
-    // This is the address of that raylet.
+    /// If this object is owned by us and stored in plasma, and reference
+    /// counting is enabled, then some raylet must be pinning the object value.
+    /// This is the address of that raylet.
     absl::optional<NodeID> pinned_at_raylet_id;
-    // If this object is owned by us and stored in plasma, this contains all
-    // object locations.
+    /// If this object is owned by us and stored in plasma, this contains all
+    /// object locations.
     absl::flat_hash_set<NodeID> locations;
+    /// A logical counter for object location updates, used for object location
+    /// subscriptions. Subscribers use -1 to indicate that they want us to
+    /// immediately send them the current location data.
+    int64_t location_version = 0;
     // Whether this object can be reconstructed via lineage. If false, then the
     // object's value will be pinned as long as it is referenced by any other
     // object's lineage.
@@ -565,7 +586,9 @@ class ReferenceCounter : public ReferenceCounterInterface,
     size_t lineage_ref_count = 0;
     /// Whether this object has been spilled to external storage.
     bool spilled = false;
-
+    /// Location subscription callbacks registered by async location get requests.
+    /// These will be invoked whenever locations or object_size are changed.
+    std::vector<LocationSubscriptionCallback> location_subscription_callbacks;
     /// Callback that will be called when this ObjectID no longer has
     /// references.
     std::function<void(const ObjectID &)> on_delete;
@@ -689,6 +712,12 @@ class ReferenceCounter : public ReferenceCounterInterface,
   void ReleaseLineageReferencesInternal(const std::vector<ObjectID> &argument_ids)
       EXCLUSIVE_LOCKS_REQUIRED(mutex_);
 
+  /// Pushes location updates to subscribers of a particular reference, invoking all
+  /// callbacks registered for the reference by GetLocationsAsync calls. This method
+  /// also increments the reference's location version counter.
+  void PushToLocationSubscribers(ReferenceTable::iterator it)
+      EXCLUSIVE_LOCKS_REQUIRED(mutex_);
+
   /// Address of our RPC server. This is used to determine whether we own a
   /// given object or not, by comparing our WorkerID with the WorkerID of the
   /// object's owner.

src/ray/object_manager/ownership_based_object_directory.cc

@@ -80,11 +80,18 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectAdded(
   request.set_node_id(node_id.Binary());
 
   rpc_client->AddObjectLocationOwner(
-      request, [worker_id, object_id](Status status,
-                                      const rpc::AddObjectLocationOwnerReply &reply) {
+      request, [worker_id, object_id, node_id](
+                   Status status, const rpc::AddObjectLocationOwnerReply &reply) {
         if (!status.ok()) {
-          RAY_LOG(ERROR) << "Worker " << worker_id << " failed to add the location for "
-                         << object_id;
+          if (status.IsObjectNotFound()) {
+            RAY_LOG(INFO) << "Worker " << worker_id << " failed to add the location "
+                          << node_id << " for " << object_id
+                          << " because the owner no longer has the object; we assume the "
+                             "object was evicted.";
+          } else {
+            RAY_LOG(INFO) << "Worker " << worker_id << " failed to add the location "
+                          << node_id << " for " << object_id << ": " << status.ToString();
+          }
         }
       });
   return Status::OK();
@@ -108,11 +115,18 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectRemoved(
   request.set_node_id(node_id.Binary());
 
   rpc_client->RemoveObjectLocationOwner(
-      request, [worker_id, object_id](Status status,
-                                      const rpc::RemoveObjectLocationOwnerReply &reply) {
+      request, [worker_id, object_id, node_id](
+                   Status status, const rpc::RemoveObjectLocationOwnerReply &reply) {
         if (!status.ok()) {
-          RAY_LOG(ERROR) << "Worker " << worker_id
-                         << " failed to remove the location for " << object_id;
+          if (status.IsObjectNotFound()) {
+            RAY_LOG(INFO) << "Worker " << worker_id << " failed to remove the location "
+                          << node_id << " for " << object_id
+                          << " because the owner no longer has the object; we assume the "
+                             "object was freed.";
+          } else {
+            RAY_LOG(INFO) << "Worker " << worker_id << " failed to remove the location "
+                          << node_id << " for " << object_id << ": " << status.ToString();
+          }
         }
       });
   return Status::OK();
@@ -121,22 +135,36 @@ ray::Status OwnershipBasedObjectDirectory::ReportObjectRemoved(
 void OwnershipBasedObjectDirectory::SubscriptionCallback(
     ObjectID object_id, WorkerID worker_id, Status status,
     const rpc::GetObjectLocationsOwnerReply &reply) {
+  // 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;
   }
-
-  if (reply.object_size() > 0) {
-    it->second.object_size = reply.object_size();
-  }
-
   std::unordered_set<NodeID> node_ids;
-  for (auto const &node_id : reply.node_ids()) {
-    node_ids.emplace(NodeID::FromBinary(node_id));
+
+  // Once this flag is set to true, it should never go back to false.
+  it->second.subscribed = true;
+
+  if (!status.ok()) {
+    RAY_LOG(INFO) << "Worker " << worker_id << " failed to return location updates to "
+                  << "subscribers  for " << object_id << ": " << status.ToString()
+                  << ", assuming that the object was freed or evicted.";
+    it->second.object_size = 0;
+  } else {
+    if (reply.object_size() > 0) {
+      it->second.object_size = reply.object_size();
+    }
+
+    for (auto const &node_id : reply.node_ids()) {
+      node_ids.emplace(NodeID::FromBinary(node_id));
+    }
+    FilterRemovedNodes(gcs_client_, &node_ids);
   }
-  FilterRemovedNodes(gcs_client_, &node_ids);
-  if (node_ids != it->second.current_object_locations) {
+  if (node_ids != it->second.current_object_locations || !status.ok()) {
     it->second.current_object_locations = std::move(node_ids);
+    // 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
@@ -154,7 +182,7 @@ void OwnershipBasedObjectDirectory::SubscriptionCallback(
   rpc::GetObjectLocationsOwnerRequest request;
   request.set_intended_worker_id(worker_id.Binary());
   request.set_object_id(object_id.Binary());
-  // TODO(zhuohan): Fix this infinite loop.
+  request.set_last_version(reply.current_version());
   worker_it->second->GetObjectLocationsOwner(
       request,
       std::bind(&OwnershipBasedObjectDirectory::SubscriptionCallback, this, object_id,
@@ -176,6 +204,7 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations(
     rpc::GetObjectLocationsOwnerRequest request;
     request.set_intended_worker_id(owner_address.worker_id());
     request.set_object_id(object_id.Binary());
+    request.set_last_version(-1);
     rpc_client->GetObjectLocationsOwner(
         request,
         std::bind(&OwnershipBasedObjectDirectory::SubscriptionCallback, this, object_id,
@@ -188,6 +217,16 @@ ray::Status OwnershipBasedObjectDirectory::SubscribeObjectLocations(
     return Status::OK();
   }
   listener_state.callbacks.emplace(callback_id, callback);
+
+  // If we previously received some notifications about the object's locations,
+  // immediately notify the caller of the current known locations.
+  if (listener_state.subscribed) {
+    auto &locations = listener_state.current_object_locations;
+    auto object_size = it->second.object_size;
+    io_service_.post([callback, locations, object_size, object_id]() {
+      callback(object_id, locations, "", NodeID::Nil(), object_size);
+    });
+  }
   return Status::OK();
 }
 
@@ -221,6 +260,7 @@ ray::Status OwnershipBasedObjectDirectory::LookupLocations(
   rpc::GetObjectLocationsOwnerRequest request;
   request.set_intended_worker_id(owner_address.worker_id());
   request.set_object_id(object_id.Binary());
+  request.set_last_version(-1);
 
   rpc_client->GetObjectLocationsOwner(
       request, [this, worker_id, object_id, callback](

src/ray/protobuf/core_worker.proto

@@ -182,11 +182,17 @@ message RemoveObjectLocationOwnerReply {
 message GetObjectLocationsOwnerRequest {
   bytes intended_worker_id = 1;
   bytes object_id = 2;
+  // The version of the last location update. Only updates more recent than this version
+  // will be returned. -1 indicates that the current location data should
+  // always be returned.
+  int64 last_version = 3;
 }
 
 message GetObjectLocationsOwnerReply {
   repeated bytes node_ids = 1;
   uint64 object_size = 2;
+  // The version of the returned location updates.
+  int64 current_version = 3;
 }
 
 message KillActorRequest {