Track an ID refcount on PyInterpreterState.

python · ericsnowcurrently · Feb 17, 2018 · Feb 16, 2018 · Feb 16, 2018 · Feb 16, 2018
commit a764edba106e7566edeb79da2b5bd9222f062cdd
diff --git a/Include/internal/pystate.h b/Include/internal/pystate.h
@@ -69,6 +69,10 @@ PyAPI_FUNC(void) _PyPathConfig_Clear(_PyPathConfig *config);
 
 PyAPI_FUNC(PyInterpreterState *) _PyInterpreterState_LookUpID(PY_INT64_T);
 
+PyAPI_FUNC(int) _PyInterpreterState_IDInitref(PyInterpreterState *);
+PyAPI_FUNC(void) _PyInterpreterState_IDIncref(PyInterpreterState *);
+PyAPI_FUNC(void) _PyInterpreterState_IDDecref(PyInterpreterState *);
+
 
 /* cross-interpreter data */
 

diff --git a/Include/pystate.h b/Include/pystate.h
@@ -8,6 +8,8 @@
 extern "C" {
 #endif
 
+#include "pythread.h"
+
 /* This limitation is for performance and simplicity. If needed it can be
 removed (with effort). */
 #define MAX_CO_EXTRA_USERS 255
@@ -111,6 +113,8 @@ typedef struct _is {
     struct _ts *tstate_head;
 
     int64_t id;
+    int64_t id_refcount;
+    PyThread_type_lock id_mutex;
 
     PyObject *modules;
     PyObject *modules_by_index;

diff --git a/Python/pystate.c b/Python/pystate.c
@@ -125,7 +125,8 @@ PyInterpreterState_New(void)
         return NULL;
     }
 
-
+    interp->id_refcount = -1;
+    interp->id_mutex = NULL;
     interp->modules = NULL;
     interp->modules_by_index = NULL;
     interp->sysdict = NULL;
@@ -247,6 +248,9 @@ PyInterpreterState_Delete(PyInterpreterState *interp)
             Py_FatalError("PyInterpreterState_Delete: remaining subinterpreters");
     }
     HEAD_UNLOCK();
+    if (interp->id_mutex != NULL) {
+        PyThread_free_lock(interp->id_mutex);
+    }
     PyMem_RawFree(interp);
 }
 
@@ -284,6 +288,58 @@ _PyInterpreterState_LookUpID(PY_INT64_T requested_id)
     return NULL;
 }
 
+
+int
+_PyInterpreterState_IDInitref(PyInterpreterState *interp)
+{
+    if (interp->id_mutex != NULL) {
+        return 0;
+    }
+    interp->id_mutex = PyThread_allocate_lock();
+    if (interp->id_mutex == NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "failed to create init interpreter ID mutex");
+        return -1;
+    }
+    interp->id_refcount = 0;
+    return 0;
+}
+
+
+void
+_PyInterpreterState_IDIncref(PyInterpreterState *interp)
+{
+    if (interp->id_mutex == NULL) {
+        return;
+    }
+    PyThread_acquire_lock(interp->id_mutex, WAIT_LOCK);
+    interp->id_refcount += 1;
+    PyThread_release_lock(interp->id_mutex);
+}
+
+
+void
+_PyInterpreterState_IDDecref(PyInterpreterState *interp)
+{
+    if (interp->id_mutex == NULL) {
+        return;
+    }
+    PyThread_acquire_lock(interp->id_mutex, WAIT_LOCK);
+    assert(interp->id_refcount != 0);
+    interp->id_refcount -= 1;
+    int64_t refcount = interp->id_refcount;
+    PyThread_release_lock(interp->id_mutex);
+
+    if (refcount == 0) {
+        PyThreadState *tstate, *save_tstate;
+        tstate = PyInterpreterState_ThreadHead(interp);
+        save_tstate = PyThreadState_Swap(tstate);
+        Py_EndInterpreter(tstate);
+        PyThreadState_Swap(save_tstate);
+    }
+}
+
+
 /* Default implementation for _PyThreadState_GetFrame */
 static struct _frame *
 threadstate_getframe(PyThreadState *self)