Use a signal-safe CM in the synchronous test case

- without this, the exception gets raised as soon as the CM's
  __exit__ method starts running
- for the async case, it turns out there's nothing we can do
  to solve this at the byte code level. Instead, we need to
  somehow push *all* pending call processing back to the
  event loop and inject a synthetic await into the current
  frame

Author: ncoghlan

Lib/test/test_with_signal_safety.py

@@ -6,6 +6,7 @@
 import asyncio
 import dis
 import sys
+import threading
 import unittest
 
 class InjectedException(Exception):
@@ -20,8 +21,9 @@ def raise_after_offset(target_function, target_offset):
     """
     target_code = target_function.__code__
     def inject_exception():
-        print("Raising injected exception")
-        raise InjectedException(f"Failing after {target_offset}")
+        exc = InjectedException(f"Failing after {target_offset}")
+        print(f"Raising injected exception: {exc}")
+        raise exc
     # This installs a trace hook that's implemented in C, and hence won't
     # trigger any of the per-bytecode processing in the eval loop
     # This means it can register the pending call that raises the exception and
@@ -51,24 +53,22 @@ def setUp(self):
         self.addCleanup(sys.settrace, old_trace)
         sys.settrace(None)
 
-    def assert_cm_exited(self, tracking_cm, target_offset, traced_operation):
-        if tracking_cm.enter_without_exit:
+    def assert_lock_released(self, test_lock, target_offset, traced_operation):
+        just_acquired = test_lock.acquire(blocking=False)
+        # Either we just acquired the lock, or the test didn't release it
+        test_lock.release()
+        if not just_acquired:
             msg = ("Context manager entered without exit due to "
                   f"exception injected at offset {target_offset} in:\n"
                   f"{dis.Bytecode(traced_operation).dis()}")
             self.fail(msg)
 
     def test_synchronous_cm(self):
-        class TrackingCM():
-            def __init__(self):
-                self.enter_without_exit = None
-            def __enter__(self):
-                self.enter_without_exit = True
-            def __exit__(self, *args):
-                self.enter_without_exit = False
-        tracking_cm = TrackingCM()
+        # Must use a signal-safe CM, otherwise __exit__ will start
+        # but then fail to actually run as the pending call gets processed
+        test_lock = threading.Lock()
         def traced_function():
-            with tracking_cm:
+            with test_lock:
                 1 + 1
             return
         target_offset = -1
@@ -80,12 +80,20 @@ def traced_function():
                 traced_function()
             except InjectedException:
                 # key invariant: if we entered the CM, we exited it
-                self.assert_cm_exited(tracking_cm, target_offset, traced_function)
+                self.assert_lock_released(test_lock, target_offset, traced_function)
             else:
                 self.fail(f"Exception wasn't raised @{target_offset}")
 
 
-    def test_asynchronous_cm(self):
+    def _test_asynchronous_cm(self):
+        # NOTE: this can't work, since asyncio is written in Python, and hence
+        # will always process pending calls at some point during the evaluation
+        # of __aenter__ and __aexit__
+        #
+        # So to handle that case, we need to some way to tell the event loop
+        # to convert pending call processing into calls to
+        # asyncio.get_event_loop().call_soon() instead of processing them
+        # immediately
         class AsyncTrackingCM():
             def __init__(self):
                 self.enter_without_exit = None
@@ -108,7 +116,7 @@ async def traced_coroutine():
                 loop.run_until_complete(traced_coroutine())
             except InjectedException:
                 # key invariant: if we entered the CM, we exited it
-                self.assert_cm_exited(tracking_cm, target_offset, traced_coroutine)
+                self.assertFalse(tracking_cm.enter_without_exit)
             else:
                 self.fail(f"Exception wasn't raised @{target_offset}")
 

Modules/_testcapimodule.c

@@ -2368,16 +2368,15 @@ error_injection_trace(PyObject *hook_args, PyFrameObject *frame,
     if (((PyObject *) frame->f_code) == target_code) {
         printf("Tracing frame of interest\n");
         frame->f_trace_opcodes = 1;
-        if (what == PyTrace_OPCODE && frame->f_lasti > target_offset) {
-            printf("Adding pending call after %d\n", frame->f_lasti);
+        if (what == PyTrace_OPCODE && frame->f_lasti >= target_offset) {
             Py_INCREF(callback);
+            PyEval_SetTrace(NULL, NULL);
+            printf("Adding pending call after %d\n", frame->f_lasti);
             if (Py_AddPendingCall(&_pending_callback, callback) < 0) {
                 printf("Failed to add pending call\n");
                 Py_DECREF(callback);
-                PyEval_SetTrace(NULL, NULL);
                 return -1;
             }
-            PyEval_SetTrace(NULL, NULL);
         }
     }
     return 0;

Python/ceval.c

@@ -419,11 +419,13 @@ Py_MakePendingCalls(void)
         PyThread_acquire_lock(pending_lock, WAIT_LOCK);
         j = pendingfirst;
         if (j == pendinglast) {
+            printf("  No pending calls remaining\n");
             func = NULL; /* Queue empty */
         } else {
             func = pendingcalls[j].func;
             arg = pendingcalls[j].arg;
             pendingfirst = (j + 1) % NPENDINGCALLS;
+            printf("  Calling %p(%p)\n", func, arg);
         }
         PyThread_release_lock(pending_lock);
         /* having released the lock, perform the callback */
@@ -982,10 +984,10 @@ _PyEval_EvalFrameDefault(PyFrameObject *f, int throwflag)
                 goto fast_next_opcode;
             }
             if (handle_pending_after > first_instr) {
-                printf("Pending calls deferred until %lu\n", DEFER_OFFSET());
+                printf("Pending calls deferred until %lu (current: %lu)\n", DEFER_OFFSET(), INSTR_OFFSET());
             }
             if (next_instr >= handle_pending_after) {
-                printf("Checking for pending calls: %lu > %lu?\n", INSTR_OFFSET(), DEFER_OFFSET());
+                printf("Checking for pending calls: %lu >= %lu?\n", INSTR_OFFSET(), DEFER_OFFSET());
                 /* Allow for subsequent jumps backwards in the bytecode */
                 handle_pending_after = first_instr;
                 if (_Py_atomic_load_relaxed(&pendingcalls_to_do)) {