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