Merge pull request #1939 from sonots/tss_api

[WIP] Add TSS (Thread Specific Storage) API in CPython 3.7+

3 files changed, 98 insertions, 1 deletions

diff --git a/Cython/Includes/cpython/pythread.pxd b/Cython/Includes/cpython/pythread.pxd
index e6c6395e1..8da456743 100644
--- a/Cython/Includes/cpython/pythread.pxd
+++ b/Cython/Includes/cpython/pythread.pxd
@@ -29,7 +29,7 @@ cdef extern from "pythread.h":
     size_t PyThread_get_stacksize()
     int PyThread_set_stacksize(size_t)
 
-    # Thread Local Storage (TLS) API
+    # Thread Local Storage (TLS) API deprecated in CPython 3.7+
     int PyThread_create_key()
     void PyThread_delete_key(int)
     int PyThread_set_key_value(int, void *)
@@ -38,3 +38,15 @@ cdef extern from "pythread.h":
 
     # Cleanup after a fork
     void PyThread_ReInitTLS()
+
+    # Thread Specific Storage (TSS) API in CPython 3.7+
+    ctypedef struct Py_tss_t:
+        pass
+    Py_tss_t Py_tss_NEEDS_INIT
+    Py_tss_t * PyThread_tss_alloc()
+    void PyThread_tss_free(Py_tss_t *key)
+    int PyThread_tss_is_created(Py_tss_t *key)
+    int PyThread_tss_create(Py_tss_t *key)
+    void PyThread_tss_delete(Py_tss_t *key)
+    int PyThread_tss_set(Py_tss_t *key, void *value)
+    void * PyThread_tss_get(Py_tss_t *key)
diff --git a/Cython/Utility/ModuleSetupCode.c b/Cython/Utility/ModuleSetupCode.c
index b4e06cabb..7deaa22af 100644
--- a/Cython/Utility/ModuleSetupCode.c
+++ b/Cython/Utility/ModuleSetupCode.c
@@ -430,6 +430,40 @@ class __Pyx_FakeReference {
   #define __Pyx_PyThreadState_Current _PyThreadState_Current
 #endif
 
+// TSS (Thread Specific Storage) API
+#if PY_VERSION_HEX < 0x03070000 && !defined(PyThread_tss_create)
+#include "pythread.h"
+#define Py_tss_NEEDS_INIT 0
+typedef int Py_tss_t;
+static CYTHON_INLINE int PyThread_tss_create(Py_tss_t *key) {
+  *key = PyThread_create_key();
+  return 0; // PyThread_create_key reports success always
+}
+static CYTHON_INLINE Py_tss_t * PyThread_tss_alloc(void) {
+  Py_tss_t *key = (Py_tss_t *)PyObject_Malloc(sizeof(Py_tss_t));
+  *key = Py_tss_NEEDS_INIT;
+  return key;
+}
+static CYTHON_INLINE void PyThread_tss_free(Py_tss_t *key) {
+  PyObject_Free(key);
+}
+static CYTHON_INLINE int PyThread_tss_is_created(Py_tss_t *key) {
+  return *key != Py_tss_NEEDS_INIT;
+}
+static CYTHON_INLINE void PyThread_tss_delete(Py_tss_t *key) {
+  PyThread_delete_key(*key);
+  *key = Py_tss_NEEDS_INIT;
+}
+static CYTHON_INLINE int PyThread_tss_set(Py_tss_t *key, void *value) {
+  return PyThread_set_key_value(*key, value);
+}
+static CYTHON_INLINE void * PyThread_tss_get(Py_tss_t *key) {
+  return PyThread_get_key_value(*key);
+}
+// PyThread_delete_key_value(key) is equalivalent to PyThread_set_key_value(key, NULL)
+// PyThread_ReInitTLS() is a no-op
+#endif // TSS (Thread Specific Storage) API
+
 #if CYTHON_COMPILING_IN_CPYTHON || defined(_PyDict_NewPresized)
 #define __Pyx_PyDict_NewPresized(n)  ((n <= 8) ? PyDict_New() : _PyDict_NewPresized(n))
 #else
diff --git a/tests/run/tss.pyx b/tests/run/tss.pyx
new file mode 100644
index 000000000..c21b869a6
--- /dev/null
+++ b/tests/run/tss.pyx
@@ -0,0 +1,51 @@
+# mode: run
+
+from cpython.pythread cimport *
+
+def tss_create_delete():
+    """
+    >>> tss_create_delete()
+    (True, False)
+    """
+    cdef Py_tss_t tss_key = Py_tss_NEEDS_INIT
+    cdef bint after_create, after_delete
+    if PyThread_tss_create(&tss_key) > 0:
+        # handle key creation failure
+        pass
+    after_create = PyThread_tss_is_created(&tss_key) != 0
+    PyThread_tss_delete(&tss_key)
+    after_delete = PyThread_tss_is_created(&tss_key) != 0
+    return (after_create, after_delete)
+
+def tss_alloc_free():
+    """
+    >>> tss_alloc_free()
+    (True, False)
+    """
+    cdef Py_tss_t *ptr_key
+    cdef bint after_alloc, after_free
+    ptr_key = PyThread_tss_alloc()
+    if ptr_key == NULL:
+        # handle key allocation failure
+        pass
+    after_alloc = PyThread_tss_is_created(ptr_key) != 0
+    PyThread_tss_free(ptr_key)
+    after_free = PyThread_tss_is_created(ptr_key) != 0
+    return (after_alloc, after_free)
+
+def tss_set_get():
+    """
+    >>> tss_set_get()
+    1
+    """
+    cdef Py_tss_t tss_key = Py_tss_NEEDS_INIT
+    cdef int the_value = 1
+    cdef int ret_value
+    if PyThread_tss_create(&tss_key) > 0:
+        # handle key creation failure
+        pass
+    if PyThread_tss_get(&tss_key) == NULL:
+        PyThread_tss_set(&tss_key, <void *>&the_value)
+    ret_value = (<int *>PyThread_tss_get(&tss_key))[0]
+    PyThread_tss_delete(&tss_key)
+    return ret_value