Propagate the use_tracing value at switch time as well on 3.10.issue256

Because it's stack based, we need to store it separately.

2 files changed, 161 insertions, 29 deletions

diff --git a/src/greenlet/greenlet.c b/src/greenlet/greenlet.c
index d743046..f47bbf8 100644
--- a/src/greenlet/greenlet.c
+++ b/src/greenlet/greenlet.c
@@ -136,6 +136,11 @@ static PyGreenlet* volatile ts_current = NULL;
 static PyObject* volatile ts_passaround_args = NULL;
 static PyObject* volatile ts_passaround_kwargs = NULL;
 
+/* Used internally in ``g_switchstack()`` */
+#if GREENLET_USE_CFRAME
+static int volatile ts__g_switchstack_use_tracing = 0;
+#endif
+
 /***********************************************************/
 /* Thread-aware routines, switching global variables when needed */
 
@@ -485,26 +490,37 @@ static int GREENLET_NOINLINE(slp_save_state)(char* stackref)
     return 0;
 }
 
+/**
+   Perform a stack switch according to some global variables
+   that must be set before calling this function. Those variables
+   are:
+
+   - ts_current: current greenlet (holds a reference)
+   - ts_target: greenlet to switch to (weak reference)
+   - ts_passaround_args: NULL if PyErr_Occurred(),
+     else a tuple of args sent to ts_target (holds a reference)
+   - ts_passaround_kwargs: switch kwargs (holds a reference)
+
+   Because the stack switch happens in this function, this function can't use
+   its own stack (local) variables, set before the switch, and then accessed after the
+   switch. Global variables beginning with ``ts__g_switchstack`` are used
+   internally instead.
+
+   On return results are passed via global variables as well:
+
+   - ts_origin: originating greenlet (holds a reference)
+   - ts_current: current greenlet (holds a reference)
+   - ts_passaround_args: NULL if PyErr_Occurred(),
+     else a tuple of args sent to ts_current (holds a reference)
+   - ts_passaround_kwargs: switch kwargs (holds a reference)
+
+   It is very important that stack switch is 'atomic', i.e. no
+   calls into other Python code allowed (except very few that
+   are safe), because global variables are very fragile.
+*/
 static int
 g_switchstack(void)
 {
-    /* Perform a stack switch according to some global variables
-       that must be set before:
-       - ts_current: current greenlet (holds a reference)
-       - ts_target: greenlet to switch to (weak reference)
-       - ts_passaround_args: NULL if PyErr_Occurred(),
-           else a tuple of args sent to ts_target (holds a reference)
-       - ts_passaround_kwargs: switch kwargs (holds a reference)
-       On return results are passed via global variables as well:
-       - ts_origin: originating greenlet (holds a reference)
-       - ts_current: current greenlet (holds a reference)
-       - ts_passaround_args: NULL if PyErr_Occurred(),
-           else a tuple of args sent to ts_current (holds a reference)
-       - ts_passaround_kwargs: switch kwargs (holds a reference)
-       It is very important that stack switch is 'atomic', i.e. no
-       calls into other Python code allowed (except very few that
-       are safe), because global variables are very fragile.
-    */
     int err;
     { /* save state */
         PyGreenlet* current = ts_current;
@@ -523,10 +539,23 @@ g_switchstack(void)
         current->exc_traceback = tstate->exc_traceback;
 #endif
 #if GREENLET_USE_CFRAME
+        /*
+         IMPORTANT: ``cframe`` is a pointer into the STACK.
+         Thus, because the call to ``slp_switch()``
+         changes the contents of the stack, you cannot read from
+         ``ts_current->cframe`` after that call and necessarily
+         get the same values you get from reading it here. Anything
+         you need to restore from now to then must be saved
+         in a global variable (because we can't use stack variables
+         here either).
+         */
         current->cframe = tstate->cframe;
+        ts__g_switchstack_use_tracing = tstate->cframe->use_tracing;
 #endif
     }
+
     err = slp_switch();
+
     if (err < 0) { /* error */
         PyGreenlet* current = ts_current;
         current->top_frame = NULL;
@@ -570,6 +599,13 @@ g_switchstack(void)
 
 #if GREENLET_USE_CFRAME
         tstate->cframe = target->cframe;
+        /*
+          If we were tracing, we need to keep tracing.
+          There should never be the possibility of hitting the
+          root_cframe here. See note above about why we can't
+          just copy this from ``origin->cframe->use_tracing``.
+        */
+        tstate->cframe->use_tracing = ts__g_switchstack_use_tracing;
 #endif
 
         assert(ts_origin == NULL);
diff --git a/src/greenlet/tests/test_tracing.py b/src/greenlet/tests/test_tracing.py
index ce059ce..2ab4d71 100644
--- a/src/greenlet/tests/test_tracing.py
+++ b/src/greenlet/tests/test_tracing.py
@@ -111,7 +111,7 @@ class TestPythonTracing(unittest.TestCase):
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
     def _trace_switch(self, glet):
@@ -130,7 +130,7 @@ class TestPythonTracing(unittest.TestCase):
             ('return', 'run'),
             ('c_return', '_trace_switch'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
     def test_trace_events_into_greenlet_func_already_set(self):
@@ -145,27 +145,123 @@ class TestPythonTracing(unittest.TestCase):
                 return tpt_callback()
         self._check_trace_events_func_already_set(X())
 
-    def test_trace_events_from_greenlet_sets_func(self):
+    def _check_trace_events_from_greenlet_sets_profiler(self, g, tracer):
+        g.switch()
+        tpt_callback()
+        tracer.__exit__()
+        self.assertEqual(tracer.actions, [
+            ('return', '__enter__'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('return', 'run'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('call', '__exit__'),
+            ('c_call', '__exit__'),
+        ])
+
+
+    def test_trace_events_from_greenlet_func_sets_profiler(self):
         tracer = PythonTracer()
         def run():
             tracer.__enter__()
             return tpt_callback()
 
-        g = greenlet.greenlet(run)
-        g.switch()
-        tpt_callback()
-        tracer.__exit__()
+        self._check_trace_events_from_greenlet_sets_profiler(greenlet.greenlet(run),
+                                                             tracer)
+
+    def test_trace_events_from_greenlet_subclass_sets_profiler(self):
+        tracer = PythonTracer()
+        class X(greenlet.greenlet):
+            def run(self):
+                tracer.__enter__()
+                return tpt_callback()
+
+        self._check_trace_events_from_greenlet_sets_profiler(X(), tracer)
+
+
+    def test_trace_events_multiple_greenlets_switching(self):
+        tracer = PythonTracer()
+
+        g1 = None
+        g2 = None
+
+        def g1_run():
+            tracer.__enter__()
+            tpt_callback()
+            g2.switch()
+            tpt_callback()
+            return 42
+
+        def g2_run():
+            tpt_callback()
+            tracer.__exit__()
+            tpt_callback()
+            g1.switch()
+
+        g1 = greenlet.greenlet(g1_run)
+        g2 = greenlet.greenlet(g2_run)
+
+        x = g1.switch()
+        self.assertEqual(x, 42)
+        tpt_callback() # ensure not in the trace
         self.assertEqual(tracer.actions, [
             ('return', '__enter__'),
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
-            ('return', 'run'),
+            ('c_call', 'g1_run'),
+            ('call', 'g2_run'),
             ('call', 'tpt_callback'),
             ('return', 'tpt_callback'),
             ('call', '__exit__'),
-            ('c_call', '__exit__')
+            ('c_call', '__exit__'),
         ])
 
-    # What about multiple greenlets switching amongst each other? Probably doesn't
-    # work either.
-    # What about with throw()? And throw() as the initial action?
+    def test_trace_events_multiple_greenlets_switching_siblings(self):
+        # Like the first version, but get both greenlets running first
+        # as "siblings" and then establish the tracing.
+        tracer = PythonTracer()
+
+        g1 = None
+        g2 = None
+
+        def g1_run():
+            greenlet.getcurrent().parent.switch()
+            tracer.__enter__()
+            tpt_callback()
+            g2.switch()
+            tpt_callback()
+            return 42
+
+        def g2_run():
+            greenlet.getcurrent().parent.switch()
+
+            tpt_callback()
+            tracer.__exit__()
+            tpt_callback()
+            g1.switch()
+
+        g1 = greenlet.greenlet(g1_run)
+        g2 = greenlet.greenlet(g2_run)
+
+        # Start g1
+        g1.switch()
+        # And it immediately returns control to us.
+        # Start g2
+        g2.switch()
+        # Which also returns. Now kick of the real part of the
+        # test.
+        x = g1.switch()
+        self.assertEqual(x, 42)
+
+        tpt_callback() # ensure not in the trace
+        self.assertEqual(tracer.actions, [
+            ('return', '__enter__'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('c_call', 'g1_run'),
+            ('call', 'tpt_callback'),
+            ('return', 'tpt_callback'),
+            ('call', '__exit__'),
+            ('c_call', '__exit__'),
+        ])