From 82a735e36776d93e53f95b0ae93e6122cf6350a7 Mon Sep 17 00:00:00 2001
From: Suke0811 <49264928+Suke0811@users.noreply.github.com>
Date: Mon, 24 Nov 2025 13:51:29 -0800
Subject: [PATCH] Await async cancellation cleanup

---
 fspin/rate_control.py     | 54 ++++++++++++++++++++++++++++++---------
 fspin/spin_context.py     |  2 +-
 tests/test_ratecontrol.py | 22 ++++++++++++++--
 3 files changed, 63 insertions(+), 15 deletions(-)

diff --git a/fspin/rate_control.py b/fspin/rate_control.py
index fa7dfc1..a819ea8 100644
--- a/fspin/rate_control.py
+++ b/fspin/rate_control.py
@@ -456,23 +456,53 @@ def start_spinning(self, func, condition_fn, *args, **kwargs):
                 raise TypeError("Expected a regular function for sync mode.")
             return self.start_spinning_sync(func, condition_fn, *args, **kwargs)
 
-    def stop_spinning(self):
+    def _close_own_loop(self):
+        if self._own_loop is not None:
+            self._own_loop.close()
+            self._own_loop = None
+
+    def _stop_thread(self):
+        if self._thread:
+            # Avoid deadlock if stop_spinning is called from within the worker thread
+            current = threading.current_thread()
+            if self._thread.is_alive() and current is not self._thread:
+                self._thread.join()
+
+    async def _cancel_task(self):
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+
+    async def stop_spinning_async(self):
         """
-        Signals the spinning loop to stop.
+        Signals the spinning loop to stop and awaits async cleanup.
         """
         self._stop_event.set()
         if self.is_coroutine:
-            if self._task:
-                self._task.cancel()
+            await self._cancel_task()
         else:
-            if self._thread:
-                # Avoid deadlock if stop_spinning is called from within the worker thread
-                current = threading.current_thread()
-                if self._thread.is_alive() and current is not self._thread:
-                    self._thread.join()
-        if self._own_loop is not None:
-            self._own_loop.close()
-            self._own_loop = None
+            self._stop_thread()
+        self._close_own_loop()
+
+    def stop_spinning(self):
+        """
+        Signals the spinning loop to stop.
+        """
+        if self.is_coroutine:
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                asyncio.run(self.stop_spinning_async())
+            else:
+                loop.create_task(self.stop_spinning_async())
+            return
+
+        self._stop_event.set()
+        self._stop_thread()
+        self._close_own_loop()
 
     def get_report(self, output=True):
         """
diff --git a/fspin/spin_context.py b/fspin/spin_context.py
index 835578a..0b96082 100644
--- a/fspin/spin_context.py
+++ b/fspin/spin_context.py
@@ -90,4 +90,4 @@ async def __aenter__(self):
         return self.rc
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self.rc.stop_spinning()
+        await self.rc.stop_spinning_async()
diff --git a/tests/test_ratecontrol.py b/tests/test_ratecontrol.py
index b6e6698..00d2796 100644
--- a/tests/test_ratecontrol.py
+++ b/tests/test_ratecontrol.py
@@ -10,6 +10,7 @@
 from fspin.reporting import ReportLogger
 from fspin.rate_control import RateControl
 from fspin.decorators import spin
+from fspin.spin_context import spin as spin_context
 from fspin.loop_context import loop
 
 def test_create_histogram():
@@ -170,8 +171,7 @@ async def awork():
     await asyncio.sleep(0.05)
 
     # Stop the task and wait for it to be cancelled
-    rc.stop_spinning()
-    await asyncio.sleep(0.05)
+    await rc.stop_spinning_async()
 
     # Check if the task is done (it might be cancelled or completed)
     assert rc._task.done(), "Task is not done after stop_spinning"
@@ -188,6 +188,24 @@ async def awork():
             pass
 
 
+@pytest.mark.asyncio
+async def test_async_context_manager_waits_for_cleanup():
+    cleanup_called = asyncio.Event()
+
+    async def awork():
+        try:
+            while True:
+                await asyncio.sleep(0.01)
+        finally:
+            cleanup_called.set()
+
+    async with spin_context(awork, freq=100) as rc:
+        await asyncio.sleep(0.03)
+
+    await asyncio.wait_for(cleanup_called.wait(), timeout=0.5)
+    assert rc._task.done()
+
+
 @pytest.mark.asyncio
 async def test_spin_async_exception_handling(caplog):
     async def awork():