diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2c127c90..08dbe489 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -66,7 +66,7 @@ repos: language: python additional_dependencies: ["./gitlint-core[trusted-deps]"] entry: gitlint - args: [--staged, --msg-filename] + args: [--staged, -c, "general.ignore=B6,T3", --msg-filename] stages: [commit-msg] - repo: https://github.com/crate-ci/typos rev: v1 diff --git a/CHANGELOG.txt b/CHANGELOG.txt index 12d00e5d..3906d9b5 100644 --- a/CHANGELOG.txt +++ b/CHANGELOG.txt @@ -1,3 +1,5 @@ +- Add argument to Session/Target send_command with_perf to return + timing information about browser write/read. - Update default chrome from 135.0.7011.0/1418433 to 144.0.7527.0/1544685 - Fix: New chrome takes longer/doesn't populate targets right away, so add a retry loop to populate targets diff --git a/pyproject.toml b/pyproject.toml index 98d3715d..f01224a9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,10 +44,6 @@ choreo_get_chrome = "choreographer.cli._cli_utils:get_chrome_cli" [dependency-groups] dev = [ - "pytest", - "pytest-asyncio; python_version < '3.14'", - "pytest-asyncio>=1.2.0; python_version >= '3.14'", - "pytest-xdist", "async-timeout", "numpy; python_version < '3.11'", "numpy>=2.3.3; python_version >= '3.11'", @@ -55,6 +51,10 @@ dev = [ "types-simplejson>=3.19.0.20241221", "poethepoet>=0.30.0", "pyright>=1.1.406", + "pytest", + "pytest-asyncio; python_version < '3.14'", + "pytest-asyncio>=1.2.0; python_version >= '3.14'", + "pytest-xdist", ] # uv doens't allow dependency groups to have separate python requirements @@ -106,6 +106,7 @@ ignore = [ ] [tool.pytest.ini_options] +asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" log_cli = false addopts = "--import-mode=append" diff --git a/src/choreographer/_brokers/_async.py b/src/choreographer/_brokers/_async.py index c567618d..66db45d7 100644 --- a/src/choreographer/_brokers/_async.py +++ b/src/choreographer/_brokers/_async.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import time import warnings from functools import partial from typing import TYPE_CHECKING @@ -22,6 +23,9 @@ _logger = logistro.getLogger(__name__) +PERFS_MAX = 5000 # maximum number of entries in the perf dicts +TRIM_SIZE = 500 # what to save after trimming it + class UnhandledMessageWarning(UserWarning): pass @@ -49,6 +53,9 @@ class Broker: ] """A mapping of session id: subscription: list[futures]""" + write_perfs: MutableMapping[protocol.MessageKey, tuple[float, float]] + read_perfs: MutableMapping[protocol.MessageKey, float] + def __init__(self, browser: Browser, channel: ChannelInterface) -> None: """ Construct a broker for a synchronous arragenment w/ both ends. @@ -66,6 +73,8 @@ def __init__(self, browser: Browser, channel: ChannelInterface) -> None: # if its a user task, can cancel self._current_read_task: asyncio.Task[Any] | None = None self.futures = {} + self.write_perfs = {} + self.read_perfs = {} self._subscriptions_futures = {} self._write_lock = asyncio.Lock() @@ -223,6 +232,14 @@ async def read_loop() -> None: # noqa: PLR0912, PLR0915, C901 raise RuntimeError(f"Couldn't find a future for key: {key}") if not future.done(): future.set_result(response) + self.read_perfs[key] = time.perf_counter() + if len(self.write_perfs) > PERFS_MAX: + self.write_perfs = dict( + list(self.write_perfs.items())[TRIM_SIZE:], + ) + self.read_perfs = dict( + list(self.read_perfs.items())[TRIM_SIZE:], + ) else: warnings.warn( f"Unhandled message type:{response!s}", @@ -237,6 +254,16 @@ async def read_loop() -> None: # noqa: PLR0912, PLR0915, C901 read_task.add_done_callback(check_read_loop_error) self._current_read_task = read_task + def get_perf( + self, + obj: protocol.BrowserCommand, + ) -> tuple[float, float, float]: + """Get the performance tuple for a certain BrowserCommand.""" + key = protocol.calculate_message_key(obj) + if not key: + return (0, 0, 0) + return (*self.write_perfs[key], self.read_perfs[key]) + async def write_json( self, obj: protocol.BrowserCommand, @@ -254,6 +281,7 @@ async def write_json( self.futures[key] = future _logger.debug(f"Created future: {key} {future}") try: + perf_start = time.perf_counter() async with self._write_lock: # this should be a queue not a lock loop = asyncio.get_running_loop() await loop.run_in_executor( @@ -261,6 +289,7 @@ async def write_json( self._channel.write_json, obj, ) + self.write_perfs[key] = (perf_start, time.perf_counter()) except (_manual_thread_pool.ExecutorClosedError, asyncio.CancelledError) as e: if not future.cancel() or not future.cancelled(): await future # it wasn't canceled, so listen to it before raising diff --git a/src/choreographer/protocol/devtools_async.py b/src/choreographer/protocol/devtools_async.py index f4d315f5..96688a1f 100644 --- a/src/choreographer/protocol/devtools_async.py +++ b/src/choreographer/protocol/devtools_async.py @@ -3,7 +3,7 @@ from __future__ import annotations import inspect -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, overload import logistro @@ -11,7 +11,7 @@ if TYPE_CHECKING: import asyncio - from typing import Any, Callable, Coroutine, MutableMapping + from typing import Any, Callable, Coroutine, Literal, MutableMapping from choreographer._brokers import Broker @@ -56,11 +56,34 @@ def __init__(self, session_id: str, broker: Broker) -> None: self.message_id = 0 self.subscriptions = {} + @overload async def send_command( self, command: str, params: MutableMapping[str, Any] | None = None, - ) -> protocol.BrowserResponse: + *, + with_perf: Literal[False] = False, + ) -> protocol.BrowserResponse: ... + + @overload + async def send_command( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + *, + with_perf: Literal[True], + ) -> tuple[protocol.BrowserResponse, tuple[float, float, float]]: ... + + async def send_command( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + *, + with_perf: bool = False, + ) -> ( + tuple[protocol.BrowserResponse, tuple[float, float, float]] + | protocol.BrowserResponse + ): """ Send a devtools command on the session. @@ -69,9 +92,12 @@ async def send_command( Args: command: devtools command to send params: the parameters to send + with_perf (bool): Return the optional tuple. Returns: A message key (session, message id) tuple or None + (Optional) A tuple[float, float, float] representing + perf_counters() for write start, end, and read end. """ current_id = self.message_id @@ -92,6 +118,11 @@ async def send_command( f"sessionId '{self.session_id}'", ) _logger.debug2(f"Full params: {str(params).replace('%', '%%')}") + if with_perf: + return ( + await self._broker.write_json(json_command), + self._broker.get_perf(json_command), + ) return await self._broker.write_json(json_command) def subscribe( @@ -201,11 +232,34 @@ def get_session(self) -> Session: session = next(iter(self.sessions.values())) return session + @overload async def send_command( self, command: str, params: MutableMapping[str, Any] | None = None, - ) -> protocol.BrowserResponse: + *, + with_perf: Literal[False] = False, + ) -> protocol.BrowserResponse: ... + + @overload + async def send_command( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + *, + with_perf: Literal[True], + ) -> tuple[protocol.BrowserResponse, tuple[float, float, float]]: ... + + async def send_command( + self, + command: str, + params: MutableMapping[str, Any] | None = None, + *, + with_perf: bool = False, + ) -> ( + protocol.BrowserResponse + | tuple[protocol.BrowserResponse, tuple[float, float, float]] + ): """ Send a command to the first session in a target. @@ -214,12 +268,13 @@ async def send_command( Args: command: devtools command to send params: the parameters to send + with_perf (bool): Also return perf tuple """ if not self.sessions.values(): raise RuntimeError("Cannot send_command without at least one valid session") session = self.get_session() - return await session.send_command(command, params) + return await session.send_command(command, params, with_perf=with_perf) async def create_session(self) -> Session: """Create a new session on this target.""" diff --git a/tests/test_browser.py b/tests/test_browser.py index ec1d56a3..3e02bd09 100644 --- a/tests/test_browser.py +++ b/tests/test_browser.py @@ -1,6 +1,9 @@ -import choreographer as choreo +import asyncio + import logistro import pytest + +import choreographer as choreo from choreographer import errors from choreographer.protocol import devtools_async @@ -115,6 +118,39 @@ async def test_browser_send_command(browser): await browser.send_command(command=12345) +@pytest.mark.asyncio +async def test_browser_send_command_with_perf(browser): + _logger.info("testing...") + perfs = [] + + # Run multiple commands and collect perf data + for _ in range(3): + response, perf = await browser.send_command( + command="Target.getTargets", + with_perf=True, + ) + assert "result" in response and "targetInfos" in response["result"] # noqa: PT018 I like this assertion + + # Validate perf is a tuple of 3 floats + assert isinstance(perf, tuple) + assert all(isinstance(t, float) for t in perf) + + # Validate timing makes sense (write_start <= write_end <= read_end) + write_start, write_end, read_end = perf + assert write_start <= write_end <= read_end + + perfs.append(perf) + await asyncio.sleep(0.1) + + # Verify each perf tuple is distinct and sequential + for i in range(len(perfs) - 1): + _, _, prev_read_end = perfs[i] + next_write_start, _, _ = perfs[i + 1] + + # Next command should start after previous command ended + assert prev_read_end <= next_write_start, "Commands should be sequential" + + @pytest.mark.asyncio async def test_populate_targets(browser): _logger.info("testing...") diff --git a/tests/test_tab.py b/tests/test_tab.py index dfcca6cf..7412fc0a 100644 --- a/tests/test_tab.py +++ b/tests/test_tab.py @@ -2,6 +2,7 @@ import logistro import pytest + from choreographer import errors from choreographer.protocol import devtools_async @@ -80,7 +81,7 @@ async def count_event(_r): assert "Page.*" in next(iter(tab.sessions.values())).subscriptions await tab.send_command("Page.enable") await tab.send_command("Page.reload") - await asyncio.sleep(0.5) + await asyncio.sleep(0.15) assert counter > old_counter tab.unsubscribe("Page.*")