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