From ebe3a8a0a82f6bc28696b14a1fb7171681a428aa Mon Sep 17 00:00:00 2001
From: Sergei Trofimov <sergei.trofimov@arm.com>
Date: Thu, 25 Feb 2016 10:28:45 +0000
Subject: [PATCH] LinuxTarget: fixing reboot sequence

When issuing a target.reboot(), the reset was immediately followed by a
boot() (for platforms that have it) and an attempt to connect. When
issuing a soft reset, it's possible the target is still shutting down
when the attempt to connect is made. This results in the connection
succeeding but being severed shortly thereafter.

This introduces a delay after the reset to the reboot sequence, giving
the target time to shutdown and improves the handling of EOF's that
result from failed reconnection attempts (while still being with the
allotted timeout period.
---
 devlib/target.py    |  7 ++++
 devlib/utils/ssh.py | 80 +++++++++++++++++++++++++++------------------
 2 files changed, 55 insertions(+), 32 deletions(-)

diff --git a/devlib/target.py b/devlib/target.py
index 5952faf..6b2f4fd 100644
--- a/devlib/target.py
+++ b/devlib/target.py
@@ -228,6 +228,13 @@ class Target(object):
                           '(in which case, a hard_reset module must be installed)'
                 raise TargetError(message)
             self.reset()
+            # Wait a fixed delay before starting polling to give the target time to
+            # shut down, otherwise, might create the connection while it's still shutting
+            # down resulting in subsequenct connection failing.
+            self.logger.debug('Waiting for target to power down...')
+            reset_delay = 20
+            time.sleep(reset_delay)
+            timeout = max(timeout - reset_delay, 10)
         if self.has('boot'):
             self.boot()  # pylint: disable=no-member
         if connect:
diff --git a/devlib/utils/ssh.py b/devlib/utils/ssh.py
index 6ad6f9d..e1b4381 100644
--- a/devlib/utils/ssh.py
+++ b/devlib/utils/ssh.py
@@ -22,6 +22,7 @@ import re
 import threading
 import tempfile
 import shutil
+import time
 
 import pexpect
 from distutils.version import StrictVersion as V
@@ -44,19 +45,28 @@ logger = logging.getLogger('ssh')
 
 def ssh_get_shell(host, username, password=None, keyfile=None, port=None, timeout=10, telnet=False):
     _check_env()
-    if telnet:
-        if keyfile:
-            raise ValueError('keyfile may not be used with a telnet connection.')
-        conn = TelnetConnection()
-    else:  # ssh
-        conn = pxssh.pxssh()
-    try:
-        if keyfile:
-            conn.login(host, username, ssh_key=keyfile, port=port, login_timeout=timeout)
-        else:
-            conn.login(host, username, password, port=port, login_timeout=timeout)
-    except EOF:
-        raise TargetError('Could not connect to {}; is the host name correct?'.format(host))
+    start_time = time.time()
+    while True:
+        if telnet:
+            if keyfile:
+                raise ValueError('keyfile may not be used with a telnet connection.')
+            conn = TelnetConnection()
+        else:  # ssh
+            conn = pxssh.pxssh()
+
+        try:
+            if keyfile:
+                conn.login(host, username, ssh_key=keyfile, port=port, login_timeout=timeout)
+            else:
+                conn.login(host, username, password, port=port, login_timeout=timeout)
+            break
+        except EOF:
+            timeout -= time.time() - start_time
+            if timeout <= 0:
+                message = 'Could not connect to {}; is the host name correct?'
+                raise TargetError(message.format(host))
+            time.sleep(5)
+
     conn.setwinsize(500,200)
     conn.sendline('')
     conn.prompt()
@@ -151,27 +161,33 @@ class SshConnection(object):
         return self._scp(source, dest, timeout)
 
     def execute(self, command, timeout=None, check_exit_code=True, as_root=False, strip_colors=True):
-        with self.lock:
-            output = self._execute_and_wait_for_prompt(command, timeout, as_root, strip_colors)
-            if check_exit_code:
-                exit_code_text = self._execute_and_wait_for_prompt('echo $?', strip_colors=strip_colors, log=False)
-                try:
-                    exit_code = int(exit_code_text.split()[0])
-                    if exit_code:
-                        message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}'
-                        raise TargetError(message.format(exit_code, command, output))
-                except (ValueError, IndexError):
-                    logger.warning('Could not get exit code for "{}",\ngot: "{}"'.format(command, exit_code_text))
-            return output
+        try:
+            with self.lock:
+                output = self._execute_and_wait_for_prompt(command, timeout, as_root, strip_colors)
+                if check_exit_code:
+                    exit_code_text = self._execute_and_wait_for_prompt('echo $?', strip_colors=strip_colors, log=False)
+                    try:
+                        exit_code = int(exit_code_text.split()[0])
+                        if exit_code:
+                            message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}'
+                            raise TargetError(message.format(exit_code, command, output))
+                    except (ValueError, IndexError):
+                        logger.warning('Could not get exit code for "{}",\ngot: "{}"'.format(command, exit_code_text))
+                return output
+        except EOF:
+            raise TargetError('Connection lost.')
 
     def background(self, command, stdout=subprocess.PIPE, stderr=subprocess.PIPE):
-        port_string = '-p {}'.format(self.port) if self.port else ''
-        keyfile_string = '-i {}'.format(self.keyfile) if self.keyfile else ''
-        command = '{} {} {} {}@{} {}'.format(ssh, keyfile_string, port_string, self.username, self.host, command)
-        logger.debug(command)
-        if self.password:
-            command = _give_password(self.password, command)
-        return subprocess.Popen(command, stdout=stdout, stderr=stderr, shell=True)
+        try:
+            port_string = '-p {}'.format(self.port) if self.port else ''
+            keyfile_string = '-i {}'.format(self.keyfile) if self.keyfile else ''
+            command = '{} {} {} {}@{} {}'.format(ssh, keyfile_string, port_string, self.username, self.host, command)
+            logger.debug(command)
+            if self.password:
+                command = _give_password(self.password, command)
+            return subprocess.Popen(command, stdout=stdout, stderr=stderr, shell=True)
+        except EOF:
+            raise TargetError('Connection lost.')
 
     def close(self):
         logger.debug('Logging out {}@{}'.format(self.username, self.host))