Merge pull request #16519 from pv/f2py-threadsafe-cb

BUG: f2py:  make callbacks threadsafe

5 files changed, 206 insertions, 47 deletions

diff --git a/numpy/f2py/cb_rules.py b/numpy/f2py/cb_rules.py
index 5c0668db1..d3e114f20 100644
--- a/numpy/f2py/cb_rules.py
+++ b/numpy/f2py/cb_rules.py
@@ -33,13 +33,61 @@ cb_routine_rules = {
     'cbtypedefs': 'typedef #rctype#(*#name#_typedef)(#optargs_td##args_td##strarglens_td##noargs#);',
     'body': """
 #begintitle#
-PyObject *#name#_capi = NULL;/*was Py_None*/
-PyTupleObject *#name#_args_capi = NULL;
-int #name#_nofargs = 0;
-jmp_buf #name#_jmpbuf;
+typedef struct {
+    PyObject *capi;
+    PyTupleObject *args_capi;
+    int nofargs;
+    jmp_buf jmpbuf;
+} #name#_t;
+
+#if defined(__GNUC__) \
+    && (__GNUC__ > 4 || (__GNUC__ == 4 && (__GNUC_MINOR__ >= 4))) \
+    && !defined(F2PY_USE_PYTHON_TLS)
+
+static __thread #name#_t *_active_#name# = NULL;
+
+static #name#_t *swap_active_#name#(#name#_t *ptr) {
+    #name#_t *prev = _active_#name#;
+    _active_#name# = ptr;
+    return prev;
+}
+
+static #name#_t *get_active_#name#(void) {
+    return _active_#name#;
+}
+
+#elif defined(_MSC_VER) && !defined(F2PY_USE_PYTHON_TLS)
+
+static __declspec(thread) #name#_t *_active_#name# = NULL;
+
+static #name#_t *swap_active_#name#(#name#_t *ptr) {
+    #name#_t *prev = _active_#name#;
+    _active_#name# = ptr;
+    return prev;
+}
+
+static #name#_t *get_active_#name#(void) {
+    return _active_#name#;
+}
+
+#else
+
+static #name#_t *swap_active_#name#(#name#_t *ptr) {
+    char *key = "__f2py_cb_#name#";
+    return (#name#_t *)F2PySwapThreadLocalCallbackPtr(key, ptr);
+}
+
+static #name#_t *get_active_#name#(void) {
+    char *key = "__f2py_cb_#name#";
+    return (#name#_t *)F2PyGetThreadLocalCallbackPtr(key);
+}
+
+#endif
+
 /*typedef #rctype#(*#name#_typedef)(#optargs_td##args_td##strarglens_td##noargs#);*/
 #static# #rctype# #callbackname# (#optargs##args##strarglens##noargs#) {
-    PyTupleObject *capi_arglist = #name#_args_capi;
+    #name#_t *cb;
+    PyTupleObject *capi_arglist = NULL;
     PyObject *capi_return = NULL;
     PyObject *capi_tmp = NULL;
     PyObject *capi_arglist_list = NULL;
@@ -49,19 +97,21 @@ jmp_buf #name#_jmpbuf;
 #ifdef F2PY_REPORT_ATEXIT
 f2py_cb_start_clock();
 #endif
+    cb = get_active_#name#();
+    capi_arglist = cb->args_capi;
     CFUNCSMESS(\"cb:Call-back function #name# (maxnofargs=#maxnofargs#(-#nofoptargs#))\\n\");
-    CFUNCSMESSPY(\"cb:#name#_capi=\",#name#_capi);
-    if (#name#_capi==NULL) {
+    CFUNCSMESSPY(\"cb:#name#_capi=\",cb->capi);
+    if (cb->capi==NULL) {
         capi_longjmp_ok = 0;
-        #name#_capi = PyObject_GetAttrString(#modulename#_module,\"#argname#\");
+        cb->capi = PyObject_GetAttrString(#modulename#_module,\"#argname#\");
     }
-    if (#name#_capi==NULL) {
+    if (cb->capi==NULL) {
         PyErr_SetString(#modulename#_error,\"cb: Callback #argname# not defined (as an argument or module #modulename# attribute).\\n\");
         goto capi_fail;
     }
-    if (F2PyCapsule_Check(#name#_capi)) {
+    if (F2PyCapsule_Check(cb->capi)) {
     #name#_typedef #name#_cptr;
-    #name#_cptr = F2PyCapsule_AsVoidPtr(#name#_capi);
+    #name#_cptr = F2PyCapsule_AsVoidPtr(cb->capi);
     #returncptr#(*#name#_cptr)(#optargs_nm##args_nm##strarglens_nm#);
     #return#
     }
@@ -103,11 +153,11 @@ f2py_cb_start_clock();
 f2py_cb_start_call_clock();
 #endif
 #ifdef PYPY_VERSION
-    capi_return = PyObject_CallObject(#name#_capi,(PyObject *)capi_arglist_list);
+    capi_return = PyObject_CallObject(cb->capi,(PyObject *)capi_arglist_list);
     Py_DECREF(capi_arglist_list);
     capi_arglist_list = NULL;
 #else
-    capi_return = PyObject_CallObject(#name#_capi,(PyObject *)capi_arglist);
+    capi_return = PyObject_CallObject(cb->capi,(PyObject *)capi_arglist);
 #endif
 #ifdef F2PY_REPORT_ATEXIT
 f2py_cb_stop_call_clock();
@@ -137,8 +187,9 @@ capi_fail:
     fprintf(stderr,\"Call-back #name# failed.\\n\");
     Py_XDECREF(capi_return);
     Py_XDECREF(capi_arglist_list);
-    if (capi_longjmp_ok)
-        longjmp(#name#_jmpbuf,-1);
+    if (capi_longjmp_ok) {
+        longjmp(cb->jmpbuf,-1);
+    }
 capi_return_pt:
     ;
 #return#
@@ -335,11 +386,11 @@ cb_arg_rules = [
         '_check': isscalar
     }, {
         'pyobjfrom': [{isintent_in: """\
-    if (#name#_nofargs>capi_i)
+    if (cb->nofargs>capi_i)
         if (CAPI_ARGLIST_SETITEM(capi_i++,pyobj_from_#ctype#1(#varname_i#)))
             goto capi_fail;"""},
                       {isintent_inout: """\
-    if (#name#_nofargs>capi_i)
+    if (cb->nofargs>capi_i)
         if (CAPI_ARGLIST_SETITEM(capi_i++,pyarr_from_p_#ctype#1(#varname_i#_cb_capi)))
             goto capi_fail;"""}],
         'need': [{isintent_in: 'pyobj_from_#ctype#1'},
@@ -360,11 +411,11 @@ cb_arg_rules = [
     }, {
         'pyobjfrom': [{debugcapi: '    fprintf(stderr,"debug-capi:cb:#varname#=\\"#showvalueformat#\\":%d:\\n",#varname_i#,#varname_i#_cb_len);'},
                       {isintent_in: """\
-    if (#name#_nofargs>capi_i)
+    if (cb->nofargs>capi_i)
         if (CAPI_ARGLIST_SETITEM(capi_i++,pyobj_from_#ctype#1size(#varname_i#,#varname_i#_cb_len)))
             goto capi_fail;"""},
                       {isintent_inout: """\
-    if (#name#_nofargs>capi_i) {
+    if (cb->nofargs>capi_i) {
         int #varname_i#_cb_dims[] = {#varname_i#_cb_len};
         if (CAPI_ARGLIST_SETITEM(capi_i++,pyarr_from_p_#ctype#1(#varname_i#,#varname_i#_cb_dims)))
             goto capi_fail;
@@ -384,13 +435,13 @@ cb_arg_rules = [
     {
         'pyobjfrom': [{debugcapi: '    fprintf(stderr,"debug-capi:cb:#varname#\\n");'},
                       {isintent_c: """\
-    if (#name#_nofargs>capi_i) {
+    if (cb->nofargs>capi_i) {
         int itemsize_ = #atype# == NPY_STRING ? 1 : 0;
         /*XXX: Hmm, what will destroy this array??? */
         PyArrayObject *tmp_arr = (PyArrayObject *)PyArray_New(&PyArray_Type,#rank#,#varname_i#_Dims,#atype#,NULL,(char*)#varname_i#,itemsize_,NPY_ARRAY_CARRAY,NULL);
 """,
                        l_not(isintent_c): """\
-    if (#name#_nofargs>capi_i) {
+    if (cb->nofargs>capi_i) {
         int itemsize_ = #atype# == NPY_STRING ? 1 : 0;
         /*XXX: Hmm, what will destroy this array??? */
         PyArrayObject *tmp_arr = (PyArrayObject *)PyArray_New(&PyArray_Type,#rank#,#varname_i#_Dims,#atype#,NULL,(char*)#varname_i#,itemsize_,NPY_ARRAY_FARRAY,NULL);
diff --git a/numpy/f2py/rules.py b/numpy/f2py/rules.py
index ad49f4590..7b25b545a 100755
--- a/numpy/f2py/rules.py
+++ b/numpy/f2py/rules.py
@@ -749,10 +749,9 @@ arg_rules = [
         'docstrcbs': '#cbdocstr#',
         'latexdocstrcbs': '\\item[] #cblatexdocstr#',
         'latexdocstropt': {isintent_nothide: '\\item[]{{}\\verb@#varname#_extra_args := () input tuple@{}} --- Extra arguments for call-back function {{}\\verb@#varname#@{}}.'},
-        'decl': ['    PyObject *#varname#_capi = Py_None;',
+        'decl': ['    #cbname#_t #varname#_cb = { Py_None, NULL, 0 };',
+                 '    #cbname#_t *#varname#_cb_ptr = &#varname#_cb;',
                  '    PyTupleObject *#varname#_xa_capi = NULL;',
-                 '    PyTupleObject *#varname#_args_capi = NULL;',
-                 '    int #varname#_nofargs_capi = 0;',
                  {l_not(isintent_callback):
                   '    #cbname#_typedef #varname#_cptr;'}
                  ],
@@ -760,25 +759,25 @@ arg_rules = [
         'argformat': {isrequired: 'O'},
         'keyformat': {isoptional: 'O'},
         'xaformat': {isintent_nothide: 'O!'},
-        'args_capi': {isrequired: ',&#varname#_capi'},
-        'keys_capi': {isoptional: ',&#varname#_capi'},
+        'args_capi': {isrequired: ',&#varname#_cb.capi'},
+        'keys_capi': {isoptional: ',&#varname#_cb.capi'},
         'keys_xa': ',&PyTuple_Type,&#varname#_xa_capi',
-        'setjmpbuf': '(setjmp(#cbname#_jmpbuf))',
+        'setjmpbuf': '(setjmp(#varname#_cb.jmpbuf))',
         'callfortran': {l_not(isintent_callback): '#varname#_cptr,'},
         'need': ['#cbname#', 'setjmp.h'],
         '_check':isexternal
     },
     {
         'frompyobj': [{l_not(isintent_callback): """\
-if(F2PyCapsule_Check(#varname#_capi)) {
-  #varname#_cptr = F2PyCapsule_AsVoidPtr(#varname#_capi);
+if(F2PyCapsule_Check(#varname#_cb.capi)) {
+  #varname#_cptr = F2PyCapsule_AsVoidPtr(#varname#_cb.capi);
 } else {
   #varname#_cptr = #cbname#;
 }
 """}, {isintent_callback: """\
-if (#varname#_capi==Py_None) {
-  #varname#_capi = PyObject_GetAttrString(#modulename#_module,\"#varname#\");
-  if (#varname#_capi) {
+if (#varname#_cb.capi==Py_None) {
+  #varname#_cb.capi = PyObject_GetAttrString(#modulename#_module,\"#varname#\");
+  if (#varname#_cb.capi) {
     if (#varname#_xa_capi==NULL) {
       if (PyObject_HasAttrString(#modulename#_module,\"#varname#_extra_args\")) {
         PyObject* capi_tmp = PyObject_GetAttrString(#modulename#_module,\"#varname#_extra_args\");
@@ -796,34 +795,28 @@ if (#varname#_capi==Py_None) {
       }
     }
   }
-  if (#varname#_capi==NULL) {
+  if (#varname#_cb.capi==NULL) {
     PyErr_SetString(#modulename#_error,\"Callback #varname# not defined (as an argument or module #modulename# attribute).\\n\");
     return NULL;
   }
 }
 """},
             """\
-    #varname#_nofargs_capi = #cbname#_nofargs;
-    if (create_cb_arglist(#varname#_capi,#varname#_xa_capi,#maxnofargs#,#nofoptargs#,&#cbname#_nofargs,&#varname#_args_capi,\"failed in processing argument list for call-back #varname#.\")) {
-        jmp_buf #varname#_jmpbuf;""",
+    if (create_cb_arglist(#varname#_cb.capi,#varname#_xa_capi,#maxnofargs#,#nofoptargs#,&#varname#_cb.nofargs,&#varname#_cb.args_capi,\"failed in processing argument list for call-back #varname#.\")) {
+""",
             {debugcapi: ["""\
-        fprintf(stderr,\"debug-capi:Assuming %d arguments; at most #maxnofargs#(-#nofoptargs#) is expected.\\n\",#cbname#_nofargs);
+        fprintf(stderr,\"debug-capi:Assuming %d arguments; at most #maxnofargs#(-#nofoptargs#) is expected.\\n\",#varname#_cb.nofargs);
         CFUNCSMESSPY(\"for #varname#=\",#cbname#_capi);""",
                          {l_not(isintent_callback): """        fprintf(stderr,\"#vardebugshowvalue# (call-back in C).\\n\",#cbname#);"""}]},
             """\
-        CFUNCSMESS(\"Saving jmpbuf for `#varname#`.\\n\");
-        SWAP(#varname#_capi,#cbname#_capi,PyObject);
-        SWAP(#varname#_args_capi,#cbname#_args_capi,PyTupleObject);
-        memcpy(&#varname#_jmpbuf,&#cbname#_jmpbuf,sizeof(jmp_buf));""",
+        CFUNCSMESS(\"Saving callback variables for `#varname#`.\\n\");
+        #varname#_cb_ptr = swap_active_#cbname#(#varname#_cb_ptr);""",
         ],
         'cleanupfrompyobj':
         """\
-        CFUNCSMESS(\"Restoring jmpbuf for `#varname#`.\\n\");
-        #cbname#_capi = #varname#_capi;
-        Py_DECREF(#cbname#_args_capi);
-        #cbname#_args_capi = #varname#_args_capi;
-        #cbname#_nofargs = #varname#_nofargs_capi;
-        memcpy(&#cbname#_jmpbuf,&#varname#_jmpbuf,sizeof(jmp_buf));
+        CFUNCSMESS(\"Restoring callback variables for `#varname#`.\\n\");
+        #varname#_cb_ptr = swap_active_#cbname#(#varname#_cb_ptr);
+        Py_DECREF(#varname#_cb.args_capi);
     }""",
         'need': ['SWAP', 'create_cb_arglist'],
         '_check':isexternal,
diff --git a/numpy/f2py/src/fortranobject.c b/numpy/f2py/src/fortranobject.c
index b3a04bcf0..aa46c57d0 100644
--- a/numpy/f2py/src/fortranobject.c
+++ b/numpy/f2py/src/fortranobject.c
@@ -30,6 +30,68 @@ F2PyDict_SetItemString(PyObject *dict, char *name, PyObject *obj)
     return PyDict_SetItemString(dict, name, obj);
 }
 
+/*
+ * Python-only fallback for thread-local callback pointers
+ */
+void *F2PySwapThreadLocalCallbackPtr(char *key, void *ptr)
+{
+    PyObject *local_dict, *value;
+    void *prev;
+
+    local_dict = PyThreadState_GetDict();
+    if (local_dict == NULL) {
+        Py_FatalError("F2PySwapThreadLocalCallbackPtr: PyThreadState_GetDict failed");
+    }
+
+    value = PyDict_GetItemString(local_dict, key);
+    if (value != NULL) {
+        prev = PyLong_AsVoidPtr(value);
+        if (PyErr_Occurred()) {
+           Py_FatalError("F2PySwapThreadLocalCallbackPtr: PyLong_AsVoidPtr failed");
+        }
+    }
+    else {
+        prev = NULL;
+    }
+
+    value = PyLong_FromVoidPtr((void *)ptr);
+    if (value == NULL) {
+        Py_FatalError("F2PySwapThreadLocalCallbackPtr: PyLong_FromVoidPtr failed");
+    }
+
+    if (PyDict_SetItemString(local_dict, key, value) != 0) {
+        Py_FatalError("F2PySwapThreadLocalCallbackPtr: PyDict_SetItemString failed");
+    }
+
+    Py_DECREF(value);
+
+    return prev;
+}
+
+void *F2PyGetThreadLocalCallbackPtr(char *key)
+{
+    PyObject *local_dict, *value;
+    void *prev;
+
+    local_dict = PyThreadState_GetDict();
+    if (local_dict == NULL) {
+        Py_FatalError("F2PyGetThreadLocalCallbackPtr: PyThreadState_GetDict failed");
+    }
+
+    value = PyDict_GetItemString(local_dict, key);
+    if (value != NULL) {
+        prev = PyLong_AsVoidPtr(value);
+        if (PyErr_Occurred()) {
+           Py_FatalError("F2PyGetThreadLocalCallbackPtr: PyLong_AsVoidPtr failed");
+        }
+    }
+    else {
+        prev = NULL;
+    }
+
+    return prev;
+}
+
 /************************* FortranObject *******************************/
 
 typedef PyObject *(*fortranfunc)(PyObject *,PyObject *,PyObject *,void *);
diff --git a/numpy/f2py/src/fortranobject.h b/numpy/f2py/src/fortranobject.h
index 5c382ab7b..d4cc10243 100644
--- a/numpy/f2py/src/fortranobject.h
+++ b/numpy/f2py/src/fortranobject.h
@@ -86,6 +86,9 @@ PyObject * F2PyCapsule_FromVoidPtr(void *ptr, void (*dtor)(PyObject *));
 void * F2PyCapsule_AsVoidPtr(PyObject *obj);
 int F2PyCapsule_Check(PyObject *ptr);
 
+extern void *F2PySwapThreadLocalCallbackPtr(char *key, void *ptr);
+extern void *F2PyGetThreadLocalCallbackPtr(char *key);
+
 #define ISCONTIGUOUS(m) (PyArray_FLAGS(m) & NPY_ARRAY_C_CONTIGUOUS)
 #define F2PY_INTENT_IN 1
 #define F2PY_INTENT_INOUT 2
diff --git a/numpy/f2py/tests/test_callback.py b/numpy/f2py/tests/test_callback.py
index 4e29ab9fc..81650a819 100644
--- a/numpy/f2py/tests/test_callback.py
+++ b/numpy/f2py/tests/test_callback.py
@@ -2,6 +2,10 @@ import math
 import textwrap
 import sys
 import pytest
+import threading
+import traceback
+import time
+import random
 
 import numpy as np
 from numpy.testing import assert_, assert_equal, IS_PYPY
@@ -161,3 +165,49 @@ cf2py  intent(out) a
         f = getattr(self.module, 'string_callback_array')
         res = f(callback, cu, len(cu))
         assert_(res == 0, repr(res))
+
+    def test_threadsafety(self):
+        # Segfaults if the callback handling is not threadsafe
+
+        errors = []
+
+        def cb():
+            # Sleep here to make it more likely for another thread
+            # to call their callback at the same time.
+            time.sleep(1e-3)
+
+            # Check reentrancy
+            r = self.module.t(lambda: 123)
+            assert_(r == 123)
+
+            return 42
+
+        def runner(name):
+            try:
+                for j in range(50):
+                    r = self.module.t(cb)
+                    assert_(r == 42)
+                    self.check_function(name)
+            except Exception:
+                errors.append(traceback.format_exc())
+
+        threads = [threading.Thread(target=runner, args=(arg,))
+                   for arg in ("t", "t2") for n in range(20)]
+
+        for t in threads:
+            t.start()
+
+        for t in threads:
+            t.join()
+
+        errors = "\n\n".join(errors)
+        if errors:
+            raise AssertionError(errors)
+
+
+class TestF77CallbackPythonTLS(TestF77Callback):
+    """
+    Callback tests using Python thread-local storage instead of
+    compiler-provided
+    """
+    options = ["-DF2PY_USE_PYTHON_TLS"]