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   # Defined in host_forwarder_main.cc
  55   _HOST_FORWARDER_LOG = '/tmp/host_forwarder_log'
  56
  57   _instance = None
  58
  59   @staticmethod
  60   def Map(port_pairs, device, tool=None):
  61     """Runs the forwarder.
  62
  63     Args:
  64       port_pairs: A list of tuples (device_port, host_port) to forward. Note
  65                  that you can specify 0 as a device_port, in which case a
  66                  port will by dynamically assigned on the device. You can
  67                  get the number of the assigned port using the
  68                  DevicePortForHostPort method.
  69       device: A DeviceUtils instance.
  70       tool: Tool class to use to get wrapper, if necessary, for executing the
  71             forwarder (see valgrind_tools.py).
  72
  73     Raises:
  74       Exception on failure to forward the port.
  75     """
  76     # TODO(jbudorick) Remove once telemetry gets switched over.
  77     if isinstance(device, pylib.android_commands.AndroidCommands):
  78       device = pylib.device.device_utils.DeviceUtils(device)
  79     if not tool:
  80       tool = valgrind_tools.CreateTool(None, device)
  81     with _FileLock(Forwarder._LOCK_PATH):
  82       instance = Forwarder._GetInstanceLocked(tool)
  83       instance._InitDeviceLocked(device, tool)
  84
  85       device_serial = str(device)
  86       redirection_commands = [
  87           ['--adb=' + constants.GetAdbPath(),
  88            '--serial-id=' + device_serial,
  89            '--map', str(device_port), str(host_port)]
  90           for device_port, host_port in port_pairs]
  91       logging.info('Forwarding using commands: %s', redirection_commands)
  92
  93       for redirection_command in redirection_commands:
  94         try:
  95           (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
  96               [instance._host_forwarder_path] + redirection_command)
  97         except OSError as e:
  98           if e.errno == 2:
  99             raise Exception('Unable to start host forwarder. Make sure you have'
 100                             ' built host_forwarder.')
 101           else: raise
 102         if exit_code != 0:
 103           Forwarder._KillDeviceLocked(device, tool)
 104           raise Exception('%s exited with %d:\n%s' % (
 105               instance._host_forwarder_path, exit_code, '\n'.join(output)))
 106         tokens = output.split(':')
 107         if len(tokens) != 2:
 108           raise Exception('Unexpected host forwarder output "%s", '
 109                           'expected "device_port:host_port"' % output)
 110         device_port = int(tokens[0])
 111         host_port = int(tokens[1])
 112         serial_with_port = (device_serial, device_port)
 113         instance._device_to_host_port_map[serial_with_port] = host_port
 114         instance._host_to_device_port_map[host_port] = serial_with_port
 115         logging.info('Forwarding device port: %d to host port: %d.',
 116                      device_port, host_port)
 117
 118   @staticmethod
 119   def UnmapDevicePort(device_port, device):
 120     """Unmaps a previously forwarded device port.
 121
 122     Args:
 123       device: A DeviceUtils instance.
 124       device_port: A previously forwarded port (through Map()).
 125     """
 126     # TODO(jbudorick) Remove once telemetry gets switched over.
 127     if isinstance(device, pylib.android_commands.AndroidCommands):
 128       device = pylib.device.device_utils.DeviceUtils(device)
 129     with _FileLock(Forwarder._LOCK_PATH):
 130       Forwarder._UnmapDevicePortLocked(device_port, device)
 131
 132   @staticmethod
 133   def UnmapAllDevicePorts(device):
 134     """Unmaps all the previously forwarded ports for the provided device.
 135
 136     Args:
 137       device: A DeviceUtils instance.
 138       port_pairs: A list of tuples (device_port, host_port) to unmap.
 139     """
 140     # TODO(jbudorick) Remove once telemetry gets switched over.
 141     if isinstance(device, pylib.android_commands.AndroidCommands):
 142       device = pylib.device.device_utils.DeviceUtils(device)
 143     with _FileLock(Forwarder._LOCK_PATH):
 144       if not Forwarder._instance:
 145         return
 146       adb_serial = str(device)
 147       if adb_serial not in Forwarder._instance._initialized_devices:
 148         return
 149       port_map = Forwarder._GetInstanceLocked(
 150           None)._device_to_host_port_map
 151       for (device_serial, device_port) in port_map.keys():
 152         if adb_serial == device_serial:
 153           Forwarder._UnmapDevicePortLocked(device_port, device)
 154       # There are no more ports mapped, kill the device_forwarder.
 155       tool = valgrind_tools.CreateTool(None, device)
 156       Forwarder._KillDeviceLocked(device, tool)
 157
 158   @staticmethod
 159   def DevicePortForHostPort(host_port):
 160     """Returns the device port that corresponds to a given host port."""
 161     with _FileLock(Forwarder._LOCK_PATH):
 162       (_device_serial, device_port) = Forwarder._GetInstanceLocked(
 163           None)._host_to_device_port_map.get(host_port)
 164       return device_port
 165
 166   @staticmethod
 167   def RemoveHostLog():
 168     if os.path.exists(Forwarder._HOST_FORWARDER_LOG):
 169       os.unlink(Forwarder._HOST_FORWARDER_LOG)
 170
 171   @staticmethod
 172   def GetHostLog():
 173     if not os.path.exists(Forwarder._HOST_FORWARDER_LOG):
 174       return ''
 175     with file(Forwarder._HOST_FORWARDER_LOG, 'r') as f:
 176       return f.read()
 177
 178   @staticmethod
 179   def _GetInstanceLocked(tool):
 180     """Returns the singleton instance.
 181
 182     Note that the global lock must be acquired before calling this method.
 183
 184     Args:
 185       tool: Tool class to use to get wrapper, if necessary, for executing the
 186             forwarder (see valgrind_tools.py).
 187     """
 188     if not Forwarder._instance:
 189       Forwarder._instance = Forwarder(tool)
 190     return Forwarder._instance
 191
 192   def __init__(self, tool):
 193     """Constructs a new instance of Forwarder.
 194
 195     Note that Forwarder is a singleton therefore this constructor should be
 196     called only once.
 197
 198     Args:
 199       tool: Tool class to use to get wrapper, if necessary, for executing the
 200             forwarder (see valgrind_tools.py).
 201     """
 202     assert not Forwarder._instance
 203     self._tool = tool
 204     self._initialized_devices = set()
 205     self._device_to_host_port_map = dict()
 206     self._host_to_device_port_map = dict()
 207     self._host_forwarder_path = os.path.join(
 208         constants.GetOutDirectory(), 'host_forwarder')
 209     assert os.path.exists(self._host_forwarder_path), 'Please build forwarder2'
 210     self._device_forwarder_path_on_host = os.path.join(
 211         constants.GetOutDirectory(), 'forwarder_dist')
 212     self._InitHostLocked()
 213
 214   @staticmethod
 215   def _UnmapDevicePortLocked(device_port, device):
 216     """Internal method used by UnmapDevicePort().
 217
 218     Note that the global lock must be acquired before calling this method.
 219     """
 220     instance = Forwarder._GetInstanceLocked(None)
 221     serial = str(device)
 222     serial_with_port = (serial, device_port)
 223     if not serial_with_port in instance._device_to_host_port_map:
 224       logging.error('Trying to unmap non-forwarded port %d' % device_port)
 225       return
 226     redirection_command = ['--adb=' + constants.GetAdbPath(),
 227                            '--serial-id=' + serial,
 228                            '--unmap', str(device_port)]
 229     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 230         [instance._host_forwarder_path] + redirection_command)
 231     if exit_code != 0:
 232       logging.error('%s exited with %d:\n%s' % (
 233           instance._host_forwarder_path, exit_code, '\n'.join(output)))
 234     host_port = instance._device_to_host_port_map[serial_with_port]
 235     del instance._device_to_host_port_map[serial_with_port]
 236     del instance._host_to_device_port_map[host_port]
 237
 238   @staticmethod
 239   def _GetPidForLock():
 240     """Returns the PID used for host_forwarder initialization.
 241
 242     The PID of the "sharder" is used to handle multiprocessing. The "sharder"
 243     is the initial process that forks that is the parent process.
 244     """
 245     return os.getpgrp()
 246
 247   def _InitHostLocked(self):
 248     """Initializes the host forwarder daemon.
 249
 250     Note that the global lock must be acquired before calling this method. This
 251     method kills any existing host_forwarder process that could be stale.
 252     """
 253     # See if the host_forwarder daemon was already initialized by a concurrent
 254     # process or thread (in case multi-process sharding is not used).
 255     pid_for_lock = Forwarder._GetPidForLock()
 256     fd = os.open(Forwarder._LOCK_PATH, os.O_RDWR | os.O_CREAT)
 257     with os.fdopen(fd, 'r+') as pid_file:
 258       pid_with_start_time = pid_file.readline()
 259       if pid_with_start_time:
 260         (pid, process_start_time) = pid_with_start_time.split(':')
 261         if pid == str(pid_for_lock):
 262           if process_start_time == str(_GetProcessStartTime(pid_for_lock)):
 263             return
 264       self._KillHostLocked()
 265       pid_file.seek(0)
 266       pid_file.write(
 267           '%s:%s' % (pid_for_lock, str(_GetProcessStartTime(pid_for_lock))))
 268       pid_file.truncate()
 269
 270   def _InitDeviceLocked(self, device, tool):
 271     """Initializes the device_forwarder daemon for a specific device (once).
 272
 273     Note that the global lock must be acquired before calling this method. This
 274     method kills any existing device_forwarder daemon on the device that could
 275     be stale, pushes the latest version of the daemon (to the device) and starts
 276     it.
 277
 278     Args:
 279       device: A DeviceUtils instance.
 280       tool: Tool class to use to get wrapper, if necessary, for executing the
 281             forwarder (see valgrind_tools.py).
 282     """
 283     device_serial = str(device)
 284     if device_serial in self._initialized_devices:
 285       return
 286     Forwarder._KillDeviceLocked(device, tool)
 287     device.PushChangedFiles([(
 288         self._device_forwarder_path_on_host,
 289         Forwarder._DEVICE_FORWARDER_FOLDER)])
 290     cmd = '%s %s' % (tool.GetUtilWrapper(), Forwarder._DEVICE_FORWARDER_PATH)
 291     device.RunShellCommand(
 292         cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER})
 293     self._initialized_devices.add(device_serial)
 294
 295   def _KillHostLocked(self):
 296     """Kills the forwarder process running on the host.
 297
 298     Note that the global lock must be acquired before calling this method.
 299     """
 300     logging.info('Killing host_forwarder.')
 301     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 302         [self._host_forwarder_path, '--kill-server'])
 303     if exit_code != 0:
 304       (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
 305           ['pkill', '-9', 'host_forwarder'])
 306       if exit_code != 0:
 307         raise Exception('%s exited with %d:\n%s' % (
 308               self._host_forwarder_path, exit_code, '\n'.join(output)))
 309
 310   @staticmethod
 311   def _KillDeviceLocked(device, tool):
 312     """Kills the forwarder process running on the device.
 313
 314     Note that the global lock must be acquired before calling this method.
 315
 316     Args:
 317       device: Instance of DeviceUtils for talking to the device.
 318       tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
 319             forwarder (see valgrind_tools.py).
 320     """
 321     logging.info('Killing device_forwarder.')
 322     Forwarder._instance._initialized_devices.discard(str(device))
 323     if not device.FileExists(Forwarder._DEVICE_FORWARDER_PATH):
 324       return
 325
 326     cmd = '%s %s --kill-server' % (tool.GetUtilWrapper(),
 327                                    Forwarder._DEVICE_FORWARDER_PATH)
 328     device.RunShellCommand(
 329         cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER})