AbstractTask.java

/*
 * Copyright (c) 2013, 2016, Oracle and/or its affiliates. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */

package toolbox;

import java.io.BufferedReader;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintStream;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.EnumMap;
import java.util.HashMap;
import java.util.Map;
import static toolbox.ToolBox.lineSeparator;

/**
 * A utility base class to simplify the implementation of tasks.
 * Provides support for running the task in a process and for
 * capturing output written by the task to stdout, stderr and
 * other writers where applicable.
 * @param <T> the implementing subclass
 */
abstract class AbstractTask<T extends AbstractTask<T>> implements Task {
    protected final ToolBox toolBox;
    protected final Mode mode;
    private final Map<OutputKind, String> redirects = new EnumMap<>(OutputKind.class);
    private final Map<String, String> envVars = new HashMap<>();
    private Expect expect = Expect.SUCCESS;
    int expectedExitCode = 0;

    /**
     * Create a task that will execute in the specified mode.
     * @param mode the mode
     */
    protected AbstractTask(ToolBox tb, Mode mode) {
        toolBox = tb;
        this.mode = mode;
    }

    /**
     * Sets the expected outcome of the task and calls {@code run()}.
     * @param expect the expected outcome
     * @return the result of calling {@code run()}
     */
    public Result run(Expect expect) {
        expect(expect, Integer.MIN_VALUE);
        return run();
    }

    /**
     * Sets the expected outcome of the task and calls {@code run()}.
     * @param expect the expected outcome
     * @param exitCode the expected exit code if the expected outcome
     *      is {@code FAIL}
     * @return the result of calling {@code run()}
     */
    public Result run(Expect expect, int exitCode) {
        expect(expect, exitCode);
        return run();
    }

    /**
     * Sets the expected outcome and expected exit code of the task.
     * The exit code will not be checked if the outcome is
     * {@code Expect.SUCCESS} or if the exit code is set to
     * {@code Integer.MIN_VALUE}.
     * @param expect the expected outcome
     * @param exitCode the expected exit code
     */
    protected void expect(Expect expect, int exitCode) {
        this.expect = expect;
        this.expectedExitCode = exitCode;
    }

    /**
     * Checks the exit code contained in a {@code Result} against the
     * expected outcome and exit value
     * @param result the result object
     * @return the result object
     * @throws TaskError if the exit code stored in the result object
     *      does not match the expected outcome and exit code.
     */
    protected Result checkExit(Result result) throws TaskError {
        switch (expect) {
            case SUCCESS:
                if (result.exitCode != 0) {
                    result.writeAll();
                    throw new TaskError("Task " + name() + " failed: rc=" + result.exitCode);
                }
                break;

            case FAIL:
                if (result.exitCode == 0) {
                    result.writeAll();
                    throw new TaskError("Task " + name() + " succeeded unexpectedly");
                }

                if (expectedExitCode != Integer.MIN_VALUE
                        && result.exitCode != expectedExitCode) {
                    result.writeAll();
                    throw new TaskError("Task " + name() + "failed with unexpected exit code "
                        + result.exitCode + ", expected " + expectedExitCode);
                }
                break;
        }
        return result;
    }

    /**
     * Sets an environment variable to be used by this task.
     * @param name the name of the environment variable
     * @param value the value for the environment variable
     * @return this task object
     * @throws IllegalStateException if the task mode is not {@code EXEC}
     */
    public T envVar(String name, String value) {
        if (mode != Mode.EXEC)
            throw new IllegalStateException();
        envVars.put(name, value);
        return (T) this;
    }

    /**
     * Redirects output from an output stream to a file.
     * @param outputKind the name of the stream to be redirected.
     * @param path the file
     * @return this task object
     * @throws IllegalStateException if the task mode is not {@code EXEC}
     */
    public T redirect(OutputKind outputKind, String path) {
        if (mode != Mode.EXEC)
            throw new IllegalStateException();
        redirects.put(outputKind, path);
        return (T) this;
    }

    /**
     * Returns a {@code ProcessBuilder} initialized with any
     * redirects and environment variables that have been set.
     * @return a {@code ProcessBuilder}
     */
    protected ProcessBuilder getProcessBuilder() {
        if (mode != Mode.EXEC)
            throw new IllegalStateException();
        ProcessBuilder pb = new ProcessBuilder();
        if (redirects.get(OutputKind.STDOUT) != null)
            pb.redirectOutput(new File(redirects.get(OutputKind.STDOUT)));
        if (redirects.get(OutputKind.STDERR) != null)
            pb.redirectError(new File(redirects.get(OutputKind.STDERR)));
        pb.environment().putAll(envVars);
        return pb;
    }

    /**
     * Collects the output from a process and saves it in a {@code Result}.
     * @param tb the {@code ToolBox} containing the task {@code t}
     * @param t the task initiating the process
     * @param p the process
     * @return a Result object containing the output from the process and its
     *      exit value.
     * @throws InterruptedException if the thread is interrupted
     */
    protected Result runProcess(ToolBox tb, Task t, Process p) throws InterruptedException {
        if (mode != Mode.EXEC)
            throw new IllegalStateException();
        ProcessOutput sysOut = new ProcessOutput(p.getInputStream()).start();
        ProcessOutput sysErr = new ProcessOutput(p.getErrorStream()).start();
        sysOut.waitUntilDone();
        sysErr.waitUntilDone();
        int rc = p.waitFor();
        Map<OutputKind, String> outputMap = new EnumMap<>(OutputKind.class);
        outputMap.put(OutputKind.STDOUT, sysOut.getOutput());
        outputMap.put(OutputKind.STDERR, sysErr.getOutput());
        return checkExit(new Result(toolBox, t, rc, outputMap));
    }

    /**
     * Thread-friendly class to read the output from a process until the stream
     * is exhausted.
     */
    static class ProcessOutput implements Runnable {
        ProcessOutput(InputStream from) {
            in = new BufferedReader(new InputStreamReader(from));
            out = new StringBuilder();
        }

        ProcessOutput start() {
            new Thread(this).start();
            return this;
        }

        @Override
        public void run() {
            try {
                String line;
                while ((line = in.readLine()) != null) {
                    out.append(line).append(lineSeparator);
                }
            } catch (IOException e) {
            }
            synchronized (this) {
                done = true;
                notifyAll();
            }
        }

        synchronized void waitUntilDone() throws InterruptedException {
            boolean interrupted = false;

            // poll interrupted flag, while waiting for copy to complete
            while (!(interrupted = Thread.interrupted()) && !done)
                wait(1000);

            if (interrupted)
                throw new InterruptedException();
        }

        String getOutput() {
            return out.toString();
        }

        private final BufferedReader in;
        private final StringBuilder out;
        private boolean done;
    }

    /**
     * Utility class to simplify the handling of temporarily setting a
     * new stream for System.out or System.err.
     */
    static class StreamOutput {
        // Functional interface to set a stream.
        // Expected use: System::setOut, System::setErr
        interface Initializer {
            void set(PrintStream s);
        }

        private final ByteArrayOutputStream baos = new ByteArrayOutputStream();
        private final PrintStream ps = new PrintStream(baos);
        private final PrintStream prev;
        private final Initializer init;

        StreamOutput(PrintStream s, Initializer init) {
            prev = s;
            init.set(ps);
            this.init = init;
        }

        /**
         * Closes the stream and returns the contents that were written to it.
         * @return the contents that were written to it.
         */
        String close() {
            init.set(prev);
            ps.close();
            return baos.toString();
        }
    }

    /**
     * Utility class to simplify the handling of creating an in-memory PrintWriter.
     */
    static class WriterOutput {
        private final StringWriter sw = new StringWriter();
        final PrintWriter pw = new PrintWriter(sw);

        /**
         * Closes the stream and returns the contents that were written to it.
         * @return the contents that were written to it.
         */
        String close() {
            pw.close();
            return sw.toString();
        }
    }
}