1# Copyright 2022 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 logging 9from typing import (TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, 10 cast) 11 12from immutabledict import immutabledict 13from ordered_set import OrderedSet 14 15from crossbench import path as pth 16from crossbench.parse import ObjectParser 17from crossbench.probes.helper import INTERNAL_NAME_PREFIX 18 19if TYPE_CHECKING: 20 from crossbench.probes.probe import Probe 21 from crossbench.runner.result_origin import ResultOrigin 22 from crossbench.types import JsonDict 23 24 25class DuplicateProbeResult(ValueError): 26 pass 27 28 29class ProbeResult(abc.ABC): 30 """ 31 Collection of result files for a given Probe. These can be URLs or any file. 32 33 We distinguish between two types of files, files that can be fed to Perfetto 34 TraceProcessor (trace) and any other file (file). Trace files will be fed to 35 the trace_processor probe if present. 36 """ 37 38 def __init__(self, 39 url: Optional[Iterable[str]] = None, 40 file: Optional[Iterable[pth.LocalPath]] = None, 41 trace: Optional[Iterable[pth.LocalPath]] = None, 42 **kwargs: Iterable[pth.LocalPath]): 43 self._url_list: Tuple[str, ...] = () 44 if url: 45 self._url_list = ObjectParser.unique_sequence( 46 tuple(url), "urls", DuplicateProbeResult) 47 self._trace_list: Tuple[pth.LocalPath, ...] = () 48 if trace: 49 self._trace_list = ObjectParser.unique_sequence( 50 tuple(trace), "traces", DuplicateProbeResult) 51 tmp_files: Dict[str, OrderedSet[pth.LocalPath]] = {} 52 if file: 53 self._extend(tmp_files, file, suffix=None, allow_duplicates=False) 54 for suffix, files in kwargs.items(): 55 self._extend(tmp_files, files, suffix=suffix, allow_duplicates=False) 56 57 # Do last and allow duplicated 58 self._extend( 59 tmp_files, self._trace_list, suffix=None, allow_duplicates=True) 60 self._files: immutabledict[str, Tuple[pth.LocalPath, ...]] = immutabledict({ 61 suffix: tuple(files) for suffix, files in tmp_files.items() 62 }) 63 # TODO: Add Metric object for keeping metrics in-memory instead of reloading 64 # them from serialized JSON files for merging. 65 self._values = None 66 self._validate() 67 68 def _append(self, 69 tmp_files: Dict[str, OrderedSet[pth.LocalPath]], 70 file: pth.LocalPath, 71 suffix: Optional[str] = None, 72 allow_duplicates: bool = False) -> None: 73 file_suffix_name = file.suffix[1:] 74 if not suffix: 75 suffix = file_suffix_name 76 elif file_suffix_name != suffix: 77 raise ValueError( 78 f"Expected '.{suffix}' suffix, but got {repr(file.suffix)} " 79 f"for {file}") 80 if files_with_suffix := tmp_files.get(suffix): 81 if file not in files_with_suffix: 82 files_with_suffix.add(file) 83 elif not allow_duplicates: 84 raise DuplicateProbeResult( 85 f"Cannot append file twice to ProbeResult: {file}") 86 else: 87 tmp_files[suffix] = OrderedSet((file,)) 88 89 def _extend(self, 90 tmp_files: Dict[str, OrderedSet[pth.LocalPath]], 91 files: Iterable[pth.LocalPath], 92 suffix: Optional[str] = None, 93 allow_duplicates=False) -> None: 94 for file in files: 95 self._append( 96 tmp_files, file, suffix=suffix, allow_duplicates=allow_duplicates) 97 98 def get(self, suffix: str) -> pth.LocalPath: 99 if files_with_suffix := self._files.get(suffix): 100 if len(files_with_suffix) != 1: 101 raise ValueError(f"Expected exactly one file with suffix {suffix}, " 102 f"but got {files_with_suffix}") 103 return files_with_suffix[0] 104 raise ValueError(f"No files with suffix '.{suffix}'. " 105 f"Options are {tuple(self._files.keys())}") 106 107 def get_all(self, suffix: str) -> List[pth.LocalPath]: 108 if files_with_suffix := self._files.get(suffix): 109 return list(files_with_suffix) 110 return [] 111 112 @property 113 def is_empty(self) -> bool: 114 return not self._url_list and not self._files 115 116 @property 117 def is_remote(self) -> bool: 118 return False 119 120 def __bool__(self) -> bool: 121 return not self.is_empty 122 123 def __eq__(self, other: Any) -> bool: 124 if not isinstance(other, ProbeResult): 125 return False 126 if self is other: 127 return True 128 if self._files != other._files: 129 return False 130 return self._url_list == other._url_list 131 132 def merge(self, other: ProbeResult) -> ProbeResult: 133 if self.is_empty: 134 return other 135 if other.is_empty: 136 return self 137 return LocalProbeResult( 138 url=self.url_list + other.url_list, 139 file=self.file_list + other.file_list, 140 trace=self.trace_list + other.trace_list) 141 142 def _validate(self) -> None: 143 for path in self.all_files(): 144 if not path.exists(): 145 raise ValueError(f"ProbeResult path does not exist: {path}") 146 147 def to_json(self) -> JsonDict: 148 result: JsonDict = {} 149 if self._url_list: 150 result["url"] = self._url_list 151 for suffix, files in self._files.items(): 152 result[suffix] = list(map(str, files)) 153 return result 154 155 @property 156 def has_files(self) -> bool: 157 return bool(self._files) 158 159 def all_files(self) -> Iterable[pth.LocalPath]: 160 for files in self._files.values(): 161 yield from files 162 163 @property 164 def url(self) -> str: 165 if len(self._url_list) != 1: 166 raise ValueError("ProbeResult has multiple URLs.") 167 return self._url_list[0] 168 169 @property 170 def url_list(self) -> List[str]: 171 return list(self._url_list) 172 173 @property 174 def file(self) -> pth.LocalPath: 175 if sum(len(files) for files in self._files.values()) > 1: 176 raise ValueError("ProbeResult has more than one file.") 177 for files in self._files.values(): 178 return files[0] 179 raise ValueError("ProbeResult has no files.") 180 181 @property 182 def file_list(self) -> List[pth.LocalPath]: 183 return list(self.all_files()) 184 185 @property 186 def trace(self) -> pth.LocalPath: 187 if len(self._trace_list) != 1: 188 raise ValueError("ProbeResult has multiple traces.") 189 return self._trace_list[0] 190 191 @property 192 def trace_list(self) -> List[pth.LocalPath]: 193 return list(self._trace_list) 194 195 @property 196 def json(self) -> pth.LocalPath: 197 return self.get("json") 198 199 @property 200 def json_list(self) -> List[pth.LocalPath]: 201 return self.get_all("json") 202 203 @property 204 def csv(self) -> pth.LocalPath: 205 return self.get("csv") 206 207 @property 208 def csv_list(self) -> List[pth.LocalPath]: 209 return self.get_all("csv") 210 211 212class EmptyProbeResult(ProbeResult): 213 214 def __init__(self) -> None: 215 super().__init__() 216 217 def __bool__(self) -> bool: 218 return False 219 220 221class LocalProbeResult(ProbeResult): 222 """LocalProbeResult can be used for files that are always available on the 223 runner/local machine.""" 224 225 226class BrowserProbeResult(ProbeResult): 227 """BrowserProbeResult are stored on the device where the browser runs. 228 Result files will be automatically transferred to the local run's results 229 folder. 230 """ 231 232 def __init__(self, 233 result_origin: ResultOrigin, 234 url: Optional[Iterable[str]] = None, 235 file: Optional[Iterable[pth.AnyPath]] = None, 236 **kwargs: Iterable[pth.AnyPath]): 237 self._browser_file = file 238 local_file: Optional[Iterable[pth.LocalPath]] = None 239 local_kwargs: Dict[str, Iterable[pth.LocalPath]] = {} 240 self._is_remote = result_origin.is_remote 241 if self._is_remote: 242 if file: 243 local_file = self._copy_files(result_origin, file) 244 for suffix_name, files in kwargs.items(): 245 local_kwargs[suffix_name] = self._copy_files(result_origin, files) 246 else: 247 # Keep local files as is. 248 local_file = cast(Iterable[pth.LocalPath], file) 249 local_kwargs = cast(Dict[str, Iterable[pth.LocalPath]], kwargs) 250 251 super().__init__(url, local_file, **local_kwargs) 252 253 @property 254 def is_remote(self) -> bool: 255 return self._is_remote 256 257 def _copy_files(self, result_origin: ResultOrigin, 258 paths: Iterable[pth.AnyPath]) -> Iterable[pth.LocalPath]: 259 assert paths, "Got no remote paths to copy." 260 # Copy result files from remote tmp dir to local results dir 261 browser_platform = result_origin.browser_platform 262 remote_tmp_dir = result_origin.browser_tmp_dir 263 out_dir = result_origin.out_dir 264 local_result_paths: List[pth.LocalPath] = [] 265 for remote_path in paths: 266 try: 267 relative_path = remote_path.relative_to(remote_tmp_dir) 268 except ValueError: 269 logging.debug( 270 "Browser result is not in browser tmp dir: " 271 "only using the name of '%s'", remote_path) 272 relative_path = result_origin.host_platform.local_path(remote_path.name) 273 local_result_path = out_dir / relative_path 274 browser_platform.pull(remote_path, local_result_path) 275 assert local_result_path.exists(), "Failed to copy result file." 276 local_result_paths.append(local_result_path) 277 return local_result_paths 278 279 280class ProbeResultDict: 281 """ 282 Maps Probes to their result files Paths. 283 """ 284 285 def __init__(self, path: pth.AnyPath) -> None: 286 self._path = path 287 self._dict: Dict[str, ProbeResult] = {} 288 289 def __setitem__(self, probe: Probe, result: ProbeResult) -> None: 290 assert isinstance(result, ProbeResult) 291 self._dict[probe.name] = result 292 293 def __getitem__(self, probe: Probe) -> ProbeResult: 294 name = probe.name 295 if name not in self._dict: 296 raise KeyError(f"No results for probe='{name}'") 297 return self._dict[name] 298 299 def __contains__(self, probe: Probe) -> bool: 300 return probe.name in self._dict 301 302 def __bool__(self) -> bool: 303 return bool(self._dict) 304 305 def __len__(self) -> int: 306 return len(self._dict) 307 308 def get(self, probe: Probe, default: Any = None) -> ProbeResult: 309 return self._dict.get(probe.name, default) 310 311 def get_by_name(self, name: str, default: Any = None) -> ProbeResult: 312 # Debug helper only. 313 # Use bracket `results[probe]` or `results.get(probe)` instead. 314 return self._dict.get(name, default) 315 316 def to_json(self) -> JsonDict: 317 data: JsonDict = {} 318 for probe_name, results in self._dict.items(): 319 if isinstance(results, (pth.AnyPath, str)): 320 data[probe_name] = str(results) 321 else: 322 if results.is_empty: 323 if not probe_name.startswith(INTERNAL_NAME_PREFIX): 324 logging.debug("probe=%s did not produce any data.", probe_name) 325 data[probe_name] = None 326 else: 327 data[probe_name] = results.to_json() 328 return data 329 330 def all_traces(self) -> Iterable[pth.LocalPath]: 331 for probe_result in self._dict.values(): 332 yield from probe_result.trace_list 333