diff --git a/tests/test_asyncio.py b/tests/test_asyncio.py index a48b6ca..9b7c8b5 100644 --- a/tests/test_asyncio.py +++ b/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(): diff --git a/yappi/_yappi.c b/yappi/_yappi.c index 8b867ef..3323de6 100644 --- a/yappi/_yappi.c +++ b/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