gh-124872: Mark the thread's default context as entered

Doc/c-api/contextvars.rst

@@ -126,7 +126,8 @@ Context object management functions:
    - ``Py_CONTEXT_SWITCHED``: The :term:`current context` has switched to a
      different context.  The object passed to the watch callback is the
      now-current :class:`contextvars.Context` object, or None if no context is
-     current.
+     current.  The thread executing the callback is guaranteed to be the thread
+     that experienced the context switch.
 
    .. versionadded:: 3.14
 

Include/cpython/context.h

@@ -31,7 +31,9 @@ typedef enum {
     /*
      * The current context has switched to a different context.  The object
      * passed to the watch callback is the now-current contextvars.Context
-     * object, or None if no context is current.
+     * object, or None if no context is current.  The thread executing the
+     * callback is guaranteed to be the thread that experienced the context
+     * switch.
      */
     Py_CONTEXT_SWITCHED = 1,
 } PyContextEvent;

Include/internal/pycore_context.h

@@ -15,6 +15,12 @@ extern PyTypeObject _PyContextTokenMissing_Type;
 
 PyStatus _PyContext_Init(PyInterpreterState *);
 
+// Exits any thread-owned contexts (see context_get) at the top of the thread's
+// context stack.  Logs a warning via PyErr_FormatUnraisable if the thread's
+// context stack is non-empty afterwards (those contexts can never be exited or
+// re-entered).
+void _PyContext_ExitThreadOwned(PyThreadState *);
+
 
 /* other API */
 
@@ -27,7 +33,11 @@ struct _pycontextobject {
     PyContext *ctx_prev;
     PyHamtObject *ctx_vars;
     PyObject *ctx_weakreflist;
-    int ctx_entered;
+    _Bool ctx_entered:1;
+    // True for the thread's default context created by context_get.  Used to
+    // safely determine whether the base context can be exited when clearing a
+    // PyThreadState.
+    _Bool ctx_owned_by_thread:1;
 };
 
 

Lib/test/test_capi/test_watchers.py

@@ -1,10 +1,11 @@
+import threading
 import unittest
 import contextvars
 
 from contextlib import contextmanager, ExitStack
 from test.support import (
     catch_unraisable_exception, import_helper,
-    gc_collect, suppress_immortalization)
+    gc_collect, suppress_immortalization, threading_helper)
 
 
 # Skip this test if the _testcapi module isn't available.
@@ -659,5 +660,75 @@ def test_exit_base_context(self):
             ctx.run(lambda: None)
         self.assertEqual(switches, [ctx, None])
 
+    def test_reenter_default_context(self):
+        _testcapi.clear_context_stack()
+        # contextvars.copy_context() creates the thread's default context (via
+        # the context_get C function).
+        ctx = contextvars.copy_context()
+        with self.context_watcher(0) as switches:
+            ctx.run(lambda: None)
+        self.assertEqual(len(switches), 2)
+        self.assertEqual(switches[0], ctx)
+        base_ctx = switches[1]
+        self.assertIsNotNone(base_ctx)
+        self.assertIsNot(base_ctx, ctx)
+        with self.assertRaisesRegex(RuntimeError, 'already entered'):
+            base_ctx.run(lambda: None)
+
+    def test_default_context_enter(self):
+        _testcapi.clear_context_stack()
+        with self.context_watcher(0) as switches:
+            ctx = contextvars.copy_context()
+            ctx.run(lambda: None)
+        self.assertEqual(len(switches), 3)
+        base_ctx = switches[0]
+        self.assertIsNotNone(base_ctx)
+        self.assertEqual(switches, [base_ctx, ctx, base_ctx])
+
+    @threading_helper.requires_working_threading()
+    def test_default_context_exit_during_thread_cleanup(self):
+        # Context watchers are per-interpreter, not per-thread.
+        with self.context_watcher(0) as switches:
+            def _thread_main():
+                _testcapi.clear_context_stack()
+                # contextvars.copy_context() creates the thread's default
+                # context (via the context_get C function).
+                contextvars.copy_context()
+                # This test only cares about the final switch that happens when
+                # exiting the thread's default context during thread cleanup.
+                switches.clear()
+
+            thread = threading.Thread(target=_thread_main)
+            thread.start()
+            threading_helper.join_thread(thread)
+        self.assertEqual(switches, [None])
+
+    @threading_helper.requires_working_threading()
+    def test_thread_cleanup_with_entered_context(self):
+        unraisables = []
+        try:
+            with catch_unraisable_exception() as cm:
+                with self.context_watcher(0) as switches:
+                    def _thread_main():
+                        _testcapi.clear_context_stack()
+                        ctx = contextvars.copy_context()
+                        _testcapi.context_enter(ctx)
+                        switches.clear()
+
+                    thread = threading.Thread(target=_thread_main)
+                    thread.start()
+                    threading_helper.join_thread(thread)
+                unraisables.append(cm.unraisable)
+            self.assertEqual(switches, [])
+            self.assertEqual(len(unraisables), 1)
+            self.assertIsNotNone(unraisables[0])
+            self.assertRegex(unraisables[0].err_msg,
+                             r'^Exception ignored during reset of thread state')
+            self.assertRegex(str(unraisables[0].exc_value), r'still entered')
+        finally:
+            # Break reference cycle
+            unraisables = None
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testcapi/watchers.c

@@ -724,6 +724,15 @@ clear_context_stack(PyObject *Py_UNUSED(self), PyObject *Py_UNUSED(args))
     Py_RETURN_NONE;
 }
 
+static PyObject *
+context_enter(PyObject *self, PyObject *ctx)
+{
+    if (PyContext_Enter(ctx)) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
 static PyObject *
 get_context_switches(PyObject *Py_UNUSED(self), PyObject *watcher_id)
 {
@@ -841,6 +850,7 @@ static PyMethodDef test_methods[] = {
     {"add_context_watcher",         add_context_watcher,        METH_O,       NULL},
     {"clear_context_watcher",       clear_context_watcher,      METH_O,       NULL},
     {"clear_context_stack",      clear_context_stack,     METH_NOARGS,  NULL},
+    {"context_enter",            context_enter,           METH_O,       NULL},
     {"get_context_switches",     get_context_switches,    METH_O,       NULL},
     {"allocate_too_many_context_watchers",
      (PyCFunction) allocate_too_many_context_watchers,       METH_NOARGS,  NULL},

Python/context.c

@@ -113,6 +113,10 @@ context_event_name(PyContextEvent event) {
 static void
 notify_context_watchers(PyThreadState *ts, PyContextEvent event, PyObject *ctx)
 {
+    // The callbacks are registered on the interpreter, not on the thread, so
+    // the only way callbacks can know which thread changed is by calling the
+    // callbacks from the affected thread.
+    assert(ts == _PyThreadState_GET());
     if (ctx == NULL) {
         // This will happen after exiting the last context in the stack, which
         // can occur if context_get was never called before entering a context
@@ -184,8 +188,9 @@ static inline void
 context_switched(PyThreadState *ts)
 {
     ts->context_ver++;
-    // ts->context is used instead of context_get() because context_get() might
-    // throw if ts->context is NULL.
+    // ts->context is used instead of context_get() because if ts->context is
+    // NULL, context_get() will either call context_switched -- causing a
+    // double notification -- or throw.
     notify_context_watchers(ts, Py_CONTEXT_SWITCHED, ts->context);
 }
 
@@ -244,6 +249,7 @@ _PyContext_Exit(PyThreadState *ts, PyObject *octx)
 
     ctx->ctx_prev = NULL;
     ctx->ctx_entered = 0;
+    ctx->ctx_owned_by_thread = 0;
     context_switched(ts);
     return 0;
 }
@@ -257,6 +263,37 @@ PyContext_Exit(PyObject *octx)
 }
 
 
+void
+_PyContext_ExitThreadOwned(PyThreadState *ts)
+{
+    assert(ts != NULL);
+    // notify_context_watchers requires the notification to come from the
+    // affected thread, so we can only exit the context(s) if ts belongs to the
+    // current thread.
+    _Bool on_thread = ts == _PyThreadState_GET();
+    while (ts->context != NULL
+           && PyContext_CheckExact(ts->context)
+           && ((PyContext *)ts->context)->ctx_owned_by_thread
+           && on_thread) {
+        if (_PyContext_Exit(ts, ts->context)) {
+            Py_UNREACHABLE();
+        }
+    }
+    if (ts->context != NULL) {
+        // This intentionally does not use tstate variants of these functions
+        // (e.g., _PyErr_GetRaisedException(ts)) because ts might not belong to
+        // the current thread.
+        PyObject *exc = PyErr_GetRaisedException();
+        PyErr_SetString(PyExc_RuntimeError,
+                        "contextvars.Context object(s) still entered during "
+                        "thread state reset");
+        PyErr_FormatUnraisable(
+            "Exception ignored during reset of thread state %p", ts);
+        PyErr_SetRaisedException(exc);
+    }
+}
+
+
 PyObject *
 PyContextVar_New(const char *name, PyObject *def)
 {
@@ -433,6 +470,7 @@ _context_alloc(void)
     ctx->ctx_vars = NULL;
     ctx->ctx_prev = NULL;
     ctx->ctx_entered = 0;
+    ctx->ctx_owned_by_thread = 0;
     ctx->ctx_weakreflist = NULL;
 
     return ctx;
@@ -478,15 +516,21 @@ context_get(void)
 {
     PyThreadState *ts = _PyThreadState_GET();
     assert(ts != NULL);
-    PyContext *current_ctx = (PyContext *)ts->context;
-    if (current_ctx == NULL) {
-        current_ctx = context_new_empty();
-        if (current_ctx == NULL) {
-            return NULL;
+    if (ts->context == NULL) {
+        PyContext *ctx = context_new_empty();
+        if (ctx != NULL) {
+            if (_PyContext_Enter(ts, (PyObject *)ctx)) {
+                Py_UNREACHABLE();
+            }
+            ctx->ctx_owned_by_thread = 1;
         }
-        ts->context = (PyObject *)current_ctx;
+        assert(ts->context == (PyObject *)ctx);
+        Py_CLEAR(ctx);  // _PyContext_Enter created its own ref.
     }
-    return current_ctx;
+    // The current context may be NULL if the above context_new_empty() call
+    // failed.
+    assert(ts->context == NULL || PyContext_CheckExact(ts->context));
+    return (PyContext *)ts->context;
 }
 
 static int

Python/pystate.c

@@ -1650,6 +1650,10 @@ PyThreadState_Clear(PyThreadState *tstate)
           "PyThreadState_Clear: warning: thread still has a frame\n");
     }
 
+    // This calls callbacks registered with PyContext_AddWatcher and can call
+    // sys.unraisablehook.
+    _PyContext_ExitThreadOwned(tstate);
+
     /* At this point tstate shouldn't be used any more,
        neither to run Python code nor for other uses.