GH-95909: Make `_PyArg_Parser` initialization thread safe (GH-95958)

2 files changed, 44 insertions, 12 deletions

diff --git a/Python/getargs.c b/Python/getargs.c
index 457dd99..f0b84b8 100644
--- a/Python/getargs.c
+++ b/Python/getargs.c
@@ -1974,15 +1974,10 @@ new_kwtuple(const char * const *keywords, int total, int pos)
 }
 
 static int
-parser_init(struct _PyArg_Parser *parser)
+_parser_init(struct _PyArg_Parser *parser)
 {
     const char * const *keywords = parser->keywords;
     assert(keywords != NULL);
-
-    if (parser->initialized) {
-        assert(parser->kwtuple != NULL);
-        return 1;
-    }
     assert(parser->pos == 0 &&
            (parser->format == NULL || parser->fname == NULL) &&
            parser->custom_msg == NULL &&
@@ -2035,6 +2030,28 @@ parser_init(struct _PyArg_Parser *parser)
     return 1;
 }
 
+static int
+parser_init(struct _PyArg_Parser *parser)
+{
+    // volatile as it can be modified by other threads
+    // and should not be optimized or reordered by compiler
+    if (*((volatile int *)&parser->initialized)) {
+        assert(parser->kwtuple != NULL);
+        return 1;
+    }
+    PyThread_acquire_lock(_PyRuntime.getargs.mutex, WAIT_LOCK);
+    // Check again if another thread initialized the parser
+    // while we were waiting for the lock.
+    if (*((volatile int *)&parser->initialized)) {
+        assert(parser->kwtuple != NULL);
+        PyThread_release_lock(_PyRuntime.getargs.mutex);
+        return 1;
+    }
+    int ret = _parser_init(parser);
+    PyThread_release_lock(_PyRuntime.getargs.mutex);
+    return ret;
+}
+
 static void
 parser_clear(struct _PyArg_Parser *parser)
 {
diff --git a/Python/pystate.c b/Python/pystate.c
index bcdb825..642d680 100644
--- a/Python/pystate.c
+++ b/Python/pystate.c
@@ -57,7 +57,7 @@ _Py_COMP_DIAG_POP
 
 static int
 alloc_for_runtime(PyThread_type_lock *plock1, PyThread_type_lock *plock2,
-                  PyThread_type_lock *plock3)
+                  PyThread_type_lock *plock3, PyThread_type_lock *plock4)
 {
     /* Force default allocator, since _PyRuntimeState_Fini() must
        use the same allocator than this function. */
@@ -82,11 +82,20 @@ alloc_for_runtime(PyThread_type_lock *plock1, PyThread_type_lock *plock2,
         return -1;
     }
 
+    PyThread_type_lock lock4 = PyThread_allocate_lock();
+    if (lock4 == NULL) {
+        PyThread_free_lock(lock1);
+        PyThread_free_lock(lock2);
+        PyThread_free_lock(lock3);
+        return -1;
+    }
+
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
     *plock1 = lock1;
     *plock2 = lock2;
     *plock3 = lock3;
+    *plock4 = lock4;
     return 0;
 }
 
@@ -97,7 +106,8 @@ init_runtime(_PyRuntimeState *runtime,
              Py_ssize_t unicode_next_index,
              PyThread_type_lock unicode_ids_mutex,
              PyThread_type_lock interpreters_mutex,
-             PyThread_type_lock xidregistry_mutex)
+             PyThread_type_lock xidregistry_mutex,
+             PyThread_type_lock getargs_mutex)
 {
     if (runtime->_initialized) {
         Py_FatalError("runtime already initialized");
@@ -119,6 +129,8 @@ init_runtime(_PyRuntimeState *runtime,
 
     runtime->xidregistry.mutex = xidregistry_mutex;
 
+    runtime->getargs.mutex = getargs_mutex;
+
     // Set it to the ID of the main thread of the main interpreter.
     runtime->main_thread = PyThread_get_thread_ident();
 
@@ -141,8 +153,8 @@ _PyRuntimeState_Init(_PyRuntimeState *runtime)
     // is called multiple times.
     Py_ssize_t unicode_next_index = runtime->unicode_ids.next_index;
 
-    PyThread_type_lock lock1, lock2, lock3;
-    if (alloc_for_runtime(&lock1, &lock2, &lock3) != 0) {
+    PyThread_type_lock lock1, lock2, lock3, lock4;
+    if (alloc_for_runtime(&lock1, &lock2, &lock3, &lock4) != 0) {
         return _PyStatus_NO_MEMORY();
     }
 
@@ -152,7 +164,7 @@ _PyRuntimeState_Init(_PyRuntimeState *runtime)
         memcpy(runtime, &initial, sizeof(*runtime));
     }
     init_runtime(runtime, open_code_hook, open_code_userdata, audit_hook_head,
-                 unicode_next_index, lock1, lock2, lock3);
+                 unicode_next_index, lock1, lock2, lock3, lock4);
 
     return _PyStatus_OK();
 }
@@ -172,6 +184,7 @@ _PyRuntimeState_Fini(_PyRuntimeState *runtime)
     FREE_LOCK(runtime->interpreters.mutex);
     FREE_LOCK(runtime->xidregistry.mutex);
     FREE_LOCK(runtime->unicode_ids.lock);
+    FREE_LOCK(runtime->getargs.mutex);
 
 #undef FREE_LOCK
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
@@ -194,6 +207,7 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
     int reinit_interp = _PyThread_at_fork_reinit(&runtime->interpreters.mutex);
     int reinit_xidregistry = _PyThread_at_fork_reinit(&runtime->xidregistry.mutex);
     int reinit_unicode_ids = _PyThread_at_fork_reinit(&runtime->unicode_ids.lock);
+    int reinit_getargs = _PyThread_at_fork_reinit(&runtime->getargs.mutex);
 
     PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
 
@@ -204,7 +218,8 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
     if (reinit_interp < 0
         || reinit_main_id < 0
         || reinit_xidregistry < 0
-        || reinit_unicode_ids < 0)
+        || reinit_unicode_ids < 0
+        || reinit_getargs < 0)
     {
         return _PyStatus_ERR("Failed to reinitialize runtime locks");