Merge pull request #218 from sumerc/fix-217

fix: decref PyFrame_GetGenerator

Author: sumerc

tests/test_asyncio.py

@@ -69,6 +69,41 @@ async def a(n):
         stats = yappi.get_func_stats()
         self.assert_traces_almost_equal(r1, stats)
 
+    def test_mem_leak(self):
+        import gc
+        import types
+
+        async def noop():
+            return
+
+        async def run_batch(n):
+            await asyncio.gather(*(noop() for _ in range(n)))
+
+        def exhausted_coroutines():
+            gc.collect()
+            return sum(
+                1
+                for o in gc.get_objects()
+                if isinstance(o, types.CoroutineType) and o.cr_frame is None
+            )
+
+        batch_size = 1000
+
+        yappi.set_clock_type("cpu")
+        yappi.start()
+        asyncio.get_event_loop().run_until_complete(run_batch(batch_size))
+        yappi.stop()
+        yappi.clear_stats()
+        baseline = exhausted_coroutines()
+
+        yappi.start()
+        asyncio.get_event_loop().run_until_complete(run_batch(batch_size))
+        yappi.stop()
+        yappi.clear_stats()
+        after = exhausted_coroutines()
+
+        self.assertLess(after - baseline, batch_size // 3)
+
     def test_basic_old_style(self):
 
         async def a():

yappi/_yappi.c

@@ -225,17 +225,20 @@ int
 IS_SUSPENDED(PyFrameObject *frame) {
 #if PY_VERSION_HEX >= 0x030B0000  // Python 3.11+
     PyGenObject *gen = (PyGenObject *)PyFrame_GetGenerator(frame);
+    int suspended;
     if (gen == NULL) {
         return 0;
     }
 #if PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION == 14
     unsigned char curr_op_code = (*frame->f_frame->instr_ptr).op.code;
-    return curr_op_code == YIELD_VALUE || curr_op_code == INSTRUMENTED_YIELD_VALUE;
+    suspended = curr_op_code == YIELD_VALUE || curr_op_code == INSTRUMENTED_YIELD_VALUE;
 #elif PY_MAJOR_VERSION >= 3 && PY_MINOR_VERSION == 13
-    return FRAME_STATE_SUSPENDED(gen->gi_frame_state);
+    suspended = FRAME_STATE_SUSPENDED(gen->gi_frame_state);
 #else
-    return gen->gi_frame_state == FRAME_SUSPENDED;
+    suspended = gen->gi_frame_state == FRAME_SUSPENDED;
 #endif
+    Py_DECREF(gen);
+    return suspended;
 #elif PY_VERSION_HEX >= 0x030A0000 // Python 3.10+
     return (frame->f_state == FRAME_SUSPENDED);
 #else