Improve forwarder2 setup/tear down in telemetry.

This CL kills the host_forwarder daemon before running the tests and also makes
Forwarder.Close() unmap the previously mapped ports.

BUG=242846
R=bulach@chromium.org

Review URL: https://codereview.chromium.org/18694002

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@210268 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 29 insertions, 14 deletions

diff --git a/build/android/bb_run_sharded_steps.py b/build/android/bb_run_sharded_steps.py
index 7a06155..afd354a 100755
--- a/build/android/bb_run_sharded_steps.py
+++ b/build/android/bb_run_sharded_steps.py
@@ -58,6 +58,7 @@ import sys
 from pylib import android_commands
 from pylib import cmd_helper
 from pylib import constants
+from pylib import forwarder
 from pylib import ports
 
 
@@ -164,6 +165,7 @@ def _PrintAllStepsOutput(steps):
 
 def _KillPendingServers():
   for retry in range(5):
+    forwarder.Forwarder.KillHost()
     for server in ['lighttpd', 'web-page-replay']:
       pids = cmd_helper.GetCmdOutput(['pgrep', '-f', server])
       pids = [pid.strip() for pid in pids.split('\n') if pid.strip()]
diff --git a/build/android/pylib/forwarder.py b/build/android/pylib/forwarder.py
index 40a5d36..ff1d710 100644
--- a/build/android/pylib/forwarder.py
+++ b/build/android/pylib/forwarder.py
@@ -45,6 +45,7 @@ class Forwarder(object):
     """
     assert build_type in ('Release', 'Debug')
     self._adb = adb
+    self._device_to_host_port_map = dict()
     self._host_to_device_port_map = dict()
     self._device_initialized = False
     self._host_adb_control_port = 0
@@ -97,6 +98,7 @@ class Forwarder(object):
                           'expected "device_port:host_port"' % output)
         device_port = int(tokens[0])
         host_port = int(tokens[1])
+        self._device_to_host_port_map[device_port] = host_port
         self._host_to_device_port_map[host_port] = device_port
         logging.info('Forwarding device port: %d to host port: %d.',
                      device_port, host_port)
@@ -137,24 +139,35 @@ class Forwarder(object):
       device_port: A previously forwarded port (through Run()).
     """
     with self._lock:
-      # Please note the minus sign below.
-      redirection_command = '%d:-%d' % (
-          self._host_adb_control_port, device_port)
-      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-          [self._host_forwarder_path, redirection_command])
-      if exit_code != 0:
-        raise Exception('%s exited with %d:\n%s' % (
-            self._host_forwarder_path, exit_code, '\n'.join(output)))
+      self._UnmapDevicePortInternalLocked(device_port)
+
+  def _UnmapDevicePortInternalLocked(self, device_port):
+    if not device_port in self._device_to_host_port_map:
+      return
+    # Please note the minus sign below.
+    redirection_command = '%d:-%d' % (
+        self._host_adb_control_port, device_port)
+    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
+        [self._host_forwarder_path, redirection_command])
+    if exit_code != 0:
+      raise Exception('%s exited with %d:\n%s' % (
+          self._host_forwarder_path, exit_code, '\n'.join(output)))
+    host_port = self._device_to_host_port_map[device_port]
+    del self._device_to_host_port_map[device_port]
+    del self._host_to_device_port_map[host_port]
 
   @staticmethod
-  def KillHost(build_type):
+  def KillHost(build_type='Debug'):
     """Kills the forwarder process running on the host.
 
     Args:
-      build_type: 'Release' or 'Debug'
+      build_type: 'Release' or 'Debug' (default='Debug')
     """
     logging.info('Killing host_forwarder.')
     host_forwarder_path = _MakeBinaryPath(build_type, 'host_forwarder')
+    if not os.path.exists(host_forwarder_path):
+      host_forwarder_path = _MakeBinaryPath(
+          'Release' if build_type == 'Debug' else 'Debug', 'host_forwarder')
     assert os.path.exists(host_forwarder_path), 'Please build forwarder2'
     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
         [host_forwarder_path, 'kill-server'])
@@ -195,8 +208,8 @@ class Forwarder(object):
     with self._lock:
       return self._host_to_device_port_map.get(host_port)
 
-  # Deprecated.
   def Close(self):
-    """Terminates the forwarder process."""
-    # TODO(pliard): Remove references in client code.
-    pass
+    """Releases the previously forwarded ports."""
+    with self._lock:
+      for device_port in self._device_to_host_port_map.copy():
+        self._UnmapDevicePortInternalLocked(device_port)