1 /*
2  * Copyright (C) 2019 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.cluster;
17 
18 import com.android.annotations.VisibleForTesting;
19 import com.android.helper.aoa.UsbDevice;
20 import com.android.helper.aoa.UsbHelper;
21 import com.android.tradefed.config.GlobalConfiguration;
22 import com.android.tradefed.config.IConfiguration;
23 import com.android.tradefed.config.IConfigurationReceiver;
24 import com.android.tradefed.config.Option;
25 import com.android.tradefed.config.OptionClass;
26 import com.android.tradefed.device.DeviceNotAvailableException;
27 import com.android.tradefed.device.ITestDevice;
28 import com.android.tradefed.invoker.IInvocationContext;
29 import com.android.tradefed.log.LogUtil.CLog;
30 import com.android.tradefed.result.ITestInvocationListener;
31 import com.android.tradefed.testtype.IInvocationContextReceiver;
32 import com.android.tradefed.testtype.IRemoteTest;
33 import com.android.tradefed.util.ArrayUtil;
34 import com.android.tradefed.util.CommandResult;
35 import com.android.tradefed.util.CommandStatus;
36 import com.android.tradefed.util.FileIdleMonitor;
37 import com.android.tradefed.util.FileUtil;
38 import com.android.tradefed.util.IRunUtil;
39 import com.android.tradefed.util.QuotationAwareTokenizer;
40 import com.android.tradefed.util.RunUtil;
41 import com.android.tradefed.util.StreamUtil;
42 import com.android.tradefed.util.StringEscapeUtils;
43 import com.android.tradefed.util.StringUtil;
44 import com.android.tradefed.util.SubprocessTestResultsParser;
45 import com.android.tradefed.util.SystemUtil;
46 
47 import java.io.File;
48 import java.io.FileOutputStream;
49 import java.io.IOException;
50 import java.time.Duration;
51 import java.util.ArrayList;
52 import java.util.LinkedHashMap;
53 import java.util.LinkedHashSet;
54 import java.util.List;
55 import java.util.Map;
56 import java.util.Map.Entry;
57 import java.util.Set;
58 
59 /**
60  * A {@link IRemoteTest} class to launch a command from TFC via a subprocess TF. FIXME: this needs
61  * to be extended to support multi-device tests.
62  */
63 @OptionClass(alias = "cluster", global_namespace = false)
64 public class ClusterCommandLauncher
65         implements IRemoteTest, IInvocationContextReceiver, IConfigurationReceiver {
66 
67     public static final String TF_JAR_DIR = "TF_JAR_DIR";
68     public static final String TF_PATH = "TF_PATH";
69     public static final String TEST_WORK_DIR = "TEST_WORK_DIR";
70 
71     private static final Duration MAX_EVENT_RECEIVER_WAIT_TIME = Duration.ofMinutes(10);
72 
73     @Option(name = "root-dir", description = "A root directory", mandatory = true)
74     private File mRootDir;
75 
76     @Option(name = "env-var", description = "Environment variables")
77     private Map<String, String> mEnvVars = new LinkedHashMap<>();
78 
79     @Option(name = "setup-script", description = "Setup scripts")
80     private List<String> mSetupScripts = new ArrayList<>();
81 
82     @Option(name = "script-timeout", description = "Script execution timeout", isTimeVal = true)
83     private long mScriptTimeout = 30 * 60 * 1000;
84 
85     @Option(name = "jvm-option", description = "JVM options")
86     private List<String> mJvmOptions = new ArrayList<>();
87 
88     @Option(name = "java-property", description = "Java properties")
89     private Map<String, String> mJavaProperties = new LinkedHashMap<>();
90 
91     @Option(name = "command-line", description = "A command line to launch.", mandatory = true)
92     private String mCommandLine = null;
93 
94     @Option(
95             name = "original-command-line",
96             description =
97                     "Original command line. It may differ from command-line in retry invocations.")
98     private String mOriginalCommandLine = null;
99 
100     @Option(name = "use-subprocess-reporting", description = "Use subprocess reporting.")
101     private boolean mUseSubprocessReporting = false;
102 
103     @Option(
104             name = "output-idle-timeout",
105             description = "Maximum time to wait for an idle subprocess",
106             isTimeVal = true)
107     private long mOutputIdleTimeout = 0L;
108 
109     private IInvocationContext mInvocationContext;
110     private IConfiguration mConfiguration;
111     private IRunUtil mRunUtil;
112 
113     @Override
setInvocationContext(IInvocationContext invocationContext)114     public void setInvocationContext(IInvocationContext invocationContext) {
115         mInvocationContext = invocationContext;
116     }
117 
118     @Override
setConfiguration(IConfiguration configuration)119     public void setConfiguration(IConfiguration configuration) {
120         mConfiguration = configuration;
121     }
122 
getEnvVar(String key)123     private String getEnvVar(String key) {
124         return getEnvVar(key, null);
125     }
126 
getEnvVar(String key, String defaultValue)127     private String getEnvVar(String key, String defaultValue) {
128         String value = mEnvVars.getOrDefault(key, defaultValue);
129         if (value != null) {
130             value = StringUtil.expand(value, mEnvVars);
131         }
132         return value;
133     }
134 
135     @Override
run(ITestInvocationListener listener)136     public void run(ITestInvocationListener listener) throws DeviceNotAvailableException {
137         // Get an expanded TF_PATH value.
138         String tfPath = getEnvVar(TF_PATH, System.getProperty(TF_JAR_DIR));
139         if (tfPath == null) {
140             throw new RuntimeException("cannot find TF path!");
141         }
142 
143         // Construct a Java class path based on TF_PATH value.
144         // This expects TF_PATH to be a colon(:) separated list of paths where each path
145         // points to a specific jar file or folder.
146         // (example: path/to/tradefed.jar:path/to/tradefed/folder:...)
147         final Set<String> jars = new LinkedHashSet<>();
148         for (final String path : tfPath.split(":")) {
149             final File jarFile = new File(path);
150             if (!jarFile.exists()) {
151                 CLog.w("TF_PATH %s doesn't exist; ignoring", path);
152                 continue;
153             }
154             if (jarFile.isFile()) {
155                 jars.add(jarFile.getAbsolutePath());
156             } else {
157                 jars.add(new File(path, "*").getAbsolutePath());
158             }
159         }
160 
161         IRunUtil runUtil = getRunUtil();
162         runUtil.setWorkingDir(mRootDir);
163         // clear the TF_GLOBAL_CONFIG env, so another tradefed will not reuse the global config file
164         runUtil.unsetEnvVariable(GlobalConfiguration.GLOBAL_CONFIG_VARIABLE);
165         for (final String key : mEnvVars.keySet()) {
166             runUtil.setEnvVariable(key, getEnvVar(key));
167         }
168 
169         final File testWorkDir = new File(getEnvVar(TEST_WORK_DIR, mRootDir.getAbsolutePath()));
170         final File logDir = new File(mRootDir, "logs");
171         logDir.mkdirs();
172         File stdoutFile = new File(logDir, "stdout.txt");
173         File stderrFile = new File(logDir, "stderr.txt");
174         FileIdleMonitor monitor = createFileMonitor(stdoutFile, stderrFile);
175 
176         SubprocessTestResultsParser subprocessEventParser = null;
177         try (FileOutputStream stdout = new FileOutputStream(stdoutFile);
178                 FileOutputStream stderr = new FileOutputStream(stderrFile)) {
179             long timeout = mScriptTimeout;
180             long startTime = System.currentTimeMillis();
181             for (String script : mSetupScripts) {
182                 script = StringUtil.expand(script, mEnvVars);
183                 CLog.i("Running a setup script: %s", script);
184                 // FIXME: Refactor command execution into a helper function.
185                 CommandResult result =
186                         runUtil.runTimedCmd(
187                                 timeout,
188                                 stdout,
189                                 stderr,
190                                 QuotationAwareTokenizer.tokenizeLine(script));
191                 if (!result.getStatus().equals(CommandStatus.SUCCESS)) {
192                     String error = null;
193                     if (result.getStatus().equals(CommandStatus.TIMED_OUT)) {
194                         error = "timeout";
195                     } else {
196                         error = FileUtil.readStringFromFile(stderrFile);
197                     }
198                     throw new RuntimeException(String.format("Script failed to run: %s", error));
199                 }
200                 timeout -= (System.currentTimeMillis() - startTime);
201                 if (timeout < 0) {
202                     throw new RuntimeException(
203                             String.format("Setup scripts failed to run in %sms", mScriptTimeout));
204                 }
205             }
206 
207             String classpath = ArrayUtil.join(":", jars);
208             String commandLine = mCommandLine;
209             if (classpath.isEmpty()) {
210                 throw new RuntimeException(
211                         String.format("cannot find any TF jars from %s!", tfPath));
212             }
213 
214             if (mOriginalCommandLine != null && !mOriginalCommandLine.equals(commandLine)) {
215                 // Make sure a wrapper XML of the original command is available because retries
216                 // try to run original commands in Q+. If the original command was run with
217                 // subprocess reporting, a recorded command would be one with .xml suffix.
218                 new SubprocessConfigBuilder()
219                         .setWorkingDir(testWorkDir)
220                         .setOriginalConfig(
221                                 QuotationAwareTokenizer.tokenizeLine(mOriginalCommandLine)[0])
222                         .build();
223             }
224             if (mUseSubprocessReporting) {
225                 SubprocessReportingHelper mHelper = new SubprocessReportingHelper();
226                 // Create standalone jar for subprocess result reporter, which is used
227                 // for pre-O cts. The created jar is put in front position of the class path to
228                 // override class with the same name.
229                 classpath =
230                         String.format(
231                                 "%s:%s",
232                                 mHelper.createSubprocessReporterJar(mRootDir).getAbsolutePath(),
233                                 classpath);
234                 subprocessEventParser =
235                         createSubprocessTestResultsParser(listener, true, mInvocationContext);
236                 String port = Integer.toString(subprocessEventParser.getSocketServerPort());
237                 commandLine = mHelper.buildNewCommandConfig(commandLine, port, testWorkDir);
238             }
239 
240             List<String> javaCommandArgs = buildJavaCommandArgs(classpath, commandLine);
241             CLog.i("Running a command line: %s", commandLine);
242             CLog.i("args = %s", javaCommandArgs);
243             CLog.i("test working directory = %s", testWorkDir);
244 
245             monitor.start();
246             runUtil.setWorkingDir(testWorkDir);
247             CommandResult result =
248                     runUtil.runTimedCmd(
249                             mConfiguration.getCommandOptions().getInvocationTimeout(),
250                             stdout,
251                             stderr,
252                             javaCommandArgs.toArray(new String[javaCommandArgs.size()]));
253             if (!result.getStatus().equals(CommandStatus.SUCCESS)) {
254                 String error = null;
255                 if (result.getStatus().equals(CommandStatus.TIMED_OUT)) {
256                     error = "timeout";
257                 } else {
258                     error = FileUtil.readStringFromFile(stderrFile);
259                 }
260                 throw new RuntimeException(String.format("Command failed to run: %s", error));
261             }
262             CLog.i("Successfully ran a command");
263 
264         } catch (IOException e) {
265             throw new RuntimeException(e);
266         } finally {
267             monitor.stop();
268             if (subprocessEventParser != null) {
269                 subprocessEventParser.joinReceiver(
270                         MAX_EVENT_RECEIVER_WAIT_TIME.toMillis(), /* wait for connection */ false);
271                 StreamUtil.close(subprocessEventParser);
272             }
273         }
274     }
275 
276     /** Build a shell command line to invoke a TF process. */
buildJavaCommandArgs(String classpath, String tfCommandLine)277     private List<String> buildJavaCommandArgs(String classpath, String tfCommandLine) {
278         // Build a command line to invoke a TF process.
279         final List<String> cmdArgs = new ArrayList<>();
280         cmdArgs.add(SystemUtil.getRunningJavaBinaryPath().getAbsolutePath());
281         cmdArgs.add("-cp");
282         cmdArgs.add(classpath);
283         cmdArgs.addAll(mJvmOptions);
284 
285         // Pass Java properties as -D options.
286         for (final Entry<String, String> entry : mJavaProperties.entrySet()) {
287             cmdArgs.add(
288                     String.format(
289                             "-D%s=%s",
290                             entry.getKey(), StringUtil.expand(entry.getValue(), mEnvVars)));
291         }
292         cmdArgs.add("com.android.tradefed.command.CommandRunner");
293         tfCommandLine = StringUtil.expand(tfCommandLine, mEnvVars);
294         cmdArgs.addAll(StringEscapeUtils.paramsToArgs(ArrayUtil.list(tfCommandLine)));
295 
296         final Integer shardCount = mConfiguration.getCommandOptions().getShardCount();
297         final Integer shardIndex = mConfiguration.getCommandOptions().getShardIndex();
298         if (shardCount != null && shardCount > 1) {
299             cmdArgs.add("--shard-count");
300             cmdArgs.add(Integer.toString(shardCount));
301             if (shardIndex != null) {
302                 cmdArgs.add("--shard-index");
303                 cmdArgs.add(Integer.toString(shardIndex));
304             }
305         }
306 
307         for (final ITestDevice device : mInvocationContext.getDevices()) {
308             // FIXME: Find a better way to support non-physical devices as well.
309             cmdArgs.add("--serial");
310             cmdArgs.add(device.getSerialNumber());
311         }
312 
313         return cmdArgs;
314     }
315 
316     /** Creates a file monitor which will perform a USB port reset if the subprocess is idle. */
createFileMonitor(File... files)317     private FileIdleMonitor createFileMonitor(File... files) {
318         // treat zero or negative timeout as infinite
319         long timeout = mOutputIdleTimeout > 0 ? mOutputIdleTimeout : Long.MAX_VALUE;
320         // reset USB ports if files are idle for too long
321         // TODO(peykov): consider making the callback customizable
322         return new FileIdleMonitor(Duration.ofMillis(timeout), this::resetUsbPorts, files);
323     }
324 
325     /** Performs a USB port reset on all devices. */
resetUsbPorts()326     private void resetUsbPorts() {
327         CLog.i("Subprocess output idle for %d ms, attempting USB port reset.", mOutputIdleTimeout);
328         try (UsbHelper usb = new UsbHelper()) {
329             for (String serial : mInvocationContext.getSerials()) {
330                 try (UsbDevice device = usb.getDevice(serial)) {
331                     if (device == null) {
332                         CLog.w("Device '%s' not found during USB reset.", serial);
333                         continue;
334                     }
335                     CLog.d("Resetting USB port for device '%s'", serial);
336                     device.reset();
337                 }
338             }
339         }
340     }
341 
342     @VisibleForTesting
getRunUtil()343     IRunUtil getRunUtil() {
344         if (mRunUtil == null) {
345             mRunUtil = new RunUtil();
346         }
347         return mRunUtil;
348     }
349 
350     @VisibleForTesting
createSubprocessTestResultsParser( ITestInvocationListener listener, boolean streaming, IInvocationContext context)351     SubprocessTestResultsParser createSubprocessTestResultsParser(
352             ITestInvocationListener listener, boolean streaming, IInvocationContext context)
353             throws IOException {
354         return new SubprocessTestResultsParser(listener, streaming, context);
355     }
356 
357     @VisibleForTesting
getEnvVars()358     Map<String, String> getEnvVars() {
359         return mEnvVars;
360     }
361 
362     @VisibleForTesting
getSetupScripts()363     List<String> getSetupScripts() {
364         return mSetupScripts;
365     }
366 
367     @VisibleForTesting
getJvmOptions()368     List<String> getJvmOptions() {
369         return mJvmOptions;
370     }
371 
372     @VisibleForTesting
getJavaProperties()373     Map<String, String> getJavaProperties() {
374         return mJavaProperties;
375     }
376 
377     @VisibleForTesting
getCommandLine()378     String getCommandLine() {
379         return mCommandLine;
380     }
381 
382     @VisibleForTesting
useSubprocessReporting()383     boolean useSubprocessReporting() {
384         return mUseSubprocessReporting;
385     }
386 }
387