• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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