1 /* 2 * Copyright (C) 2017 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 package com.android.tradefed.invoker.shard; 17 18 import com.android.annotations.VisibleForTesting; 19 import com.android.tradefed.config.Configuration; 20 import com.android.tradefed.config.ConfigurationDescriptor; 21 import com.android.tradefed.config.ConfigurationException; 22 import com.android.tradefed.config.ConfigurationFactory; 23 import com.android.tradefed.config.GlobalConfiguration; 24 import com.android.tradefed.config.IConfiguration; 25 import com.android.tradefed.config.IGlobalConfiguration; 26 import com.android.tradefed.invoker.IInvocationContext; 27 import com.android.tradefed.invoker.IRescheduler; 28 import com.android.tradefed.invoker.ShardListener; 29 import com.android.tradefed.invoker.ShardMasterResultForwarder; 30 import com.android.tradefed.invoker.shard.token.ITokenRequest; 31 import com.android.tradefed.log.LogUtil.CLog; 32 import com.android.tradefed.result.IShardableListener; 33 import com.android.tradefed.result.ITestInvocationListener; 34 import com.android.tradefed.suite.checker.ISystemStatusChecker; 35 import com.android.tradefed.testtype.IBuildReceiver; 36 import com.android.tradefed.testtype.IDeviceTest; 37 import com.android.tradefed.testtype.IInvocationContextReceiver; 38 import com.android.tradefed.testtype.IMultiDeviceTest; 39 import com.android.tradefed.testtype.IRemoteTest; 40 import com.android.tradefed.testtype.IShardableTest; 41 import com.android.tradefed.util.QuotationAwareTokenizer; 42 import com.android.tradefed.util.keystore.IKeyStoreClient; 43 import com.android.tradefed.util.keystore.KeyStoreException; 44 45 import java.util.ArrayList; 46 import java.util.Collection; 47 import java.util.Collections; 48 import java.util.Iterator; 49 import java.util.List; 50 import java.util.concurrent.CountDownLatch; 51 52 /** Helper class that handles creating the shards and scheduling them for an invocation. */ 53 public class ShardHelper implements IShardHelper { 54 55 /** 56 * List of the list configuration obj that should be clone to each shard in order to avoid state 57 * issues. 58 */ 59 private static final List<String> CONFIG_OBJ_TO_CLONE = new ArrayList<>(); 60 61 static { 62 CONFIG_OBJ_TO_CLONE.add(Configuration.SYSTEM_STATUS_CHECKER_TYPE_NAME); 63 CONFIG_OBJ_TO_CLONE.add(Configuration.DEVICE_METRICS_COLLECTOR_TYPE_NAME); 64 CONFIG_OBJ_TO_CLONE.add(Configuration.TARGET_PREPARER_TYPE_NAME); 65 CONFIG_OBJ_TO_CLONE.add(Configuration.MULTI_PREPARER_TYPE_NAME); 66 CONFIG_OBJ_TO_CLONE.add(Configuration.CMD_OPTIONS_TYPE_NAME); 67 CONFIG_OBJ_TO_CLONE.add(Configuration.LOGGER_TYPE_NAME); 68 // Deep clone of log_saver to ensure each shard manages its own logs 69 CONFIG_OBJ_TO_CLONE.add(Configuration.LOG_SAVER_TYPE_NAME); 70 } 71 72 /** 73 * Attempt to shard the configuration into sub-configurations, to be re-scheduled to run on 74 * multiple resources in parallel. 75 * 76 * <p>A successful shard action renders the current config empty, and invocation should not 77 * proceed. 78 * 79 * @see IShardableTest 80 * @see IRescheduler 81 * @param config the current {@link IConfiguration}. 82 * @param context the {@link IInvocationContext} holding the tests information. 83 * @param rescheduler the {@link IRescheduler} 84 * @return true if test was sharded. Otherwise return <code>false</code> 85 */ 86 @Override shardConfig( IConfiguration config, IInvocationContext context, IRescheduler rescheduler)87 public boolean shardConfig( 88 IConfiguration config, IInvocationContext context, IRescheduler rescheduler) { 89 List<IRemoteTest> shardableTests = new ArrayList<IRemoteTest>(); 90 boolean isSharded = false; 91 Integer shardCount = config.getCommandOptions().getShardCount(); 92 for (IRemoteTest test : config.getTests()) { 93 isSharded |= shardTest(shardableTests, test, shardCount, context); 94 } 95 if (!isSharded) { 96 return false; 97 } 98 // shard this invocation! 99 // create the TestInvocationListener that will collect results from all the shards, 100 // and forward them to the original set of listeners (minus any ISharddableListeners) 101 // once all shards complete 102 int expectedShard = shardableTests.size(); 103 if (shardCount != null) { 104 expectedShard = Math.min(shardCount, shardableTests.size()); 105 } 106 ShardMasterResultForwarder resultCollector = 107 new ShardMasterResultForwarder(buildMasterShardListeners(config), expectedShard); 108 109 config.getLogSaver().invocationStarted(context); 110 resultCollector.invocationStarted(context); 111 synchronized (shardableTests) { 112 // When shardCount is available only create 1 poller per shard 113 // TODO: consider aggregating both case by picking a predefined shardCount if not 114 // available (like 4) for autosharding. 115 if (shardCount != null) { 116 // We shuffle the tests for best results: avoid having the same module sub-tests 117 // contiguously in the list. 118 Collections.shuffle(shardableTests); 119 int maxShard = Math.min(shardCount, shardableTests.size()); 120 CountDownLatch tracker = new CountDownLatch(maxShard); 121 Collection<ITokenRequest> tokenPool = null; 122 if (config.getCommandOptions().shouldUseTokenSharding()) { 123 tokenPool = extractTokenTests(shardableTests); 124 } 125 for (int i = 0; i < maxShard; i++) { 126 IConfiguration shardConfig = config.clone(); 127 TestsPoolPoller poller = 128 new TestsPoolPoller(shardableTests, tokenPool, tracker); 129 shardConfig.setTest(poller); 130 rescheduleConfig(shardConfig, config, context, rescheduler, resultCollector, i); 131 } 132 } else { 133 CountDownLatch tracker = new CountDownLatch(shardableTests.size()); 134 Collection<ITokenRequest> tokenPool = null; 135 if (config.getCommandOptions().shouldUseTokenSharding()) { 136 tokenPool = extractTokenTests(shardableTests); 137 } 138 int i = 0; 139 for (IRemoteTest testShard : shardableTests) { 140 CLog.d("Rescheduling sharded config..."); 141 IConfiguration shardConfig = config.clone(); 142 if (config.getCommandOptions().shouldUseDynamicSharding()) { 143 TestsPoolPoller poller = 144 new TestsPoolPoller(shardableTests, tokenPool, tracker); 145 shardConfig.setTest(poller); 146 } else { 147 shardConfig.setTest(testShard); 148 } 149 rescheduleConfig(shardConfig, config, context, rescheduler, resultCollector, i); 150 i++; 151 } 152 } 153 } 154 // clean up original builds 155 for (String deviceName : context.getDeviceConfigNames()) { 156 config.getDeviceConfigByName(deviceName) 157 .getBuildProvider() 158 .cleanUp(context.getBuildInfo(deviceName)); 159 } 160 return true; 161 } 162 rescheduleConfig( IConfiguration shardConfig, IConfiguration config, IInvocationContext context, IRescheduler rescheduler, ShardMasterResultForwarder resultCollector, int index)163 private void rescheduleConfig( 164 IConfiguration shardConfig, 165 IConfiguration config, 166 IInvocationContext context, 167 IRescheduler rescheduler, 168 ShardMasterResultForwarder resultCollector, 169 int index) { 170 cloneConfigObject(config, shardConfig); 171 ShardBuildCloner.cloneBuildInfos(config, shardConfig, context); 172 173 shardConfig.setTestInvocationListeners( 174 buildShardListeners(resultCollector, config.getTestInvocationListeners())); 175 176 // Set the host_log suffix to avoid similar names 177 String suffix = String.format("_shard_index_%s", index); 178 if (shardConfig.getCommandOptions().getHostLogSuffix() != null) { 179 suffix = shardConfig.getCommandOptions().getHostLogSuffix() + suffix; 180 } 181 shardConfig.getCommandOptions().setHostLogSuffix(suffix); 182 183 // Use the same {@link ITargetPreparer}, {@link IDeviceRecovery} etc as original config 184 // Make sure we don't run as sandboxed in shards, only parent invocation needs to 185 // run as sandboxed 186 shardConfig.getConfigurationDescription().setSandboxed(false); 187 rescheduler.scheduleConfig(shardConfig); 188 } 189 190 /** Returns the current global configuration. */ 191 @VisibleForTesting getGlobalConfiguration()192 protected IGlobalConfiguration getGlobalConfiguration() { 193 return GlobalConfiguration.getInstance(); 194 } 195 196 /** Runs the {@link IConfiguration#validateOptions(boolean)} on the config. */ 197 @VisibleForTesting validateOptions(IConfiguration config)198 protected void validateOptions(IConfiguration config) throws ConfigurationException { 199 config.validateOptions(true); 200 } 201 202 /** 203 * Helper to clone {@link ISystemStatusChecker}s from the original config to the clonedConfig. 204 */ cloneConfigObject(IConfiguration oriConfig, IConfiguration clonedConfig)205 private void cloneConfigObject(IConfiguration oriConfig, IConfiguration clonedConfig) { 206 IKeyStoreClient client = null; 207 try { 208 client = getGlobalConfiguration().getKeyStoreFactory().createKeyStoreClient(); 209 } catch (KeyStoreException e) { 210 throw new RuntimeException( 211 String.format( 212 "failed to load keystore client when sharding: %s", e.getMessage()), 213 e); 214 } 215 try { 216 IConfiguration deepCopy = 217 ConfigurationFactory.getInstance() 218 .createConfigurationFromArgs( 219 QuotationAwareTokenizer.tokenizeLine( 220 oriConfig.getCommandLine()), 221 null, 222 client); 223 for (String objType : CONFIG_OBJ_TO_CLONE) { 224 clonedConfig.setConfigurationObjectList( 225 objType, deepCopy.getConfigurationObjectList(objType)); 226 } 227 // Sharding was done, no need for children to look into it. 228 clonedConfig.getCommandOptions().setShardCount(null); 229 clonedConfig 230 .getConfigurationDescription() 231 .addMetadata(ConfigurationDescriptor.LOCAL_SHARDED_KEY, "true"); 232 // Validate and download the dynamic options 233 validateOptions(clonedConfig); 234 } catch (ConfigurationException e) { 235 // should not happen 236 throw new RuntimeException( 237 String.format("failed to deep copy a configuration: %s", e.getMessage()), e); 238 } 239 } 240 241 /** 242 * Attempt to shard given {@link IRemoteTest}. 243 * 244 * @param shardableTests the list of {@link IRemoteTest}s to add to 245 * @param test the {@link IRemoteTest} to shard 246 * @param shardCount attempted number of shard, can be null. 247 * @param context the {@link IInvocationContext} of the current invocation. 248 * @return <code>true</code> if test was sharded 249 */ shardTest( List<IRemoteTest> shardableTests, IRemoteTest test, Integer shardCount, IInvocationContext context)250 private static boolean shardTest( 251 List<IRemoteTest> shardableTests, 252 IRemoteTest test, 253 Integer shardCount, 254 IInvocationContext context) { 255 boolean isSharded = false; 256 if (test instanceof IShardableTest) { 257 // inject device and build since they might be required to shard. 258 if (test instanceof IBuildReceiver) { 259 ((IBuildReceiver) test).setBuild(context.getBuildInfos().get(0)); 260 } 261 if (test instanceof IDeviceTest) { 262 ((IDeviceTest) test).setDevice(context.getDevices().get(0)); 263 } 264 if (test instanceof IMultiDeviceTest) { 265 ((IMultiDeviceTest) test).setDeviceInfos(context.getDeviceBuildMap()); 266 } 267 if (test instanceof IInvocationContextReceiver) { 268 ((IInvocationContextReceiver) test).setInvocationContext(context); 269 } 270 271 IShardableTest shardableTest = (IShardableTest) test; 272 Collection<IRemoteTest> shards = null; 273 // Give the shardCount hint to tests if they need it. 274 if (shardCount != null) { 275 shards = shardableTest.split(shardCount); 276 } else { 277 shards = shardableTest.split(); 278 } 279 if (shards != null) { 280 shardableTests.addAll(shards); 281 isSharded = true; 282 } 283 } 284 if (!isSharded) { 285 shardableTests.add(test); 286 } 287 return isSharded; 288 } 289 290 /** 291 * Builds the {@link ITestInvocationListener} listeners that will collect the results from all 292 * shards. Currently excludes {@link IShardableListener}s. 293 */ buildMasterShardListeners(IConfiguration config)294 private static List<ITestInvocationListener> buildMasterShardListeners(IConfiguration config) { 295 List<ITestInvocationListener> newListeners = new ArrayList<ITestInvocationListener>(); 296 for (ITestInvocationListener l : config.getTestInvocationListeners()) { 297 if (!(l instanceof IShardableListener)) { 298 newListeners.add(l); 299 } 300 } 301 return newListeners; 302 } 303 304 /** 305 * Builds the list of {@link ITestInvocationListener}s for each shard. Currently includes any 306 * {@link IShardableListener}, plus a single listener that will forward results to the master 307 * shard collector. 308 */ buildShardListeners( ITestInvocationListener resultCollector, List<ITestInvocationListener> origListeners)309 private static List<ITestInvocationListener> buildShardListeners( 310 ITestInvocationListener resultCollector, List<ITestInvocationListener> origListeners) { 311 List<ITestInvocationListener> shardListeners = new ArrayList<ITestInvocationListener>(); 312 for (ITestInvocationListener l : origListeners) { 313 if (l instanceof IShardableListener) { 314 shardListeners.add(((IShardableListener) l).clone()); 315 } 316 } 317 ShardListener origConfigListener = new ShardListener(resultCollector); 318 shardListeners.add(origConfigListener); 319 return shardListeners; 320 } 321 extractTokenTests(Collection<IRemoteTest> shardableTests)322 private Collection<ITokenRequest> extractTokenTests(Collection<IRemoteTest> shardableTests) { 323 List<ITokenRequest> tokenPool = new ArrayList<>(); 324 Iterator<IRemoteTest> itr = new ArrayList<>(shardableTests).iterator(); 325 326 while (itr.hasNext()) { 327 IRemoteTest test = itr.next(); 328 if (test instanceof ITokenRequest) { 329 tokenPool.add((ITokenRequest) test); 330 shardableTests.remove(test); 331 } 332 } 333 return tokenPool; 334 } 335 } 336