From baa4b7cae341146c0ab371b80aba3bc2e9d2805b Mon Sep 17 00:00:00 2001
From: Robert Nishihara <robertnishihara@gmail.com>
Date: Fri, 22 Jul 2016 14:15:02 -0700
Subject: [PATCH] fix bug in reference counting None object (#286)

---
 src/raylib.cc   |  8 ++++++--
 test/runtest.py | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/src/raylib.cc b/src/raylib.cc
index bf0a3f7dd..1cbf3a36a 100644
--- a/src/raylib.cc
+++ b/src/raylib.cc
@@ -714,9 +714,13 @@ static PyObject* wait_for_next_message(PyObject* self, PyObject* args) {
   Worker* worker;
   PyObjectToWorker(worker_capsule, &worker);
   if (std::unique_ptr<WorkerMessage> message = worker->receive_next_message()) {
+    // The tuple constructed below will take ownership of some None objects.
+    // When the tuple goes out of scope, the reference count for None will be
+    // decremented. Therefore, we need to increment the reference count for None
+    // every time we put a None in the tuple.
     PyObject* t = PyTuple_New(2); // We set the items of the tuple using PyTuple_SetItem, because that transfers ownership to the tuple.
-    PyTuple_SetItem(t, 0, message->task.name().empty() ? Py_None : deserialize_task(worker_capsule, &message->task));
-    PyTuple_SetItem(t, 1, message->function.empty() ? Py_None : PyString_FromStringAndSize(message->function.data(), static_cast<ssize_t>(message->function.size())));
+    PyTuple_SetItem(t, 0, message->task.name().empty() ? Py_INCREF(Py_None), Py_None : deserialize_task(worker_capsule, &message->task));
+    PyTuple_SetItem(t, 1, message->function.empty() ? Py_INCREF(Py_None), Py_None : PyString_FromStringAndSize(message->function.data(), static_cast<ssize_t>(message->function.size())));
     return t;
   }
   Py_RETURN_NONE;
diff --git a/test/runtest.py b/test/runtest.py
index 7c91c940b..f558fd330 100644
--- a/test/runtest.py
+++ b/test/runtest.py
@@ -4,6 +4,7 @@ import numpy as np
 import time
 import subprocess32 as subprocess
 import os
+import sys
 
 import test_functions
 import ray.array.remote as ra
@@ -489,5 +490,22 @@ class PythonModeTest(unittest.TestCase):
 
     ray.services.cleanup()
 
+class PythonCExtensionTest(unittest.TestCase):
+
+  def testReferenceCountNone(self):
+    worker_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "test_worker.py")
+    ray.services.start_ray_local(num_workers=1, worker_path=worker_path)
+
+    # Make sure that we aren't accidentally messing up Python's reference counts.
+    for obj in [None, True, False]:
+      @ray.remote([], [int])
+      def f():
+        return sys.getrefcount(obj)
+      first_count = ray.get(f())
+      second_count = ray.get(f())
+      self.assertEqual(first_count, second_count)
+
+    ray.services.cleanup()
+
 if __name__ == "__main__":
     unittest.main()