gh-122581: Avoid data races when collecting parser statistics (#122694)

python · Aug 6, 2024 · ce0d66c · ce0d66c
1 parent a8be8fc
commit ce0d66c
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 0 deletions.
diff --git a/Include/internal/pycore_parser.h b/Include/internal/pycore_parser.h
@@ -21,15 +21,20 @@ extern "C" {
 struct _parser_runtime_state {
 #ifdef Py_DEBUG
     long memo_statistics[_PYPEGEN_NSTATISTICS];
+#ifdef Py_GIL_DISABLED
+    PyMutex mutex;
+#endif
 #else
     int _not_used;
 #endif
     struct _expr dummy_name;
 };
 
 _Py_DECLARE_STR(empty, "")
+#if defined(Py_DEBUG) && defined(Py_GIL_DISABLED)
 #define _parser_runtime_state_INIT \
     { \
+        .mutex = {0}, \
         .dummy_name = { \
             .kind = Name_kind, \
             .v.Name.id = &_Py_STR(empty), \
@@ -40,6 +45,20 @@ _Py_DECLARE_STR(empty, "")
             .end_col_offset = 0, \
         }, \
     }
+#else
+#define _parser_runtime_state_INIT \
+    { \
+        .dummy_name = { \
+            .kind = Name_kind, \
+            .v.Name.id = &_Py_STR(empty), \
+            .v.Name.ctx = Load, \
+            .lineno = 1, \
+            .col_offset = 0, \
+            .end_lineno = 1, \
+            .end_col_offset = 0, \
+        }, \
+    }
+#endif
 
 extern struct _mod* _PyParser_ASTFromString(
     const char *str,

diff --git a/Parser/pegen.c b/Parser/pegen.c
@@ -296,12 +296,22 @@ _PyPegen_fill_token(Parser *p)
 #define NSTATISTICS _PYPEGEN_NSTATISTICS
 #define memo_statistics _PyRuntime.parser.memo_statistics
 
+#ifdef Py_GIL_DISABLED
+#define MUTEX_LOCK() PyMutex_Lock(&_PyRuntime.parser.mutex)
+#define MUTEX_UNLOCK() PyMutex_Unlock(&_PyRuntime.parser.mutex)
+#else
+#define MUTEX_LOCK()
+#define MUTEX_UNLOCK()
+#endif
+
 void
 _PyPegen_clear_memo_statistics(void)
 {
+    MUTEX_LOCK();
     for (int i = 0; i < NSTATISTICS; i++) {
         memo_statistics[i] = 0;
     }
+    MUTEX_UNLOCK();
 }
 
 PyObject *
@@ -311,18 +321,23 @@ _PyPegen_get_memo_statistics(void)
     if (ret == NULL) {
         return NULL;
     }
+
+    MUTEX_LOCK();
     for (int i = 0; i < NSTATISTICS; i++) {
         PyObject *value = PyLong_FromLong(memo_statistics[i]);
         if (value == NULL) {
+            MUTEX_UNLOCK();
             Py_DECREF(ret);
             return NULL;
         }
         // PyList_SetItem borrows a reference to value.
         if (PyList_SetItem(ret, i, value) < 0) {
+            MUTEX_UNLOCK();
             Py_DECREF(ret);
             return NULL;
         }
     }
+    MUTEX_UNLOCK();
     return ret;
 }
 #endif
@@ -348,7 +363,9 @@ _PyPegen_is_memoized(Parser *p, int type, void *pres)
                 if (count <= 0) {
                     count = 1;
                 }
+                MUTEX_LOCK();
                 memo_statistics[type] += count;
+                MUTEX_UNLOCK();
             }
 #endif
             p->mark = m->mark;