From 603cba646ad9fcd2b432f7ebd18ed3317b7d7a11 Mon Sep 17 00:00:00 2001
From: Jiajun Yao <jeromeyjj@gmail.com>
Date: Tue, 24 May 2022 10:25:44 -0700
Subject: [PATCH] Use OBOD report as the source of truth of in-memory locations
 of object (#25004)

* hang

* update

* up

* up

* comment
---
 src/ray/core_worker/reference_count.cc        | 11 +++++++----
 .../core_worker/test/reference_count_test.cc  | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc
index 1ee295b64..4ea492091 100644
--- a/src/ray/core_worker/reference_count.cc
+++ b/src/ray/core_worker/reference_count.cc
@@ -695,8 +695,6 @@ void ReferenceCounter::UpdateObjectPinnedAtRaylet(const ObjectID &object_id,
     if (!it->second.OutOfScope(lineage_pinning_enabled_)) {
       if (check_node_alive_(raylet_id)) {
         it->second.pinned_at_raylet_id = raylet_id;
-        // We eagerly add the pinned location to the set of object locations.
-        AddObjectLocationInternal(it, raylet_id);
       } else {
         ReleasePlasmaObject(it);
         objects_to_recover_.push_back(object_id);
@@ -1281,11 +1279,16 @@ absl::optional<LocalityData> ReferenceCounter::GetLocalityData(
   //   locations.
   // - If we don't own this object, this will contain a snapshot of the object locations
   //   at future resolution time.
-  const auto &node_ids = it->second.locations;
+  auto node_ids = it->second.locations;
+  // Add location of the primary copy since the object must be there: either in memory or
+  // spilled.
+  if (it->second.pinned_at_raylet_id.has_value()) {
+    node_ids.emplace(it->second.pinned_at_raylet_id.value());
+  }
 
   // We should only reach here if we have valid locality data to return.
   absl::optional<LocalityData> locality_data(
-      {static_cast<uint64_t>(object_size), node_ids});
+      {static_cast<uint64_t>(object_size), std::move(node_ids)});
   return locality_data;
 }
 
diff --git a/src/ray/core_worker/test/reference_count_test.cc b/src/ray/core_worker/test/reference_count_test.cc
index c0e86ccea..1fc5ac611 100644
--- a/src/ray/core_worker/test/reference_count_test.cc
+++ b/src/ray/core_worker/test/reference_count_test.cc
@@ -648,6 +648,7 @@ TEST_F(ReferenceCountTest, TestHandleObjectSpilled) {
 TEST_F(ReferenceCountTest, TestGetLocalityData) {
   ObjectID obj1 = ObjectID::FromRandom();
   ObjectID obj2 = ObjectID::FromRandom();
+  ObjectID obj3 = ObjectID::FromRandom();
   NodeID node1 = NodeID::FromRandom();
   NodeID node2 = NodeID::FromRandom();
   rpc::Address address;
@@ -735,8 +736,25 @@ TEST_F(ReferenceCountTest, TestGetLocalityData) {
   auto locality_data_obj2_no_object_size = rc->GetLocalityData(obj2);
   ASSERT_FALSE(locality_data_obj2_no_object_size.has_value());
 
+  // Primary copy location is always returned
+  // even if it's not in-memory (i.e. spilled).
+  rc->AddOwnedObject(obj3,
+                     {},
+                     address,
+                     "file2.py:43",
+                     -1,
+                     false,
+                     /*add_local_ref=*/true);
+  rc->UpdateObjectSize(obj3, 101);
+  rc->UpdateObjectPinnedAtRaylet(obj3, node1);
+  auto locality_data_obj3 = rc->GetLocalityData(obj3);
+  ASSERT_TRUE(locality_data_obj3.has_value());
+  ASSERT_EQ(locality_data_obj3->nodes_containing_object,
+            absl::flat_hash_set<NodeID>({node1}));
+
   rc->RemoveLocalReference(obj1, nullptr);
   rc->RemoveLocalReference(obj2, nullptr);
+  rc->RemoveLocalReference(obj3, nullptr);
 }
 
 // Tests that we can get the owner address correctly for objects that we own,
@@ -2625,6 +2643,7 @@ TEST_F(ReferenceCountLineageEnabledTest, TestPlasmaLocation) {
   ASSERT_TRUE(rc->IsPlasmaObjectPinnedOrSpilled(id, &owned_by_us, &pinned_at, &spilled));
   ASSERT_TRUE(owned_by_us);
   ASSERT_FALSE(pinned_at.IsNil());
+  ASSERT_TRUE(rc->GetObjectLocations(id)->empty());
 
   rc->RemoveLocalReference(id, nullptr);
   ASSERT_FALSE(rc->IsPlasmaObjectPinnedOrSpilled(id, &owned_by_us, &pinned_at, &spilled));