Skip to content

Commit

Permalink
Allow connections to be used as context managers
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
setrofim committed Jan 23, 2020
1 parent 29438c8 commit 252a4ff
Show file tree
Hide file tree
Showing 7 changed files with 56 additions and 5 deletions.
13 changes: 13 additions & 0 deletions devlib/connection.py
Original file line number Diff line number Diff line change
@@ -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)

4 changes: 3 additions & 1 deletion devlib/host.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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'
Expand All @@ -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
Expand Down
4 changes: 3 additions & 1 deletion devlib/target.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
4 changes: 3 additions & 1 deletion devlib/utils/android.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
4 changes: 3 additions & 1 deletion devlib/utils/ssh.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
29 changes: 29 additions & 0 deletions doc/connection.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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')
```
3 changes: 2 additions & 1 deletion doc/target.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down

0 comments on commit 252a4ff

Please sign in to comment.