gh-111956: Add thread-safe one-time initialization. (gh-111960)

Author: colesbury

Include/internal/pycore_ast_state.h

@@ -46,6 +46,7 @@ typedef struct _PyMutex PyMutex;
 #define _Py_UNLOCKED    0
 #define _Py_LOCKED      1
 #define _Py_HAS_PARKED  2
+#define _Py_ONCE_INITIALIZED 4
 
 // (private) slow path for locking the mutex
 PyAPI_FUNC(void) _PyMutex_LockSlow(PyMutex *m);
@@ -166,6 +167,35 @@ _PyRawMutex_Unlock(_PyRawMutex *m)
     _PyRawMutex_UnlockSlow(m);
 }
 
+// A data structure that can be used to run initialization code once in a
+// thread-safe manner. The C++11 equivalent is std::call_once.
+typedef struct {
+    uint8_t v;
+} _PyOnceFlag;
+
+// Type signature for one-time initialization functions. The function should
+// return 0 on success and -1 on failure.
+typedef int _Py_once_fn_t(void *arg);
+
+// (private) slow path for one time initialization
+PyAPI_FUNC(int)
+_PyOnceFlag_CallOnceSlow(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg);
+
+// Calls `fn` once using `flag`. The `arg` is passed to the call to `fn`.
+//
+// Returns 0 on success and -1 on failure.
+//
+// If `fn` returns 0 (success), then subsequent calls immediately return 0.
+// If `fn` returns -1 (failure), then subsequent calls will retry the call.
+static inline int
+_PyOnceFlag_CallOnce(_PyOnceFlag *flag, _Py_once_fn_t *fn, void *arg)
+{
+    if (_Py_atomic_load_uint8(&flag->v) == _Py_ONCE_INITIALIZED) {
+        return 0;
+    }
+    return _PyOnceFlag_CallOnceSlow(flag, fn, arg);
+}
+
 #ifdef __cplusplus
 }
 #endif

Include/internal/pycore_lock.h

@@ -1,5 +1,8 @@
 #ifndef Py_INTERNAL_MODSUPPORT_H
 #define Py_INTERNAL_MODSUPPORT_H
+
+#include "pycore_lock.h"    // _PyOnceFlag
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -65,15 +68,16 @@ PyAPI_FUNC(void) _PyArg_BadArgument(
 // --- _PyArg_Parser API ---------------------------------------------------
 
 typedef struct _PyArg_Parser {
-    int initialized;
     const char *format;
     const char * const *keywords;
     const char *fname;
     const char *custom_msg;
-    int pos;            /* number of positional-only arguments */
-    int min;            /* minimal number of arguments */
-    int max;            /* maximal number of positional arguments */
-    PyObject *kwtuple;  /* tuple of keyword parameter names */
+    _PyOnceFlag once;       /* atomic one-time initialization flag */
+    int is_kwtuple_owned;   /* does this parser own the kwtuple object? */
+    int pos;                /* number of positional-only arguments */
+    int min;                /* minimal number of arguments */
+    int max;                /* maximal number of positional arguments */
+    PyObject *kwtuple;      /* tuple of keyword parameter names */
     struct _PyArg_Parser *next;
 } _PyArg_Parser;
 

Include/internal/pycore_modsupport.h

@@ -27,7 +27,6 @@ extern "C" {
 #include "pycore_unicodeobject.h"   // struct _Py_unicode_runtime_state
 
 struct _getargs_runtime_state {
-    PyThread_type_lock mutex;
     struct _PyArg_Parser *static_parsers;
 };
 

Include/internal/pycore_runtime.h

@@ -0,0 +1,2 @@
+Add internal-only one-time initialization API: ``_PyOnceFlag`` and
+``_PyOnceFlag_CallOnce``.

Misc/NEWS.d/next/C API/2023-11-10-10-24-28.gh-issue-111956.ImE6Cx.rst

@@ -341,13 +341,45 @@ test_lock_benchmark(PyObject *module, PyObject *obj)
     Py_RETURN_NONE;
 }
 
+static int
+init_maybe_fail(void *arg)
+{
+    int *counter = (int *)arg;
+    (*counter)++;
+    if (*counter < 5) {
+        // failure
+        return -1;
+    }
+    assert(*counter == 5);
+    return 0;
+}
+
+static PyObject *
+test_lock_once(PyObject *self, PyObject *obj)
+{
+    _PyOnceFlag once = {0};
+    int counter = 0;
+    for (int i = 0; i < 10; i++) {
+        int res = _PyOnceFlag_CallOnce(&once, init_maybe_fail, &counter);
+        if (i < 4) {
+            assert(res == -1);
+        }
+        else {
+            assert(res == 0);
+            assert(counter == 5);
+        }
+    }
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef test_methods[] = {
     {"test_lock_basic", test_lock_basic, METH_NOARGS},
     {"test_lock_two_threads", test_lock_two_threads, METH_NOARGS},
     {"test_lock_counter", test_lock_counter, METH_NOARGS},
     {"test_lock_counter_slow", test_lock_counter_slow, METH_NOARGS},
     _TESTINTERNALCAPI_BENCHMARK_LOCKS_METHODDEF
     {"test_lock_benchmark", test_lock_benchmark, METH_NOARGS},
+    {"test_lock_once", test_lock_once, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };