Plumb timeout through ReadAndDispatch, Shell(), use in long-running cmds

This fixes a bug in which ReadAndDispatch hits a timeout for a
long-running command (eg. 'pm list packages') and assumes that the
connection was broken. When the command finally finishes, we get errors
when reading output from subsequent commands.

Author: erock2112

adb/adb_protocol.py

@@ -175,7 +175,6 @@ def __init__(self, header, data=''):
 
   def Write(self, usb, timeout_ms=None):
     """Send this message over USB."""
-    timeout_ms = usb.Timeout(timeout_ms)
     # We can't merge these 2 writes, the device wouldn't be able to read the
     # packet.
     try:
@@ -190,7 +189,6 @@ def Read(cls, usb, timeout_ms=None):
 
     Returns None if it failed to read the header with a ReadFailedError.
     """
-    timeout_ms = usb.Timeout(timeout_ms)
     packet = usb.BulkRead(24, timeout_ms)
     hdr = _AdbMessageHeader.Unpack(packet)
     if hdr.data_length:
@@ -228,9 +226,10 @@ def __str__(self):
 class _AdbConnection(object):
   """One logical ADB connection to a service."""
   class _MessageQueue(object):
-    def __init__(self, manager):
+    def __init__(self, manager, timeout_ms=None):
       self._queue = Queue.Queue()
       self._manager = manager
+      self._timeout_ms = timeout_ms
 
     def __iter__(self):
       return self
@@ -241,7 +240,7 @@ def next(self):
           i = self._queue.get_nowait()
         except Queue.Empty:
           # Will reentrantly call self._Add() via parent._OnRead()
-          if not self._manager.ReadAndDispatch():
+          if not self._manager.ReadAndDispatch(timeout_ms=self._timeout_ms):
             # Failed to read from the device, the connection likely dropped.
             raise StopIteration()
           continue
@@ -255,14 +254,14 @@ def _Add(self, message):
     def _Close(self):
       self._queue.put(StopIteration())
 
-  def __init__(self, manager, local_id, service_name):
+  def __init__(self, manager, local_id, service_name, timeout_ms=None):
     # ID as given by the remote device.
     self.remote_id = 0
     # Service requested on the remote device.
     self.service_name = service_name
     # Self assigned local ID.
     self._local_id = local_id
-    self._yielder = self._MessageQueue(manager)
+    self._yielder = self._MessageQueue(manager, timeout_ms=timeout_ms)
     self._manager = manager
 
   @property
@@ -423,7 +422,7 @@ def Open(self, destination, timeout_ms=None):
       next_id = self._next_local_id
       self._next_local_id += 1
 
-    conn = _AdbConnection(self, next_id, destination)
+    conn = _AdbConnection(self, next_id, destination, timeout_ms=timeout_ms)
     conn._Write('OPEN', destination + '\0')
     with self._lock:
       self._connections[conn.local_id] = conn

adb/common.py

@@ -63,7 +63,7 @@ def __init__(self, serial=None, timeout_ms=None):
     self._timeout_ms = timeout_ms or DEFAULT_TIMEOUT_MS
 
   def Timeout(self, timeout_ms):
-    return timeout_ms if timeout_ms is not None else self._timeout_ms
+    return max(timeout_ms, self._timeout_ms)
 
   # TODO(bpastene) Remove all dependencies on a non UsbHandle needing port_path
   @property

adb/contrib/adb_commands_safe.py

@@ -382,7 +382,7 @@ def Remount(self):
             break
     return False
 
-  def Shell(self, cmd):
+  def Shell(self, cmd, timeout_ms=None):
     """Runs a command on an Android device while swallowing exceptions.
 
     Traps all kinds of USB errors so callers do not have to handle this.
@@ -395,7 +395,7 @@ def Shell(self, cmd):
     if self._adb_cmd:
       for _ in self._Loop():
         try:
-          return self.ShellRaw(cmd)
+          return self.ShellRaw(cmd, timeout_ms=timeout_ms)
         except self._ERRORS as e:
           if not self._Reset('(%s): %s', cmd, e):
             break
@@ -413,7 +413,7 @@ def IsShellOk(self, cmd):
     # Has to keep one byte for trailing nul byte.
     return cmd_size < pkt_size
 
-  def ShellRaw(self, cmd):
+  def ShellRaw(self, cmd, timeout_ms=None):
     """Runs a command on an Android device.
 
     It is expected that the user quote cmd properly.
@@ -433,8 +433,10 @@ def ShellRaw(self, cmd):
     # The adb protocol doesn't return the exit code, so embed it inside the
     # command.
     assert self.IsShellOk(cmd), 'Command is too long: %r' % cmd
-    out = self._adb_cmd.Shell(cmd + self._SHELL_SUFFIX).decode(
-        'utf-8', 'replace')
+    timeout_ms = max(timeout_ms, self._default_timeout_ms)
+    out = self._adb_cmd.Shell(
+        cmd + self._SHELL_SUFFIX,
+        timeout_ms=timeout_ms).decode('utf-8', 'replace')
     # Protect against & or other bash conditional execution that wouldn't make
     # the 'echo $?' command to run.
     if not out:

adb/contrib/high.py

@@ -444,12 +444,12 @@ def Remount(self):
   def Root(self):
     return self._device.Root()
 
-  def Shell(self, cmd):
+  def Shell(self, cmd, timeout_ms=None):
     """Automatically uses WrappedShell() when necessary."""
     if self._device.IsShellOk(cmd):
-      return self._device.Shell(cmd)
+      return self._device.Shell(cmd, timeout_ms=timeout_ms)
     else:
-      return self.WrappedShell([cmd])
+      return self.WrappedShell([cmd], timeout_ms=timeout_ms)
 
   def ShellRaw(self, cmd):
     return self._device.ShellRaw(cmd)
@@ -729,7 +729,9 @@ def GetLastUID(self):
 
   def GetPackages(self):
     """Returns the list of packages installed."""
-    out, _ = self.Shell('pm list packages')
+    # pm can be very slow at times. Use a longer timeout to prevent
+    # confusing a long-running command with an interrupted connection.
+    out, _ = self.Shell('pm list packages', timeout_ms=30000)
     if not out:
       return None
     return [l.split(':', 1)[1] for l in out.strip().splitlines() if ':' in l]
@@ -818,7 +820,9 @@ def WaitUntilFullyBooted(self, timeout=300):
             'time',
             self.port_path)
         return False
-      out, _ = self.Shell('pm path')
+      # pm can be very slow at times. Use a longer timeout to prevent
+      # confusing a long-running command with an interrupted connection.
+      out, _ = self.Shell('pm path', timeout_ms=30000)
       if out == 'Error: no package specified\n':
         # It's up!
         break
@@ -880,7 +884,7 @@ def Mkstemp(self, content, prefix='python-adb', suffix=''):
       if self.PushContent(content, name):
         return name
 
-  def WrappedShell(self, commands):
+  def WrappedShell(self, commands, timeout_ms=None):
     """Creates a temporary shell script, runs it then return the data.
 
     This is needed when:
@@ -899,7 +903,8 @@ def WrappedShell(self, commands):
       if not outfile:
         return False
       try:
-        _, exit_code = self.Shell('sh %s &> %s' % (script, outfile))
+        _, exit_code = self.Shell('sh %s &> %s' % (script, outfile),
+                                  timeout_ms=timeout_ms)
         out = self.PullContent(outfile)
         return out, exit_code
       finally: