1# Copyright 2024 The Chromium Authors 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5from __future__ import annotations 6 7import abc 8import contextlib 9import datetime as dt 10from typing import (TYPE_CHECKING, Generic, Iterable, Iterator, Optional, 11 TypeVar) 12 13from crossbench import plt 14from crossbench.probes.results import (BrowserProbeResult, EmptyProbeResult, 15 LocalProbeResult, ProbeResult) 16 17if TYPE_CHECKING: 18 from selenium.webdriver.common.options import BaseOptions 19 20 from crossbench.browsers.browser import Browser 21 from crossbench.path import AnyPath, LocalPath 22 from crossbench.probes.probe import Probe 23 from crossbench.runner.groups.session import BrowserSessionRunGroup 24 from crossbench.runner.result_origin import ResultOrigin 25 from crossbench.runner.run import Run 26 from crossbench.runner.runner import Runner 27 28# Redefine here to avoid circular imports 29ProbeT = TypeVar("ProbeT", bound="Probe") 30 31 32class BaseProbeContext(Generic[ProbeT], metaclass=abc.ABCMeta): 33 """ 34 Base class for an activation of a probe where active data collection 35 happens. See specific subclasses for implementations that can be used 36 for data collection during runs or whole sessions. 37 Override in Probe subclasses to implement actual performance data 38 collection. 39 - The data should be written to self.result_path. 40 - A file / list / dict of result file Paths should be returned by the 41 override teardown() method 42 """ 43 44 def __init__(self, probe: ProbeT, result_origin: ResultOrigin) -> None: 45 self._probe: ProbeT = probe 46 self._result_origin = result_origin 47 self._is_active: bool = False 48 self._is_success: bool = False 49 self._start_time: Optional[dt.datetime] = None 50 self._stop_time: Optional[dt.datetime] = None 51 52 def set_start_time(self, start_datetime: dt.datetime) -> None: 53 assert self._start_time is None 54 self._start_time = start_datetime 55 56 @contextlib.contextmanager 57 def open(self) -> Iterator[None]: 58 assert self._start_time 59 assert not self._is_active 60 assert not self._is_success 61 62 with self.result_origin.exception_handler(f"Probe {self.name} start"): 63 self._is_active = True 64 self.start() 65 66 try: 67 yield 68 finally: 69 with self.result_origin.exception_handler(f"Probe {self.name} stop"): 70 self.stop() 71 self._is_success = True 72 assert self._stop_time is None 73 self._stop_time = dt.datetime.now() 74 75 @property 76 def probe(self) -> ProbeT: 77 return self._probe 78 79 @property 80 def result_origin(self) -> ResultOrigin: 81 return self._result_origin 82 83 @property 84 def browser_platform(self) -> plt.Platform: 85 return self.browser.platform 86 87 @property 88 def host_platform(self) -> plt.Platform: 89 return self.browser.host_platform 90 91 @property 92 @abc.abstractmethod 93 def browser(self) -> Browser: 94 pass 95 96 @property 97 @abc.abstractmethod 98 def runner(self) -> Runner: 99 pass 100 101 @property 102 @abc.abstractmethod 103 def session(self) -> BrowserSessionRunGroup: 104 pass 105 106 @property 107 def start_time(self) -> dt.datetime: 108 """ 109 Returns a unified start time that is the same for all probe contexts 110 within a run. This can be used to account for startup delays caused by other 111 Probes. 112 """ 113 assert self._start_time 114 return self._start_time 115 116 @property 117 def duration(self) -> dt.timedelta: 118 assert self._start_time and self._stop_time 119 return self._stop_time - self._start_time 120 121 @property 122 def is_success(self) -> bool: 123 return self._is_success 124 125 @property 126 @abc.abstractmethod 127 def result_path(self) -> AnyPath: 128 pass 129 130 @property 131 @abc.abstractmethod 132 def local_result_path(self) -> LocalPath: 133 pass 134 135 @property 136 def name(self) -> str: 137 return self.probe.name 138 139 @property 140 def browser_pid(self) -> int: 141 maybe_pid = self.browser.pid 142 assert maybe_pid, "Browser is not runner or does not provide a pid." 143 return maybe_pid 144 145 def browser_result(self, 146 url: Optional[Iterable[str]] = None, 147 file: Optional[Iterable[AnyPath]] = None, 148 **kwargs: Iterable[AnyPath]) -> BrowserProbeResult: 149 """Helper to create BrowserProbeResult that might be stored on a remote 150 browser/device and need to be copied over to the local machine.""" 151 return BrowserProbeResult(self.result_origin, url=url, file=file, **kwargs) 152 153 def local_result(self, 154 url: Optional[Iterable[str]] = None, 155 file: Optional[Iterable[LocalPath]] = None, 156 **kwargs: Iterable[LocalPath]) -> LocalProbeResult: 157 """Helper to create LocalProbeResult.""" 158 return LocalProbeResult(url=url, file=file, **kwargs) 159 160 def setup(self) -> None: 161 """ 162 Called before starting the browser, typically used to set run-specific 163 browser flags. 164 """ 165 166 @abc.abstractmethod 167 def start(self) -> None: 168 pass 169 170 @abc.abstractmethod 171 def stop(self) -> None: 172 pass 173 174 @abc.abstractmethod 175 def teardown(self) -> ProbeResult: 176 pass 177 178 179class ProbeContext(BaseProbeContext[ProbeT], metaclass=abc.ABCMeta): 180 """ 181 A scope during which a probe is actively collecting data during a Run. 182 See BaseProbeContext additional usage. 183 """ 184 185 def __init__(self, probe: ProbeT, run: Run) -> None: 186 super().__init__(probe, run) 187 self._run: Run = run 188 self._default_result_path: AnyPath = self.get_default_result_path() 189 190 def get_default_result_path(self) -> AnyPath: 191 return self._run.get_default_probe_result_path(self._probe) 192 193 @property 194 def run(self) -> Run: 195 return self._run 196 197 @property 198 def result_origin(self) -> ResultOrigin: 199 return self._run 200 201 @property 202 def session(self) -> BrowserSessionRunGroup: 203 return self._run.session 204 205 @property 206 def browser(self) -> Browser: 207 return self._run.browser 208 209 @property 210 def runner(self) -> Runner: 211 return self._run.runner 212 213 @property 214 def result_path(self) -> AnyPath: 215 return self._default_result_path 216 217 @property 218 def local_result_path(self) -> LocalPath: 219 return self.host_platform.local_path(self.result_path) 220 221 def setup_selenium_options(self, options: BaseOptions) -> None: 222 """ 223 Custom hook to change selenium options before starting the browser. 224 """ 225 # TODO: move to SessionContext 226 del options 227 228 @abc.abstractmethod 229 def start(self) -> None: 230 """ 231 Called immediately before starting the given Run, after the browser started. 232 This method should have as little overhead as possible. If possible, 233 delegate heavy computation to the "SetUp" method. 234 """ 235 236 def start_story_run(self) -> None: 237 """ 238 Called before running a Story's core workload (Story.run) 239 and after running Story.setup. 240 """ 241 242 def stop_story_run(self) -> None: 243 """ 244 Called after running a Story's core workload (Story.run) and before running 245 Story.teardown. 246 """ 247 248 @abc.abstractmethod 249 def stop(self) -> None: 250 """ 251 Called immediately after finishing the given Run with the browser still 252 running. 253 This method should have as little overhead as possible. If possible, 254 delegate heavy computation to the "teardown" method. 255 """ 256 return None 257 258 @abc.abstractmethod 259 def teardown(self) -> ProbeResult: 260 """ 261 Called after stopping all probes and shutting down the browser. 262 Returns 263 - None if no data was collected 264 - If Data was collected: 265 - Either a path (or list of paths) to results file 266 - Directly a primitive json-serializable object containing the data 267 """ 268 return EmptyProbeResult() 269 270 271class ProbeSessionContext(BaseProbeContext[ProbeT], metaclass=abc.ABCMeta): 272 """ 273 A scope during which a probe is actively collecting data during an active 274 browser session, which might span several runs. 275 See BaseProbeContext additional usage. 276 """ 277 278 def __init__(self, probe: ProbeT, session: BrowserSessionRunGroup) -> None: 279 super().__init__(probe, session) 280 self._session: BrowserSessionRunGroup = session 281 self._default_result_path: AnyPath = self.get_default_result_path() 282 283 def get_default_result_path(self) -> AnyPath: 284 return self._session.get_default_probe_result_path(self._probe) 285 286 @property 287 def session(self) -> BrowserSessionRunGroup: 288 return self._session 289 290 @property 291 def result_origin(self) -> ResultOrigin: 292 return self._session 293 294 @property 295 def browser(self) -> Browser: 296 return self._session.browser 297 298 @property 299 def result_path(self) -> AnyPath: 300 return self._default_result_path 301