1/* 2 * Copyright (c) 2013, 2016, Oracle and/or its affiliates. All rights reserved. 3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER. 4 * 5 * This code is free software; you can redistribute it and/or modify it 6 * under the terms of the GNU General Public License version 2 only, as 7 * published by the Free Software Foundation. 8 * 9 * This code is distributed in the hope that it will be useful, but WITHOUT 10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 11 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License 12 * version 2 for more details (a copy is included in the LICENSE file that 13 * accompanied this code). 14 * 15 * You should have received a copy of the GNU General Public License version 16 * 2 along with this work; if not, write to the Free Software Foundation, 17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA. 18 * 19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA 20 * or visit www.oracle.com if you need additional information or have any 21 * questions. 22 */ 23 24package toolbox; 25 26import java.io.BufferedReader; 27import java.io.ByteArrayOutputStream; 28import java.io.File; 29import java.io.IOException; 30import java.io.InputStream; 31import java.io.InputStreamReader; 32import java.io.PrintStream; 33import java.io.PrintWriter; 34import java.io.StringWriter; 35import java.util.EnumMap; 36import java.util.HashMap; 37import java.util.Map; 38import static toolbox.ToolBox.lineSeparator; 39 40/** 41 * A utility base class to simplify the implementation of tasks. 42 * Provides support for running the task in a process and for 43 * capturing output written by the task to stdout, stderr and 44 * other writers where applicable. 45 * @param <T> the implementing subclass 46 */ 47abstract class AbstractTask<T extends AbstractTask<T>> implements Task { 48 protected final ToolBox toolBox; 49 protected final Mode mode; 50 private final Map<OutputKind, String> redirects = new EnumMap<>(OutputKind.class); 51 private final Map<String, String> envVars = new HashMap<>(); 52 private Expect expect = Expect.SUCCESS; 53 int expectedExitCode = 0; 54 55 /** 56 * Create a task that will execute in the specified mode. 57 * @param mode the mode 58 */ 59 protected AbstractTask(ToolBox tb, Mode mode) { 60 toolBox = tb; 61 this.mode = mode; 62 } 63 64 /** 65 * Sets the expected outcome of the task and calls {@code run()}. 66 * @param expect the expected outcome 67 * @return the result of calling {@code run()} 68 */ 69 public Result run(Expect expect) { 70 expect(expect, Integer.MIN_VALUE); 71 return run(); 72 } 73 74 /** 75 * Sets the expected outcome of the task and calls {@code run()}. 76 * @param expect the expected outcome 77 * @param exitCode the expected exit code if the expected outcome 78 * is {@code FAIL} 79 * @return the result of calling {@code run()} 80 */ 81 public Result run(Expect expect, int exitCode) { 82 expect(expect, exitCode); 83 return run(); 84 } 85 86 /** 87 * Sets the expected outcome and expected exit code of the task. 88 * The exit code will not be checked if the outcome is 89 * {@code Expect.SUCCESS} or if the exit code is set to 90 * {@code Integer.MIN_VALUE}. 91 * @param expect the expected outcome 92 * @param exitCode the expected exit code 93 */ 94 protected void expect(Expect expect, int exitCode) { 95 this.expect = expect; 96 this.expectedExitCode = exitCode; 97 } 98 99 /** 100 * Checks the exit code contained in a {@code Result} against the 101 * expected outcome and exit value 102 * @param result the result object 103 * @return the result object 104 * @throws TaskError if the exit code stored in the result object 105 * does not match the expected outcome and exit code. 106 */ 107 protected Result checkExit(Result result) throws TaskError { 108 switch (expect) { 109 case SUCCESS: 110 if (result.exitCode != 0) { 111 result.writeAll(); 112 throw new TaskError("Task " + name() + " failed: rc=" + result.exitCode); 113 } 114 break; 115 116 case FAIL: 117 if (result.exitCode == 0) { 118 result.writeAll(); 119 throw new TaskError("Task " + name() + " succeeded unexpectedly"); 120 } 121 122 if (expectedExitCode != Integer.MIN_VALUE 123 && result.exitCode != expectedExitCode) { 124 result.writeAll(); 125 throw new TaskError("Task " + name() + "failed with unexpected exit code " 126 + result.exitCode + ", expected " + expectedExitCode); 127 } 128 break; 129 } 130 return result; 131 } 132 133 /** 134 * Sets an environment variable to be used by this task. 135 * @param name the name of the environment variable 136 * @param value the value for the environment variable 137 * @return this task object 138 * @throws IllegalStateException if the task mode is not {@code EXEC} 139 */ 140 public T envVar(String name, String value) { 141 if (mode != Mode.EXEC) 142 throw new IllegalStateException(); 143 envVars.put(name, value); 144 return (T) this; 145 } 146 147 /** 148 * Redirects output from an output stream to a file. 149 * @param outputKind the name of the stream to be redirected. 150 * @param path the file 151 * @return this task object 152 * @throws IllegalStateException if the task mode is not {@code EXEC} 153 */ 154 public T redirect(OutputKind outputKind, String path) { 155 if (mode != Mode.EXEC) 156 throw new IllegalStateException(); 157 redirects.put(outputKind, path); 158 return (T) this; 159 } 160 161 /** 162 * Returns a {@code ProcessBuilder} initialized with any 163 * redirects and environment variables that have been set. 164 * @return a {@code ProcessBuilder} 165 */ 166 protected ProcessBuilder getProcessBuilder() { 167 if (mode != Mode.EXEC) 168 throw new IllegalStateException(); 169 ProcessBuilder pb = new ProcessBuilder(); 170 if (redirects.get(OutputKind.STDOUT) != null) 171 pb.redirectOutput(new File(redirects.get(OutputKind.STDOUT))); 172 if (redirects.get(OutputKind.STDERR) != null) 173 pb.redirectError(new File(redirects.get(OutputKind.STDERR))); 174 pb.environment().putAll(envVars); 175 return pb; 176 } 177 178 /** 179 * Collects the output from a process and saves it in a {@code Result}. 180 * @param tb the {@code ToolBox} containing the task {@code t} 181 * @param t the task initiating the process 182 * @param p the process 183 * @return a Result object containing the output from the process and its 184 * exit value. 185 * @throws InterruptedException if the thread is interrupted 186 */ 187 protected Result runProcess(ToolBox tb, Task t, Process p) throws InterruptedException { 188 if (mode != Mode.EXEC) 189 throw new IllegalStateException(); 190 ProcessOutput sysOut = new ProcessOutput(p.getInputStream()).start(); 191 ProcessOutput sysErr = new ProcessOutput(p.getErrorStream()).start(); 192 sysOut.waitUntilDone(); 193 sysErr.waitUntilDone(); 194 int rc = p.waitFor(); 195 Map<OutputKind, String> outputMap = new EnumMap<>(OutputKind.class); 196 outputMap.put(OutputKind.STDOUT, sysOut.getOutput()); 197 outputMap.put(OutputKind.STDERR, sysErr.getOutput()); 198 return checkExit(new Result(toolBox, t, rc, outputMap)); 199 } 200 201 /** 202 * Thread-friendly class to read the output from a process until the stream 203 * is exhausted. 204 */ 205 static class ProcessOutput implements Runnable { 206 ProcessOutput(InputStream from) { 207 in = new BufferedReader(new InputStreamReader(from)); 208 out = new StringBuilder(); 209 } 210 211 ProcessOutput start() { 212 new Thread(this).start(); 213 return this; 214 } 215 216 @Override 217 public void run() { 218 try { 219 String line; 220 while ((line = in.readLine()) != null) { 221 out.append(line).append(lineSeparator); 222 } 223 } catch (IOException e) { 224 } 225 synchronized (this) { 226 done = true; 227 notifyAll(); 228 } 229 } 230 231 synchronized void waitUntilDone() throws InterruptedException { 232 boolean interrupted = false; 233 234 // poll interrupted flag, while waiting for copy to complete 235 while (!(interrupted = Thread.interrupted()) && !done) 236 wait(1000); 237 238 if (interrupted) 239 throw new InterruptedException(); 240 } 241 242 String getOutput() { 243 return out.toString(); 244 } 245 246 private final BufferedReader in; 247 private final StringBuilder out; 248 private boolean done; 249 } 250 251 /** 252 * Utility class to simplify the handling of temporarily setting a 253 * new stream for System.out or System.err. 254 */ 255 static class StreamOutput { 256 // Functional interface to set a stream. 257 // Expected use: System::setOut, System::setErr 258 interface Initializer { 259 void set(PrintStream s); 260 } 261 262 private final ByteArrayOutputStream baos = new ByteArrayOutputStream(); 263 private final PrintStream ps = new PrintStream(baos); 264 private final PrintStream prev; 265 private final Initializer init; 266 267 StreamOutput(PrintStream s, Initializer init) { 268 prev = s; 269 init.set(ps); 270 this.init = init; 271 } 272 273 /** 274 * Closes the stream and returns the contents that were written to it. 275 * @return the contents that were written to it. 276 */ 277 String close() { 278 init.set(prev); 279 ps.close(); 280 return baos.toString(); 281 } 282 } 283 284 /** 285 * Utility class to simplify the handling of creating an in-memory PrintWriter. 286 */ 287 static class WriterOutput { 288 private final StringWriter sw = new StringWriter(); 289 final PrintWriter pw = new PrintWriter(sw); 290 291 /** 292 * Closes the stream and returns the contents that were written to it. 293 * @return the contents that were written to it. 294 */ 295 String close() { 296 pw.close(); 297 return sw.toString(); 298 } 299 } 300} 301