Merge pull request #271 from plotly/andrew/add_profiler

andrew/add profiler

Author: ayjayt

.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

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

pyproject.toml

@@ -44,17 +44,17 @@ 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'",
   "mypy>=1.14.1",
   "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"

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,13 +281,15 @@ 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(
                     self._executor,
                     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

src/choreographer/protocol/devtools_async.py

@@ -3,15 +3,15 @@
 from __future__ import annotations
 
 import inspect
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 import logistro
 
 from choreographer import protocol
 
 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."""

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...")

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.*")