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