Garbage collect greenlets after threads exit

Thanks to Armin Ronacher and Armin Rigo

greenlet.c

@@ -245,8 +245,13 @@ static int green_updatecurrent(void)
 
 static PyObject* green_statedict(PyGreenlet* g)
 {
-	while (!PyGreenlet_STARTED(g))
+	while (!PyGreenlet_STARTED(g)) {
 		g = g->parent;
+		if (g == NULL) {
+			/* garbage collected greenlet in chain */
+			return NULL;
+		}
+	}
 	return g->run_info;
 }
 
@@ -465,16 +470,19 @@ g_switch(PyGreenlet* target, PyObject* args, PyObject* kwargs)
 {
 	/* _consumes_ a reference to the args tuple and kwargs dict,
 	   and return a new tuple reference */
+	PyObject* run_info;
 
 	/* check ts_current */
 	if (!STATE_OK) {
 		Py_DECREF(args);
 		Py_XDECREF(kwargs);
 		return NULL;
 	}
-	if (green_statedict(target) != ts_current->run_info) {
-		PyErr_SetString(PyExc_GreenletError,
-				"cannot switch to a different thread");
+	run_info = green_statedict(target);
+	if (run_info == NULL || run_info != ts_current->run_info) {
+		PyErr_SetString(PyExc_GreenletError, run_info
+				? "cannot switch to a different thread"
+				: "cannot switch to a garbage collected greenlet");
 		Py_DECREF(args);
 		Py_XDECREF(kwargs);
 		return NULL;
@@ -734,15 +742,28 @@ green_traverse(PyGreenlet *self, visitproc visit, void *arg)
 
 static int green_is_gc(PyGreenlet* self)
 {
-	int rval;
-	/* Main and alive greenlets are not garbage collectable */
-	rval = (self->stack_stop == (char *)-1 || self->stack_start != NULL) ? 0 : 1;
-	return rval;
+	/* Main greenlet can be garbage collected since it can only
+	   become unreachable if the underlying thread exited.
+	   Active greenlet cannot be garbage collected, however. */
+	if (PyGreenlet_MAIN(self) || !PyGreenlet_ACTIVE(self))
+		return 1;
+	return 0;
 }
 
 static int green_clear(PyGreenlet* self)
 {
-	return 0; /* greenlet is not alive, so there's nothing to clear */
+	/* Greenlet is only cleared if it is about to be collected.
+	   Since active greenlets are not garbage collectable, we can
+	   be sure that, even if they are deallocated during clear,
+	   nothing they reference is in unreachable or finalizers,
+	   so even if it switches we are relatively safe. */
+	Py_CLEAR(self->parent);
+	Py_CLEAR(self->run_info);
+	Py_CLEAR(self->exc_type);
+	Py_CLEAR(self->exc_value);
+	Py_CLEAR(self->exc_traceback);
+	Py_CLEAR(self->dict);
+	return 0;
 }
 #endif
 
@@ -754,7 +775,7 @@ static void green_dealloc(PyGreenlet* self)
 	PyObject_GC_UnTrack((PyObject *)self);
 	Py_TRASHCAN_SAFE_BEGIN(self);
 #endif /* GREENLET_USE_GC */
-	if (PyGreenlet_ACTIVE(self)) {
+	if (PyGreenlet_ACTIVE(self) && self->run_info != NULL) {
 		/* Hacks hacks hacks copied from instance_dealloc() */
 		/* Temporarily resurrect the greenlet. */
 		assert(Py_REFCNT(self) == 0);
@@ -1057,6 +1078,7 @@ static PyObject* green_getparent(PyGreenlet* self, void* c)
 static int green_setparent(PyGreenlet* self, PyObject* nparent, void* c)
 {
 	PyGreenlet* p;
+	PyObject* run_info = NULL;
 	if (nparent == NULL) {
 		PyErr_SetString(PyExc_AttributeError, "can't delete attribute");
 		return -1;
@@ -1070,6 +1092,15 @@ static int green_setparent(PyGreenlet* self, PyObject* nparent, void* c)
 			PyErr_SetString(PyExc_ValueError, "cyclic parent chain");
 			return -1;
 		}
+		run_info = PyGreenlet_ACTIVE(p) ? p->run_info : NULL;
+	}
+	if (run_info == NULL) {
+		PyErr_SetString(PyExc_ValueError, "parent must not be garbage collected");
+		return -1;
+	}
+	if (PyGreenlet_STARTED(self) && self->run_info != run_info) {
+		PyErr_SetString(PyExc_ValueError, "parent cannot be on a different thread");
+		return -1;
 	}
 	p = self->parent;
 	self->parent = (PyGreenlet*) nparent;

greenlet.h

@@ -32,6 +32,7 @@ typedef struct _greenlet {
 } PyGreenlet;
 
 #define PyGreenlet_Check(op)      PyObject_TypeCheck(op, &PyGreenlet_Type)
+#define PyGreenlet_MAIN(op)       (((PyGreenlet*)(op))->stack_stop == (char*) -1)
 #define PyGreenlet_STARTED(op)    (((PyGreenlet*)(op))->stack_stop != NULL)
 #define PyGreenlet_ACTIVE(op)     (((PyGreenlet*)(op))->stack_start != NULL)
 #define PyGreenlet_GET_PARENT(op) (((PyGreenlet*)(op))->parent)

tests/test_greenlet.py

@@ -262,3 +262,28 @@ def setdict(g, value):
         self.assertEqual(g.__dict__, {'test': 42})
         self.assertRaises(TypeError, deldict, g)
         self.assertRaises(TypeError, setdict, g, 42)
+
+    def test_threaded_reparent(self):
+        data = {}
+        created_event = threading.Event()
+        done_event = threading.Event()
+
+        def foo():
+            data['g'] = greenlet(lambda: None)
+            created_event.set()
+            done_event.wait()
+
+        def blank():
+            greenlet.getcurrent().parent.switch()
+
+        def setparent(g, value):
+            g.parent = value
+
+        thread = threading.Thread(target=foo)
+        thread.start()
+        created_event.wait()
+        g = greenlet(blank)
+        g.switch()
+        self.assertRaises(ValueError, setparent, g, data['g'])
+        done_event.set()
+        thread.join()

tests/test_leaks.py

@@ -1,7 +1,10 @@
 import unittest
 import sys
+import gc
 
+import weakref
 import greenlet
+import threading
 
 
 class ArgRefcountTests(unittest.TestCase):
@@ -21,3 +24,39 @@ def test_kwarg_refs(self):
         for i in range(100):
             g.switch(**kwargs)
         self.assertEqual(sys.getrefcount(kwargs), 2)
+
+    if greenlet.GREENLET_USE_GC:
+        # These only work with greenlet gc support
+
+        def test_threaded_leak(self):
+            gg = []
+            def worker():
+                # only main greenlet present
+                gg.append(weakref.ref(greenlet.getcurrent()))
+            for i in range(2):
+                t = threading.Thread(target=worker)
+                t.start()
+                t.join()
+            greenlet.getcurrent() # update ts_current
+            gc.collect()
+            for g in gg:
+                self.assertTrue(g() is None)
+
+        def test_threaded_adv_leak(self):
+            gg = []
+            def worker():
+                # main and additional *finished* greenlets
+                ll = greenlet.getcurrent().ll = []
+                def additional():
+                    ll.append(greenlet.getcurrent())
+                for i in range(2):
+                    greenlet.greenlet(additional).switch()
+                gg.append(weakref.ref(greenlet.getcurrent()))
+            for i in range(2):
+                t = threading.Thread(target=worker)
+                t.start()
+                t.join()
+            greenlet.getcurrent() # update ts_current
+            gc.collect()
+            for g in gg:
+                self.assertTrue(g() is None)