crosperf: find sshwatcher relative to chromeos_root

The root of the ChromiumOS source tree is not always in the same
location relative to where the crosperf scripts are. This makes
the logic for finding sshwatcher.go not work everywhere. To fix
this, this change uses the label.chromeos_root path that is
present in crosperf runs to find sshwatcher.go.

BUG=b:235119397
TEST=run crosperf in standalone toolchain-utils as well as
  inside ChromiumOS tree to see that sshwatcher.go is found
  in both cases.

Change-Id: I981a55d3cfd2638da0bf320b86ebc5601101ce3f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/3935652
Tested-by: Bob Haarman <inglorion@chromium.org>
Reviewed-by: Denis Nikitin <denik@chromium.org>
Commit-Queue: Bob Haarman <inglorion@chromium.org>

2 files changed, 19 insertions, 16 deletions

diff --git a/crosperf/suite_runner.py b/crosperf/suite_runner.py
index 0e4ba045..e777a57f 100644
--- a/crosperf/suite_runner.py
+++ b/crosperf/suite_runner.py
@@ -19,16 +19,8 @@ import time
 from cros_utils import command_executer
 
 
-SSHWATCHER = [
-    "go",
-    "run",
-    str(
-        Path(
-            __file__,
-            "../../../../platform/dev/contrib/sshwatcher/sshwatcher.go",
-        ).resolve()
-    ),
-]
+# sshwatcher path, relative to ChromiumOS source root.
+SSHWATCHER = "src/platform/dev/contrib/sshwatcher/sshwatcher.go"
 TEST_THAT_PATH = "/usr/bin/test_that"
 TAST_PATH = "/usr/bin/tast"
 CROSFLEET_PATH = "crosfleet"
@@ -68,25 +60,34 @@ def GetDutConfigArgs(dut_config):
 
 
 @contextlib.contextmanager
-def ssh_tunnel(machinename):
+def ssh_tunnel(sshwatcher: "os.PathLike", machinename: str) -> str:
     """Context manager that forwards a TCP port over SSH while active.
 
     This class is used to set up port forwarding before entering the
     chroot, so that the forwarded port can be used from inside
     the chroot.
 
-    The value yielded by ssh_tunnel is a host:port string.
+    Args:
+        sshwatcher: Path to sshwatcher.go
+        machinename: Hostname of the machine to connect to.
+
+    Returns:
+        host:port string that can be passed to tast
     """
     # We have to tell sshwatcher which port we want to use.
     # We pick a port that is likely to be available.
     port = random.randrange(4096, 32768)
-    cmd = SSHWATCHER + [machinename, str(port)]
+    cmd = ["go", "run", str(sshwatcher), machinename, str(port)]
     # Pylint wants us to use subprocess.Popen as a context manager,
     # but we don't, so that we can ask sshwatcher to terminate and
     # limit the time we wait for it to do so.
     # pylint: disable=consider-using-with
-    proc = subprocess.Popen(cmd)
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE)
     try:
+        # sshwatcher takes a few seconds before it binds to the port,
+        # presumably due to SSH handshaking taking a while.
+        # Give it 12 seconds before we ask the client to connect.
+        time.sleep(12)
         yield f"localhost:{port}"
     finally:
         proc.terminate()
@@ -123,7 +124,9 @@ class SuiteRunner(object):
                 )
             else:
                 if benchmark.suite == "tast":
-                    with ssh_tunnel(machine_name) as hostport:
+                    with ssh_tunnel(
+                        Path(label.chromeos_root, SSHWATCHER), machine_name
+                    ) as hostport:
                         ret_tup = self.Tast_Run(hostport, label, benchmark)
                 else:
                     ret_tup = self.Test_That_Run(
diff --git a/crosperf/suite_runner_unittest.py b/crosperf/suite_runner_unittest.py
index cc96ee4a..c936a074 100755
--- a/crosperf/suite_runner_unittest.py
+++ b/crosperf/suite_runner_unittest.py
@@ -122,7 +122,7 @@ class SuiteRunnerTest(unittest.TestCase):
     @mock.patch("suite_runner.ssh_tunnel")
     def test_run(self, ssh_tunnel):
         @contextlib.contextmanager
-        def mock_ssh_tunnel(_host):
+        def mock_ssh_tunnel(_watcher, _host):
             yield "fakelocalhost:1234"
 
         ssh_tunnel.side_effect = mock_ssh_tunnel