From e6a188b2e07683ad77e183d8131d1cb8f3fd0d20 Mon Sep 17 00:00:00 2001
From: Philipp Moritz <pcmoritz@gmail.com>
Date: Mon, 20 Jun 2016 12:05:54 -0700
Subject: [PATCH] Fix memory leaks for Python dicts in C extension

---
 src/raylib.cc                   | 30 +++++++++++++++++++++++-------
 test/memory_leak_deserialize.py | 17 +++++++++++++++++
 2 files changed, 40 insertions(+), 7 deletions(-)
 create mode 100644 test/memory_leak_deserialize.py

diff --git a/src/raylib.cc b/src/raylib.cc
index c727e67d4..87ca0d787 100644
--- a/src/raylib.cc
+++ b/src/raylib.cc
@@ -183,6 +183,20 @@ void TaskCapsule_Destructor(PyObject* capsule) {
   delete obj;
 }
 
+// Helper methods
+
+// Pass ownership of both the key and the value to the PyDict.
+// This is only required for PyDicts, not for PyLists or PyTuples, compare
+// https://docs.python.org/2/c-api/dict.html
+// https://docs.python.org/2/c-api/list.html
+// https://docs.python.org/2/c-api/tuple.html
+
+void set_dict_item_and_transfer_ownership(PyObject* dict, PyObject* key, PyObject* val) {
+  PyDict_SetItem(dict, key, val);
+  Py_XDECREF(key);
+  Py_XDECREF(val);
+}
+
 // Serialization
 
 // serialize will serialize the python object val into the protocol buffer
@@ -364,7 +378,9 @@ PyObject* deserialize(PyObject* worker_capsule, const Obj& obj, std::vector<ObjR
     PyObject* dict = PyDict_New();
     size_t size = data.elem_size();
     for (size_t i = 0; i < size; ++i) {
-      PyDict_SetItem(dict, deserialize(worker_capsule, data.elem(i).key(), objrefs), deserialize(worker_capsule, data.elem(i).value(), objrefs));
+      PyObject* pykey = deserialize(worker_capsule, data.elem(i).key(), objrefs);
+      PyObject* pyval = deserialize(worker_capsule, data.elem(i).value(), objrefs);
+      set_dict_item_and_transfer_ownership(dict, pykey, pyval);
     }
     return dict;
   } else if (obj.has_string_data()) {
@@ -785,8 +801,8 @@ PyObject* scheduler_info(PyObject* self, PyObject* args) {
   }
 
   PyObject* dict = PyDict_New();
-  PyDict_SetItem(dict, PyString_FromString("target_objrefs"), target_objref_list);
-  PyDict_SetItem(dict, PyString_FromString("reference_counts"), reference_count_list);
+  set_dict_item_and_transfer_ownership(dict, PyString_FromString("target_objrefs"), target_objref_list);
+  set_dict_item_and_transfer_ownership(dict, PyString_FromString("reference_counts"), reference_count_list);
   return dict;
 }
 
@@ -804,14 +820,14 @@ PyObject* task_info(PyObject* self, PyObject* args) {
   for (size_t i = 0; i < reply.failed_task_size(); ++i) {
     const TaskStatus& info = reply.failed_task(i);
     PyObject* info_dict = PyDict_New();
-    PyDict_SetItem(info_dict, PyString_FromString("worker_address"), PyString_FromStringAndSize(info.worker_address().data(), info.worker_address().size()));
-    PyDict_SetItem(info_dict, PyString_FromString("operationid"), PyInt_FromLong(info.operationid()));
-    PyDict_SetItem(info_dict, PyString_FromString("error_message"), PyString_FromStringAndSize(info.error_message().data(), info.error_message().size()));
+    set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("worker_address"), PyString_FromStringAndSize(info.worker_address().data(), info.worker_address().size()));
+    set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("operationid"), PyInt_FromLong(info.operationid()));
+    set_dict_item_and_transfer_ownership(info_dict, PyString_FromString("error_message"), PyString_FromStringAndSize(info.error_message().data(), info.error_message().size()));
     PyList_SetItem(failed_tasks_list, i, info_dict);
   }
 
   PyObject* dict = PyDict_New();
-  PyDict_SetItem(dict, PyString_FromString("failed_tasks"), failed_tasks_list);
+  set_dict_item_and_transfer_ownership(dict, PyString_FromString("failed_tasks"), failed_tasks_list);
   return dict;
 }
 
diff --git a/test/memory_leak_deserialize.py b/test/memory_leak_deserialize.py
new file mode 100644
index 000000000..5ff87095d
--- /dev/null
+++ b/test/memory_leak_deserialize.py
@@ -0,0 +1,17 @@
+# This code reproduces a memory leak we had in the past
+
+import os
+import numpy as np
+import ray
+import ray.worker
+import ray.services as services
+
+worker_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_worker.py")
+services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=1, worker_path=worker_path)
+
+d = {"w": np.zeros(1000000)}
+
+obj_capsule, contained_objrefs = ray.lib.serialize_object(ray.worker.global_worker.handle, d)
+
+while True:
+  ray.lib.deserialize_object(ray.worker.global_worker.handle, obj_capsule)