Fix module-level preconditions sharding

When a device fails a module-level precondition during a sharded run,
there is no mechanism to end the invocation, so any other devices that
have passed their preconditions continue to run their assigned
tests. Reimplement ModuleRepo to determine that all devices have been
properly prepared before any devices begin running tests.

Bug:30022329
Change-Id: Ia1b14ac9227370640ac5bd396965409f8e3f61ce
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 e75131f..5510036 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
@@ -342,13 +342,23 @@
             }
 
             // Set values and run preconditions
+            boolean isPrepared = true; // whether the device has been successfully prepared
             for (int i = 0; i < moduleCount; i++) {
                 IModuleDef module = modules.get(i);
                 module.setBuild(mBuildHelper.getBuildInfo());
                 module.setDevice(mDevice);
                 module.setPreparerWhitelist(mPreparerWhitelist);
-                module.prepare(mSkipPreconditions);
+                isPrepared &= (module.prepare(mSkipPreconditions));
             }
+            mModuleRepo.setPrepared(isPrepared);
+
+            if (!mModuleRepo.isPrepared()) {
+                CLog.logAndDisplay(LogLevel.ERROR,
+                        "Incorrect preparation detected, exiting test run from %s",
+                        mDevice.getSerialNumber());
+                return;
+            }
+
             // Run the tests
             for (int i = 0; i < moduleCount; i++) {
                 IModuleDef module = modules.get(i);
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleDef.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleDef.java
index 45ef438..14427ca 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleDef.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/IModuleDef.java
@@ -66,7 +66,8 @@
 
     /**
      * Runs the module's precondition checks and setup tasks.
+     * @return whether preparation succeeded.
      */
-    void prepare(boolean skipPrep) throws DeviceNotAvailableException;
+    boolean prepare(boolean skipPrep) throws DeviceNotAvailableException;
 
 }
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 cf52f35..7cbf7a3 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
@@ -29,6 +29,16 @@
 public interface IModuleRepo {
 
     /**
+     * @return true after each shard has prepared successfully.
+     */
+    boolean isPrepared();
+
+    /**
+     * Indicates to the repo whether a shard is prepared to run.
+     */
+    void setPrepared(boolean isPrepared);
+
+    /**
      * @return true if this repository has been initialized.
      */
     boolean isInitialized();
diff --git a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleDef.java b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleDef.java
index 0844a6e..2c3f96b 100644
--- a/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleDef.java
+++ b/common/host-side/tradefed/src/com/android/compatibility/common/tradefed/testtype/ModuleDef.java
@@ -268,7 +268,7 @@
      * {@inheritDoc}
      */
     @Override
-    public void prepare(boolean skipPrep) throws DeviceNotAvailableException {
+    public boolean prepare(boolean skipPrep) throws DeviceNotAvailableException {
         for (ITargetPreparer preparer : mPreconditions) {
             CLog.d("Preparer: %s", preparer.getClass().getSimpleName());
             if (preparer instanceof IAbiReceiver) {
@@ -282,13 +282,16 @@
                 // This should only happen for flashing new build
                 CLog.e("Unexpected BuildError from precondition: %s",
                         preparer.getClass().getCanonicalName());
+                return false;
             } catch (TargetSetupError e) {
                 // log precondition class then rethrow & let caller handle
                 CLog.e("TargetSetupError in precondition: %s",
                         preparer.getClass().getCanonicalName());
-                throw new RuntimeException(e);
+                e.printStackTrace();
+                return false;
             }
         }
+        return true;
     }
 
     private void setOption(Object target, String option, String value) {
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 de509c8..c4ff1b0 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
@@ -47,6 +47,7 @@
 import java.util.Map;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
 /**
@@ -78,6 +79,10 @@
     private IConfigurationFactory mConfigFactory = ConfigurationFactory.getInstance();
 
     private volatile boolean mInitialized = false;
+    // Whether the modules in this repo are ready to run on their assigned devices.
+    // True until explicitly set false in setPrepared().
+    private volatile boolean mPrepared = true;
+    private CountDownLatch mPreparedLatch;
 
     // Holds all the small tests waiting to be run.
     private List<IModuleDef> mSmallModules = new ArrayList<>();
@@ -197,6 +202,28 @@
      * {@inheritDoc}
      */
     @Override
+    public boolean isPrepared() {
+        try {
+            mPreparedLatch.await();
+        } catch (InterruptedException e) {
+            return false;
+        }
+        return mPrepared;
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
+    public void setPrepared(boolean isPrepared) {
+        mPrepared &= isPrepared;
+        mPreparedLatch.countDown();
+    }
+
+    /**
+     * {@inheritDoc}
+     */
+    @Override
     public boolean isInitialized() {
         return mInitialized;
     }
@@ -214,6 +241,7 @@
                 includeFilters, excludeFilters);
         mInitialized = true;
         mShards = shards;
+        mPreparedLatch = new CountDownLatch(shards);
         for (String line : deviceTokens) {
             String[] parts = line.split(":");
             if (parts.length == 2) {
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 24c7e3d..3b8e4cb 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
@@ -262,6 +262,26 @@
         assertArrayEquals(EXPECTED_MODULE_IDS, mRepo.getModuleIds());
     }
 
+    public void testIsPrepared() {
+        mRepo.initialize(3, mTestsDir, ABIS, DEVICE_TOKENS, TEST_ARGS, MODULE_ARGS, INCLUDES,
+                EXCLUDES, mBuild);
+        assertTrue("Should be initialized", mRepo.isInitialized());
+        mRepo.setPrepared(true);
+        mRepo.setPrepared(true);
+        mRepo.setPrepared(true); // each shard should call setPrepared() once
+        assertTrue(mRepo.isPrepared());
+    }
+
+    public void testIsNotPrepared() {
+        mRepo.initialize(3, mTestsDir, ABIS, DEVICE_TOKENS, TEST_ARGS, MODULE_ARGS, INCLUDES,
+                EXCLUDES, mBuild);
+        assertTrue("Should be initialized", mRepo.isInitialized());
+        mRepo.setPrepared(true);
+        mRepo.setPrepared(false); // mRepo should return false for setPrepared() after third call
+        mRepo.setPrepared(true);
+        assertFalse(mRepo.isPrepared());
+    }
+
     private void assertArrayEquals(Object[] expected, Object[] actual) {
         assertEquals(Arrays.asList(expected), Arrays.asList(actual));
     }