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