src/build/android/pylib/forwarder.py

   1 # Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 # Use of this source code is governed by a BSD-style license that can be
   3 # found in the LICENSE file.
   4
   5 # pylint: disable=W0212
   6
   7 import fcntl
   8 import logging
   9 import os
  10 import psutil
  11
  12 from pylib import cmd_helper
  13 from pylib import constants
  14 from pylib import valgrind_tools
  15
  16
  17 def _GetProcessStartTime(pid):
  18   return psutil.Process(pid).create_time
  19
  20
  21 class _FileLock(object):
  22   """With statement-aware implementation of a file lock.
  23
  24   File locks are needed for cross-process synchronization when the
  25   multiprocessing Python module is used.
  26   """
  27   def __init__(self, path):
  28     self._fd = -1
  29     self._path = path
  30
  31   def __enter__(self):
  32     self._fd = os.open(self._path, os.O_RDONLY | os.O_CREAT)
  33     if self._fd < 0:
  34       raise Exception('Could not open file %s for reading' % self._path)
  35     fcntl.flock(self._fd, fcntl.LOCK_EX)
  36
  37   def __exit__(self, _exception_type, _exception_value, traceback):
  38     fcntl.flock(self._fd, fcntl.LOCK_UN)
  39     os.close(self._fd)
  40
  41
  42 class Forwarder(object):
  43   """Thread-safe class to manage port forwards from the device to the host."""
  44
  45   _DEVICE_FORWARDER_FOLDER = (constants.TEST_EXECUTABLE_DIR +
  46                               '/forwarder/')
  47   _DEVICE_FORWARDER_PATH = (constants.TEST_EXECUTABLE_DIR +
  48                             '/forwarder/device_forwarder')
  49   _LD_LIBRARY_PATH = 'LD_LIBRARY_PATH=%s' % _DEVICE_FORWARDER_FOLDER
  50   _LOCK_PATH = '/tmp/chrome.forwarder.lock'
  51   _MULTIPROCESSING_ENV_VAR = 'CHROME_FORWARDER_USE_MULTIPROCESSING'
  52   # Defined in host_forwarder_main.cc
  53   _HOST_FORWARDER_LOG = '/tmp/host_forwarder_log'
  54
  55   _instance = None
  56
  57   @staticmethod
  58   def UseMultiprocessing():
  59     """Tells the forwarder that multiprocessing is used."""
  60     os.environ[Forwarder._MULTIPROCESSING_ENV_VAR] = '1'
  61
  62   @staticmethod
  63   def Map(port_pairs, adb, tool=None):
  64     """Runs the forwarder.
  65
  66     Args:
  67       port_pairs: A list of tuples (device_port, host_port) to forward. Note
  68                  that you can specify 0 as a device_port, in which case a
  69                  port will by dynamically assigned on the device. You can
  70                  get the number of the assigned port using the
  71                  DevicePortForHostPort method.
  72       adb: An AndroidCommands instance.
  73       tool: Tool class to use to get wrapper, if necessary, for executing the
  74             forwarder (see valgrind_tools.py).
  75
  76     Raises:
  77       Exception on failure to forward the port.
  78     """
  79     if not tool:
  80       tool = valgrind_tools.CreateTool(None, adb)
  81     with _FileLock(Forwarder._LOCK_PATH):
  82       instance = Forwarder._GetInstanceLocked(tool)
  83       instance._InitDeviceLocked(adb, tool)
  84
  85       device_serial = adb.Adb().GetSerialNumber()
  86       redirection_commands = [
  87           ['--serial-id=' + device_serial, '--map', str(device),
  88            str(host)] for device, host in port_pairs]
  89       logging.info('Forwarding using commands: %s', redirection_commands)
  90
  91       for redirection_command in redirection_commands:
  92         try:
  93           (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
  94               [instance._host_forwarder_path] + redirection_command)
  95         except OSError as e:
  96           if e.errno == 2:
  97             raise Exception('Unable to start host forwarder. Make sure you have'
  98                             ' built host_forwarder.')
  99           else: raise
 100         if exit_code != 0:
 101           raise Exception('%s exited with %d:\n%s' % (
 102               instance._host_forwarder_path, exit_code, '\n'.join(output)))
 103         tokens = output.split(':')
 104         if len(tokens) != 2:
 105           raise Exception(('Unexpected host forwarder output "%s", ' +
 106                           'expected "device_port:host_port"') % output)
 107         device_port = int(tokens[0])
 108         host_port = int(tokens[1])
 109         serial_with_port = (device_serial, device_port)
 110         instance._device_to_host_port_map[serial_with_port] = host_port
 111         instance._host_to_device_port_map[host_port] = serial_with_port
 112         logging.info('Forwarding device port: %d to host port: %d.',
 113                      device_port, host_port)
 114
 115   @staticmethod
 116   def UnmapDevicePort(device_port, adb):
 117     """Unmaps a previously forwarded device port.
 118
 119     Args:
 120       adb: An AndroidCommands instance.
 121       device_port: A previously forwarded port (through Map()).
 122     """
 123     with _FileLock(Forwarder._LOCK_PATH):
 124       Forwarder._UnmapDevicePortLocked(device_port, adb)
 125
 126   @staticmethod
 127   def UnmapAllDevicePorts(adb):
 128     """Unmaps all the previously forwarded ports for the provided device.
 129
 130     Args:
 131       adb: An AndroidCommands instance.
 132       port_pairs: A list of tuples (device_port, host_port) to unmap.
 133     """
 134     with _FileLock(Forwarder._LOCK_PATH):
 135       if not Forwarder._instance:
 136         return
 137       adb_serial = adb.Adb().GetSerialNumber()
 138       if adb_serial not in Forwarder._instance._initialized_devices:
 139         return
 140       port_map = Forwarder._GetInstanceLocked(
 141           None)._device_to_host_port_map
 142       for (device_serial, device_port) in port_map.keys():
 143         if adb_serial == device_serial:
 144           Forwarder._UnmapDevicePortLocked(device_port, adb)
 145       # There are no more ports mapped, kill the device_forwarder.
 146       tool = valgrind_tools.CreateTool(None, adb)
 147       Forwarder._KillDeviceLocked(adb, tool)
 148       Forwarder._instance._initialized_devices.remove(adb_serial)
 149
 150
 151   @staticmethod
 152   def DevicePortForHostPort(host_port):
 153     """Returns the device port that corresponds to a given host port."""
 154     with _FileLock(Forwarder._LOCK_PATH):
 155       (_device_serial, device_port) = Forwarder._GetInstanceLocked(
 156           None)._host_to_device_port_map.get(host_port)
 157       return device_port
 158
 159   @staticmethod
 160   def RemoveHostLog():
 161     if os.path.exists(Forwarder._HOST_FORWARDER_LOG):
 162       os.unlink(Forwarder._HOST_FORWARDER_LOG)
 163
 164   @staticmethod
 165   def GetHostLog():
 166     if not os.path.exists(Forwarder._HOST_FORWARDER_LOG):
 167       return ''
 168     with file(Forwarder._HOST_FORWARDER_LOG, 'r') as f:
 169       return f.read()
 170
 171   @staticmethod
 172   def _GetInstanceLocked(tool):
 173     """Returns the singleton instance.
 174
 175     Note that the global lock must be acquired before calling this method.
 176
 177     Args:
 178       tool: Tool class to use to get wrapper, if necessary, for executing the
 179             forwarder (see valgrind_tools.py).
 180     """
 181     if not Forwarder._instance:
 182       Forwarder._instance = Forwarder(tool)
 183     return Forwarder._instance
 184
 185   def __init__(self, tool):
 186     """Constructs a new instance of Forwarder.
 187
 188     Note that Forwarder is a singleton therefore this constructor should be
 189     called only once.
 190
 191     Args:
 192       tool: Tool class to use to get wrapper, if necessary, for executing the
 193             forwarder (see valgrind_tools.py).
 194     """
 195     assert not Forwarder._instance
 196     self._tool = tool
 197     self._initialized_devices = set()
 198     self._device_to_host_port_map = dict()
 199     self._host_to_device_port_map = dict()
 200     self._host_forwarder_path = os.path.join(
 201         constants.GetOutDirectory(), 'host_forwarder')
 202     assert os.path.exists(self._host_forwarder_path), 'Please build forwarder2'
 203     self._device_forwarder_path_on_host = os.path.join(
 204         constants.GetOutDirectory(), 'forwarder_dist')
 205     self._InitHostLocked()
 206
 207   @staticmethod
 208   def _UnmapDevicePortLocked(device_port, adb):
 209     """Internal method used by UnmapDevicePort().
 210
 211     Note that the global lock must be acquired before calling this method.
 212     """
 213     instance = Forwarder._GetInstanceLocked(None)
 214     serial = adb.Adb().GetSerialNumber()
 215     serial_with_port = (serial, device_port)
 216     if not serial_with_port in instance._device_to_host_port_map:
 217       logging.error('Trying to unmap non-forwarded port %d' % device_port)
 218       return
 219     redirection_command = ['--serial-id=' + serial, '--unmap', str(device_port)]
 220     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 221         [instance._host_forwarder_path] + redirection_command)
 222     if exit_code != 0:
 223       logging.error('%s exited with %d:\n%s' % (
 224           instance._host_forwarder_path, exit_code, '\n'.join(output)))
 225     host_port = instance._device_to_host_port_map[serial_with_port]
 226     del instance._device_to_host_port_map[serial_with_port]
 227     del instance._host_to_device_port_map[host_port]
 228
 229   @staticmethod
 230   def _GetPidForLock():
 231     """Returns the PID used for host_forwarder initialization.
 232
 233     In case multi-process sharding is used, the PID of the "sharder" is used.
 234     The "sharder" is the initial process that forks that is the parent process.
 235     By default, multi-processing is not used. In that case the PID of the
 236     current process is returned.
 237     """
 238     use_multiprocessing = Forwarder._MULTIPROCESSING_ENV_VAR in os.environ
 239     return os.getppid() if use_multiprocessing else os.getpid()
 240
 241   def _InitHostLocked(self):
 242     """Initializes the host forwarder daemon.
 243
 244     Note that the global lock must be acquired before calling this method. This
 245     method kills any existing host_forwarder process that could be stale.
 246     """
 247     # See if the host_forwarder daemon was already initialized by a concurrent
 248     # process or thread (in case multi-process sharding is not used).
 249     pid_for_lock = Forwarder._GetPidForLock()
 250     fd = os.open(Forwarder._LOCK_PATH, os.O_RDWR | os.O_CREAT)
 251     with os.fdopen(fd, 'r+') as pid_file:
 252       pid_with_start_time = pid_file.readline()
 253       if pid_with_start_time:
 254         (pid, process_start_time) = pid_with_start_time.split(':')
 255         if pid == str(pid_for_lock):
 256           if process_start_time == str(_GetProcessStartTime(pid_for_lock)):
 257             return
 258       self._KillHostLocked()
 259       pid_file.seek(0)
 260       pid_file.write(
 261           '%s:%s' % (pid_for_lock, str(_GetProcessStartTime(pid_for_lock))))
 262
 263   def _InitDeviceLocked(self, adb, tool):
 264     """Initializes the device_forwarder daemon for a specific device (once).
 265
 266     Note that the global lock must be acquired before calling this method. This
 267     method kills any existing device_forwarder daemon on the device that could
 268     be stale, pushes the latest version of the daemon (to the device) and starts
 269     it.
 270
 271     Args:
 272       adb: An AndroidCommands instance.
 273       tool: Tool class to use to get wrapper, if necessary, for executing the
 274             forwarder (see valgrind_tools.py).
 275     """
 276     device_serial = adb.Adb().GetSerialNumber()
 277     if device_serial in self._initialized_devices:
 278       return
 279     Forwarder._KillDeviceLocked(adb, tool)
 280     adb.PushIfNeeded(
 281         self._device_forwarder_path_on_host,
 282         Forwarder._DEVICE_FORWARDER_FOLDER)
 283     (exit_code, output) = adb.GetShellCommandStatusAndOutput(
 284         '%s %s %s' % (Forwarder._LD_LIBRARY_PATH, tool.GetUtilWrapper(),
 285                       Forwarder._DEVICE_FORWARDER_PATH))
 286     if exit_code != 0:
 287       raise Exception(
 288           'Failed to start device forwarder:\n%s' % '\n'.join(output))
 289     self._initialized_devices.add(device_serial)
 290
 291   def _KillHostLocked(self):
 292     """Kills the forwarder process running on the host.
 293
 294     Note that the global lock must be acquired before calling this method.
 295     """
 296     logging.info('Killing host_forwarder.')
 297     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 298         [self._host_forwarder_path, '--kill-server'])
 299     if exit_code != 0:
 300       (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 301           ['pkill', '-9', 'host_forwarder'])
 302       if exit_code != 0:
 303         raise Exception('%s exited with %d:\n%s' % (
 304               self._host_forwarder_path, exit_code, '\n'.join(output)))
 305
 306   @staticmethod
 307   def _KillDeviceLocked(adb, tool):
 308     """Kills the forwarder process running on the device.
 309
 310     Note that the global lock must be acquired before calling this method.
 311
 312     Args:
 313       adb: Instance of AndroidCommands for talking to the device.
 314       tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
 315             forwarder (see valgrind_tools.py).
 316     """
 317     logging.info('Killing device_forwarder.')
 318     if not adb.FileExistsOnDevice(Forwarder._DEVICE_FORWARDER_PATH):
 319       return
 320     adb.GetShellCommandStatusAndOutput(
 321         '%s %s --kill-server' % (tool.GetUtilWrapper(),
 322                                  Forwarder._DEVICE_FORWARDER_PATH))
 323     # TODO(pliard): Remove the following call to KillAllBlocking() when we are
 324     # sure that the old version of device_forwarder (not supporting
 325     # 'kill-server') is not running on the bots anymore.
 326     timeout_sec = 5
 327     processes_killed = adb.KillAllBlocking('device_forwarder', timeout_sec)
 328     if not processes_killed:
 329       pids = adb.ExtractPid('device_forwarder')
 330       if pids:
 331         raise Exception('Timed out while killing device_forwarder')