Ensure metric collectors are cloned per shard

In order to avoid state issue when doing local sharding
with device metric collectors, ensure a copy is given to
each shard.

Test: unit tests
Bug: 36351803
Change-Id: Ib756efde7a25c69b694a2bec2b02e0e9145419e8
diff --git a/src/com/android/tradefed/invoker/shard/ShardHelper.java b/src/com/android/tradefed/invoker/shard/ShardHelper.java
index d599718..d31b9cc 100644
--- a/src/com/android/tradefed/invoker/shard/ShardHelper.java
+++ b/src/com/android/tradefed/invoker/shard/ShardHelper.java
@@ -15,6 +15,7 @@
  */
 package com.android.tradefed.invoker.shard;
 
+import com.android.tradefed.config.Configuration;
 import com.android.tradefed.config.ConfigurationException;
 import com.android.tradefed.config.ConfigurationFactory;
 import com.android.tradefed.config.IConfiguration;
@@ -42,6 +43,17 @@
 public class ShardHelper implements IShardHelper {
 
     /**
+     * List of the list configuration obj that should be clone to each shard in order to avoid state
+     * issues.
+     */
+    private static final List<String> CONFIG_OBJ_TO_CLONE = new ArrayList<>();
+
+    static {
+        CONFIG_OBJ_TO_CLONE.add(Configuration.SYSTEM_STATUS_CHECKER_TYPE_NAME);
+        CONFIG_OBJ_TO_CLONE.add(Configuration.DEVICE_METRICS_COLLECTOR_TYPE_NAME);
+    }
+
+    /**
      * Attempt to shard the configuration into sub-configurations, to be re-scheduled to run on
      * multiple resources in parallel.
      *
@@ -124,7 +136,7 @@
             IInvocationContext context,
             IRescheduler rescheduler,
             ShardMasterResultForwarder resultCollector) {
-        cloneStatusChecker(config, shardConfig);
+        cloneConfigObject(config, shardConfig);
         ShardBuildCloner.cloneBuildInfos(config, shardConfig, context);
 
         shardConfig.setTestInvocationListeners(
@@ -138,14 +150,17 @@
     /**
      * Helper to clone {@link ISystemStatusChecker}s from the original config to the clonedConfig.
      */
-    private static void cloneStatusChecker(IConfiguration oriConfig, IConfiguration clonedConfig) {
+    private static void cloneConfigObject(IConfiguration oriConfig, IConfiguration clonedConfig) {
         try {
             IConfiguration deepCopy =
                     ConfigurationFactory.getInstance()
                             .createConfigurationFromArgs(
                                     QuotationAwareTokenizer.tokenizeLine(
                                             oriConfig.getCommandLine()));
-            clonedConfig.setSystemStatusCheckers(deepCopy.getSystemStatusCheckers());
+            for (String objType : CONFIG_OBJ_TO_CLONE) {
+                clonedConfig.setConfigurationObjectList(
+                        objType, deepCopy.getConfigurationObjectList(objType));
+            }
         } catch (ConfigurationException e) {
             // should not happen
             throw new RuntimeException("failed to deep copy a configuration", e);
diff --git a/tests/src/com/android/tradefed/invoker/shard/ShardHelperTest.java b/tests/src/com/android/tradefed/invoker/shard/ShardHelperTest.java
index 2038c03..4f3fad1 100644
--- a/tests/src/com/android/tradefed/invoker/shard/ShardHelperTest.java
+++ b/tests/src/com/android/tradefed/invoker/shard/ShardHelperTest.java
@@ -16,6 +16,7 @@
 package com.android.tradefed.invoker.shard;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertTrue;
 
 import com.android.tradefed.build.BuildInfo;
@@ -24,11 +25,14 @@
 import com.android.tradefed.config.Configuration;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.config.OptionSetter;
+import com.android.tradefed.device.metric.BaseDeviceMetricCollector;
 import com.android.tradefed.invoker.IInvocationContext;
 import com.android.tradefed.invoker.IRescheduler;
 import com.android.tradefed.invoker.InvocationContext;
 import com.android.tradefed.result.ILogSaver;
+import com.android.tradefed.suite.checker.KeyguardStatusChecker;
 import com.android.tradefed.testtype.StubTest;
+import com.android.tradefed.util.FileUtil;
 
 import org.junit.After;
 import org.junit.Before;
@@ -38,10 +42,19 @@
 import org.mockito.ArgumentMatcher;
 import org.mockito.Mockito;
 
+import java.io.File;
+import java.io.IOException;
+import java.util.Arrays;
+
 /** Unit tests for {@link ShardHelper}. */
 @RunWith(JUnit4.class)
 public class ShardHelperTest {
 
+    private static final String TEST_CONFIG =
+            "<configuration description=\"shard config test\">\n"
+                    + "    <%s class=\"%s\" />\n"
+                    + "</configuration>";
+
     private ShardHelper mHelper;
     private IConfiguration mConfig;
     private ILogSaver mMockLogSaver;
@@ -158,4 +171,103 @@
                                     }
                                 }));
     }
+
+    private File createTmpConfig(String objType, Object obj) throws IOException {
+        File configFile = FileUtil.createTempFile("shard-helper-test", ".xml");
+        String content = String.format(TEST_CONFIG, objType, obj.getClass().getCanonicalName());
+        FileUtil.writeToFile(content, configFile);
+        return configFile;
+    }
+
+    /**
+     * Test that some objects are being cloned to the shards in order to avoid shared state issues.
+     */
+    @Test
+    public void testCloneStatusChecker() throws Exception {
+        KeyguardStatusChecker checker = new KeyguardStatusChecker();
+        File configFile = createTmpConfig(Configuration.SYSTEM_STATUS_CHECKER_TYPE_NAME, checker);
+        try {
+            CommandOptions options = new CommandOptions();
+            OptionSetter setter = new OptionSetter(options);
+            // shard-count is the number of shards we are requesting
+            setter.setOptionValue("shard-count", "3");
+            mConfig.setCommandOptions(options);
+            mConfig.setCommandLine(new String[] {configFile.getAbsolutePath()});
+            mConfig.setSystemStatusChecker(checker);
+            StubTest test = new StubTest();
+            setter = new OptionSetter(test);
+            // num-shards is a {@link StubTest} option that specify how many tests can stubtest split
+            // into.
+            setter.setOptionValue("num-shards", "5");
+            mConfig.setTest(test);
+            assertEquals(1, mConfig.getTests().size());
+            assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler));
+            // Ensure that we did split 1 tests per shard rescheduled.
+            Mockito.verify(mRescheduler, Mockito.times(3))
+                    .scheduleConfig(
+                            Mockito.argThat(
+                                    new ArgumentMatcher<IConfiguration>() {
+                                        @Override
+                                        public boolean matches(IConfiguration argument) {
+                                            assertEquals(1, argument.getTests().size());
+                                            // Status checker is in all shard and a new one has been
+                                            // created
+                                            assertEquals(
+                                                    1, argument.getSystemStatusCheckers().size());
+                                            assertNotSame(
+                                                    checker,
+                                                    argument.getSystemStatusCheckers().get(0));
+                                            return true;
+                                        }
+                                    }));
+        } finally {
+            FileUtil.deleteFile(configFile);
+        }
+    }
+
+    /**
+     * Test that some objects are being cloned to the shards in order to avoid shared state issues.
+     */
+    @Test
+    public void testCloneMetricCollector() throws Exception {
+        BaseDeviceMetricCollector collector = new BaseDeviceMetricCollector();
+        File configFile =
+                createTmpConfig(Configuration.DEVICE_METRICS_COLLECTOR_TYPE_NAME, collector);
+        try {
+            CommandOptions options = new CommandOptions();
+            OptionSetter setter = new OptionSetter(options);
+            // shard-count is the number of shards we are requesting
+            setter.setOptionValue("shard-count", "3");
+            mConfig.setCommandOptions(options);
+            mConfig.setCommandLine(new String[] {configFile.getAbsolutePath()});
+            mConfig.setDeviceMetricCollectors(Arrays.asList(collector));
+            StubTest test = new StubTest();
+            setter = new OptionSetter(test);
+            // num-shards is a {@link StubTest} option that specify how many tests can stubtest split
+            // into.
+            setter.setOptionValue("num-shards", "5");
+            mConfig.setTest(test);
+            assertEquals(1, mConfig.getTests().size());
+            assertTrue(mHelper.shardConfig(mConfig, mContext, mRescheduler));
+            // Ensure that we did split 1 tests per shard rescheduled.
+            Mockito.verify(mRescheduler, Mockito.times(3))
+                    .scheduleConfig(
+                            Mockito.argThat(
+                                    new ArgumentMatcher<IConfiguration>() {
+                                        @Override
+                                        public boolean matches(IConfiguration argument) {
+                                            assertEquals(1, argument.getTests().size());
+                                            // Status checker is in all shard and a new one has been
+                                            // created
+                                            assertEquals(1, argument.getMetricCollectors().size());
+                                            assertNotSame(
+                                                    collector,
+                                                    argument.getMetricCollectors().get(0));
+                                            return true;
+                                        }
+                                    }));
+        } finally {
+            FileUtil.deleteFile(configFile);
+        }
+    }
 }