From 8297522580e6cc3788d126a61f19848ab516a2f4 Mon Sep 17 00:00:00 2001
From: Robert Nishihara <robertnishihara@gmail.com>
Date: Tue, 5 Jul 2016 19:13:57 -0700
Subject: [PATCH] don't destruct worker object before all of the object
 references go out of scope (#212)

---
 lib/python/ray/services.py | 14 +++++++-------
 src/raylib.cc              | 24 ++++++++++++++++++------
 2 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/lib/python/ray/services.py b/lib/python/ray/services.py
index 31a83547a..6845e5487 100644
--- a/lib/python/ray/services.py
+++ b/lib/python/ray/services.py
@@ -39,6 +39,13 @@ def new_objstore_port():
   return 20000 + objstore_port_counter
 
 def cleanup():
+  global drivers
+  for driver in drivers:
+    ray.disconnect(driver)
+  if len(drivers) == 0:
+    ray.disconnect()
+  drivers = []
+
   global all_processes
   for p, address in all_processes:
     if p.poll() is not None: # process has already terminated
@@ -59,13 +66,6 @@ def cleanup():
     print "Termination attempt failed, giving up."
   all_processes = []
 
-  global drivers
-  for driver in drivers:
-    ray.disconnect(driver)
-  if len(drivers) == 0:
-    ray.disconnect()
-  drivers = []
-
 # atexit.register(cleanup)
 
 def start_scheduler(scheduler_address):
diff --git a/src/raylib.cc b/src/raylib.cc
index 4f05204c6..70539a761 100644
--- a/src/raylib.cc
+++ b/src/raylib.cc
@@ -22,16 +22,25 @@ static int PyObjectToWorker(PyObject* object, Worker **worker);
 // Object references
 
 typedef struct {
-    PyObject_HEAD
-    ObjRef val;
-    Worker* worker;
+  PyObject_HEAD
+  ObjRef val;
+  // We give the PyObjRef object a reference to the worker capsule object to
+  // make sure that the worker capsule does not go out of scope until all of the
+  // object references have gone out of scope. The reason for this is that the
+  // worker capsule destructor destroys the worker object. If the worker object
+  // has been destroyed, then when the object reference tries to call
+  // worker->decrement_reference_count, we can get a segfault.
+  PyObject* worker_capsule;
 } PyObjRef;
 
 static void PyObjRef_dealloc(PyObjRef *self) {
+  Worker* worker;
+  PyObjectToWorker(self->worker_capsule, &worker);
   std::vector<ObjRef> objrefs;
   objrefs.push_back(self->val);
-  self->worker->decrement_reference_count(objrefs);
+  worker->decrement_reference_count(objrefs);
   self->ob_type->tp_free((PyObject*) self);
+  Py_DECREF(self->worker_capsule); // The corresponding increment happens in PyObjRef_init.
 }
 
 static PyObject* PyObjRef_new(PyTypeObject *type, PyObject *args, PyObject *kwds) {
@@ -43,13 +52,16 @@ static PyObject* PyObjRef_new(PyTypeObject *type, PyObject *args, PyObject *kwds
 }
 
 static int PyObjRef_init(PyObjRef *self, PyObject *args, PyObject *kwds) {
-  if (!PyArg_ParseTuple(args, "iO&", &self->val, &PyObjectToWorker, &self->worker)) {
+  if (!PyArg_ParseTuple(args, "iO", &self->val, &self->worker_capsule)) {
     return -1;
   }
+  Worker* worker;
+  PyObjectToWorker(self->worker_capsule, &worker);
+  Py_INCREF(self->worker_capsule); // The corresponding decrement happens in PyObjRef_dealloc.
   std::vector<ObjRef> objrefs;
   objrefs.push_back(self->val);
   RAY_LOG(RAY_REFCOUNT, "In PyObjRef_init, calling increment_reference_count for objref " << objrefs[0]);
-  self->worker->increment_reference_count(objrefs);
+  worker->increment_reference_count(objrefs);
   return 0;
 };