Merge pull request #2758 from Textualize/win-sleep-bug

Fix win sleep bug

src/textual/_time.py

@@ -1,5 +1,5 @@
+import asyncio
 import platform
-from asyncio import get_running_loop
 from asyncio import sleep as asyncio_sleep
 from time import monotonic, perf_counter
 
@@ -26,7 +26,7 @@ async def sleep(secs: float) -> None:
         Args:
             secs: Number of seconds to sleep for.
         """
-        await get_running_loop().run_in_executor(None, win_sleep, secs)
+        await asyncio.create_task(win_sleep(secs))
 
 else:
 

src/textual/_win_sleep.py

@@ -4,7 +4,11 @@
 This should only be imported on Windows.
 """
 
+from __future__ import annotations
+
+import asyncio
 from time import sleep as time_sleep
+from typing import Coroutine
 
 __all__ = ["sleep"]
 
@@ -14,21 +18,40 @@
 CREATE_WAITABLE_TIMER_HIGH_RESOLUTION = 0x00000002
 TIMER_ALL_ACCESS = 0x1F0003
 
+
+async def time_sleep_coro(secs: float):
+    """Coroutine wrapper around `time.sleep`."""
+    await asyncio.sleep(secs)
+
+
 try:
     import ctypes
-    from ctypes.wintypes import LARGE_INTEGER
+    from ctypes.wintypes import HANDLE, LARGE_INTEGER
 
     kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]
 except Exception:
-    sleep = time_sleep
+    print("Exception found")
+
+    def sleep(secs: float) -> Coroutine[None, None, None]:
+        """Wrapper around `time.sleep` to match the signature of the main case below."""
+        return time_sleep_coro(secs)
+
 else:
 
-    def sleep(secs: float) -> None:
+    async def no_sleep_coro():
+        """Creates a coroutine that does nothing for when no sleep is needed."""
+        pass
+
+    def sleep(secs: float) -> Coroutine[None, None, None]:
         """A replacement sleep for Windows.
 
         Note that unlike `time.sleep` this *may* sleep for slightly less than the
         specified time. This is generally not an issue for Textual's use case.
 
+        In order to create a timer that _can_ be cancelled on Windows, we need to
+        create a timer and a separate event, and then we wait for either of the two
+        things. When Textual wants to quit, we set the cancel event.
+
         Args:
             secs: Seconds to sleep for.
         """
@@ -37,32 +60,63 @@ def sleep(secs: float) -> None:
         sleep_for = max(0, secs - 0.001)
         if sleep_for < 0.0005:
             # Less than 0.5ms and its not worth doing the sleep
-            return
+            return no_sleep_coro()
 
-        handle = kernel32.CreateWaitableTimerExW(
+        timer = kernel32.CreateWaitableTimerExW(
             None,
             None,
             CREATE_WAITABLE_TIMER_HIGH_RESOLUTION,
             TIMER_ALL_ACCESS,
         )
-        if not handle:
-            time_sleep(sleep_for)
-            return
-
-        try:
-            if not kernel32.SetWaitableTimer(
-                handle,
-                ctypes.byref(LARGE_INTEGER(int(sleep_for * -10_000_000))),
-                0,
-                None,
-                None,
-                0,
+        if not timer:
+            return time_sleep_coro(sleep_for)
+
+        if not kernel32.SetWaitableTimer(
+            timer,
+            ctypes.byref(LARGE_INTEGER(int(sleep_for * -10_000_000))),
+            0,
+            None,
+            None,
+            0,
+        ):
+            kernel32.CloseHandle(timer)
+            return time_sleep_coro(sleep_for)
+
+        cancel_event = kernel32.CreateEventExW(None, None, 0, TIMER_ALL_ACCESS)
+        if not cancel_event:
+            kernel32.CloseHandle(timer)
+            return time_sleep_coro(sleep_for)
+
+        def cancel_inner():
+            """Sets the cancel event so we know we can stop waiting for the timer."""
+            kernel32.SetEvent(cancel_event)
+
+        async def cancel():
+            """Cancels the timer by setting the cancel event."""
+            await asyncio.get_running_loop().run_in_executor(None, cancel_inner)
+
+        def wait_inner():
+            """Function responsible for waiting for the timer or the cancel event."""
+            if (
+                kernel32.WaitForMultipleObjects(
+                    2,
+                    ctypes.pointer((HANDLE * 2)(cancel_event, timer)),
+                    False,
+                    INFINITE,
+                )
+                == WAIT_FAILED
             ):
-                kernel32.CloseHandle(handle)
                 time_sleep(sleep_for)
-                return
 
-            if kernel32.WaitForSingleObject(handle, INFINITE) == WAIT_FAILED:
-                time_sleep(sleep_for)
-        finally:
-            kernel32.CloseHandle(handle)
+        async def wait():
+            """Wraps the actual sleeping so we can detect if the thread was cancelled."""
+            try:
+                await asyncio.get_running_loop().run_in_executor(None, wait_inner)
+            except asyncio.CancelledError:
+                await cancel()
+                raise
+            finally:
+                kernel32.CloseHandle(timer)
+                kernel32.CloseHandle(cancel_event)
+
+        return wait()

tests/test_win_sleep.py

@@ -0,0 +1,41 @@
+import asyncio
+import time
+import sys
+
+import pytest
+
+from textual.app import App
+
+pytestmark = pytest.mark.skipif(
+    sys.platform != "win32", reason="We only need to test this on Windows."
+)
+
+
+def test_win_sleep_timer_is_cancellable():
+    """Regression test for https://github.com/Textualize/textual/issues/2711.
+
+    When we exit an app with a "long" timer, everything asyncio-related
+    should shutdown quickly. So, we create an app with a timer that triggers
+    every SLEEP_FOR seconds and we shut the app down immediately after creating
+    it. `asyncio` should be done quickly (i.e., the timer was cancelled) and
+    thus the total time this takes should be considerably lesser than the time
+    we originally set the timer for.
+    """
+
+    SLEEP_FOR = 10
+
+    class WindowsIntervalBugApp(App[None]):
+        def on_mount(self) -> None:
+            self.set_interval(SLEEP_FOR, lambda: None)
+
+        def key_e(self):
+            self.exit()
+
+    async def actual_test():
+        async with WindowsIntervalBugApp().run_test() as pilot:
+            await pilot.press("e")
+
+    start = time.perf_counter()
+    asyncio.run(actual_test())
+    end = time.perf_counter()
+    assert end - start < 1