Manually merge preconditions sharding fix in mr1

Add fix from go/aog/288805 that was unable to merge into nyc-mr1-dev

Test: turn off screen lock for one of multiple devices, execute
"cts-tradefed run cts --shards 2"
Bug:30022329
Bug:31150255

Change-Id: Ic57eb0eb730202e26df91e8ce49c8220f3dd8b75
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/build/CompatibilityBuildHelper.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/build/CompatibilityBuildHelper.java
index 8196aea..ed9150e 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/build/CompatibilityBuildHelper.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/build/CompatibilityBuildHelper.java
@@ -258,6 +258,13 @@
     }
 
     /**
+     * @return a {@link File} in the resultDir for logging invocation failures
+     */
+    public File getInvocationFailureFile() throws FileNotFoundException {
+        return new File(getResultDir(), "invocation_failure.txt");
+    }
+
+    /**
      * @return a {@link String} to use for directory suffixes created from the given time.
      */
     public static String getDirSuffix(long millis) {
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/InvocationFailureHandler.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/InvocationFailureHandler.java
new file mode 100644
index 0000000..e588fbb
--- /dev/null
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/InvocationFailureHandler.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.compatibility.common.tradefed.result;
+
+import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.tradefed.log.LogUtil.CLog;
+import com.android.tradefed.util.FileUtil;
+
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.IOException;
+
+
+/**
+ * A helper class for setting and checking whether an invocation has failed.
+ */
+public class InvocationFailureHandler {
+
+    /**
+     * Determine whether the invocation for this session has previously failed.
+     *
+     * @param buildHelper the {@link CompatibilityBuildHelper} from which to retrieve invocation
+     * failure file
+     * @return if invocation has previously failed
+     */
+    public static boolean hasFailed(final CompatibilityBuildHelper buildHelper) {
+        try {
+            File f = buildHelper.getInvocationFailureFile();
+            return (f.exists() && f.length() != 0);
+        } catch (FileNotFoundException e) {
+            CLog.e("Could not find invocation failure file for session %s",
+                buildHelper.getDirSuffix(buildHelper.getStartTime()));
+            CLog.e(e);
+            return false;
+        }
+    }
+
+    /**
+     * Write the cause of invocation failure to the result's invocation failure file.
+     *
+     * @param buildHelper the {@link CompatibilityBuildHelper} from which to retrieve the
+     * invocation failure file
+     * @param cause the throwable responsible for invocation failure
+     */
+    public static void setFailed(final CompatibilityBuildHelper buildHelper, Throwable cause) {
+        try {
+            File f = buildHelper.getInvocationFailureFile();
+            if (!f.exists()) {
+                f.createNewFile();
+                FileUtil.writeToFile(cause.toString(), f);
+            }
+        } catch (IOException e) {
+            CLog.e("Exception while writing invocation failure file.");
+            CLog.e(e);
+        }
+    }
+}
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/ResultReporter.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/ResultReporter.java
index a88f223..6a9a9ab 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/ResultReporter.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/result/ResultReporter.java
@@ -16,6 +16,7 @@
 package com.android.compatibility.common.tradefed.result;
 
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.compatibility.common.tradefed.result.InvocationFailureHandler;
 import com.android.compatibility.common.tradefed.testtype.CompatibilityTest;
 import com.android.compatibility.common.util.ICaseResult;
 import com.android.compatibility.common.util.IInvocationResult;
@@ -488,6 +489,7 @@
     @Override
     public void invocationFailed(Throwable cause) {
         warn("Invocation failed: %s", cause);
+        InvocationFailureHandler.setFailed(mBuildHelper, cause);
     }
 
     /**
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/CompatibilityTest.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/CompatibilityTest.java
index 9bd8647..0397a42 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/CompatibilityTest.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/CompatibilityTest.java
@@ -18,6 +18,7 @@
 
 import com.android.compatibility.SuiteInfo;
 import com.android.compatibility.common.tradefed.build.CompatibilityBuildHelper;
+import com.android.compatibility.common.tradefed.result.InvocationFailureHandler;
 import com.android.compatibility.common.tradefed.result.SubPlanCreator;
 import com.android.compatibility.common.tradefed.targetprep.NetworkConnectivityChecker;
 import com.android.compatibility.common.tradefed.targetprep.SystemStatusChecker;
@@ -73,6 +74,7 @@
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 /**
  * A Test for running Compatibility Suites
@@ -99,6 +101,10 @@
     public static final String LOGCAT_ON_FAILURE_SIZE_OPTION = "logcat-on-failure-size";
     private static final String URL = "dynamic-config-url";
 
+    // Constants for checking invocation or preconditions preparation failure
+    private static final int NUM_PREP_ATTEMPTS = 10;
+    private static final int MINUTES_PER_PREP_ATTEMPT = 2;
+
     /* API Key for compatibility test project, used for dynamic configuration */
     private static final String API_KEY = "AIzaSyAbwX5JRlmsLeygY2WWihpIJPXFLueOQ3U";
 
@@ -380,11 +386,20 @@
             }
             mModuleRepo.setPrepared(isPrepared);
 
-            if (!mModuleRepo.isPrepared()) {
-                CLog.logAndDisplay(LogLevel.ERROR,
-                        "Incorrect preparation detected, exiting test run from %s",
-                        mDevice.getSerialNumber());
-                return;
+            int prepAttempt = 1;
+            while (!mModuleRepo.isPrepared(MINUTES_PER_PREP_ATTEMPT, TimeUnit.MINUTES)) {
+                if (prepAttempt >= NUM_PREP_ATTEMPTS
+                        || InvocationFailureHandler.hasFailed(mBuildHelper)) {
+                    CLog.logAndDisplay(LogLevel.ERROR,
+                            "Incorrect preparation detected, exiting test run from %s",
+                            mDevice.getSerialNumber());
+                    return;
+                } else {
+                    CLog.logAndDisplay(LogLevel.INFO,
+                            "Device %s on standby while all shards complete preparation",
+                            mDevice.getSerialNumber());
+                }
+                prepAttempt++;
             }
 
             // Run the tests
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleRepo.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleRepo.java
index 540373b..e689a99 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleRepo.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleRepo.java
@@ -22,6 +22,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Interface for accessing tests from the Compatibility repository.
@@ -31,7 +32,7 @@
     /**
      * @return true after each shard has prepared successfully.
      */
-    boolean isPrepared();
+    boolean isPrepared(long timeout, TimeUnit unit);
 
     /**
      * Indicates to the repo whether a shard is prepared to run.
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleRepo.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleRepo.java
index 3d1b234..a2287d4 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleRepo.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleRepo.java
@@ -202,13 +202,13 @@
      * {@inheritDoc}
      */
     @Override
-    public boolean isPrepared() {
+    public boolean isPrepared(long timeout, TimeUnit unit) {
+        // returns true only if CountDownLatch reaches zero && no shards have setPrepared to false
         try {
-            mPreparedLatch.await();
+            return (mPreparedLatch.await(timeout, unit)) ? mPrepared : false;
         } catch (InterruptedException e) {
             return false;
         }
-        return mPrepared;
     }
 
     /**
diff --git a/common/host-side/tradefed/tests/src/com/android/compatibility/common/tradefed/testtype/ModuleRepoTest.java b/common/host-side/tradefed/tests/src/com/android/compatibility/common/tradefed/testtype/ModuleRepoTest.java
index 5a0ebed..d54277b 100644
--- a/common/host-side/tradefed/tests/src/com/android/compatibility/common/tradefed/testtype/ModuleRepoTest.java
+++ b/common/host-side/tradefed/tests/src/com/android/compatibility/common/tradefed/testtype/ModuleRepoTest.java
@@ -36,6 +36,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 
 public class ModuleRepoTest extends TestCase {
 
@@ -269,7 +270,7 @@
         mRepo.setPrepared(true);
         mRepo.setPrepared(true);
         mRepo.setPrepared(true); // each shard should call setPrepared() once
-        assertTrue(mRepo.isPrepared());
+        assertTrue(mRepo.isPrepared(0, TimeUnit.MINUTES));
     }
 
     public void testIsNotPrepared() {
@@ -279,7 +280,7 @@
         mRepo.setPrepared(true);
         mRepo.setPrepared(false); // mRepo should return false for setPrepared() after third call
         mRepo.setPrepared(true);
-        assertFalse(mRepo.isPrepared());
+        assertFalse(mRepo.isPrepared(0, TimeUnit.MINUTES));
     }
 
     private void assertArrayEquals(Object[] expected, Object[] actual) {