gh-124872: Replace enter/exit events with "switched" (#125532)

Users want to know when the current context switches to a different
context object.  Right now this happens when and only when a context
is entered or exited, so the enter and exit events are synonymous with
"switched".  However, if the changes proposed for gh-99633 are
implemented, the current context will also switch for reasons other
than context enter or exit.  Since users actually care about context
switches and not enter or exit, replace the enter and exit events with
a single switched event.

The former exit event was emitted just before exiting the context.
The new switched event is emitted after the context is exited to match
the semantics users expect of an event with a past-tense name.  If
users need the ability to clean up before the switch takes effect,
another event type can be added in the future.  It is not added here
because YAGNI.

I skipped 0 in the enum as a matter of practice.  Skipping 0 makes it
easier to troubleshoot when code forgets to set zeroed memory, and it
aligns with best practices for other tools (e.g.,
https://protobuf.dev/programming-guides/dos-donts/#unspecified-enum).

Co-authored-by: Richard Hansen <rhansen@rhansen.org>
Co-authored-by: Victor Stinner <vstinner@python.org>

1 files changed, 41 insertions, 38 deletions

diff --git a/Modules/_testcapi/watchers.c b/Modules/_testcapi/watchers.c
index b4233d0..321d3ae 100644
--- a/Modules/_testcapi/watchers.c
+++ b/Modules/_testcapi/watchers.c
@@ -626,16 +626,12 @@ allocate_too_many_func_watchers(PyObject *self, PyObject *args)
 // Test contexct object watchers
 #define NUM_CONTEXT_WATCHERS 2
 static int context_watcher_ids[NUM_CONTEXT_WATCHERS] = {-1, -1};
-static int num_context_object_enter_events[NUM_CONTEXT_WATCHERS] = {0, 0};
-static int num_context_object_exit_events[NUM_CONTEXT_WATCHERS] = {0, 0};
+static PyObject *context_switches[NUM_CONTEXT_WATCHERS];
 
 static int
 handle_context_watcher_event(int which_watcher, PyContextEvent event, PyObject *ctx) {
-    if (event == Py_CONTEXT_EVENT_ENTER) {
-        num_context_object_enter_events[which_watcher]++;
-    }
-    else if (event == Py_CONTEXT_EVENT_EXIT)  {
-        num_context_object_exit_events[which_watcher]++;
+    if (event == Py_CONTEXT_SWITCHED) {
+        PyList_Append(context_switches[which_watcher], ctx);
     }
     else {
         return -1;
@@ -667,31 +663,28 @@ error_context_event_handler(PyContextEvent event, PyObject *ctx) {
 static PyObject *
 add_context_watcher(PyObject *self, PyObject *which_watcher)
 {
-    int watcher_id;
+    static const PyContext_WatchCallback callbacks[] = {
+        &first_context_watcher_callback,
+        &second_context_watcher_callback,
+        &error_context_event_handler,
+    };
     assert(PyLong_Check(which_watcher));
     long which_l = PyLong_AsLong(which_watcher);
-    if (which_l == 0) {
-        watcher_id = PyContext_AddWatcher(first_context_watcher_callback);
-        context_watcher_ids[0] = watcher_id;
-        num_context_object_enter_events[0] = 0;
-        num_context_object_exit_events[0] = 0;
-    }
-    else if (which_l == 1) {
-        watcher_id = PyContext_AddWatcher(second_context_watcher_callback);
-        context_watcher_ids[1] = watcher_id;
-        num_context_object_enter_events[1] = 0;
-        num_context_object_exit_events[1] = 0;
-    }
-    else if (which_l == 2) {
-        watcher_id = PyContext_AddWatcher(error_context_event_handler);
-    }
-    else {
+    if (which_l < 0 || which_l >= (long)Py_ARRAY_LENGTH(callbacks)) {
         PyErr_Format(PyExc_ValueError, "invalid watcher %d", which_l);
         return NULL;
     }
+    int watcher_id = PyContext_AddWatcher(callbacks[which_l]);
     if (watcher_id < 0) {
         return NULL;
     }
+    if (which_l >= 0 && which_l < NUM_CONTEXT_WATCHERS) {
+        context_watcher_ids[which_l] = watcher_id;
+        Py_XSETREF(context_switches[which_l], PyList_New(0));
+        if (context_switches[which_l] == NULL) {
+            return NULL;
+        }
+    }
     return PyLong_FromLong(watcher_id);
 }
 
@@ -708,8 +701,7 @@ clear_context_watcher(PyObject *self, PyObject *watcher_id)
         for (int i = 0; i < NUM_CONTEXT_WATCHERS; i++) {
             if (watcher_id_l == context_watcher_ids[i]) {
                 context_watcher_ids[i] = -1;
-                num_context_object_enter_events[i] = 0;
-                num_context_object_exit_events[i] = 0;
+                Py_CLEAR(context_switches[i]);
             }
         }
     }
@@ -717,21 +709,34 @@ clear_context_watcher(PyObject *self, PyObject *watcher_id)
 }
 
 static PyObject *
-get_context_watcher_num_enter_events(PyObject *self, PyObject *watcher_id)
+clear_context_stack(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
 {
-    assert(PyLong_Check(watcher_id));
-    long watcher_id_l = PyLong_AsLong(watcher_id);
-    assert(watcher_id_l >= 0 && watcher_id_l < NUM_CONTEXT_WATCHERS);
-    return PyLong_FromLong(num_context_object_enter_events[watcher_id_l]);
+    PyThreadState *tstate = PyThreadState_Get();
+    if (tstate->context == NULL) {
+        Py_RETURN_NONE;
+    }
+    if (((PyContext *)tstate->context)->ctx_prev != NULL) {
+        PyErr_SetString(PyExc_RuntimeError,
+                        "must first exit all non-base contexts");
+        return NULL;
+    }
+    Py_CLEAR(tstate->context);
+    Py_RETURN_NONE;
 }
 
 static PyObject *
-get_context_watcher_num_exit_events(PyObject *self, PyObject *watcher_id)
+get_context_switches(PyObject *Py_UNUSED(self), PyObject *watcher_id)
 {
     assert(PyLong_Check(watcher_id));
     long watcher_id_l = PyLong_AsLong(watcher_id);
-    assert(watcher_id_l >= 0 && watcher_id_l < NUM_CONTEXT_WATCHERS);
-    return PyLong_FromLong(num_context_object_exit_events[watcher_id_l]);
+    if (watcher_id_l < 0 || watcher_id_l >= NUM_CONTEXT_WATCHERS) {
+        PyErr_Format(PyExc_ValueError, "invalid watcher %ld", watcher_id_l);
+        return NULL;
+    }
+    if (context_switches[watcher_id_l] == NULL) {
+        return PyList_New(0);
+    }
+    return Py_NewRef(context_switches[watcher_id_l]);
 }
 
 static PyObject *
@@ -835,10 +840,8 @@ static PyMethodDef test_methods[] = {
     // Code object watchers.
     {"add_context_watcher",         add_context_watcher,        METH_O,       NULL},
     {"clear_context_watcher",       clear_context_watcher,      METH_O,       NULL},
-    {"get_context_watcher_num_enter_events",
-     get_context_watcher_num_enter_events,                 METH_O,       NULL},
-    {"get_context_watcher_num_exit_events",
-     get_context_watcher_num_exit_events,               METH_O,       NULL},
+    {"clear_context_stack",      clear_context_stack,     METH_NOARGS,  NULL},
+    {"get_context_switches",     get_context_switches,    METH_O,       NULL},
     {"allocate_too_many_context_watchers",
      (PyCFunction) allocate_too_many_context_watchers,       METH_NOARGS,  NULL},
     {NULL},