Merge pull request #71 from DiamondLightSource/py3.11_fixes

Clean up the created threadstate

Author: Araneidae

context/_coroutine.c

@@ -107,7 +107,11 @@ static void *coroutine_wrapper(void *action_, void *arg_)
     Py_DECREF(arg);
 
 
-#if PY_VERSION_HEX < 0x30B0000
+#if PY_VERSION_HEX >= 0x30B0000
+    new_threadstate = PyThreadState_Swap(thread_state);
+    PyThreadState_Clear(new_threadstate);
+    PyThreadState_Delete(new_threadstate);
+#else
     /* Some of the stuff we've initialised can leak through, so far I've only
      * seen exc_type still set at this point, but maybe other fields can also
      * leak.  Avoid a memory leak by making sure we're not holding onto these.

pyproject.toml

@@ -39,6 +39,7 @@ dev = [
     "sphinx-design",
     "tox-direct",
     "types-mock",
+    "psutil",
 ]
 
 [project.scripts]

tests/test_cothread.py

@@ -3,6 +3,8 @@
 import multiprocessing as mp
 import time
 
+import psutil
+
 # Add cothread onto file and import
 import sys
 import os
@@ -169,6 +171,33 @@ def target():
     after = time.time()
     assert float(after - before) > TIMEOUT
 
+def test_threadstate_memory_leak():
+    """Test that the memory leak reported in issue #70 is fixed and does not
+    reoccur"""
+
+    process = psutil.Process()
+
+    vms_start = process.memory_info().vms
+
+    def test():
+        pass
+
+    # Arbitrary large number of spawns. On my machine with the memory leak
+    # active,this results in a leak of approximately 2GiB.
+    for i in range(500000):
+        cothread.Spawn(test)
+        cothread.Yield()
+
+    vms_end = process.memory_info().vms
+
+    print(f"VMS start {vms_start} end {vms_end} diff {vms_end-vms_start}")
+
+    memory_increase = vms_end - vms_start
+
+    # With the memory leak fixed, the memory increase on my machine is only
+    # 16384 bytes. Added an order of magnitude for future safety.
+    assert memory_increase < 200000
+
 
 if __name__ == '__main__':
     unittest.main(verbosity=2)