Make tracing state exception safe.object_ownership

2 files changed, 77 insertions, 32 deletions

diff --git a/src/greenlet/greenlet.cpp b/src/greenlet/greenlet.cpp
index 2e5ccb4..36636fd 100644
--- a/src/greenlet/greenlet.cpp
+++ b/src/greenlet/greenlet.cpp
@@ -1547,6 +1547,50 @@ static void GREENLET_NOINLINE(slp_restore_state_trampoline)()
 
 /***********************************************************/
 
+class TracingGuard
+{
+private:
+    PyThreadState* tstate;
+public:
+    TracingGuard()
+        : tstate(PyThreadState_GET())
+    {
+        PyThreadState_EnterTracing(this->tstate);
+    }
+
+    ~TracingGuard()
+    {
+        PyThreadState_LeaveTracing(this->tstate);
+        this->tstate = nullptr;
+    }
+
+    inline void CallTraceFunction(const OwnedObject& tracefunc,
+                                  const ImmortalEventName& event,
+                                  const BorrowedGreenlet& origin,
+                                  const BorrowedGreenlet& target)
+    {
+        // TODO: This calls tracefunc(event, (origin, target)). Add a shortcut
+        // function for that that's specialized to avoid the Py_BuildValue
+        // string parsing, or start with just using "ON" format with PyTuple_Pack(2,
+        // origin, target). That seems like what the N format is meant
+        // for.
+        // XXX: Why does event not automatically cast back to a PyObject?
+        // It tries to call the "deleted constructor ImmortalEventName
+        // const" instead.
+        assert(tracefunc);
+        assert(event);
+        assert(origin);
+        assert(target);
+        NewReference retval(PyObject_CallFunction(tracefunc.borrow(),
+                                             "O(OO)",
+                                             event.borrow(),
+                                             origin.borrow(),
+                                             target.borrow()));
+        if (!retval) {
+            throw PyErrOccurred();
+        }
+    }
+};
 
 static void
 g_calltrace(const OwnedObject& tracefunc,
@@ -1554,43 +1598,17 @@ g_calltrace(const OwnedObject& tracefunc,
             const BorrowedGreenlet& origin,
             const BorrowedGreenlet& target)
 {
-    PyThreadState* tstate;
     PyErrPieces saved_exc;
-    tstate = PyThreadState_GET();
-    tstate->tracing++;
-    TSTATE_USE_TRACING(tstate) = 0;
-    // TODO: This calls tracefunc(event, (origin, target)). Add a shortcut
-    // function for that that's specialized to avoid the Py_BuildValue
-    // string parsing, or start with just using "ON" format with PyTuple_Pack(2,
-    // origin, target). That seems like what the N format is meant
-    // for.
-    // XXX: Why does event not automatically cast back to a PyObject?
-    // It tries to call the "deleted constructor ImmortalEventName
-    // const" instead.
-    assert(event);
-    assert(origin);
-    assert(target);
-    NewReference retval(PyObject_CallFunction(tracefunc.borrow(),
-                                             "O(OO)",
-                                             event.borrow(),
-                                             origin.borrow(),
-                                             target.borrow()));
-    tstate->tracing--;
-    TSTATE_USE_TRACING(tstate) =
-        (tstate->tracing <= 0 &&
-         ((tstate->c_tracefunc != NULL) || (tstate->c_profilefunc != NULL)));
-
-    if (!retval) {
+    try {
+        TracingGuard tracing_guard;
+        tracing_guard.CallTraceFunction(tracefunc, event, origin, target);
+    }
+    catch (const PyErrOccurred&) {
         // In case of exceptions trace function is removed,
         // and any existing exception is replaced with the tracing
         // exception.
-        assert(PyErr_Occurred());
-        // In principle, the DECREF of the tracefunc that happens here
-        // could run arbitrary Python code and maybe even switch
-        // greenlets. Our caller is holding another reference to the
-        // tracefunc, though, so the decref won't happen until they unwind.
         GET_THREAD_STATE().state().set_tracefunc(Py_None);
-        throw PyErrOccurred();
+        throw;
     }
 
     saved_exc.PyErrRestore();
diff --git a/src/greenlet/greenlet_cpython_compat.hpp b/src/greenlet/greenlet_cpython_compat.hpp
index 1b04cd8..740955c 100644
--- a/src/greenlet/greenlet_cpython_compat.hpp
+++ b/src/greenlet/greenlet_cpython_compat.hpp
@@ -98,4 +98,31 @@ PyObject* PyModule_Create(PyModuleDef* m)
 }
 #endif
 
+// bpo-43760 added PyThreadState_EnterTracing() to Python 3.11.0a2
+// bpo-43760 added PyThreadState_LeaveTracing() to Python 3.11.0a2
+#if PY_VERSION_HEX < 0x030B00A2
+static inline void PyThreadState_EnterTracing(PyThreadState *tstate)
+{
+    tstate->tracing++;
+#if GREENLET_USE_CFRAME
+    tstate->cframe->use_tracing = 0;
+#else
+    tstate->use_tracing = 0;
+#endif
+}
+
+
+static inline void PyThreadState_LeaveTracing(PyThreadState *tstate)
+{
+    tstate->tracing--;
+    int use_tracing = (tstate->c_tracefunc != NULL
+                       || tstate->c_profilefunc != NULL);
+#if GREENLET_USE_CFRAME
+    tstate->cframe->use_tracing = use_tracing;
+#else
+    tstate->use_tracing = use_tracing;
+#endif
+}
+#endif
+
 #endif /* GREENLET_CPYTHON_COMPAT_H */