1 /*
2  * Copyright (C) 2009 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 
17 package vogar.commands;
18 
19 import com.google.common.annotations.VisibleForTesting;
20 import com.google.common.collect.ImmutableList;
21 import java.io.BufferedReader;
22 import java.io.File;
23 import java.io.IOException;
24 import java.io.InputStream;
25 import java.io.InputStreamReader;
26 import java.io.PrintStream;
27 import java.lang.reflect.Field;
28 import java.util.ArrayList;
29 import java.util.Arrays;
30 import java.util.Collection;
31 import java.util.Collections;
32 import java.util.LinkedHashMap;
33 import java.util.List;
34 import java.util.Map;
35 import java.util.concurrent.Executors;
36 import java.util.concurrent.ScheduledExecutorService;
37 import java.util.concurrent.TimeUnit;
38 import java.util.concurrent.TimeoutException;
39 import vogar.Log;
40 import vogar.util.Strings;
41 
42 /**
43  * An out of process executable.
44  */
45 public final class Command {
46     private static final ScheduledExecutorService timer
47             = Executors.newSingleThreadScheduledExecutor();
48 
49     private final Log log;
50     private final File workingDir;
51     private final List<String> args;
52     private final Map<String, String> env;
53     private final boolean permitNonZeroExitStatus;
54     private final PrintStream tee;
55 
56     private volatile Process process;
57     private volatile boolean destroyed;
58     private volatile long timeoutNanoTime;
59 
Command(Log log, String... args)60     public Command(Log log, String... args) {
61         this.log = log;
62         this.workingDir = null;
63         this.args = ImmutableList.copyOf(args);
64         this.env = Collections.emptyMap();
65         this.permitNonZeroExitStatus = false;
66         this.tee = null;
67     }
68 
Command(Builder builder)69     private Command(Builder builder) {
70         this.log = builder.log;
71         this.workingDir = builder.workingDir;
72         this.args = ImmutableList.copyOf(builder.args);
73         this.env = builder.env;
74         this.permitNonZeroExitStatus = builder.permitNonZeroExitStatus;
75         this.tee = builder.tee;
76         if (builder.maxLength != -1) {
77             String string = toString();
78             if (string.length() > builder.maxLength) {
79                 throw new IllegalStateException("Maximum command length " + builder.maxLength
80                                                 + " exceeded by: " + string);
81             }
82         }
83     }
84 
start()85     public void start() throws IOException {
86         if (isStarted()) {
87             throw new IllegalStateException("Already started!");
88         }
89 
90         log.verbose("executing " + args + (workingDir != null ? " in " + workingDir : ""));
91 
92         ProcessBuilder processBuilder = new ProcessBuilder()
93                 .directory(workingDir)
94                 .command(args)
95                 .redirectErrorStream(true);
96 
97         processBuilder.environment().putAll(env);
98 
99         process = processBuilder.start();
100     }
101 
isStarted()102     public boolean isStarted() {
103         return process != null;
104     }
105 
getInputStream()106     public InputStream getInputStream() {
107         if (!isStarted()) {
108             throw new IllegalStateException("Not started!");
109         }
110 
111         return process.getInputStream();
112     }
113 
gatherOutput()114     public List<String> gatherOutput()
115             throws IOException, InterruptedException {
116         if (!isStarted()) {
117             throw new IllegalStateException("Not started!");
118         }
119 
120         BufferedReader in = new BufferedReader(
121                 new InputStreamReader(getInputStream(), "UTF-8"));
122         List<String> outputLines = new ArrayList<String>();
123         String outputLine;
124         while ((outputLine = in.readLine()) != null) {
125             if (tee != null) {
126                 tee.println(outputLine);
127             }
128             outputLines.add(outputLine);
129         }
130 
131         int exitValue = process.waitFor();
132         destroyed = true;
133         if (exitValue != 0 && !permitNonZeroExitStatus) {
134             throw new CommandFailedException(args, outputLines);
135         }
136 
137         return outputLines;
138     }
139 
execute()140     public List<String> execute() {
141         try {
142             start();
143             return gatherOutput();
144         } catch (IOException e) {
145             throw new RuntimeException("Failed to execute process: " + args, e);
146         } catch (InterruptedException e) {
147             throw new RuntimeException("Interrupted while executing process: " + args, e);
148         }
149     }
150 
151     /**
152      * Executes a command with a specified timeout. If the process does not
153      * complete normally before the timeout has elapsed, it will be destroyed.
154      *
155      * @param timeoutSeconds how long to wait, or 0 to wait indefinitely
156      * @return the command's output, or null if the command timed out
157      */
executeWithTimeout(int timeoutSeconds)158     public List<String> executeWithTimeout(int timeoutSeconds) throws TimeoutException {
159         if (timeoutSeconds == 0) {
160             return execute();
161         }
162 
163         scheduleTimeout(timeoutSeconds);
164         return execute();
165     }
166 
167     /**
168      * Destroys the underlying process and closes its associated streams.
169      */
destroy()170     public void destroy() {
171         Process process = this.process;
172         if (process == null) {
173             throw new IllegalStateException();
174         }
175         if (destroyed) {
176             return;
177         }
178 
179         destroyed = true;
180         process.destroy();
181         try {
182             process.waitFor();
183             int exitValue = process.exitValue();
184             log.verbose("received exit value " + exitValue + " from destroyed command " + this);
185         } catch (IllegalThreadStateException | InterruptedException destroyUnsuccessful) {
186             log.warn("couldn't destroy " + this);
187         }
188     }
189 
toString()190     @Override public String toString() {
191         String envString = !env.isEmpty() ? (Strings.join(env.entrySet(), " ") + " ") : "";
192         return envString + Strings.join(args, " ");
193     }
194 
195     /**
196      * Sets the time at which this process will be killed. If a timeout has
197      * already been scheduled, it will be rescheduled.
198      */
scheduleTimeout(int timeoutSeconds)199     public void scheduleTimeout(int timeoutSeconds) {
200         timeoutNanoTime = System.nanoTime() + TimeUnit.SECONDS.toNanos(timeoutSeconds);
201 
202         new TimeoutTask() {
203             @Override protected void onTimeout(Process process) {
204                 // send a quit signal immediately
205                 log.verbose("sending quit signal to command " + Command.this);
206                 sendQuitSignal(process);
207 
208                 // hard kill in 2 seconds
209                 timeoutNanoTime = System.nanoTime() + TimeUnit.SECONDS.toNanos(2);
210                 new TimeoutTask() {
211                     @Override protected void onTimeout(Process process) {
212                         log.verbose("killing timed out command " + Command.this);
213                         destroy();
214                     }
215                 }.schedule();
216             }
217         }.schedule();
218     }
219 
sendQuitSignal(Process process)220     private void sendQuitSignal(Process process) {
221         // TODO: 'adb shell kill' to kill on processes running on Androids
222         new Command(log, "kill", "-3", Integer.toString(getPid(process))).execute();
223     }
224 
225     /**
226      * Return the PID of this command's process.
227      */
getPid(Process process)228     private int getPid(Process process) {
229         try {
230             // See org.openqa.selenium.ProcessUtils.getProcessId()
231             Field field = process.getClass().getDeclaredField("pid");
232             field.setAccessible(true);
233             return (Integer) field.get(process);
234         } catch (Exception e) {
235             throw new RuntimeException(e);
236         }
237     }
238 
timedOut()239     public boolean timedOut() {
240         return System.nanoTime() >= timeoutNanoTime;
241     }
242 
243     @VisibleForTesting
getArgs()244     public List<String> getArgs() {
245         return args;
246     }
247 
248     public static class Builder {
249         private final Log log;
250         private final List<String> args = new ArrayList<String>();
251         private final Map<String, String> env = new LinkedHashMap<String, String>();
252         private boolean permitNonZeroExitStatus = false;
253         private PrintStream tee = null;
254         private int maxLength = -1;
255         private File workingDir;
256 
Builder(Log log)257         public Builder(Log log) {
258             this.log = log;
259         }
260 
Builder(Builder other)261         public Builder(Builder other) {
262             this.log = other.log;
263             this.workingDir = other.workingDir;
264             this.args.addAll(other.args);
265             this.env.putAll(other.env);
266             this.permitNonZeroExitStatus = other.permitNonZeroExitStatus;
267             this.tee = other.tee;
268             this.maxLength = other.maxLength;
269         }
270 
args(Object... args)271         public Builder args(Object... args) {
272             return args(Arrays.asList(args));
273         }
274 
args(Collection<?> args)275         public Builder args(Collection<?> args) {
276             for (Object object : args) {
277                 this.args.add(object.toString());
278             }
279             return this;
280         }
281 
env(String key, String value)282         public Builder env(String key, String value) {
283             env.put(key, value);
284             return this;
285         }
286 
287         /**
288          * Controls whether execute() throws if the invoked process returns a
289          * nonzero exit code.
290          */
permitNonZeroExitStatus(boolean value)291         public Builder permitNonZeroExitStatus(boolean value) {
292             this.permitNonZeroExitStatus = value;
293             return this;
294         }
295 
tee(PrintStream printStream)296         public Builder tee(PrintStream printStream) {
297             tee = printStream;
298             return this;
299         }
300 
maxLength(int maxLength)301         public Builder maxLength(int maxLength) {
302             this.maxLength = maxLength;
303             return this;
304         }
305 
workingDir(File workingDir)306         public Builder workingDir(File workingDir) {
307             this.workingDir = workingDir;
308             return this;
309         }
310 
build()311         public Command build() {
312             return new Command(this);
313         }
314 
execute()315         public List<String> execute() {
316             return build().execute();
317         }
318     }
319 
320     /**
321      * Runs some code when the command times out.
322      */
323     private abstract class TimeoutTask implements Runnable {
schedule()324         public final void schedule() {
325             timer.schedule(this, System.nanoTime() - timeoutNanoTime, TimeUnit.NANOSECONDS);
326         }
327 
onTimeout(Process process)328         protected abstract void onTimeout(Process process);
329 
run()330         @Override public final void run() {
331             // don't destroy commands that have already been destroyed
332             Process process = Command.this.process;
333             if (destroyed) {
334                 return;
335             }
336 
337             if (timedOut()) {
338                 onTimeout(process);
339             } else {
340                 // if the kill time has been pushed back, reschedule
341                 timer.schedule(this, System.nanoTime() - timeoutNanoTime, TimeUnit.NANOSECONDS);
342             }
343         }
344     }
345 }
346