From 2c1f4fd82c4eeddf473319091f9f2da556b6dc49 Mon Sep 17 00:00:00 2001
From: Stephanie Wang <swang@cs.berkeley.edu>
Date: Mon, 24 Feb 2020 17:58:42 -0800
Subject: [PATCH] [core] Add long running regression test for distributed ref
 counting and fix memory leak (#7302)

* Add long running test for serialized IDs and fix mem leak

* comment
---
 .../ray-project/project.yaml                  |  1 +
 .../workloads/many_tasks_serialized_ids.py    | 72 +++++++++++++++++++
 src/ray/core_worker/reference_count.cc        |  9 +--
 3 files changed, 78 insertions(+), 4 deletions(-)
 create mode 100644 ci/long_running_tests/workloads/many_tasks_serialized_ids.py

diff --git a/ci/long_running_tests/ray-project/project.yaml b/ci/long_running_tests/ray-project/project.yaml
index da9bed221..fc37e2b07 100644
--- a/ci/long_running_tests/ray-project/project.yaml
+++ b/ci/long_running_tests/ray-project/project.yaml
@@ -29,6 +29,7 @@ commands:
             "node_failures",
             "pbt",
             "serve",
+            "many_tasks_serialized_ids",
           ]
     config:
       tmux: true
diff --git a/ci/long_running_tests/workloads/many_tasks_serialized_ids.py b/ci/long_running_tests/workloads/many_tasks_serialized_ids.py
new file mode 100644
index 000000000..bb300428f
--- /dev/null
+++ b/ci/long_running_tests/workloads/many_tasks_serialized_ids.py
@@ -0,0 +1,72 @@
+# This workload stresses distributed reference counting by passing and
+# returning serialized ObjectIDs.
+
+import time
+
+import numpy as np
+
+import ray
+from ray.cluster_utils import Cluster
+
+num_redis_shards = 5
+redis_max_memory = 10**8
+object_store_memory = 10**8
+num_nodes = 10
+
+message = ("Make sure there is enough memory on this machine to run this "
+           "workload. We divide the system memory by 2 to provide a buffer.")
+assert (num_nodes * object_store_memory + num_redis_shards * redis_max_memory <
+        ray.utils.get_system_memory() / 2)
+
+# Simulate a cluster on one machine.
+
+cluster = Cluster()
+for i in range(num_nodes):
+    cluster.add_node(
+        redis_port=6379 if i == 0 else None,
+        num_redis_shards=num_redis_shards if i == 0 else None,
+        num_cpus=2,
+        num_gpus=0,
+        resources={str(i): 2},
+        object_store_memory=object_store_memory,
+        redis_max_memory=redis_max_memory,
+        webui_host="0.0.0.0")
+ray.init(address=cluster.address)
+
+# Run the workload.
+
+
+@ray.remote(max_retries=0)
+def child(*xs):
+    return np.zeros(1024, dtype=np.uint8)
+
+
+@ray.remote(max_retries=0)
+def f(xs):
+    return child.remote(*xs)
+
+
+iteration = 0
+ids = []
+start_time = time.time()
+previous_time = start_time
+while True:
+    for _ in range(50):
+        new_constrained_ids = [
+            f._remote(args=[[ids]], resources={str(i % num_nodes): 1})
+            for i in range(25)
+        ]
+        new_unconstrained_ids = [f.remote([ids]) for _ in range(25)]
+        ids = new_constrained_ids + new_unconstrained_ids
+
+    ray.get(ids)
+
+    new_time = time.time()
+    print("Iteration {}:\n"
+          "  - Iteration time: {}.\n"
+          "  - Absolute time: {}.\n"
+          "  - Total elapsed time: {}.".format(
+              iteration, new_time - previous_time, new_time,
+              new_time - start_time))
+    previous_time = new_time
+    iteration += 1
diff --git a/src/ray/core_worker/reference_count.cc b/src/ray/core_worker/reference_count.cc
index 04f42b107..e852bcf99 100644
--- a/src/ray/core_worker/reference_count.cc
+++ b/src/ray/core_worker/reference_count.cc
@@ -218,8 +218,7 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it,
                                                std::vector<ObjectID> *deleted) {
   const ObjectID id = it->first;
   RAY_LOG(DEBUG) << "Attempting to delete object " << id;
-  if (distributed_ref_counting_enabled_ && it->second.RefCount() == 0 &&
-      it->second.on_ref_removed) {
+  if (it->second.RefCount() == 0 && it->second.on_ref_removed) {
     RAY_LOG(DEBUG) << "Calling on_ref_removed for object " << id;
     it->second.on_ref_removed(id);
     it->second.on_ref_removed = nullptr;
@@ -257,7 +256,8 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it,
           // If this object ID was nested in a borrowed object, make sure that
           // we have already returned this information through a previous
           // GetAndClearLocalBorrowers call.
-          RAY_CHECK(!inner_it->second.contained_in_borrowed_id.has_value());
+          RAY_CHECK(!inner_it->second.contained_in_borrowed_id.has_value())
+              << "Outer object " << id << ", inner object " << inner_id;
         }
         DeleteReferenceInternal(inner_it, deleted);
       }
@@ -266,8 +266,8 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it,
 
   // Perform the deletion.
   if (should_delete_value) {
-    RAY_LOG(DEBUG) << "Deleting object " << id;
     if (it->second.on_delete) {
+      RAY_LOG(DEBUG) << "Calling on_delete for object " << id;
       it->second.on_delete(id);
       it->second.on_delete = nullptr;
     }
@@ -276,6 +276,7 @@ void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it,
     }
   }
   if (should_delete_reference) {
+    RAY_LOG(DEBUG) << "Deleting Reference to object " << id;
     object_id_refs_.erase(it);
   }
 }