gh-119333: Back up exception before calling PyContext_WatchCallback

rhansen · rhansen · commit 9e4ea98b2f45 · 2024-09-29T19:15:00.000-04:00
I believe that the value of a simpler API (and defense against poorly
written callbacks) outweighs the cost of backing up and restoring the
thread's exception state.
diff --git a/Doc/c-api/contextvars.rst b/Doc/c-api/contextvars.rst
@@ -139,11 +139,8 @@ Context object management functions:
    exception will be printed as an unraisable exception using
    :c:func:`PyErr_FormatUnraisable`. Otherwise it should return ``0``.
 
-   There may already be a pending exception set on entry to the callback. In
-   this case, the callback should return ``0`` with the same exception still
-   set. This means the callback may not call any other API that can set an
-   exception unless it saves and clears the exception state first, and restores
-   it before returning.
+   Any pending exception is cleared before the callback is called and restored
+   after the callback returns.
 
    .. versionadded:: 3.14
 
diff --git a/Include/cpython/context.h b/Include/cpython/context.h
@@ -40,6 +40,9 @@ typedef enum {
  *
  * if the callback returns with an exception set, it must return -1. Otherwise
  * it should return 0
+ *
+ * Any pending exception is cleared before the callback is called and restored
+ * after the callback returns.
  */
 typedef int (*PyContext_WatchCallback)(PyContextEvent, PyContext *);
 
diff --git a/Lib/test/test_capi/test_watchers.py b/Lib/test/test_capi/test_watchers.py
@@ -640,6 +640,16 @@ def _in_context(stack):
                 ctx.run(_in_context, stack)
             self.assertEqual(str(cm.unraisable.exc_value), "boom!")
 
+    def test_exception_save(self):
+        with self.context_watcher(2):
+            with catch_unraisable_exception() as cm:
+                def _in_context():
+                    raise RuntimeError("test")
+
+                with self.assertRaisesRegex(RuntimeError, "test"):
+                    contextvars.copy_context().run(_in_context)
+                self.assertEqual(str(cm.unraisable.exc_value), "boom!")
+
     def test_clear_out_of_range_watcher_id(self):
         with self.assertRaisesRegex(ValueError, r"Invalid context watcher ID -1"):
             _testcapi.clear_context_watcher(-1)
diff --git a/Python/context.c b/Python/context.c
@@ -119,6 +119,7 @@ static void notify_context_watchers(PyContextEvent event, PyContext *ctx, PyThre
     assert(interp->_initialized);
     uint8_t bits = interp->active_context_watchers;
     int i = 0;
+    PyObject *exc = _PyErr_GetRaisedException(ts);
     while (bits) {
         assert(i < CONTEXT_MAX_WATCHERS);
         if (bits & 1) {
@@ -133,6 +134,7 @@ static void notify_context_watchers(PyContextEvent event, PyContext *ctx, PyThre
         i++;
         bits >>= 1;
     }
+    _PyErr_SetRaisedException(ts, exc);
 }
 
 

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,9 @@ typedef enum {`
`40`	`40`	`*`
`41`	`41`	`* if the callback returns with an exception set, it must return -1. Otherwise`
`42`	`42`	`* it should return 0`
	`43`	`+ *`
	`44`	`+ * Any pending exception is cleared before the callback is called and restored`
	`45`	`+ * after the callback returns.`
`43`	`46`	`*/`
`44`	`47`	`typedef int (PyContext_WatchCallback)(PyContextEvent, PyContext );`
`45`	`48`
Original file line number	Diff line number	Diff line change
`@@ -119,6 +119,7 @@ static void notify_context_watchers(PyContextEvent event, PyContext *ctx, PyThre`
`119`	`119`	`assert(interp->_initialized);`
`120`	`120`	`uint8_t bits = interp->active_context_watchers;`
`121`	`121`	`int i = 0;`
	`122`	`+ PyObject *exc = _PyErr_GetRaisedException(ts);`
`122`	`123`	`while (bits) {`
`123`	`124`	`assert(i < CONTEXT_MAX_WATCHERS);`
`124`	`125`	`if (bits & 1) {`
`@@ -133,6 +134,7 @@ static void notify_context_watchers(PyContextEvent event, PyContext *ctx, PyThre`
`133`	`134`	`i++;`
`134`	`135`	`bits >>= 1;`
`135`	`136`	`}`
	`137`	`+ _PyErr_SetRaisedException(ts, exc);`
`136`	`138`	`}`
`137`	`139`
`138`	`140`