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.tradefed.config.ConfigurationException; 19 import com.android.tradefed.config.ConfigurationFactory; 20 import com.android.tradefed.config.IConfiguration; 21 import com.android.tradefed.invoker.IInvocationContext; 22 import com.android.tradefed.invoker.IRescheduler; 23 import com.android.tradefed.invoker.ShardListener; 24 import com.android.tradefed.invoker.ShardMasterResultForwarder; 25 import com.android.tradefed.log.LogUtil.CLog; 26 import com.android.tradefed.result.IShardableListener; 27 import com.android.tradefed.result.ITestInvocationListener; 28 import com.android.tradefed.suite.checker.ISystemStatusChecker; 29 import com.android.tradefed.testtype.IBuildReceiver; 30 import com.android.tradefed.testtype.IDeviceTest; 31 import com.android.tradefed.testtype.IRemoteTest; 32 import com.android.tradefed.testtype.IShardableTest; 33 import com.android.tradefed.util.QuotationAwareTokenizer; 34 35 import java.util.ArrayList; 36 import java.util.Collection; 37 import java.util.Collections; 38 import java.util.List; 39 import java.util.concurrent.CountDownLatch; 40 41 /** Helper class that handles creating the shards and scheduling them for an invocation. */ 42 public class ShardHelper implements IShardHelper { 43 44 /** 45 * Attempt to shard the configuration into sub-configurations, to be re-scheduled to run on 46 * multiple resources in parallel. 47 * 48 * <p>A successful shard action renders the current config empty, and invocation should not 49 * proceed. 50 * 51 * @see IShardableTest 52 * @see IRescheduler 53 * @param config the current {@link IConfiguration}. 54 * @param context the {@link IInvocationContext} holding the tests information. 55 * @param rescheduler the {@link IRescheduler} 56 * @return true if test was sharded. Otherwise return <code>false</code> 57 */ 58 @Override 59 public boolean shardConfig( 60 IConfiguration config, IInvocationContext context, IRescheduler rescheduler) { 61 List<IRemoteTest> shardableTests = new ArrayList<IRemoteTest>(); 62 boolean isSharded = false; 63 Integer shardCount = config.getCommandOptions().getShardCount(); 64 for (IRemoteTest test : config.getTests()) { 65 isSharded |= shardTest(shardableTests, test, shardCount, context); 66 } 67 if (!isSharded) { 68 return false; 69 } 70 // shard this invocation! 71 // create the TestInvocationListener that will collect results from all the shards, 72 // and forward them to the original set of listeners (minus any ISharddableListeners) 73 // once all shards complete 74 int expectedShard = shardableTests.size(); 75 if (shardCount != null) { 76 expectedShard = Math.min(shardCount, shardableTests.size()); 77 } 78 ShardMasterResultForwarder resultCollector = 79 new ShardMasterResultForwarder( 80 config.getLogSaver(), buildMasterShardListeners(config), expectedShard); 81 82 resultCollector.invocationStarted(context); 83 synchronized (shardableTests) { 84 // When shardCount is available only create 1 poller per shard 85 // TODO: consider aggregating both case by picking a predefined shardCount if not 86 // available (like 4) for autosharding. 87 if (shardCount != null) { 88 // We shuffle the tests for best results: avoid having the same module sub-tests 89 // contiguously in the list. 90 Collections.shuffle(shardableTests); 91 int maxShard = Math.min(shardCount, shardableTests.size()); 92 CountDownLatch tracker = new CountDownLatch(maxShard); 93 for (int i = 0; i < maxShard; i++) { 94 IConfiguration shardConfig = config.clone(); 95 shardConfig.setTest(new TestsPoolPoller(shardableTests, tracker)); 96 rescheduleConfig(shardConfig, config, context, rescheduler, resultCollector); 97 } 98 } else { 99 CountDownLatch tracker = new CountDownLatch(shardableTests.size()); 100 for (IRemoteTest testShard : shardableTests) { 101 CLog.i("Rescheduling sharded config..."); 102 IConfiguration shardConfig = config.clone(); 103 if (config.getCommandOptions().shouldUseDynamicSharding()) { 104 shardConfig.setTest(new TestsPoolPoller(shardableTests, tracker)); 105 } else { 106 shardConfig.setTest(testShard); 107 } 108 rescheduleConfig(shardConfig, config, context, rescheduler, resultCollector); 109 } 110 } 111 } 112 // clean up original builds 113 for (String deviceName : context.getDeviceConfigNames()) { 114 config.getDeviceConfigByName(deviceName) 115 .getBuildProvider() 116 .cleanUp(context.getBuildInfo(deviceName)); 117 } 118 return true; 119 } 120 121 public void rescheduleConfig( 122 IConfiguration shardConfig, 123 IConfiguration config, 124 IInvocationContext context, 125 IRescheduler rescheduler, 126 ShardMasterResultForwarder resultCollector) { 127 cloneStatusChecker(config, shardConfig); 128 ShardBuildCloner.cloneBuildInfos(config, shardConfig, context); 129 130 shardConfig.setTestInvocationListeners( 131 buildShardListeners(resultCollector, config.getTestInvocationListeners())); 132 shardConfig.setLogOutput(config.getLogOutput().clone()); 133 shardConfig.setCommandOptions(config.getCommandOptions().clone()); 134 // use the same {@link ITargetPreparer}, {@link IDeviceRecovery} etc as original config 135 rescheduler.scheduleConfig(shardConfig); 136 } 137 138 /** 139 * Helper to clone {@link ISystemStatusChecker}s from the original config to the clonedConfig. 140 */ 141 private static void cloneStatusChecker(IConfiguration oriConfig, IConfiguration clonedConfig) { 142 try { 143 IConfiguration deepCopy = 144 ConfigurationFactory.getInstance() 145 .createConfigurationFromArgs( 146 QuotationAwareTokenizer.tokenizeLine( 147 oriConfig.getCommandLine())); 148 clonedConfig.setSystemStatusCheckers(deepCopy.getSystemStatusCheckers()); 149 } catch (ConfigurationException e) { 150 // should not happen 151 throw new RuntimeException("failed to deep copy a configuration", e); 152 } 153 } 154 155 /** 156 * Attempt to shard given {@link IRemoteTest}. 157 * 158 * @param shardableTests the list of {@link IRemoteTest}s to add to 159 * @param test the {@link IRemoteTest} to shard 160 * @param shardCount attempted number of shard, can be null. 161 * @param context the {@link IInvocationContext} of the current invocation. 162 * @return <code>true</code> if test was sharded 163 */ 164 private static boolean shardTest( 165 List<IRemoteTest> shardableTests, 166 IRemoteTest test, 167 Integer shardCount, 168 IInvocationContext context) { 169 boolean isSharded = false; 170 if (test instanceof IShardableTest) { 171 // inject device and build since they might be required to shard. 172 if (test instanceof IBuildReceiver) { 173 ((IBuildReceiver) test).setBuild(context.getBuildInfos().get(0)); 174 } 175 if (test instanceof IDeviceTest) { 176 ((IDeviceTest) test).setDevice(context.getDevices().get(0)); 177 } 178 IShardableTest shardableTest = (IShardableTest) test; 179 Collection<IRemoteTest> shards = null; 180 // Give the shardCount hint to tests if they need it. 181 if (shardCount != null) { 182 shards = shardableTest.split(shardCount); 183 } else { 184 shards = shardableTest.split(); 185 } 186 if (shards != null) { 187 shardableTests.addAll(shards); 188 isSharded = true; 189 } 190 } 191 if (!isSharded) { 192 shardableTests.add(test); 193 } 194 return isSharded; 195 } 196 197 /** 198 * Builds the {@link ITestInvocationListener} listeners that will collect the results from all 199 * shards. Currently excludes {@link IShardableListener}s. 200 */ 201 private static List<ITestInvocationListener> buildMasterShardListeners(IConfiguration config) { 202 List<ITestInvocationListener> newListeners = new ArrayList<ITestInvocationListener>(); 203 for (ITestInvocationListener l : config.getTestInvocationListeners()) { 204 if (!(l instanceof IShardableListener)) { 205 newListeners.add(l); 206 } 207 } 208 return newListeners; 209 } 210 211 /** 212 * Builds the list of {@link ITestInvocationListener}s for each shard. Currently includes any 213 * {@link IShardableListener}, plus a single listener that will forward results to the master 214 * shard collector. 215 */ 216 private static List<ITestInvocationListener> buildShardListeners( 217 ITestInvocationListener resultCollector, List<ITestInvocationListener> origListeners) { 218 List<ITestInvocationListener> shardListeners = new ArrayList<ITestInvocationListener>(); 219 for (ITestInvocationListener l : origListeners) { 220 if (l instanceof IShardableListener) { 221 shardListeners.add(((IShardableListener) l).clone()); 222 } 223 } 224 ShardListener origConfigListener = new ShardListener(resultCollector); 225 shardListeners.add(origConfigListener); 226 return shardListeners; 227 } 228 229 } 230
