Reland r210268: "Improve forwarder2 setup/tear down in telemetry."

build/android/pylib/forwarder.py

 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import logging
 import os
 import re
 import sys
 import threading
 import time
@@ skipping 22 matching lines @@
     """Forwards TCP ports on the device back to the host.
 
     Works like adb forward, but in reverse.
 
     Args:
       adb: Instance of AndroidCommands for talking to the device.
       build_type: 'Release' or 'Debug'.
     """
     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
     self._lock = threading.Lock()
     self._host_forwarder_path = _MakeBinaryPath(build_type, 'host_forwarder')
     self._device_forwarder_path_on_host = os.path.join(
         cmd_helper.OutDirectory.get(), build_type, 'forwarder_dist')
 
   def Run(self, port_pairs, tool):
     """Runs the forwarder.
@@ skipping 29 matching lines @@
           else: raise
         if exit_code != 0:
           raise Exception('%s exited with %d:\n%s' % (
               self._host_forwarder_path, exit_code, '\n'.join(output)))
         tokens = output.split(':')
         if len(tokens) != 2:
           raise Exception(('Unexpected host forwarder output "%s", ' +
                           '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)
 
   def _InitDeviceLocked(self, tool):
     """Initializes the device forwarder process (only once)."""
     if self._device_initialized:
       return
     self._adb.PushIfNeeded(
         self._device_forwarder_path_on_host,
         Forwarder._DEVICE_FORWARDER_FOLDER)
     (exit_code, output) = self._adb.GetShellCommandStatusAndOutput(
         '%s %s %s' % (Forwarder._LD_LIBRARY_PATH, tool.GetUtilWrapper(),
                       Forwarder._DEVICE_FORWARDER_PATH))
     if exit_code != 0:
       raise Exception(
           'Failed to start device forwarder:\n%s' % '\n'.join(output))
     self._device_initialized = True
 
   def UnmapDevicePort(self, device_port):
     """Unmaps a previously forwarded device port.
 
     Args:
       device_port: A previously forwarded port (through Run()).
     """
     with self._lock:
-      redirection_command = [
-          '--serial-id=' + self._adb.Adb().GetSerialNumber(), '--unmap',
-          str(device_port)]
-      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-          [self._host_forwarder_path] + redirection_command)
-      if exit_code != 0:
-        logging.error('%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
+    redirection_command = [
+        '--serial-id=' + self._adb.Adb().GetSerialNumber(), '--unmap',
+        str(device_port)]
+    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
+        [self._host_forwarder_path] + redirection_command)
+    if exit_code != 0:
+      logging.error('%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'])
     if exit_code != 0:
       (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
           ['pkill', 'host_forwarder'])
       if exit_code != 0:
         raise Exception('%s exited with %d:\n%s' % (
               host_forwarder_path, exit_code, '\n'.join(output)))
 
@@ skipping 20 matching lines @@
     if not processes_killed:
       pids = adb.ExtractPid('device_forwarder')
       if pids:
         raise Exception('Timed out while killing device_forwarder')
 
   def DevicePortForHostPort(self, host_port):
     """Returns the device port that corresponds to a given host port."""
     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)