Allow connections to be used as context managers

Add a base class for connection objects that implements the context
manager interface and keeps track of the connection's target (which gets
set by Target.get_connection). When used as a context manager, the connection
substitutes itself for the target's default connection for the thread
within the context using Target.set_connect(). The old connection is
restored upon exiting the context.

Author: setrofim

devlib/connection.py

@@ -0,0 +1,13 @@
+class ConnectionBase:
+
+    def __init__(self):
+        self.target = None
+        self._old_conn = None
+
+    def __enter__(self):
+        self._old_conn = self.target.set_connection(self)
+        return self.target
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.target.set_connection(self._old_conn)
+

devlib/host.py

@@ -22,6 +22,7 @@
 from getpass import getpass
 from pipes import quote
 
+from devlib.connection import ConnectionBase
 from devlib.exception import TargetTransientError, TargetStableError
 from devlib.utils.misc import check_output
 
@@ -37,7 +38,7 @@ def kill_children(pid, signal=signal.SIGKILL):
             os.kill(cpid, signal)
 
 
-class LocalConnection(object):
+class LocalConnection(ConnectionBase):
 
     name = 'local'
     host = 'localhost'
@@ -56,6 +57,7 @@ def connected_as_root(self, state):
     # pylint: disable=unused-argument
     def __init__(self, platform=None, keep_password=True, unrooted=False,
                  password=None, timeout=None):
+        super(LocalConnection, self).__init__()
         self._connected_as_root = None
         self.logger = logging.getLogger('local_connection')
         self.keep_password = keep_password

devlib/target.py

@@ -311,7 +311,9 @@ def disconnect(self):
     def get_connection(self, timeout=None):
         if self.conn_cls is None:
             raise ValueError('Connection class not specified on Target creation.')
-        return self.conn_cls(timeout=timeout, **self.connection_settings)  # pylint: disable=not-callable
+        conn = self.conn_cls(timeout=timeout, **self.connection_settings)  # pylint: disable=not-callable
+        conn.target = self
+        return conn
 
     def set_connection(self, conn):
         tid = id(threading.current_thread())

devlib/utils/android.py

@@ -36,6 +36,7 @@
 except ImportError:
     from pipes import quote
 
+from devlib.connection import ConnectionBase
 from devlib.exception import TargetTransientError, TargetStableError, HostError
 from devlib.utils.misc import check_output, which, ABI_MAP
 
@@ -233,7 +234,7 @@ def _run(self, command):
         return output
 
 
-class AdbConnection(object):
+class AdbConnection(ConnectionBase):
 
     # maintains the count of parallel active connections to a device, so that
     # adb disconnect is not invoked untill all connections are closed
@@ -263,6 +264,7 @@ def connected_as_root(self, state):
     # pylint: disable=unused-argument
     def __init__(self, device=None, timeout=None, platform=None, adb_server=None,
                  adb_as_root=False):
+        super(AdbConnection, self).__init__()
         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)

devlib/utils/ssh.py

@@ -39,6 +39,7 @@
 from pexpect import EOF, TIMEOUT, spawn
 
 # pylint: disable=redefined-builtin,wrong-import-position
+from devlib.connection import ConnectionBase
 from devlib.exception import (HostError, TargetStableError, TargetNotRespondingError,
                               TimeoutError, TargetTransientError)
 from devlib.utils.misc import (which, strip_bash_colors, check_output,
@@ -157,7 +158,7 @@ def check_keyfile(keyfile):
         return keyfile
 
 
-class SshConnection(object):
+class SshConnection(ConnectionBase):
 
     default_password_prompt = '[sudo] password'
     max_cancel_attempts = 5
@@ -194,6 +195,7 @@ def __init__(self,
                  sudo_cmd="sudo -- sh -c {}",
                  options=None
                  ):
+        super(SshConnection, self).__init__()
         self._connected_as_root = None
         self.host = host
         self.username = username

doc/connection.rst

@@ -249,3 +249,32 @@ The only methods discussed below are those that will be overwritten by the
 .. method:: _wait_for_boot(self)
 
     Wait for the gem5 simulated system to have booted and finished the booting animation.
+
+
+.. _multiple-connections:
+
+Multipe Connections With One Target
+-----------------------------------
+
+A ``Target`` will automatically maintain one connection per active thread and
+will seemlessly switch between them, so that target commands being executed from
+parallel threads won't block each other.
+
+It is also possible to create additional connection objects within the same
+thread. You can then use these connections as context managers to execute
+target commands using them rather than the default conection for the thread:
+
+```
+conn = target.get_connection()
+with conn:
+    target.execute('ls')  # uses conn rather than the default connection.
+```
+
+If the connection object is being uses withn another function, you do not need
+to pass the target into that function as well, as the target will be returned on
+entering the connection's context:
+
+```
+with conn as target:
+   target.execute('ls')
+```

doc/target.rst

@@ -182,7 +182,8 @@ Target
              thread, so you don't normally need to use this explicitly in
              threaded code. This is generally useful if you want to perform a
              blocking operation (e.g. using ``background()``) while at the same
-             time doing something else in the same host-side thread.
+             time doing something else in the same host-side thread. See also
+             :ref:`multple-connections`.
 
 .. method:: Target.set_connection(conn)