crosperf: use sshwatcher to set up port forwarding to DUTs

Some DUTs (notably those on gTransit) require port forwarding to be
accessible from inside the chroot. This CL makes crosperf's SuitRunner
use sshwatcher to set up the required forwarding.

BUG=b:235119397
TEST=suite_runner_unittest, then try it out for real

Change-Id: Idad2abf40713e2d0aa3b81fdb4762a22b638400e
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/3928354
Reviewed-by: Denis Nikitin <denik@chromium.org>
Reviewed-by: Jordan Abrahams-Whitehead <ajordanr@google.com>
Commit-Queue: Bob Haarman <inglorion@chromium.org>
Tested-by: Bob Haarman <inglorion@chromium.org>

2 files changed, 53 insertions, 3 deletions

diff --git a/crosperf/suite_runner.py b/crosperf/suite_runner.py
index dfd14b21..0e4ba045 100644
--- a/crosperf/suite_runner.py
+++ b/crosperf/suite_runner.py
@@ -6,15 +6,29 @@
 """SuiteRunner defines the interface from crosperf to test script."""
 
 
+import contextlib
 import json
 import os
+from pathlib import Path
 import pipes
+import random
 import shlex
+import subprocess
 import time
 
 from cros_utils import command_executer
 
 
+SSHWATCHER = [
+    "go",
+    "run",
+    str(
+        Path(
+            __file__,
+            "../../../../platform/dev/contrib/sshwatcher/sshwatcher.go",
+        ).resolve()
+    ),
+]
 TEST_THAT_PATH = "/usr/bin/test_that"
 TAST_PATH = "/usr/bin/tast"
 CROSFLEET_PATH = "crosfleet"
@@ -53,6 +67,32 @@ def GetDutConfigArgs(dut_config):
     return f"dut_config={pipes.quote(json.dumps(dut_config))}"
 
 
+@contextlib.contextmanager
+def ssh_tunnel(machinename):
+    """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.
+    """
+    # 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)]
+    # 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)
+    try:
+        yield f"localhost:{port}"
+    finally:
+        proc.terminate()
+        proc.wait(timeout=5)
+
+
 class SuiteRunner(object):
     """This defines the interface from crosperf to test script."""
 
@@ -83,7 +123,8 @@ class SuiteRunner(object):
                 )
             else:
                 if benchmark.suite == "tast":
-                    ret_tup = self.Tast_Run(machine_name, label, benchmark)
+                    with ssh_tunnel(machine_name) as hostport:
+                        ret_tup = self.Tast_Run(hostport, label, benchmark)
                 else:
                     ret_tup = self.Test_That_Run(
                         machine_name, label, benchmark, test_args, profiler_args
diff --git a/crosperf/suite_runner_unittest.py b/crosperf/suite_runner_unittest.py
index 69476f37..cc96ee4a 100755
--- a/crosperf/suite_runner_unittest.py
+++ b/crosperf/suite_runner_unittest.py
@@ -8,6 +8,7 @@
 """Unittest for suite_runner."""
 
 
+import contextlib
 import json
 import unittest
 import unittest.mock as mock
@@ -118,7 +119,14 @@ class SuiteRunnerTest(unittest.TestCase):
         res = suite_runner.GetDutConfigArgs(dut_config)
         self.assertEqual(res, output_str)
 
-    def test_run(self):
+    @mock.patch("suite_runner.ssh_tunnel")
+    def test_run(self, ssh_tunnel):
+        @contextlib.contextmanager
+        def mock_ssh_tunnel(_host):
+            yield "fakelocalhost:1234"
+
+        ssh_tunnel.side_effect = mock_ssh_tunnel
+
         def reset():
             self.test_that_args = []
             self.crosfleet_run_args = []
@@ -254,7 +262,8 @@ class SuiteRunnerTest(unittest.TestCase):
         self.assertFalse(self.call_test_that_run)
         self.assertFalse(self.call_crosfleet_run)
         self.assertEqual(
-            self.tast_args, ["fake_machine", self.mock_label, self.tast_bench]
+            self.tast_args,
+            ["fakelocalhost:1234", self.mock_label, self.tast_bench],
         )
 
     def test_gen_test_args(self):