diff --git a/news/700.bugfix.rst b/news/700.bugfix.rst new file mode 100644 index 0000000000..169dbe912a --- /dev/null +++ b/news/700.bugfix.rst @@ -0,0 +1 @@ +Fix a crash in greenlet when a greenlet switch happens after Memray's tracking function is deactivated. diff --git a/requirements-test.txt b/requirements-test.txt index 1030ae6375..cb80c11b40 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -1,6 +1,6 @@ Cython coverage[toml] -greenlet; python_version < '3.13' +greenlet; python_version < '3.14' pytest pytest-cov ipython diff --git a/src/memray/_memray/tracking_api.cpp b/src/memray/_memray/tracking_api.cpp index a6fa414e3c..a8b37a1425 100644 --- a/src/memray/_memray/tracking_api.cpp +++ b/src/memray/_memray/tracking_api.cpp @@ -1219,6 +1219,16 @@ Tracker::handleGreenletSwitch(PyObject* from, PyObject* to) // Grab the Tracker lock, as this may need to write pushes/pops. std::unique_lock lock(*s_mutex); RecursionGuard guard; + + // Check if the trace function is still installed in the current thread as + // it's possible that the thread is dying and the trace function has been + // uninstalled and therefore the stack will not be properly tracked so we + // should not recreate the current thread stack. + PyThreadState* ts = PyThreadState_Get(); + if (ts->c_profilefunc != PyTraceFunction) { + return; + } + PythonStackTracker::get().handleGreenletSwitch(from, to); } diff --git a/tests/integration/test_greenlet.py b/tests/integration/test_greenlet.py index 4090aae279..4c899b4718 100644 --- a/tests/integration/test_greenlet.py +++ b/tests/integration/test_greenlet.py @@ -10,7 +10,7 @@ from tests.utils import filter_relevant_allocations pytestmark = pytest.mark.skipif( - sys.version_info >= (3, 12), reason="Greenlet does not yet support Python 3.12" + sys.version_info >= (3, 14), reason="Greenlet does not yet support Python 3.14" ) @@ -194,3 +194,73 @@ def stack(alloc): assert vallocs[0].tid != vallocs[1].tid != vallocs[6].tid assert vallocs[0].tid == vallocs[2].tid assert vallocs[1].tid == vallocs[3].tid == vallocs[4].tid == vallocs[5].tid + + +def test_uninstall_profile_in_greenlet(tmpdir): + """Verify that memray handles profile function changes in greenlets correctly.""" + # GIVEN + output = Path(tmpdir) / "test.bin" + subprocess_code = textwrap.dedent( + f""" + import greenlet + import sys + + from memray import Tracker + from memray._test import MemoryAllocator + + def foo(): + bar() + allocator.valloc(1024 * 10) + + def bar(): + baz() + + def baz(): + sys.setprofile(None) + other.switch() + + def test(): + allocator.valloc(1024 * 70) + main_greenlet.switch() + + + allocator = MemoryAllocator() + output = "{output}" + + with Tracker(output): + main_greenlet = greenlet.getcurrent() + other = greenlet.greenlet(test) + foo() + + """ + ) + + # WHEN + subprocess.run([sys.executable, "-Xdev", "-c", subprocess_code], timeout=5) + + # THEN + reader = FileReader(output) + records = list(reader.get_allocation_records()) + vallocs = [ + record + for record in filter_relevant_allocations(records) + if record.allocator == AllocatorType.VALLOC + ] + + def stack(alloc): + return [frame[0] for frame in alloc.stack_trace()] + + # Verify allocations and their stack traces (which should be empty + # because we remove the tracking function) + assert len(vallocs) == 2 + + assert stack(vallocs[0]) == [] + assert vallocs[0].size == 70 * 1024 + + assert stack(vallocs[1]) == [] + assert vallocs[1].size == 10 * 1024 + + # Verify thread IDs + main_tid = vallocs[0].tid # inner greenlet + outer_tid = vallocs[1].tid # outer greenlet + assert main_tid == outer_tid