GH-113710: Fix updating of dict version tag and add watched dict stats (GH-115221)

6 files changed, 22 insertions, 23 deletions

diff --git a/Include/cpython/pystats.h b/Include/cpython/pystats.h
index bf0cfe4..0f50439 100644
--- a/Include/cpython/pystats.h
+++ b/Include/cpython/pystats.h
@@ -133,6 +133,9 @@ typedef struct _rare_event_stats {
     uint64_t builtin_dict;
     /* Modifying a function, e.g. func.__defaults__ = ..., etc. */
     uint64_t func_modification;
+    /* Modifying a dict that is being watched */
+    uint64_t watched_dict_modification;
+    uint64_t watched_globals_modification;
 } RareEventStats;
 
 typedef struct _stats {
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
index 233da05..0ebe701 100644
--- a/Include/internal/pycore_dict.h
+++ b/Include/internal/pycore_dict.h
@@ -209,6 +209,7 @@ static inline PyDictUnicodeEntry* DK_UNICODE_ENTRIES(PyDictKeysObject *dk) {
 
 #define DICT_VERSION_INCREMENT (1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS))
 #define DICT_WATCHER_MASK ((1 << DICT_MAX_WATCHERS) - 1)
+#define DICT_WATCHER_AND_MODIFICATION_MASK ((1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS)) - 1)
 
 #ifdef Py_GIL_DISABLED
 #define DICT_NEXT_VERSION(INTERP) \
@@ -236,10 +237,10 @@ _PyDict_NotifyEvent(PyInterpreterState *interp,
     assert(Py_REFCNT((PyObject*)mp) > 0);
     int watcher_bits = mp->ma_version_tag & DICT_WATCHER_MASK;
     if (watcher_bits) {
+        RARE_EVENT_STAT_INC(watched_dict_modification);
         _PyDict_SendEvent(watcher_bits, event, mp, key, value);
-        return DICT_NEXT_VERSION(interp) | watcher_bits;
     }
-    return DICT_NEXT_VERSION(interp);
+    return DICT_NEXT_VERSION(interp) | (mp->ma_version_tag & DICT_WATCHER_AND_MODIFICATION_MASK);
 }
 
 extern PyObject *_PyObject_MakeDictFromInstanceAttributes(PyObject *obj, PyDictValues *values);
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 2cfbf4b..b14e695 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -28,25 +28,23 @@ increment_mutations(PyObject* dict) {
     d->ma_version_tag += (1 << DICT_MAX_WATCHERS);
 }
 
+/* The first two dict watcher IDs are reserved for CPython,
+ * so we don't need to check that they haven't been used */
+#define BUILTINS_WATCHER_ID 0
+#define GLOBALS_WATCHER_ID  1
+
 static int
 globals_watcher_callback(PyDict_WatchEvent event, PyObject* dict,
                          PyObject* key, PyObject* new_value)
 {
-    if (event == PyDict_EVENT_CLONED) {
-        return 0;
-    }
-    uint64_t watched_mutations = get_mutations(dict);
-    if (watched_mutations < _Py_MAX_ALLOWED_GLOBALS_MODIFICATIONS) {
-        _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), dict);
-        increment_mutations(dict);
-    }
-    else {
-        PyDict_Unwatch(1, dict);
-    }
+    RARE_EVENT_STAT_INC(watched_globals_modification);
+    assert(get_mutations(dict) < _Py_MAX_ALLOWED_GLOBALS_MODIFICATIONS);
+    _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), dict);
+    increment_mutations(dict);
+    PyDict_Unwatch(GLOBALS_WATCHER_ID, dict);
     return 0;
 }
 
-
 static void
 global_to_const(_PyUOpInstruction *inst, PyObject *obj)
 {
@@ -82,11 +80,6 @@ incorrect_keys(_PyUOpInstruction *inst, PyObject *obj)
     return 0;
 }
 
-/* The first two dict watcher IDs are reserved for CPython,
- * so we don't need to check that they haven't been used */
-#define BUILTINS_WATCHER_ID 0
-#define GLOBALS_WATCHER_ID  1
-
 /* Returns 1 if successfully optimized
  *         0 if the trace is not suitable for optimization (yet)
  *        -1 if there was an error. */
@@ -117,8 +110,8 @@ remove_globals(_PyInterpreterFrame *frame, _PyUOpInstruction *buffer,
     uint32_t builtins_watched = 0;
     uint32_t globals_checked = 0;
     uint32_t globals_watched = 0;
-    if (interp->dict_state.watchers[1] == NULL) {
-        interp->dict_state.watchers[1] = globals_watcher_callback;
+    if (interp->dict_state.watchers[GLOBALS_WATCHER_ID] == NULL) {
+        interp->dict_state.watchers[GLOBALS_WATCHER_ID] = globals_watcher_callback;
     }
     for (int pc = 0; pc < buffer_size; pc++) {
         _PyUOpInstruction *inst = &buffer[pc];
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index 61c9d4f..2300180 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -611,7 +611,7 @@ static int
 builtins_dict_watcher(PyDict_WatchEvent event, PyObject *dict, PyObject *key, PyObject *new_value)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
-    if (event != PyDict_EVENT_CLONED && interp->rare_events.builtin_dict < _Py_MAX_ALLOWED_BUILTINS_MODIFICATIONS) {
+    if (interp->rare_events.builtin_dict < _Py_MAX_ALLOWED_BUILTINS_MODIFICATIONS) {
         _Py_Executors_InvalidateAll(interp);
     }
     RARE_EVENT_INTERP_INC(interp, builtin_dict);
diff --git a/Python/specialize.c b/Python/specialize.c
index e38e355..ea26385 100644
--- a/Python/specialize.c
+++ b/Python/specialize.c
@@ -275,6 +275,8 @@ print_rare_event_stats(FILE *out, RareEventStats *stats)
     fprintf(out, "Rare event (set_eval_frame_func): %" PRIu64 "\n", stats->set_eval_frame_func);
     fprintf(out, "Rare event (builtin_dict): %" PRIu64 "\n", stats->builtin_dict);
     fprintf(out, "Rare event (func_modification): %" PRIu64 "\n", stats->func_modification);
+    fprintf(out, "Rare event (watched_dict_modification): %" PRIu64 "\n", stats->watched_dict_modification);
+    fprintf(out, "Rare event (watched_globals_modification): %" PRIu64 "\n", stats->watched_globals_modification);
 }
 
 static void
diff --git a/Tools/scripts/summarize_stats.py b/Tools/scripts/summarize_stats.py
index 9b7e7b9..7891b9c 100644
--- a/Tools/scripts/summarize_stats.py
+++ b/Tools/scripts/summarize_stats.py
@@ -415,7 +415,7 @@ class Stats:
     def get_rare_events(self) -> list[tuple[str, int]]:
         prefix = "Rare event "
         return [
-            (key[len(prefix) + 1:-1], val)
+            (key[len(prefix) + 1:-1].replace("_", " "), val)
             for key, val in self._data.items()
             if key.startswith(prefix)
         ]