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 7from typing import TYPE_CHECKING, Iterable, List, Optional 8 9from crossbench.parse import ObjectParser 10from crossbench.probes.probe import Probe, ProbeConfigParser, ProbeKeyT 11from crossbench.probes.probe_context import ProbeContext 12from crossbench.probes.result_location import ResultLocation 13from crossbench.probes.results import LocalProbeResult, ProbeResult 14 15if TYPE_CHECKING: 16 from crossbench import path as pth 17 from crossbench.env import HostEnvironment 18 from crossbench.plt.base import CmdArg, TupleCmdArgs 19 from crossbench.runner.run import Run 20 21 22class ShellProbe(Probe): 23 """ 24 Run an arbitrary shell command on the browser platform and store the 25 stdout and stderr of the command as a result file. 26 """ 27 NAME = "shell" 28 IS_GENERAL_PURPOSE = True 29 RESULT_LOCATION = ResultLocation.LOCAL 30 31 @classmethod 32 def config_parser(cls) -> ProbeConfigParser: 33 parser = super().config_parser() 34 parser.add_argument( 35 "setup_cmd", 36 aliases=("setup",), 37 type=ObjectParser.sh_cmd, 38 required=False, 39 help="CMD is run before the browser is started.") 40 parser.add_argument( 41 "start_cmd", 42 type=ObjectParser.sh_cmd, 43 aliases=("start",), 44 required=False, 45 help=("CMD is run right before each story is started " 46 "and the browser is already running.")) 47 parser.add_argument( 48 "start_story_run_cmd", 49 aliases=("start-story",), 50 type=ObjectParser.sh_cmd, 51 required=False, 52 help=("CMD is run right before the measurement phase " 53 "of a story is started.")) 54 parser.add_argument( 55 "stop_story_run_cmd", 56 aliases=("stop-story",), 57 type=ObjectParser.sh_cmd, 58 required=False, 59 help=("CMD is run right after the measurement phase " 60 "of a story has ended.")) 61 parser.add_argument( 62 "stop_cmd", 63 aliases=("cmd", "stop"), 64 type=ObjectParser.sh_cmd, 65 required=True, 66 help=("CMD is run right after the workload ended and the browser " 67 "is still running.")) 68 parser.add_argument( 69 "teardown_cmd", 70 aliases=("teardown",), 71 type=ObjectParser.sh_cmd, 72 required=False, 73 help="CMD is run after the browser is stopped.") 74 return parser 75 76 def __init__(self, 77 setup_cmd: Optional[Iterable[CmdArg]] = None, 78 start_cmd: Optional[Iterable[CmdArg]] = None, 79 start_story_run_cmd: Optional[Iterable[CmdArg]] = None, 80 stop_story_run_cmd: Optional[Iterable[CmdArg]] = None, 81 stop_cmd: Optional[Iterable[CmdArg]] = None, 82 teardown_cmd: Optional[Iterable[CmdArg]] = None) -> None: 83 super().__init__() 84 self._setup_cmd: TupleCmdArgs = tuple(setup_cmd) if setup_cmd else () 85 self._start_cmd: TupleCmdArgs = tuple(start_cmd) if start_cmd else () 86 self._start_story_run_cmd: TupleCmdArgs = ( 87 tuple(start_story_run_cmd) if start_story_run_cmd else ()) 88 self._stop_story_run_cmd: TupleCmdArgs = ( 89 tuple(stop_story_run_cmd) if stop_story_run_cmd else ()) 90 self._stop_cmd: TupleCmdArgs = tuple(stop_cmd) if stop_cmd else () 91 self._teardown_cmd: TupleCmdArgs = ( 92 tuple(teardown_cmd) if teardown_cmd else ()) 93 94 @property 95 def key(self) -> ProbeKeyT: 96 return super().key + ( 97 ("setup_cmd", tuple(map(str, self.stop_cmd))), 98 ("start_cmd", tuple(map(str, self.start_cmd))), 99 ("start_story_run_cmd", tuple(map(str, self.start_story_run_cmd))), 100 ("stop_story_run_cmd", tuple(map(str, self.stop_story_run_cmd))), 101 ("stop_cmd", tuple(map(str, self.stop_cmd))), 102 ("teardown_cmd", tuple(map(str, self.teardown_cmd))), 103 ) 104 105 @property 106 def setup_cmd(self) -> TupleCmdArgs: 107 return self._setup_cmd 108 109 @property 110 def start_cmd(self) -> TupleCmdArgs: 111 return self._start_cmd 112 113 @property 114 def start_story_run_cmd(self) -> TupleCmdArgs: 115 return self._start_story_run_cmd 116 117 @property 118 def stop_story_run_cmd(self) -> TupleCmdArgs: 119 return self._stop_story_run_cmd 120 121 @property 122 def stop_cmd(self) -> TupleCmdArgs: 123 return self._stop_cmd 124 125 @property 126 def teardown_cmd(self) -> TupleCmdArgs: 127 return self._teardown_cmd 128 129 def validate_env(self, env: HostEnvironment) -> None: 130 super().validate_env(env) 131 if env.repetitions != 1: 132 env.handle_warning(f"Probe={self.NAME} cannot merge data over multiple " 133 f"repetitions={env.repetitions}.") 134 135 def get_context(self, run: Run) -> ShellProbeContext: 136 return ShellProbeContext(self, run) 137 138 139class ShellProbeContext(ProbeContext[ShellProbe]): 140 141 def __init__(self, probe: ShellProbe, run: Run) -> None: 142 super().__init__(probe, run) 143 self._result_files: List[pth.LocalPath] = [] 144 145 def _maybe_run_cmd(self, name: str, cmd: TupleCmdArgs) -> None: 146 if not cmd: 147 return 148 stdout_path = self.local_result_path / f"{name}.stdout.txt" 149 self.host_platform.touch(stdout_path) 150 self._result_files.append(stdout_path) 151 stderr_path = self.local_result_path / f"{name}.stderr.txt" 152 self.host_platform.touch(stderr_path) 153 self._result_files.append(stderr_path) 154 with stdout_path.open("w") as stdout, stderr_path.open("w") as stderr: 155 self.browser_platform.sh(*cmd, shell=True, stdout=stdout, stderr=stderr) 156 157 def setup(self) -> None: 158 self.host_platform.mkdir(self.local_result_path) 159 self._maybe_run_cmd("setup", self.probe.setup_cmd) 160 161 def start(self) -> None: 162 self._maybe_run_cmd("start", self.probe.start_cmd) 163 164 def start_story_run(self) -> None: 165 self._maybe_run_cmd("start_story_run", self.probe.start_story_run_cmd) 166 167 def stop_story_run(self) -> None: 168 self._maybe_run_cmd("stop_story_run", self.probe.stop_story_run_cmd) 169 170 def stop(self) -> None: 171 self._maybe_run_cmd("stop", self.probe.stop_cmd) 172 173 def teardown(self) -> ProbeResult: 174 self._maybe_run_cmd("teardown", self.probe.teardown_cmd) 175 return LocalProbeResult(file=tuple(self._result_files)) 176