connection: Rework TransferManager

* Split TransferManager and TransferHandle:
    * TransferManager deals with the generic monitoring. To abort a
      transfer, it simply cancels the transfer and raises an exception
      from manage().
    * TransferHandle provides a way for the manager to query the state
      of the transfer and cancel it. It is backend-specific.

* Remove most of the state in TransferManager, along with the associated
  background command leak etc

* Use a daemonic monitor thread to behave as excpected on interpreter
  shutdown.

* Ensure a transfer manager _always_ exists. When no management is
  desired, a noop object is used. This avoids using a None sentinel,
  which is invariably mishandled by some code leading to crashes.

* Try to merge more paths in the code to uncover as many issues as
  possible in testing.

Author: douglas-raillard-arm

devlib/connection.py

@@ -14,12 +14,11 @@
 #
 
 from abc import ABC, abstractmethod
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from datetime import datetime
 from functools import partial
 from weakref import WeakSet
 from shlex import quote
-from time import monotonic
 import os
 import signal
 import socket
@@ -61,12 +60,26 @@ class ConnectionBase(InitCheckpoint):
     """
     Base class for all connections.
     """
-    def __init__(self):
+    def __init__(
+        self,
+        poll_transfers=False,
+        start_transfer_poll_delay=30,
+        total_transfer_timeout=3600,
+        transfer_poll_period=30,
+    ):
         self._current_bg_cmds = WeakSet()
         self._closed = False
         self._close_lock = threading.Lock()
         self.busybox = None
 
+        self.transfer_mgr = TransferManager(
+            self,
+            start_transfer_poll_delay=start_transfer_poll_delay,
+            total_transfer_timeout=total_transfer_timeout,
+            transfer_poll_period=transfer_poll_period,
+        ) if poll_transfers else NoopTransferManager()
+
+
     def cancel_running_command(self):
         bg_cmds = set(self._current_bg_cmds)
         for bg_cmd in bg_cmds:
@@ -500,161 +513,147 @@ def __exit__(self, *args, **kwargs):
         self.adb_popen.__exit__(*args, **kwargs)
 
 
-class TransferManagerBase(ABC):
-
-    def _pull_dest_size(self, dest):
-        if os.path.isdir(dest):
-            return sum(
-                os.stat(os.path.join(dirpath, f)).st_size
-	            for dirpath, _, fnames in os.walk(dest)
-	            for f in fnames
-            )
-        else:
-            return os.stat(dest).st_size
-
-    def _push_dest_size(self, dest):
-        cmd = '{} du -s {}'.format(quote(self.conn.busybox), quote(dest))
-        out = self.conn.execute(cmd)
-        try:
-            return int(out.split()[0])
-        except ValueError:
-            return 0
-
-    def __init__(self, conn, poll_period, start_transfer_poll_delay, total_timeout):
+class TransferManager:
+    def __init__(self, conn, poll_period=30, start_transfer_poll_delay=30, total_timeout=3600):
         self.conn = conn
         self.poll_period = poll_period
         self.total_timeout = total_timeout
         self.start_transfer_poll_delay = start_transfer_poll_delay
 
         self.logger = logging.getLogger('FileTransfer')
-        self.managing = threading.Event()
-        self.transfer_started = threading.Event()
-        self.transfer_completed = threading.Event()
-        self.transfer_aborted = threading.Event()
 
-        self.monitor_thread = None
-        self.sources = None
-        self.dest = None
-        self.direction = None
+    @contextmanager
+    def manage(self, sources, dest, direction, handle):
+        excep = None
+        stop_thread = threading.Event()
+
+        def monitor():
+            nonlocal excep
+
+            def cancel(reason):
+                self.logger.warning(
+                    f'Cancelling file transfer {sources} -> {dest} due to: {reason}'
+                )
+                handle.cancel()
+
+            start_t = time.monotonic()
+            stop_thread.wait(self.start_transfer_poll_delay)
+            while not stop_thread.wait(self.poll_period):
+                if not handle.isactive():
+                    cancel(reason='transfer inactive')
+                elif monotonic() - start_t > self.total_timeout:
+                    cancel(reason='transfer timed out')
+                    excep = TimeoutError(f'{direction}: {sources} -> {dest}')
+
+        m_thread = threading.Thread(target=monitor, daemon=True)
+        try:
+            m_thread.start()
+            yield self
+        finally:
+            stop_thread.set()
+            m_thread.join()
+            if excep is not None:
+                raise excep
 
-    @abstractmethod
-    def _cancel(self):
-        pass
 
-    def cancel(self, reason=None):
-        msg = 'Cancelling file transfer {} -> {}'.format(self.sources, self.dest)
-        if reason is not None:
-            msg += ' due to \'{}\''.format(reason)
-        self.logger.warning(msg)
-        self.transfer_aborted.set()
-        self._cancel()
+class NoopTransferManager:
+    def manage(self, *args, **kwargs):
+        return nullcontext(self)
+
+
+class TransferHandleBase(ABC):
+    def __init__(self, mgr):
+        self.mgr = mgr
 
     @abstractmethod
     def isactive(self):
         pass
 
-    @contextmanager
-    def manage(self, sources, dest, direction):
-        try:
-            self.sources, self.dest, self.direction = sources, dest, direction
-            m_thread = threading.Thread(target=self._monitor)
+    @abstractmethod
+    def cancel(self):
+        pass
 
-            self.transfer_completed.clear()
-            self.transfer_aborted.clear()
-            self.transfer_started.set()
 
-            m_thread.start()
-            yield self
-        except BaseException:
-            self.cancel(reason='exception during transfer')
-            raise
-        finally:
-            self.transfer_completed.set()
-            self.transfer_started.set()
-            m_thread.join()
-            self.transfer_started.clear()
-            self.transfer_completed.clear()
-            self.transfer_aborted.clear()
+class PopenTransferHandle(TransferHandleBase):
+    def __init__(self, bg_cmd, dest, direction, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
-    def _monitor(self):
-        start_t = monotonic()
-        self.transfer_completed.wait(self.start_transfer_poll_delay)
-        while not self.transfer_completed.wait(self.poll_period):
-            if not self.isactive():
-                self.cancel(reason='transfer inactive')
-            elif monotonic() - start_t > self.total_timeout:
-                self.cancel(reason='transfer timed out')
+        if direction == 'push':
+            sample_size = self._push_dest_size
+        elif direction == 'pull':
+            sample_size = self._pull_dest_size
+        else:
+            raise ValueError(f'Unknown direction: {direction}')
 
+        self.sample_size = lambda: sample_size(dest)
 
-class PopenTransferManager(TransferManagerBase):
+        self.bg_cmd = bg_cmd
+        self.last_sample = 0
 
-    def __init__(self, conn, poll_period=30, start_transfer_poll_delay=30, total_timeout=3600):
-        super().__init__(conn, poll_period, start_transfer_poll_delay, total_timeout)
-        self.transfer = None
-        self.last_sample = None
+    @staticmethod
+    def _pull_dest_size(dest):
+        if os.path.isdir(dest):
+            return sum(
+                os.stat(os.path.join(dirpath, f)).st_size
+	            for dirpath, _, fnames in os.walk(dest)
+	            for f in fnames
+            )
+        else:
+            return os.stat(dest).st_size
 
-    def _cancel(self):
-        if self.transfer:
-            self.transfer.cancel()
-            self.transfer = None
-            self.last_sample = None
+    def _push_dest_size(self, dest):
+        conn = self.mgr.conn
+        cmd = '{} du -s -- {}'.format(quote(conn.busybox), quote(dest))
+        out = conn.execute(cmd)
+        return int(out.split()[0])
+
+    def cancel(self):
+        self.bg_cmd.cancel()
 
     def isactive(self):
-        size_fn = self._push_dest_size if self.direction == 'push' else self._pull_dest_size
-        curr_size = size_fn(self.dest)
-        self.logger.debug('Polled file transfer, destination size {}'.format(curr_size))
-        active = True if self.last_sample is None else curr_size > self.last_sample
-        self.last_sample = curr_size
-        return active
+        try:
+            curr_size = self.sample_size()
+        except Exception as e:
+            self.logger.debug(f'File size polling failed: {e}')
+            return True
+        else:
+            self.logger.debug(f'Polled file transfer, destination size: {curr_size}')
+            if curr_size:
+                active = curr_size > self.last_sample
+                self.last_sample = curr_size
+                return active
+            # If the file is empty it will never grow in size, so we assume
+            # everything is going well.
+            else:
+                return True
 
-    def set_transfer_and_wait(self, popen_bg_cmd):
-        self.transfer = popen_bg_cmd
-        self.last_sample = None
-        ret = self.transfer.wait()
 
-        if ret and not self.transfer_aborted.is_set():
-            raise subprocess.CalledProcessError(ret, self.transfer.popen.args)
-        elif self.transfer_aborted.is_set():
-            raise TimeoutError(self.transfer.popen.args)
+class SSHTransferHandle(TransferHandleBase):
 
+    def __init__(self, handle, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
-class SSHTransferManager(TransferManagerBase):
+        # SFTPClient or SSHClient
+        self.handle = handle
 
-    def __init__(self, conn, poll_period=30, start_transfer_poll_delay=30, total_timeout=3600):
-        super().__init__(conn, poll_period, start_transfer_poll_delay, total_timeout)
-        self.transferer = None
         self.progressed = False
-        self.transferred = None
-        self.to_transfer = None
+        self.transferred = 0
+        self.to_transfer = 0
 
-    def _cancel(self):
-        self.transferer.close()
+    def cancel(self):
+        self.handle.close()
 
     def isactive(self):
         progressed = self.progressed
-        self.progressed = False
-        msg = 'Polled transfer: {}% [{}B/{}B]'
-        pc = format((self.transferred / self.to_transfer) * 100, '.2f')
-        self.logger.debug(msg.format(pc, self.transferred, self.to_transfer))
+        if progressed:
+            self.progressed = False
+            pc = self.transferred / self.to_transfer * 100
+            self.logger.debug(
+                f'Polled transfer: {pc:.2f}% [{self.transferred}B/{self.to_transfer}B]'
+            )
         return progressed
 
-    @contextmanager
-    def manage(self, sources, dest, direction, transferer):
-        with super().manage(sources, dest, direction):
-            try:
-                self.progressed = False
-                self.transferer = transferer  # SFTPClient or SCPClient
-                yield self
-            except socket.error as e:
-                if self.transfer_aborted.is_set():
-                    self.transfer_aborted.clear()
-                    method = 'SCP' if self.conn.use_scp else 'SFTP'
-                    raise TimeoutError('{} {}: {} -> {}'.format(method, self.direction, sources, self.dest))
-                else:
-                    raise e
-
     def progress_cb(self, to_transfer, transferred):
-        if self.transfer_started.is_set():
-            self.progressed = True
-            self.transferred = transferred
-            self.to_transfer = to_transfer
+        self.progressed = True
+        self.transferred = transferred
+        self.to_transfer = to_transfer

devlib/utils/android.py

@@ -39,7 +39,7 @@
 
 from devlib.exception import TargetTransientError, TargetStableError, HostError, TargetTransientCalledProcessError, TargetStableCalledProcessError, AdbRootError
 from devlib.utils.misc import check_output, which, ABI_MAP, redirect_streams, get_subprocess
-from devlib.connection import ConnectionBase, AdbBackgroundCommand, PopenBackgroundCommand, PopenTransferManager
+from devlib.connection import ConnectionBase, AdbBackgroundCommand, PopenBackgroundCommand, PopenTransferManager, NoopTransferManager
 
 
 logger = logging.getLogger('android')
@@ -278,26 +278,32 @@ def connected_as_root(self, state):
         self._connected_as_root[self.device] = state
 
     # pylint: disable=unused-argument
-    def __init__(self, device=None, timeout=None, platform=None, adb_server=None,
-                 adb_as_root=False, connection_attempts=MAX_ATTEMPTS,
-                 poll_transfers=False,
-                 start_transfer_poll_delay=30,
-                 total_transfer_timeout=3600,
-                 transfer_poll_period=30,):
-        super().__init__()
+    def __init__(
+        self,
+        device=None,
+        timeout=None,
+        platform=None,
+        adb_server=None,
+        adb_as_root=False,
+        connection_attempts=MAX_ATTEMPTS,
+
+        poll_transfers=False,
+        start_transfer_poll_delay=30,
+        total_transfer_timeout=3600,
+        transfer_poll_period=30,
+    ):
+        super().__init__(
+            poll_transfers=poll_transfers,
+            start_transfer_poll_delay=start_transfer_poll_delay,
+            total_transfer_timeout=total_transfer_timeout,
+            transfer_poll_period=transfer_poll_period,
+        )
         self.timeout = timeout if timeout is not None else self.default_timeout
         if device is None:
             device = adb_get_device(timeout=timeout, adb_server=adb_server)
         self.device = device
         self.adb_server = adb_server
         self.adb_as_root = adb_as_root
-        self.poll_transfers = poll_transfers
-        if poll_transfers:
-            transfer_opts = {'start_transfer_poll_delay': start_transfer_poll_delay,
-                            'total_timeout': total_transfer_timeout,
-                            'poll_period': transfer_poll_period,
-                            }
-        self.transfer_mgr = PopenTransferManager(self, **transfer_opts) if poll_transfers else None
         lock, nr_active = AdbConnection.active_connections
         with lock:
             nr_active[self.device] += 1
@@ -330,12 +336,18 @@ def _push_pull(self, action, sources, dest, timeout):
         paths = ' '.join(map(do_quote, paths))
 
         command = "{} {}".format(action, paths)
-        if timeout or not self.poll_transfers:
+        if timeout:
             adb_command(self.device, command, timeout=timeout, adb_server=self.adb_server)
         else:
-            with self.transfer_mgr.manage(sources, dest, action):
-                bg_cmd = adb_command_background(self.device, command, adb_server=self.adb_server)
-                self.transfer_mgr.set_transfer_and_wait(bg_cmd)
+            bg_cmd = adb_command_background(self.device, command, adb_server=self.adb_server)
+            handle = PopenTransferHandle(
+                mgr=self.transfer_mgr,
+                bg_cmd=bg_cmd,
+                dest=dest,
+                direction=action
+            )
+            with bg_cmd, self.transfer_mgr.manage(sources, dest, action, handle):
+                bg_cmd.communicate()
 
     # pylint: disable=unused-argument
     def execute(self, command, timeout=None, check_exit_code=False,