fix: cancel protocol calls on connection / driver close (microsoft#1897)

mxschmitt · web-flow · commit b913f4830f06 · 2023-05-08T18:21:19.000+02:00
diff --git a/playwright/_impl/_browser.py b/playwright/_impl/_browser.py
@@ -29,6 +29,7 @@
 from playwright._impl._cdp_session import CDPSession
 from playwright._impl._connection import ChannelOwner, from_channel
 from playwright._impl._helper import (
+    BROWSER_CLOSED_ERROR,
     ColorScheme,
     ForcedColors,
     HarContentPolicy,
@@ -182,7 +183,7 @@ async def close(self) -> None:
             if not is_safe_close_error(e):
                 raise e
         if self._should_close_connection_on_close:
-            await self._connection.stop_async()
+            await self._connection.stop_async(BROWSER_CLOSED_ERROR)
 
     @property
     def version(self) -> str:
diff --git a/playwright/_impl/_browser_type.py b/playwright/_impl/_browser_type.py
@@ -33,6 +33,7 @@
     from_nullable_channel,
 )
 from playwright._impl._helper import (
+    BROWSER_CLOSED_ERROR,
     ColorScheme,
     Env,
     ForcedColors,
@@ -218,7 +219,7 @@ async def connect(
 
         timeout_future = throw_on_timeout(timeout, Error("Connection timed out"))
         done, pending = await asyncio.wait(
-            {transport.on_error_future, playwright_future, timeout_future},
+            {playwright_future, timeout_future},
             return_when=asyncio.FIRST_COMPLETED,
         )
         if not playwright_future.done():
@@ -234,13 +235,13 @@ async def connect(
         self._did_launch_browser(browser)
         browser._should_close_connection_on_close = True
 
-        def handle_transport_close() -> None:
+        def handle_transport_close(transport_exception: str) -> None:
             for context in browser.contexts:
                 for page in context.pages:
                     page._on_close()
                 context._on_close()
             browser._on_close()
-            connection.cleanup()
+            connection.cleanup(transport_exception or BROWSER_CLOSED_ERROR)
 
         transport.once("close", handle_transport_close)
 
diff --git a/playwright/_impl/_connection.py b/playwright/_impl/_connection.py
@@ -36,7 +36,7 @@
 from pyee.asyncio import AsyncIOEventEmitter
 
 import playwright
-from playwright._impl._helper import ParsedMessagePayload, parse_error
+from playwright._impl._helper import Error, ParsedMessagePayload, parse_error
 from playwright._impl._transport import Transport
 
 if TYPE_CHECKING:
@@ -242,6 +242,7 @@ def __init__(
         ] = contextvars.ContextVar("ApiZone", default=None)
         self._local_utils: Optional["LocalUtils"] = local_utils
         self._tracing_count = 0
+        self._closed_error_message: Optional[str] = None
 
     @property
     def local_utils(self) -> "LocalUtils":
@@ -272,16 +273,24 @@ def stop_sync(self) -> None:
         self._loop.run_until_complete(self._transport.wait_until_stopped())
         self.cleanup()
 
-    async def stop_async(self) -> None:
+    async def stop_async(self, error_message: str = None) -> None:
         self._transport.request_stop()
         await self._transport.wait_until_stopped()
-        self.cleanup()
+        self.cleanup(error_message)
 
-    def cleanup(self) -> None:
+    def cleanup(self, error_message: str = None) -> None:
+        if not error_message:
+            error_message = "Connection closed"
+        self._closed_error_message = error_message
         if self._init_task and not self._init_task.done():
             self._init_task.cancel()
         for ws_connection in self._child_ws_connections:
             ws_connection._transport.dispose()
+        for callback in self._callbacks.values():
+            callback.future.set_exception(Error(error_message))
+            # Prevent 'Task exception was never retrieved'
+            callback.future.exception()
+        self._callbacks.clear()
         self.emit("close")
 
     def call_on_object_with_known_name(
@@ -298,6 +307,8 @@ def set_in_tracing(self, is_tracing: bool) -> None:
     def _send_message_to_server(
         self, guid: str, method: str, params: Dict
     ) -> ProtocolCallback:
+        if self._closed_error_message:
+            raise Error(self._closed_error_message)
         self._last_id += 1
         id = self._last_id
         callback = ProtocolCallback(self._loop)
@@ -339,6 +350,8 @@ def _send_message_to_server(
         return callback
 
     def dispatch(self, msg: ParsedMessagePayload) -> None:
+        if self._closed_error_message:
+            return
         id = msg.get("id")
         if id:
             callback = self._callbacks.pop(id)
diff --git a/playwright/_impl/_helper.py b/playwright/_impl/_helper.py
@@ -314,10 +314,14 @@ def prepare_interception_patterns(
         return patterns
 
 
+BROWSER_CLOSED_ERROR = "Browser has been closed"
+BROWSER_OR_CONTEXT_CLOSED_ERROR = "Target page, context or browser has been closed"
+
+
 def is_safe_close_error(error: Exception) -> bool:
     message = str(error)
-    return message.endswith("Browser has been closed") or message.endswith(
-        "Target page, context or browser has been closed"
+    return message.endswith(BROWSER_CLOSED_ERROR) or message.endswith(
+        BROWSER_OR_CONTEXT_CLOSED_ERROR
     )
 
 
diff --git a/playwright/_impl/_json_pipe.py b/playwright/_impl/_json_pipe.py
@@ -17,9 +17,8 @@
 
 from pyee.asyncio import AsyncIOEventEmitter
 
-from playwright._impl._api_types import Error
 from playwright._impl._connection import Channel
-from playwright._impl._helper import ParsedMessagePayload, parse_error
+from playwright._impl._helper import ParsedMessagePayload
 from playwright._impl._transport import Transport
 
 
@@ -37,7 +36,7 @@ def __init__(
 
     def request_stop(self) -> None:
         self._stop_requested = True
-        self._loop.create_task(self._pipe_channel.send("close", {}))
+        self._pipe_channel.send_no_reply("close", {})
 
     def dispose(self) -> None:
         self.on_error_future.cancel()
@@ -49,17 +48,17 @@ async def wait_until_stopped(self) -> None:
     async def connect(self) -> None:
         self._stopped_future: asyncio.Future = asyncio.Future()
 
+        close_error: Optional[str] = None
+
         def handle_message(message: Dict) -> None:
-            if not self._stop_requested:
+            try:
                 self.on_message(cast(ParsedMessagePayload, message))
+            except Exception as e:
+                nonlocal close_error
+                close_error = str(e)
 
-        def handle_closed(error: Optional[Dict]) -> None:
-            self.emit("close")
-            self.on_error_future.set_exception(
-                parse_error(error["error"])
-                if error
-                else Error("Playwright connection closed")
-            )
+        def handle_closed() -> None:
+            self.emit("close", close_error)
             self._stopped_future.set_result(None)
 
         self._pipe_channel.on(
@@ -68,13 +67,11 @@ def handle_closed(error: Optional[Dict]) -> None:
         )
         self._pipe_channel.on(
             "closed",
-            lambda params: handle_closed(params.get("error")),
+            lambda _: handle_closed(),
         )
 
     async def run(self) -> None:
         await self._stopped_future
 
     def send(self, message: Dict) -> None:
-        if self._stop_requested:
-            raise Error("Playwright connection closed")
-        self._loop.create_task(self._pipe_channel.send("send", {"message": message}))
+        self._pipe_channel.send_no_reply("send", {"message": message})
diff --git a/tests/async/test_asyncio.py b/tests/async/test_asyncio.py
@@ -19,6 +19,8 @@
 
 from playwright.async_api import async_playwright
 
+from ..server import Server
+
 
 async def test_should_cancel_underlying_protocol_calls(
     browser_name: str, launch_arguments: Dict
@@ -54,3 +56,14 @@ async def test_async_playwright_stop_multiple_times() -> None:
     playwright = await async_playwright().start()
     await playwright.stop()
     await playwright.stop()
+
+
+async def test_cancel_pending_protocol_call_on_playwright_stop(server: Server) -> None:
+    server.set_route("/hang", lambda _: None)
+    playwright = await async_playwright().start()
+    api_request_context = await playwright.request.new_context()
+    pending_task = asyncio.create_task(api_request_context.get(server.PREFIX + "/hang"))
+    await playwright.stop()
+    with pytest.raises(Exception) as exc_info:
+        await pending_task
+    assert "Connection closed" in str(exc_info.value)
diff --git a/tests/async/test_browsertype_connect.py b/tests/async/test_browsertype_connect.py
@@ -145,7 +145,7 @@ async def test_browser_type_connect_should_throw_when_used_after_is_connected_re
 
     with pytest.raises(Error) as exc_info:
         await page.evaluate("1 + 1")
-    assert "Playwright connection closed" == exc_info.value.message
+    assert "has been closed" in exc_info.value.message
     assert browser.is_connected() is False
 
 
@@ -159,7 +159,7 @@ async def test_browser_type_connect_should_reject_navigation_when_browser_closes
 
     with pytest.raises(Error) as exc_info:
         await page.goto(server.PREFIX + "/one-style.html")
-    assert "Playwright connection closed" in exc_info.value.message
+    assert "has been closed" in exc_info.value.message
 
 
 async def test_should_not_allow_getting_the_path(
diff --git a/tests/sync/test_browsertype_connect.py b/tests/sync/test_browsertype_connect.py
@@ -140,7 +140,7 @@ def test_browser_type_connect_should_throw_when_used_after_is_connected_returns_
 
     with pytest.raises(Error) as exc_info:
         page.evaluate("1 + 1")
-    assert "Playwright connection closed" == exc_info.value.message
+    assert "has been closed" in exc_info.value.message
     assert browser.is_connected() is False
 
 
