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.build.BuildRetrievalError; 20 import com.android.tradefed.config.Configuration; 21 import com.android.tradefed.config.ConfigurationDescriptor; 22 import com.android.tradefed.config.ConfigurationException; 23 import com.android.tradefed.config.DynamicRemoteFileResolver; 24 import com.android.tradefed.config.GlobalConfiguration; 25 import com.android.tradefed.config.IConfiguration; 26 import com.android.tradefed.config.IGlobalConfiguration; 27 import com.android.tradefed.invoker.IInvocationContext; 28 import com.android.tradefed.invoker.IRescheduler; 29 import com.android.tradefed.invoker.ShardListener; 30 import com.android.tradefed.invoker.ShardMainResultForwarder; 31 import com.android.tradefed.invoker.TestInformation; 32 import com.android.tradefed.invoker.shard.token.ITokenRequest; 33 import com.android.tradefed.log.ITestLogger; 34 import com.android.tradefed.log.LogUtil.CLog; 35 import com.android.tradefed.result.IShardableListener; 36 import com.android.tradefed.result.ITestInvocationListener; 37 import com.android.tradefed.result.ITestLoggerReceiver; 38 import com.android.tradefed.retry.IRetryDecision; 39 import com.android.tradefed.suite.checker.ISystemStatusChecker; 40 import com.android.tradefed.testtype.IBuildReceiver; 41 import com.android.tradefed.testtype.IDeviceTest; 42 import com.android.tradefed.testtype.IInvocationContextReceiver; 43 import com.android.tradefed.testtype.IRemoteTest; 44 import com.android.tradefed.testtype.IShardableTest; 45 import com.android.tradefed.util.keystore.IKeyStoreClient; 46 import com.android.tradefed.util.keystore.KeyStoreException; 47 48 import java.util.ArrayList; 49 import java.util.Collection; 50 import java.util.Collections; 51 import java.util.Iterator; 52 import java.util.List; 53 import java.util.concurrent.CountDownLatch; 54 55 /** Helper class that handles creating the shards and scheduling them for an invocation. */ 56 public class ShardHelper implements IShardHelper { 57 58 public static final String LAST_SHARD_DETECTOR = "last_shard_detector"; 59 public static final String SHARED_TEST_INFORMATION = "shared_test_information"; 60 61 /** 62 * List of the list configuration obj that should be clone to each shard in order to avoid state 63 * issues. 64 */ 65 private static final List<String> CONFIG_OBJ_TO_CLONE = new ArrayList<>(); 66 67 static { 68 CONFIG_OBJ_TO_CLONE.add(Configuration.SYSTEM_STATUS_CHECKER_TYPE_NAME); 69 CONFIG_OBJ_TO_CLONE.add(Configuration.DEVICE_METRICS_COLLECTOR_TYPE_NAME); 70 // Copy all the objects under the <device> tag from 71 // {@link Configuration#getMultiDeviceSupportedTag()} except DEVICE_REQUIREMENTS_TYPE_NAME 72 // which should be shared since all shards should have the same requirements. 73 CONFIG_OBJ_TO_CLONE.add(Configuration.BUILD_PROVIDER_TYPE_NAME); 74 CONFIG_OBJ_TO_CLONE.add(Configuration.TARGET_PREPARER_TYPE_NAME); 75 CONFIG_OBJ_TO_CLONE.add(Configuration.DEVICE_RECOVERY_TYPE_NAME); 76 CONFIG_OBJ_TO_CLONE.add(Configuration.DEVICE_OPTIONS_TYPE_NAME); 77 78 CONFIG_OBJ_TO_CLONE.add(Configuration.MULTI_PREPARER_TYPE_NAME); 79 CONFIG_OBJ_TO_CLONE.add(Configuration.CMD_OPTIONS_TYPE_NAME); 80 CONFIG_OBJ_TO_CLONE.add(Configuration.LOGGER_TYPE_NAME); 81 // Deep clone of log_saver to ensure each shard manages its own logs 82 CONFIG_OBJ_TO_CLONE.add(Configuration.LOG_SAVER_TYPE_NAME); 83 // Deep clone RetryDecision to ensure each shard retry independently 84 CONFIG_OBJ_TO_CLONE.add(Configuration.RETRY_DECISION_TYPE_NAME); 85 } 86 87 /** 88 * Attempt to shard the configuration into sub-configurations, to be re-scheduled to run on 89 * multiple resources in parallel. 90 * 91 * <p>A successful shard action renders the current config empty, and invocation should not 92 * proceed. 93 * 94 * @see IShardableTest 95 * @see IRescheduler 96 * @param config the current {@link IConfiguration}. 97 * @param testInfo the {@link TestInformation} holding the tests information. 98 * @param rescheduler the {@link IRescheduler} 99 * @return true if test was sharded. Otherwise return <code>false</code> 100 */ 101 @Override shardConfig( IConfiguration config, TestInformation testInfo, IRescheduler rescheduler, ITestLogger logger)102 public boolean shardConfig( 103 IConfiguration config, 104 TestInformation testInfo, 105 IRescheduler rescheduler, 106 ITestLogger logger) { 107 IInvocationContext context = testInfo.getContext(); 108 List<IRemoteTest> shardableTests = new ArrayList<IRemoteTest>(); 109 boolean isSharded = false; 110 Integer shardCount = config.getCommandOptions().getShardCount(); 111 for (IRemoteTest test : config.getTests()) { 112 isSharded |= shardTest(shardableTests, test, shardCount, testInfo, logger); 113 } 114 if (!isSharded) { 115 return false; 116 } 117 // shard this invocation! 118 // create the TestInvocationListener that will collect results from all the shards, 119 // and forward them to the original set of listeners (minus any ISharddableListeners) 120 // once all shards complete 121 int expectedShard = shardableTests.size(); 122 if (shardCount != null) { 123 expectedShard = Math.min(shardCount, shardableTests.size()); 124 } 125 // Add a tracker so we know in invocation if the last shard is done running. 126 LastShardDetector lastShard = new LastShardDetector(); 127 ShardMainResultForwarder resultCollector = 128 new ShardMainResultForwarder( 129 buildMainShardListeners(config, lastShard), expectedShard); 130 131 config.getLogSaver().invocationStarted(context); 132 resultCollector.invocationStarted(context); 133 synchronized (shardableTests) { 134 // When shardCount is available only create 1 poller per shard 135 // TODO: consider aggregating both case by picking a predefined shardCount if not 136 // available (like 4) for autosharding. 137 if (shardCount != null) { 138 // We shuffle the tests for best results: avoid having the same module sub-tests 139 // contiguously in the list. 140 Collections.shuffle(shardableTests); 141 int maxShard = Math.min(shardCount, shardableTests.size()); 142 CountDownLatch tracker = new CountDownLatch(maxShard); 143 Collection<ITokenRequest> tokenPool = null; 144 if (config.getCommandOptions().shouldUseTokenSharding()) { 145 tokenPool = extractTokenTests(shardableTests); 146 } 147 for (int i = 0; i < maxShard; i++) { 148 IConfiguration shardConfig = cloneConfigObject(config); 149 try { 150 shardConfig.setConfigurationObject(LAST_SHARD_DETECTOR, lastShard); 151 } catch (ConfigurationException e) { 152 throw new RuntimeException(e); 153 } 154 TestsPoolPoller poller = 155 new TestsPoolPoller(shardableTests, tokenPool, tracker); 156 shardConfig.setTest(poller); 157 rescheduleConfig( 158 shardConfig, config, testInfo, rescheduler, resultCollector, i); 159 } 160 } else { 161 CountDownLatch tracker = new CountDownLatch(shardableTests.size()); 162 Collection<ITokenRequest> tokenPool = null; 163 if (config.getCommandOptions().shouldUseTokenSharding()) { 164 tokenPool = extractTokenTests(shardableTests); 165 } 166 int i = 0; 167 for (IRemoteTest testShard : shardableTests) { 168 CLog.d("Rescheduling sharded config..."); 169 IConfiguration shardConfig = cloneConfigObject(config); 170 try { 171 shardConfig.setConfigurationObject(LAST_SHARD_DETECTOR, lastShard); 172 } catch (ConfigurationException e) { 173 throw new RuntimeException(e); 174 } 175 if (config.getCommandOptions().shouldUseDynamicSharding()) { 176 TestsPoolPoller poller = 177 new TestsPoolPoller(shardableTests, tokenPool, tracker); 178 shardConfig.setTest(poller); 179 } else { 180 shardConfig.setTest(testShard); 181 } 182 rescheduleConfig( 183 shardConfig, config, testInfo, rescheduler, resultCollector, i); 184 i++; 185 } 186 } 187 } 188 // clean up original builds 189 for (String deviceName : context.getDeviceConfigNames()) { 190 config.getDeviceConfigByName(deviceName) 191 .getBuildProvider() 192 .cleanUp(context.getBuildInfo(deviceName)); 193 } 194 return true; 195 } 196 rescheduleConfig( IConfiguration shardConfig, IConfiguration config, TestInformation testInfo, IRescheduler rescheduler, ShardMainResultForwarder resultCollector, int index)197 private void rescheduleConfig( 198 IConfiguration shardConfig, 199 IConfiguration config, 200 TestInformation testInfo, 201 IRescheduler rescheduler, 202 ShardMainResultForwarder resultCollector, 203 int index) { 204 validateOptions(testInfo, shardConfig); 205 ShardBuildCloner.cloneBuildInfos(config, shardConfig, testInfo); 206 207 shardConfig.setTestInvocationListeners( 208 buildShardListeners(resultCollector, config, config.getTestInvocationListeners())); 209 210 // Set the host_log suffix to avoid similar names 211 String suffix = String.format("_shard_index_%s", index); 212 if (shardConfig.getCommandOptions().getHostLogSuffix() != null) { 213 suffix = shardConfig.getCommandOptions().getHostLogSuffix() + suffix; 214 } 215 shardConfig.getCommandOptions().setHostLogSuffix(suffix); 216 217 // Use the same {@link ITargetPreparer}, {@link IDeviceRecovery} etc as original config 218 // Make sure we don't run as sandboxed in shards, only parent invocation needs to 219 // run as sandboxed 220 shardConfig.getConfigurationDescription().setSandboxed(false); 221 rescheduler.scheduleConfig(shardConfig); 222 } 223 224 /** Returns the current global configuration. */ 225 @VisibleForTesting getGlobalConfiguration()226 protected IGlobalConfiguration getGlobalConfiguration() { 227 return GlobalConfiguration.getInstance(); 228 } 229 230 /** Runs the {@link IConfiguration#validateOptions()} on the config. */ 231 @VisibleForTesting validateOptions(TestInformation testInfo, IConfiguration config)232 protected void validateOptions(TestInformation testInfo, IConfiguration config) { 233 try { 234 config.validateOptions(); 235 DynamicRemoteFileResolver resolver = new DynamicRemoteFileResolver(); 236 resolver.setDevice(testInfo.getDevice()); 237 resolver.addExtraArgs(config.getCommandOptions().getDynamicDownloadArgs()); 238 config.resolveDynamicOptions(resolver); 239 } catch (ConfigurationException | BuildRetrievalError e) { 240 throw new RuntimeException(e); 241 } 242 } 243 244 /** 245 * Helper to clone {@link ISystemStatusChecker}s from the original config to the clonedConfig. 246 */ cloneConfigObject(IConfiguration origConfig)247 private IConfiguration cloneConfigObject(IConfiguration origConfig) { 248 IKeyStoreClient client = null; 249 try { 250 client = getGlobalConfiguration().getKeyStoreFactory().createKeyStoreClient(); 251 } catch (KeyStoreException e) { 252 throw new RuntimeException( 253 String.format( 254 "failed to load keystore client when sharding: %s", e.getMessage()), 255 e); 256 } 257 258 try { 259 IConfiguration deepCopy = origConfig.partialDeepClone(CONFIG_OBJ_TO_CLONE, client); 260 // Sharding was done, no need for children to look into it. 261 deepCopy.getCommandOptions().setShardCount(null); 262 deepCopy.getConfigurationDescription() 263 .addMetadata(ConfigurationDescriptor.LOCAL_SHARDED_KEY, "true"); 264 return deepCopy; 265 } catch (ConfigurationException e) { 266 throw new RuntimeException( 267 String.format("failed to deep copy a configuration: %s", e.getMessage()), e); 268 } 269 } 270 271 /** 272 * Attempt to shard given {@link IRemoteTest}. 273 * 274 * @param shardableTests the list of {@link IRemoteTest}s to add to 275 * @param test the {@link IRemoteTest} to shard 276 * @param shardCount attempted number of shard, can be null. 277 * @param testInfo the {@link TestInformation} of the current invocation. 278 * @return <code>true</code> if test was sharded 279 */ shardTest( List<IRemoteTest> shardableTests, IRemoteTest test, Integer shardCount, TestInformation testInfo, ITestLogger logger)280 private static boolean shardTest( 281 List<IRemoteTest> shardableTests, 282 IRemoteTest test, 283 Integer shardCount, 284 TestInformation testInfo, 285 ITestLogger logger) { 286 boolean isSharded = false; 287 if (test instanceof IShardableTest) { 288 // inject device and build since they might be required to shard. 289 if (test instanceof IBuildReceiver) { 290 ((IBuildReceiver) test).setBuild(testInfo.getBuildInfo()); 291 } 292 if (test instanceof IDeviceTest) { 293 ((IDeviceTest) test).setDevice(testInfo.getDevice()); 294 } 295 if (test instanceof IInvocationContextReceiver) { 296 ((IInvocationContextReceiver) test).setInvocationContext(testInfo.getContext()); 297 } 298 if (test instanceof ITestLoggerReceiver) { 299 ((ITestLoggerReceiver) test).setTestLogger(logger); 300 } 301 302 IShardableTest shardableTest = (IShardableTest) test; 303 Collection<IRemoteTest> shards = null; 304 // Give the shardCount hint to tests if they need it. 305 shards = shardableTest.split(shardCount, testInfo); 306 if (shards != null) { 307 shardableTests.addAll(shards); 308 isSharded = true; 309 } 310 } 311 if (!isSharded) { 312 shardableTests.add(test); 313 } 314 return isSharded; 315 } 316 317 /** 318 * Builds the {@link ITestInvocationListener} listeners that will collect the results from all 319 * shards. Currently excludes {@link IShardableListener}s. 320 */ buildMainShardListeners( IConfiguration config, LastShardDetector lastShardDetector)321 private static List<ITestInvocationListener> buildMainShardListeners( 322 IConfiguration config, LastShardDetector lastShardDetector) { 323 List<ITestInvocationListener> newListeners = new ArrayList<ITestInvocationListener>(); 324 for (ITestInvocationListener l : config.getTestInvocationListeners()) { 325 if (!(l instanceof IShardableListener)) { 326 newListeners.add(l); 327 } 328 } 329 newListeners.add(lastShardDetector); 330 return newListeners; 331 } 332 333 /** 334 * Builds the list of {@link ITestInvocationListener}s for each shard. Currently includes any 335 * {@link IShardableListener}, plus a single listener that will forward results to the main 336 * shard collector. 337 */ buildShardListeners( ITestInvocationListener resultCollector, IConfiguration config, List<ITestInvocationListener> origListeners)338 private static List<ITestInvocationListener> buildShardListeners( 339 ITestInvocationListener resultCollector, 340 IConfiguration config, 341 List<ITestInvocationListener> origListeners) { 342 List<ITestInvocationListener> shardListeners = new ArrayList<ITestInvocationListener>(); 343 for (ITestInvocationListener l : origListeners) { 344 if (l instanceof IShardableListener) { 345 shardListeners.add(((IShardableListener) l).clone()); 346 } 347 } 348 ShardListener origConfigListener = new ShardListener(resultCollector); 349 origConfigListener.setSupportGranularResults(isAutoRetryEnabled(config)); 350 shardListeners.add(origConfigListener); 351 return shardListeners; 352 } 353 isAutoRetryEnabled(IConfiguration config)354 private static boolean isAutoRetryEnabled(IConfiguration config) { 355 IRetryDecision decision = config.getRetryDecision(); 356 if (decision.isAutoRetryEnabled() && decision.getMaxRetryCount() > 0) { 357 return true; 358 } 359 return false; 360 } 361 extractTokenTests(Collection<IRemoteTest> shardableTests)362 private Collection<ITokenRequest> extractTokenTests(Collection<IRemoteTest> shardableTests) { 363 List<ITokenRequest> tokenPool = new ArrayList<>(); 364 Iterator<IRemoteTest> itr = new ArrayList<>(shardableTests).iterator(); 365 366 while (itr.hasNext()) { 367 IRemoteTest test = itr.next(); 368 if (test instanceof ITokenRequest) { 369 tokenPool.add((ITokenRequest) test); 370 shardableTests.remove(test); 371 } 372 } 373 return tokenPool; 374 } 375 } 376