From dbcab733ed98e9ce2fd9bf7f5e179a967e3ce71f Mon Sep 17 00:00:00 2001
From: AlexWells <alex.wells@diamond.ac.uk>
Date: Mon, 19 May 2025 08:59:43 +0100
Subject: [PATCH 1/2] Clean up the created threadstate

This avoids a memory leak, as reported in #70.
---
 context/_coroutine.c | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/context/_coroutine.c b/context/_coroutine.c
index 3a0c3e77..10b350cc 100644
--- a/context/_coroutine.c
+++ b/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.

From e83c1155efb48a48051086f4811473c7310a5463 Mon Sep 17 00:00:00 2001
From: AlexWells <alex.wells@diamond.ac.uk>
Date: Mon, 19 May 2025 11:19:51 +0100
Subject: [PATCH 2/2] Add test for memory leak

---
 pyproject.toml         |  1 +
 tests/test_cothread.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/pyproject.toml b/pyproject.toml
index 79ccdf8f..07044b85 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,6 +39,7 @@ dev = [
     "sphinx-design",
     "tox-direct",
     "types-mock",
+    "psutil",
 ]
 
 [project.scripts]
diff --git a/tests/test_cothread.py b/tests/test_cothread.py
index 47e54689..41211e6b 100644
--- a/tests/test_cothread.py
+++ b/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)