fix(sync): survive an exception raised by a signal handler during a sync call

Fixes: #3222

Author: pavelfeldman

playwright/_impl/_disposable.py

@@ -15,8 +15,6 @@
 import traceback
 from typing import Any, Awaitable, Callable, Coroutine, Dict
 
-import greenlet
-
 from playwright._impl._connection import ChannelOwner, _capture_stack_trace
 from playwright._impl._errors import Error, is_target_closed_error
 from playwright._impl._greenlets import connection_closed_error, wait_for_future
@@ -78,15 +76,13 @@ def _sync(self, coro: Coroutine[Any, Any, Any]) -> object:
         if self._dispatcher_fiber.dead:
             coro.close()
             raise connection_closed_error()
-        g_self = greenlet.getcurrent()
         task = self._loop.create_task(coro)
         setattr(
             task,
             "__pw_stack__",
             _capture_stack_trace(),
         )
         setattr(task, "__pw_stack_trace__", traceback.extract_stack(limit=10))
-        task.add_done_callback(lambda _: g_self.switch())
         wait_for_future(self._loop, self._dispatcher_fiber, task)
         return task.result()
 

playwright/_impl/_greenlets.py

@@ -62,14 +62,33 @@ def wait_for_future(
     future: "asyncio.Future",
 ) -> None:
     __tracebackhide__ = True
-    while not future.done():
-        # The dispatcher fiber exits once the connection to the driver ends, e.g.
-        # when the driver process dies. Nothing can settle the future after that,
-        # and switching to a dead greenlet returns right away, so we would spin.
-        if dispatcher_fiber.dead:
-            future.cancel()
-            raise connection_closed_error()
-        dispatcher_fiber.switch()
-    # The loop is only running for as long as the dispatcher fiber is alive.
-    if not dispatcher_fiber.dead:
-        asyncio._set_running_loop(loop)
+    g_self = greenlet.getcurrent()
+
+    def resume(_: "asyncio.Future") -> None:
+        g_self.switch()
+
+    # The wait owns the callback that resumes this greenlet, so that it can
+    # remove it again when the wait is abandoned. A callback left behind would
+    # resume this greenlet at an arbitrary later point.
+    future.add_done_callback(resume)
+    try:
+        while not future.done():
+            # The dispatcher fiber exits once the connection to the driver ends,
+            # e.g. when the driver process dies. Nothing can settle the future
+            # after that, and switching to a dead greenlet returns right away,
+            # so we would spin.
+            if dispatcher_fiber.dead:
+                raise connection_closed_error()
+            # Raises when a signal handler interrupted the event loop while it
+            # waited for the driver, see greenlet_main in _context_manager.py.
+            dispatcher_fiber.switch()
+    except BaseException:
+        future.remove_done_callback(resume)
+        future.cancel()
+        raise
+    finally:
+        # The loop is only running for as long as the dispatcher fiber is alive.
+        # Also after a signal handler interrupted it: run_forever() cleared this
+        # on its way out, but the dispatcher is going to run the loop again.
+        if not dispatcher_fiber.dead:
+            asyncio._set_running_loop(loop)

playwright/_impl/_sync_base.py

@@ -29,8 +29,6 @@
     cast,
 )
 
-import greenlet
-
 from playwright._impl._connection import _capture_stack_trace
 from playwright._impl._greenlets import connection_closed_error, wait_for_future
 from playwright._impl._helper import Error
@@ -47,8 +45,6 @@ class EventInfo(Generic[T]):
     def __init__(self, sync_base: "SyncBase", future: "asyncio.Future[T]") -> None:
         self._sync_base = sync_base
         self._future = future
-        g_self = greenlet.getcurrent()
-        self._future.add_done_callback(lambda _: g_self.switch())
 
     @property
     def value(self) -> T:
@@ -107,12 +103,10 @@ def _sync(
             coro.close()
             raise connection_closed_error()
 
-        g_self = greenlet.getcurrent()
         task: asyncio.tasks.Task[Any] = self._loop.create_task(coro)
         setattr(task, "__pw_stack__", _capture_stack_trace())
         setattr(task, "__pw_stack_trace__", traceback.extract_stack(limit=10))
 
-        task.add_done_callback(lambda _: g_self.switch())
         wait_for_future(self._loop, self._dispatcher_fiber, task)
         return task.result()
 

playwright/sync_api/_context_manager.py

@@ -15,7 +15,7 @@
 import asyncio
 from typing import TYPE_CHECKING, Any, Optional
 
-from greenlet import greenlet
+from greenlet import GreenletExit, greenlet
 
 from playwright._impl._connection import Connection
 from playwright._impl._errors import Error
@@ -48,11 +48,33 @@ def __enter__(self) -> SyncPlaywright:
 Please use the Async API instead."""
             )
 
+        g_self = greenlet.getcurrent()
+
         # Create a new fiber for the protocol dispatcher. It will be pumping events
         # until the end of times. We will pass control to that fiber every time we
         # block while waiting for a response.
         def greenlet_main() -> None:
-            self._loop.run_until_complete(self._connection.run_as_sync())
+            task = self._loop.create_task(self._connection.run_as_sync())
+            while not task.done():
+                try:
+                    self._loop.run_until_complete(task)
+                except GreenletExit:
+                    raise
+                except BaseException as exc:
+                    if task.done():
+                        raise
+                    # asyncio contains exceptions raised in tasks and callbacks.
+                    # An exception that escapes the loop while the connection
+                    # still runs comes from a signal handler, e.g. pytest-timeout
+                    # or Ctrl+C, that fired while the loop waited for the driver.
+                    # The user code is suspended in a sync call in the greenlet
+                    # that created us: raise it there, and stay alive. The loop
+                    # and its tasks are intact, so the next sync call resumes
+                    # them.
+                    g_self.throw(exc)
+                    # Between sync calls the user code marks the loop as running,
+                    # see wait_for_future. run_forever() refuses to start then.
+                    asyncio._set_running_loop(None)
 
         dispatcher_fiber = MainGreenlet(greenlet_main)
 
@@ -63,8 +85,6 @@ def greenlet_main() -> None:
             self._loop,
         )
 
-        g_self = greenlet.getcurrent()
-
         # Wait until initialize completes (not just Playwright __create__), matching async.
         self._connection.playwright_future.add_done_callback(lambda _: g_self.switch())
         dispatcher_fiber.switch()

tests/sync/test_sync.py

@@ -12,20 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import contextlib
 import multiprocessing
 import os
+import signal
 import subprocess
 import sys
 import textwrap
+import time
 from datetime import timedelta
 from pathlib import Path
-from typing import Any, Callable, Dict
+from typing import Any, Callable, Dict, Generator, List
 
 import pytest
 
 from playwright.sync_api import (
     Browser,
     BrowserContext,
+    ConsoleMessage,
     Dialog,
     Error,
     Page,
@@ -448,6 +452,77 @@ def dispose_route():
     assert result.returncode == 0, result.stderr
 
 
+class Interrupted(BaseException):
+    """Raised from a signal handler, like pytest-timeout's Failed."""
+
+
+@contextlib.contextmanager
+def interrupt_after(seconds: float) -> Generator[None, None, None]:
+    # Python runs the handler in whatever code executes when the alarm fires.
+    # During a blocking sync call that is the dispatcher's event loop, not the
+    # test. The previous handler and timer belong to pytest-timeout.
+    def handler(signum: int, frame: Any) -> None:
+        raise Interrupted()
+
+    previous_handler = signal.signal(signal.SIGALRM, handler)
+    previous_timer = signal.setitimer(signal.ITIMER_REAL, seconds)
+    try:
+        yield
+    finally:
+        signal.setitimer(signal.ITIMER_REAL, *previous_timer)
+        signal.signal(signal.SIGALRM, previous_handler)
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGALRM is not available")
+def test_sync_should_survive_exception_raised_by_signal_handler(page: Page) -> None:
+    with interrupt_after(0.2):
+        with pytest.raises(Interrupted):
+            page.wait_for_timeout(10_000)
+    # The connection is still usable, and the next call does not wait for the
+    # abandoned one to finish.
+    start = time.monotonic()
+    assert page.evaluate("1 + 1") == 2
+    assert time.monotonic() - start < 5
+    assert page.title() == ""
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGALRM is not available")
+def test_sync_expect_event_should_survive_exception_raised_by_signal_handler(
+    page: Page,
+) -> None:
+    with interrupt_after(0.2):
+        with pytest.raises(Interrupted):
+            with page.expect_event("console"):
+                pass
+    with page.expect_console_message() as console_info:
+        page.evaluate("console.log('hello')")
+    assert console_info.value.text == "hello"
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason="SIGALRM is not available")
+def test_sync_should_survive_exception_raised_by_signal_handler_while_event_handler_blocks(
+    page: Page,
+) -> None:
+    handled: List[str] = []
+
+    def handler(message: ConsoleMessage) -> None:
+        page.wait_for_timeout(500)
+        handled.append(message.text)
+
+    page.on("console", handler)
+    with interrupt_after(0.2):
+        with pytest.raises(Interrupted):
+            page.evaluate("console.log('hello')")
+            page.wait_for_timeout(10_000)
+    assert page.evaluate("1 + 1") == 2
+    # The handler was suspended in its own sync call. It finishes once the
+    # dispatcher runs the event loop again.
+    deadline = time.monotonic() + 10
+    while not handled and time.monotonic() < deadline:
+        page.wait_for_timeout(100)
+    assert handled == ["hello"]
+
+
 def test_click_should_accept_timedelta_for_timeout(page: Page) -> None:
     with pytest.raises(TimeoutError, match="Timeout 1ms exceeded"):
         page.click("does-not-exist", timeout=timedelta(milliseconds=1))