diff --git a/playwright/_impl/_browser_type.py b/playwright/_impl/_browser_type.py index f78ba82c0..451179030 100644 --- a/playwright/_impl/_browser_type.py +++ b/playwright/_impl/_browser_type.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio from pathlib import Path from typing import Dict, List, Optional, Union, cast @@ -21,6 +22,7 @@ ProxySettings, ViewportSize, ) +from playwright._impl._api_types import Error from playwright._impl._browser import Browser, normalize_context_params from playwright._impl._browser_context import BrowserContext from playwright._impl._connection import ( @@ -37,6 +39,7 @@ not_installed_error, ) from playwright._impl._transport import WebSocketTransport +from playwright._impl._wait_helper import throw_on_timeout class BrowserType(ChannelOwner): @@ -172,8 +175,10 @@ async def connect( slow_mo: float = None, headers: Dict[str, str] = None, ) -> Browser: - transport = WebSocketTransport(ws_endpoint, timeout, headers) + if timeout is None: + timeout = 30000 + transport = WebSocketTransport(self._connection._loop, ws_endpoint, headers) connection = Connection( self._connection._dispatcher_fiber, self._connection._object_factory, @@ -182,8 +187,20 @@ async def connect( connection._is_sync = self._connection._is_sync connection._loop = self._connection._loop connection._loop.create_task(connection.run()) + future = connection._loop.create_task( + connection.wait_for_object_with_known_name("Playwright") + ) + timeout_future = throw_on_timeout(timeout, Error("Connection timed out")) + done, pending = await asyncio.wait( + {transport.on_error_future, future, timeout_future}, + return_when=asyncio.FIRST_COMPLETED, + ) + if not future.done(): + future.cancel() + if not timeout_future.done(): + timeout_future.cancel() + playwright = next(iter(done)).result() self._connection._child_ws_connections.append(connection) - playwright = await connection.wait_for_object_with_known_name("Playwright") pre_launched_browser = playwright._initializer.get("preLaunchedBrowser") assert pre_launched_browser browser = cast(Browser, from_channel(pre_launched_browser)) diff --git a/playwright/_impl/_transport.py b/playwright/_impl/_transport.py index b315318b3..b4dfb8d8e 100644 --- a/playwright/_impl/_transport.py +++ b/playwright/_impl/_transport.py @@ -41,8 +41,10 @@ def _get_stderr_fileno() -> Optional[int]: class Transport(ABC): - def __init__(self) -> None: + def __init__(self, loop: asyncio.AbstractEventLoop) -> None: + self._loop = loop self.on_message = lambda _: None + self.on_error_future: asyncio.Future = loop.create_future() @abstractmethod def request_stop(self) -> None: @@ -55,9 +57,9 @@ def dispose(self) -> None: async def wait_until_stopped(self) -> None: pass + @abstractmethod async def run(self) -> None: - self._loop = asyncio.get_running_loop() - self.on_error_future: asyncio.Future = asyncio.Future() + pass @abstractmethod def send(self, message: Dict) -> None: @@ -78,11 +80,12 @@ def deserialize_message(self, data: bytes) -> Any: class PipeTransport(Transport): - def __init__(self, driver_executable: Path) -> None: - super().__init__() + def __init__( + self, loop: asyncio.AbstractEventLoop, driver_executable: Path + ) -> None: + super().__init__(loop) self._stopped = False self._driver_executable = driver_executable - self._loop: asyncio.AbstractEventLoop def request_stop(self) -> None: self._stopped = True @@ -93,17 +96,21 @@ async def wait_until_stopped(self) -> None: await self._proc.wait() async def run(self) -> None: - await super().run() self._stopped_future: asyncio.Future = asyncio.Future() - self._proc = proc = await asyncio.create_subprocess_exec( - str(self._driver_executable), - "run-driver", - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=_get_stderr_fileno(), - limit=32768, - ) + try: + self._proc = proc = await asyncio.create_subprocess_exec( + str(self._driver_executable), + "run-driver", + stdin=asyncio.subprocess.PIPE, + stdout=asyncio.subprocess.PIPE, + stderr=_get_stderr_fileno(), + limit=32768, + ) + except Exception as exc: + self.on_error_future.set_exception(exc) + return + assert proc.stdout assert proc.stdin self._output = proc.stdin @@ -138,16 +145,17 @@ def send(self, message: Dict) -> None: class WebSocketTransport(AsyncIOEventEmitter, Transport): def __init__( - self, ws_endpoint: str, timeout: float = None, headers: Dict[str, str] = None + self, + loop: asyncio.AbstractEventLoop, + ws_endpoint: str, + headers: Dict[str, str] = None, ) -> None: - super().__init__() - Transport.__init__(self) + super().__init__(loop) + Transport.__init__(self, loop) self._stopped = False self.ws_endpoint = ws_endpoint - self.timeout = timeout self.headers = headers - self._loop: asyncio.AbstractEventLoop def request_stop(self) -> None: self._stopped = True @@ -160,15 +168,13 @@ async def wait_until_stopped(self) -> None: await self._connection.wait_closed() async def run(self) -> None: - await super().run() - - options: Dict[str, Any] = {} - if self.timeout is not None: - options["close_timeout"] = self.timeout / 1000 - options["ping_timeout"] = self.timeout / 1000 - if self.headers is not None: - options["extra_headers"] = self.headers - self._connection = await websockets.connect(self.ws_endpoint, **options) + try: + self._connection = await websockets.connect( + self.ws_endpoint, extra_headers=self.headers + ) + except Exception as exc: + self.on_error_future.set_exception(Error(f"websocket.connect: {str(exc)}")) + return while not self._stopped: try: @@ -188,8 +194,8 @@ async def run(self) -> None: ) break except Exception as exc: - print(f"Received unhandled exception: {exc}") self.on_error_future.set_exception(exc) + break def send(self, message: Dict) -> None: if self._stopped or self._connection.closed: diff --git a/playwright/_impl/_wait_helper.py b/playwright/_impl/_wait_helper.py index a5ca59c41..a57724b10 100644 --- a/playwright/_impl/_wait_helper.py +++ b/playwright/_impl/_wait_helper.py @@ -91,3 +91,11 @@ def listener(event_data: Any = None) -> None: def result(self) -> asyncio.Future: return self._result + + +def throw_on_timeout(timeout: float, exception: Exception) -> asyncio.Task: + async def throw() -> None: + await asyncio.sleep(timeout / 1000) + raise exception + + return asyncio.create_task(throw()) diff --git a/playwright/async_api/_context_manager.py b/playwright/async_api/_context_manager.py index 51800dd48..7e9cae8db 100644 --- a/playwright/async_api/_context_manager.py +++ b/playwright/async_api/_context_manager.py @@ -28,14 +28,23 @@ def __init__(self) -> None: async def __aenter__(self) -> AsyncPlaywright: self._connection = Connection( - None, create_remote_object, PipeTransport(compute_driver_executable()) + None, + create_remote_object, + PipeTransport(asyncio.get_event_loop(), compute_driver_executable()), ) loop = asyncio.get_running_loop() self._connection._loop = loop loop.create_task(self._connection.run()) - playwright = AsyncPlaywright( - await self._connection.wait_for_object_with_known_name("Playwright") + playwright_future = asyncio.create_task( + self._connection.wait_for_object_with_known_name("Playwright") ) + done, pending = await asyncio.wait( + {self._connection._transport.on_error_future, playwright_future}, + return_when=asyncio.FIRST_COMPLETED, + ) + if not playwright_future.done(): + playwright_future.cancel() + playwright = AsyncPlaywright(next(iter(done)).result()) playwright.stop = self.__aexit__ # type: ignore return playwright diff --git a/playwright/sync_api/_context_manager.py b/playwright/sync_api/_context_manager.py index df12d3af8..289483bb9 100644 --- a/playwright/sync_api/_context_manager.py +++ b/playwright/sync_api/_context_manager.py @@ -31,21 +31,20 @@ def __init__(self) -> None: self._playwright: SyncPlaywright def __enter__(self) -> SyncPlaywright: - def greenlet_main() -> None: - loop = None - own_loop = None - try: - loop = asyncio.get_running_loop() - except RuntimeError: - loop = asyncio.new_event_loop() - own_loop = loop - - if loop.is_running(): - raise Error( - """It looks like you are using Playwright Sync API inside the asyncio loop. + loop: asyncio.AbstractEventLoop + own_loop = None + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + own_loop = loop + if loop.is_running(): + raise Error( + """It looks like you are using Playwright Sync API inside the asyncio loop. Please use the Async API instead.""" - ) + ) + def greenlet_main() -> None: loop.run_until_complete(self._connection.run_as_sync()) if own_loop: @@ -56,7 +55,7 @@ def greenlet_main() -> None: self._connection = Connection( dispatcher_fiber, create_remote_object, - PipeTransport(compute_driver_executable()), + PipeTransport(loop, compute_driver_executable()), ) g_self = greenlet.getcurrent() diff --git a/tests/async/test_browsertype_connect.py b/tests/async/test_browsertype_connect.py index 842b6c1ee..ed16a00e9 100644 --- a/tests/async/test_browsertype_connect.py +++ b/tests/async/test_browsertype_connect.py @@ -182,3 +182,13 @@ async def test_prevent_getting_video_path( == "Path is not available when using browserType.connect(). Use save_as() to save a local copy." ) remote_server.kill() + + +async def test_connect_to_closed_server_without_hangs( + browser_type: BrowserType, launch_server +): + remote_server = launch_server() + remote_server.kill() + with pytest.raises(Error) as exc: + await browser_type.connect(remote_server.ws_endpoint) + assert "websocket.connect: " in exc.value.message diff --git a/tests/sync/test_browsertype_connect.py b/tests/sync/test_browsertype_connect.py index d698d8da1..620d129c5 100644 --- a/tests/sync/test_browsertype_connect.py +++ b/tests/sync/test_browsertype_connect.py @@ -145,3 +145,13 @@ def test_browser_type_connect_should_forward_close_events_to_pages( assert events == ["page::close", "context::close", "browser::disconnected"] remote.kill() assert events == ["page::close", "context::close", "browser::disconnected"] + + +def test_connect_to_closed_server_without_hangs( + browser_type: BrowserType, launch_server +): + remote_server = launch_server() + remote_server.kill() + with pytest.raises(Error) as exc: + browser_type.connect(remote_server.ws_endpoint) + assert "websocket.connect: " in exc.value.message