Move Python setup/tear down logic into Forwarder itself.

build/android/pylib/forwarder.py

Index: build/android/pylib/forwarder.py
diff --git a/build/android/pylib/forwarder.py b/build/android/pylib/forwarder.py
index 5bdb1cd4ce92b2598e9cff38295f19fd3950c572..6646f7c44611de16bd29ba8d2118616841603252 100644
--- a/build/android/pylib/forwarder.py
+++ b/build/android/pylib/forwarder.py
@@ -2,24 +2,46 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import fcntl
 import logging
 import os
 import re
 import sys
-import threading
 import time
 
 import android_commands
 import cmd_helper
 import constants
 
-from pylib import pexpect
+from pylib import valgrind_tools
 
 
 def _MakeBinaryPath(build_type, binary_name):
   return os.path.join(cmd_helper.OutDirectory.get(), build_type, binary_name)
 
 
+def _GetProcessStartTime(pid):
+  return open('/proc/%s/stat' % pid, 'r').readline().split(' ')[21]

[frankf, 2013/07/15 22:23:22]

Can you use psutil: https://code.google.com/p/psutil/wiki/Documentation

[Philippe, 2013/07/16 12:31:21]

Yes, thanks. I thought that this was a third-party library.

+
+
+class _FileLock(object):
+  """With statement-aware implementation of a file lock.
+
+  File locks are needed for cross-process synchronization when the
+  multiprocessing Python module is used.
+  """
+  def __init__(self, path):
+    self._path = path
+
+  def __enter__(self):
+    self._file = open(self._path, 'a+')
+    fcntl.flock(self._file, fcntl.LOCK_EX)
+
+  def __exit__(self, type, value, traceback):
+    fcntl.flock(self._file, fcntl.LOCK_UN)
+    self._file.close()
+
+
 class Forwarder(object):
   """Thread-safe class to manage port forwards from the device to the host."""
 
@@ -28,28 +50,19 @@ class Forwarder(object):
   _DEVICE_FORWARDER_PATH = (constants.TEST_EXECUTABLE_DIR +
                             '/forwarder/device_forwarder')
   _LD_LIBRARY_PATH = 'LD_LIBRARY_PATH=%s' % _DEVICE_FORWARDER_FOLDER
+  _LOCK_PATH = '/tmp/chrome.forwarder.lock'
+  _MULTIPROCESSING_ENV_VAR = 'CHROME_FORWARDER_USE_MULTIPROCESSING'
 
-  def __init__(self, adb, build_type):
-    """Forwards TCP ports on the device back to the host.
-
-    Works like adb forward, but in reverse.
+  _lock = _FileLock(_LOCK_PATH)
+  _instance = None
 
-    Args:
-      adb: Instance of AndroidCommands for talking to the device.
-      build_type: 'Release' or 'Debug'.
-    """
-    assert build_type in ('Release', 'Debug')
-    self._adb = adb
-    self._device_to_host_port_map = dict()
-    self._host_to_device_port_map = dict()
-    self._device_initialized = False
-    self._host_adb_control_port = 0
-    self._lock = threading.Lock()
-    self._host_forwarder_path = _MakeBinaryPath(build_type, 'host_forwarder')
-    self._device_forwarder_path_on_host = os.path.join(
-        cmd_helper.OutDirectory.get(), build_type, 'forwarder_dist')
+  @staticmethod
+  def use_multiprocessing():
+    """Tells the forwarder that multiprocessing is used."""
+    os.environ[Forwarder._MULTIPROCESSING_ENV_VAR] = '1'
 
-  def Run(self, port_pairs, tool):
+  @staticmethod
+  def Map(port_pairs, adb, build_type='Debug', tool=None):
     """Runs the forwarder.
 
     Args:
@@ -58,24 +71,28 @@ class Forwarder(object):
                  port will by dynamically assigned on the device. You can
                  get the number of the assigned port using the
                  DevicePortForHostPort method.
+      adb: An AndroidCommands instance.
       tool: Tool class to use to get wrapper, if necessary, for executing the
             forwarder (see valgrind_tools.py).
 
     Raises:
       Exception on failure to forward the port.
     """
-    with self._lock:
-      self._InitDeviceLocked(tool)
-      host_name = '127.0.0.1'
+    if not tool:
+      tool = valgrind_tools.CreateTool(None, adb)
+    with Forwarder._lock:
+      instance = Forwarder._GetInstanceLocked(build_type, tool)
+      instance._InitDeviceOnceLocked(adb, tool)
+
       redirection_commands = [
-          ['--serial-id=' + self._adb.Adb().GetSerialNumber(), '--map',
+          ['--serial-id=' + adb.Adb().GetSerialNumber(), '--map',
            str(device), str(host)] for device, host in port_pairs]
       logging.info('Forwarding using commands: %s', redirection_commands)
 
       for redirection_command in redirection_commands:
         try:
           (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-              [self._host_forwarder_path] + redirection_command)
+              [instance._host_forwarder_path] + redirection_command)
         except OSError as e:
           if e.errno == 2:
             raise Exception('Unable to start host forwarder. Make sure you have'
@@ -83,83 +100,189 @@ class Forwarder(object):
           else: raise
         if exit_code != 0:
           raise Exception('%s exited with %d:\n%s' % (
-              self._host_forwarder_path, exit_code, '\n'.join(output)))
+              instance._host_forwarder_path, exit_code, '\n'.join(output)))
         tokens = output.split(':')
         if len(tokens) != 2:
           raise Exception(('Unexpected host forwarder output "%s", ' +
                           'expected "device_port:host_port"') % output)
         device_port = int(tokens[0])
         host_port = int(tokens[1])
-        self._device_to_host_port_map[device_port] = host_port
-        self._host_to_device_port_map[host_port] = device_port
+        instance._device_to_host_port_map[device_port] = host_port
+        instance._host_to_device_port_map[host_port] = device_port
         logging.info('Forwarding device port: %d to host port: %d.',
                      device_port, host_port)
 
-  def _InitDeviceLocked(self, tool):
-    """Initializes the device forwarder process (only once)."""
-    if self._device_initialized:
+  @staticmethod
+  def UnmapDevicePort(device_port, adb):
+    """Unmaps a previously forwarded device port.
+
+    Args:
+      adb: An AndroidCommands instance.
+      device_port: A previously forwarded port (through Map()).
+    """
+    with Forwarder._lock:
+      instance = Forwarder._GetInstanceLocked(None, None)
+      if not device_port in instance._device_to_host_port_map:
+        return
+      redirection_command = [
+          '--serial-id=' + adb.Adb().GetSerialNumber(), '--unmap',
+          str(device_port)]
+      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
+          [instance._host_forwarder_path] + redirection_command)
+      if exit_code != 0:
+        logging.error('%s exited with %d:\n%s' % (
+            instance._host_forwarder_path, exit_code, '\n'.join(output)))
+      host_port = instance._device_to_host_port_map[device_port]
+      del instance._device_to_host_port_map[device_port]
+      del instance._host_to_device_port_map[host_port]
+
+  @staticmethod
+  def UnmapPorts(port_pairs, adb):

[frankf, 2013/07/15 22:23:22]

To clarify my earlier comment: In many cases, you just want to unmap all the
ports, so it makes sense to have a method that doesn't take the port as a
parmater, and just uses the _device_to_host_port_map dict.

[Philippe, 2013/07/16 12:31:21]

Done.

+    """Unmaps previously forwarded ports.
+
+    Args:
+      adb: An AndroidCommands instance.
+      port_pairs: A list of tuples (device_port, host_port) to unmap.
+    """
+    for (device_port, _) in port_pairs:
+      Forwarder.UnmapDevicePort(device_port, adb)
+
+  @staticmethod
+  def DevicePortForHostPort(host_port):
+    """Returns the device port that corresponds to a given host port."""
+    with Forwarder._lock:
+      return Forwarder._GetInstanceLocked(
+          None, None)._host_to_device_port_map.get(host_port)
+
+  @staticmethod
+  def _GetInstanceLocked(build_type, tool):
+    """Returns the singleton instance.
+
+    Note that the global lock must be acquired before calling this method.
+
+    Args:
+      build_type: 'Release' or 'Debug'
+      tool: Tool class to use to get wrapper, if necessary, for executing the
+            forwarder (see valgrind_tools.py).
+    """
+    if not Forwarder._instance:
+      Forwarder._instance = Forwarder(build_type, tool)
+    return Forwarder._instance
+
+  def __init__(self, build_type, tool):
+    """Constructs a new instance of Forwarder.
+
+    Note that Forwarder is a singleton therefore this constructor should be
+    called only once.
+
+    Args:
+      build_type: 'Release' or 'Debug'
+      tool: Tool class to use to get wrapper, if necessary, for executing the
+            forwarder (see valgrind_tools.py).
+    """
+    assert not Forwarder._instance
+    self._build_type = build_type
+    self._tool = tool
+    self._initialized_devices = set()
+    self._device_to_host_port_map = dict()
+    self._host_to_device_port_map = dict()
+    self._InitHostOnceLocked()

[bulach, 2013/07/15 17:51:43]

nit: maybe move 189 below, i.e., do all member initialization then start calling
functions?

[Philippe, 2013/07/16 12:31:21]

Good idea.

+    self._device_forwarder_path_on_host = os.path.join(
+        cmd_helper.OutDirectory.get(), self._build_type, 'forwarder_dist')
+
+  @staticmethod
+  def _GetPid():

[bulach, 2013/07/15 17:51:43]

nit: maybe GetPidForLock() ?

[Philippe, 2013/07/16 12:31:21]

Done.

+    """Returns the PID used for host_forwarder initialization.
+
+    In case multi-process sharding is used, the PID of the "sharder" is used.
+    The "sharder" is the initial process that forks that is the parent process.
+    By default, multi-processing is not used. In that case the PID of the
+    current process is returned.
+    """
+    use_multiprocessing = Forwarder._MULTIPROCESSING_ENV_VAR in os.environ
+    return os.getppid() if use_multiprocessing else os.getpid()
+
+  def _InitHostOnceLocked(self):

[frankf, 2013/07/15 22:23:22]

To be consistent remove the 'Once' from method names.

[Philippe, 2013/07/16 12:31:21]

Done. Note that I meant 'only one time' by 'once' rather than 'after the lock
was acquired' :) But I removed it anyway since it was ambiguous.

+    """Initializes the host forwarder daemon.
+
+    Note that the global lock must be acquired before calling this method. This
+    method determines the host_forwarder binary location and kills any existing
+    host_forwarder process that could be stale.
+    """
+    # See if the host_forwarder daemon was already initialized by a concurrent
+    # process or thread (in case multi-process sharding is not used).
+    current_pid = str(Forwarder._GetPid())

[bulach, 2013/07/15 17:51:43]

nit: maybe pid_for_lock? (current_pid is a big confusing..)

[Philippe, 2013/07/16 12:31:21]

Done.

+    with open(Forwarder._LOCK_PATH, 'r') as pid_file:

[bulach, 2013/07/15 17:51:43]

w+ ?

[Philippe, 2013/07/16 12:31:21]

I had a fun debugging session with these file modes :) 'w+' would overwrite the
file if it already exists which would remove its content. I used the open
syscall directly with explicit flags and then passed the file descriptor to
fdopen.

+      pid_with_start_time = pid_file.readline()
+      if pid_with_start_time:
+        (pid, process_start_time) = pid_with_start_time.split(':')

[frankf, 2013/07/15 22:23:22]

where's process_start_time used?

[Philippe, 2013/07/16 12:31:21]

Great catch! This was a bug.

+        if pid == current_pid:
+          if _GetProcessStartTime(pid) == _GetProcessStartTime(current_pid):
+            return
+      self._host_forwarder_path = _MakeBinaryPath(
+          self._build_type, 'host_forwarder')
+      if not os.path.exists(self._host_forwarder_path):
+        self._build_type = 'Release' if self._build_type == 'Debug' else 'Debug'
+        self._host_forwarder_path = _MakeBinaryPath(
+            self._build_type, 'host_forwarder')
+        assert os.path.exists(
+            self._host_forwarder_path), 'Please build forwarder2'

[bulach, 2013/07/15 17:51:43]

nit: I suppose 222-229 could be moved to the ctor itself? it'd be simpler to
read this..

[Philippe, 2013/07/16 12:31:21]

Yes, good idea.

+      self._KillHostLocked()
+      pid_file.close()

[bulach, 2013/07/15 17:51:43]

pid_file.seek(0), and avoid re-opening it below.

[Philippe, 2013/07/16 12:31:21]

Done.

+      pid_file = open(Forwarder._LOCK_PATH, 'w+')
+      pid_file.write(
+          '%s:%s\n' % (current_pid, _GetProcessStartTime(current_pid)))
+      pid_file.flush()

[bulach, 2013/07/15 17:51:43]

doesn't "with" take care of flushing?

[Philippe, 2013/07/16 12:31:21]

Yeah, indeed. I used not to have this 'with' statement.

+
+  def _InitDeviceOnceLocked(self, adb, tool):
+    """Initializes the device_forwarder daemon for a specific device (once).
+
+    Note that the global lock must be acquired before calling this method. This
+    method kills any existing device_forwarder daemon on the device that could
+    be stale, pushes the latest version of the daemon (to the device) and starts
+    it.
+
+    Args:
+      adb: An AndroidCommands instance.
+      tool: Tool class to use to get wrapper, if necessary, for executing the
+            forwarder (see valgrind_tools.py).
+    """
+    device_serial = adb.Adb().GetSerialNumber()
+    if device_serial in self._initialized_devices:
       return
-    self._adb.PushIfNeeded(
+    Forwarder._KillDeviceLocked(adb, tool)
+    adb.PushIfNeeded(
         self._device_forwarder_path_on_host,
         Forwarder._DEVICE_FORWARDER_FOLDER)
-    (exit_code, output) = self._adb.GetShellCommandStatusAndOutput(
+    (exit_code, output) = adb.GetShellCommandStatusAndOutput(
         '%s %s %s' % (Forwarder._LD_LIBRARY_PATH, tool.GetUtilWrapper(),
                       Forwarder._DEVICE_FORWARDER_PATH))
     if exit_code != 0:
       raise Exception(
           'Failed to start device forwarder:\n%s' % '\n'.join(output))
-    self._device_initialized = True
-
-  def UnmapDevicePort(self, device_port):
-    """Unmaps a previously forwarded device port.
-
-    Args:
-      device_port: A previously forwarded port (through Run()).
-    """
-    with self._lock:
-      self._UnmapDevicePortInternalLocked(device_port)
+    self._initialized_devices.add(device_serial)
 
-  def _UnmapDevicePortInternalLocked(self, device_port):
-    if not device_port in self._device_to_host_port_map:
-      return
-    redirection_command = [
-        '--serial-id=' + self._adb.Adb().GetSerialNumber(), '--unmap',
-        str(device_port)]
-    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-        [self._host_forwarder_path] + redirection_command)
-    if exit_code != 0:
-      logging.error('%s exited with %d:\n%s' % (
-          self._host_forwarder_path, exit_code, '\n'.join(output)))
-    host_port = self._device_to_host_port_map[device_port]
-    del self._device_to_host_port_map[device_port]
-    del self._host_to_device_port_map[host_port]
-
-  @staticmethod
-  def KillHost(build_type='Debug'):
+  def _KillHostLocked(self):
     """Kills the forwarder process running on the host.
 
-    Args:
-      build_type: 'Release' or 'Debug' (default='Debug')
+    Note that the global lock must be acquired before calling this method.
     """
     logging.info('Killing host_forwarder.')
-    host_forwarder_path = _MakeBinaryPath(build_type, 'host_forwarder')
-    if not os.path.exists(host_forwarder_path):
-      host_forwarder_path = _MakeBinaryPath(
-          'Release' if build_type == 'Debug' else 'Debug', 'host_forwarder')
-    assert os.path.exists(host_forwarder_path), 'Please build forwarder2'
     (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-        [host_forwarder_path, '--kill-server'])
+        [self._host_forwarder_path, '--kill-server'])
     if exit_code != 0:
       (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
-          ['pkill', 'host_forwarder'])
+          ['pkill', '-9', 'host_forwarder'])
       if exit_code != 0:
         raise Exception('%s exited with %d:\n%s' % (
-              host_forwarder_path, exit_code, '\n'.join(output)))
+              self._host_forwarder_path, exit_code, '\n'.join(output)))
 
   @staticmethod
-  def KillDevice(adb, tool):
+  def _KillDeviceLocked(adb, tool):
     """Kills the forwarder process running on the device.
 
+    Note that the global lock must be acquired before calling this method.
+
     Args:
       adb: Instance of AndroidCommands for talking to the device.
       tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
@@ -180,14 +303,3 @@ class Forwarder(object):
       pids = adb.ExtractPid('device_forwarder')
       if pids:
         raise Exception('Timed out while killing device_forwarder')
-
-  def DevicePortForHostPort(self, host_port):
-    """Returns the device port that corresponds to a given host port."""
-    with self._lock:
-      return self._host_to_device_port_map.get(host_port)
-
-  def Close(self):
-    """Releases the previously forwarded ports."""
-    with self._lock:
-      for device_port in self._device_to_host_port_map.copy():
-        self._UnmapDevicePortInternalLocked(device_port)