cheets: Add a script to keep adb connections alive

This change wraps the GTS invocation around a context that continuously
polls adb and reconnects any lost connections.

BUG=b:34280955
TEST=test_that minnie cheets_GTS.4.1_r1.GtsGmscoreHostTestCases

Change-Id: I5cd8b070a85bd3cd240ed99ec5adc987cfa5b1bc
Reviewed-on: https://chromium-review.googlesource.com/434161
Commit-Ready: Luis Hector Chavez <lhchavez@chromium.org>
Tested-by: Luis Hector Chavez <lhchavez@chromium.org>
Reviewed-by: Luis Hector Chavez <lhchavez@chromium.org>
(cherry picked from commit 554c6f8f1ca1384aaa6f541d29c4e12d2e527cc4)
Reviewed-on: https://chromium-review.googlesource.com/437753
Reviewed-by: David Haddock <dhaddock@chromium.org>
Commit-Queue: David Haddock <dhaddock@chromium.org>
Tested-by: David Haddock <dhaddock@chromium.org>
diff --git a/client/common_lib/cros/adb_keepalive.py b/client/common_lib/cros/adb_keepalive.py
new file mode 100755
index 0000000..ddcfec2
--- /dev/null
+++ b/client/common_lib/cros/adb_keepalive.py
@@ -0,0 +1,50 @@
+#!/usr/bin/python
+
+# Copyright 2017 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import argparse
+import logging
+import pipes
+import time
+
+import common
+from autotest_lib.client.bin import utils
+
+_ADB_POLLING_INTERVAL_SECONDS = 10
+_ADB_CONNECT_INTERVAL_SECONDS = 1
+
+
+def _is_adb_connected():
+    """Return true if adb is connected to the container."""
+    output = utils.system_output('adb get-state', ignore_status=True)
+    logging.debug('adb get-state: %s', output)
+    return output.strip() == 'device'
+
+
+def _ensure_adb_connected(target):
+    """Ensures adb is connected to the container, reconnects otherwise."""
+    while not _is_adb_connected():
+        logging.info('adb not connected. attempting to reconnect')
+        output = utils.system_output('adb connect %s' % pipes.quote(target),
+                                     ignore_status=True)
+        logging.debug('adb connect %s: %s', target, output)
+        time.sleep(_ADB_CONNECT_INTERVAL_SECONDS)
+
+
+if __name__ == '__main__':
+    logging.basicConfig(level=logging.DEBUG)
+    parser = argparse.ArgumentParser(description='ensure adb is connected')
+    parser.add_argument('target', help='Device to connect to')
+    args = parser.parse_args()
+
+    logging.info('Starting adb_keepalive for target %s', args.target)
+
+    while True:
+        try:
+            time.sleep(_ADB_POLLING_INTERVAL_SECONDS)
+            _ensure_adb_connected(args.target)
+        except KeyboardInterrupt:
+            logging.info('Shutting down')
+            break
diff --git a/server/cros/tradefed_test.py b/server/cros/tradefed_test.py
index 186e0f1..fde0229 100644
--- a/server/cros/tradefed_test.py
+++ b/server/cros/tradefed_test.py
@@ -34,6 +34,7 @@
 import urlparse
 
 from autotest_lib.client.bin import utils as client_utils
+from autotest_lib.client.common_lib import base_utils
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import dev_server
 from autotest_lib.server import afe_utils
@@ -140,6 +141,38 @@
         logging.info('Released cache lock.')
 
 
+@contextlib.contextmanager
+def adb_keepalive(target, extra_paths):
+    """A context manager that keeps the adb connection alive.
+
+    AdbKeepalive will spin off a new process that will continuously poll for
+    adb's connected state, and will attempt to reconnect if it ever goes down.
+    This is the only way we can currently recover safely from (intentional)
+    reboots.
+
+    @param target: the hostname and port of the DUT.
+    @param extra_paths: any additional components to the PATH environment
+                        variable.
+    """
+    from autotest_lib.client.common_lib.cros import adb_keepalive as module
+    # |__file__| returns the absolute path of the compiled bytecode of the
+    # module. We want to run the original .py file, so we need to change the
+    # extension back.
+    script_filename = module.__file__.replace('.pyc', '.py')
+    job = base_utils.BgJob([script_filename, target],
+                           nickname='adb_keepalive', stderr_level=logging.DEBUG,
+                           stdout_tee=base_utils.TEE_TO_LOGS,
+                           stderr_tee=base_utils.TEE_TO_LOGS,
+                           extra_paths=extra_paths)
+
+    try:
+        yield
+    finally:
+        # The adb_keepalive.py script runs forever until SIGTERM is sent.
+        base_utils.nuke_subprocess(job.sp)
+        base_utils.join_bg_jobs([job])
+
+
 class TradefedTest(test.test):
     """Base class to prepare DUT to run tests via tradefed."""
     version = 1
@@ -192,6 +225,9 @@
         """
         return _ChromeLogin(self._host)
 
+    def _get_adb_target(self):
+        return '{}:{}'.format(self._host.hostname, self._host.port)
+
     def _try_adb_connect(self):
         """Attempts to connect to adb on the DUT.
 
@@ -200,7 +236,7 @@
         # This may fail return failure due to a race condition in adb connect
         # (b/29370989). If adb is already connected, this command will
         # immediately return success.
-        hostport = '{}:{}'.format(self._host.hostname, self._host.port)
+        hostport = self._get_adb_target()
         result = self._run(
                 'adb',
                 args=('connect', hostport),
@@ -261,6 +297,9 @@
         # This starts adbd.
         self._android_shell('setprop sys.usb.config mtp,adb')
 
+        # Also let it be automatically started upon reboot.
+        self._android_shell('setprop persist.sys.usb.config mtp,adb')
+
         # adbd may take some time to come up. Repeatedly try to connect to adb.
         utils.poll_for_condition(lambda: self._try_adb_connect(),
                                  exception=error.TestFail(
diff --git a/server/site_tests/cheets_GTS/cheets_GTS.py b/server/site_tests/cheets_GTS/cheets_GTS.py
index 3c4bfa3..3d14954 100644
--- a/server/site_tests/cheets_GTS/cheets_GTS.py
+++ b/server/site_tests/cheets_GTS/cheets_GTS.py
@@ -60,14 +60,16 @@
                              '--skip-device-info', '--module',
                              self._target_package]
         # Run GTS via tradefed and obtain stdout, sterr as output.
-        output = self._run(
-                gts_tradefed,
-                args=gts_tradefed_args,
-                verbose=True,
-                # Make sure to tee tradefed stdout/stderr to autotest logs
-                # already during the test run.
-                stdout_tee=utils.TEE_TO_LOGS,
-                stderr_tee=utils.TEE_TO_LOGS)
+        with tradefed_test.adb_keepalive(self._get_adb_target(),
+                                         self._install_paths):
+            output = self._run(
+                    gts_tradefed,
+                    args=gts_tradefed_args,
+                    verbose=True,
+                    # Make sure to tee tradefed stdout/stderr to autotest logs
+                    # already during the test run.
+                    stdout_tee=utils.TEE_TO_LOGS,
+                    stderr_tee=utils.TEE_TO_LOGS)
         # Parse stdout to obtain datetime IDs of directories into which tradefed
         # wrote result xml files and logs.
         datetime_id = self._parse_tradefed_datetime(output)