org.apache.drill.yarn.scripts.ScriptUtils.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.drill.yarn.scripts.ScriptUtils.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.drill.yarn.scripts;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;

import org.apache.commons.io.FileUtils;

public class ScriptUtils {

    private static ScriptUtils instance = new ScriptUtils();
    public File distribDir;
    public File javaHome;
    public File testDir;
    public File testDrillHome;
    public File testSiteDir;
    public File testLogDir;
    public boolean externalLogDir;

    /**
     * Out-of-the-box command-line arguments when launching sqlline.
     * Order is not important here (though it is to Java.)
     */

    public static String sqlLineArgs[] = makeSqlLineArgs();

    private static String[] makeSqlLineArgs() {
        String args[] = { "-Dlog.path=/.*/drill/log/sqlline\\.log",
                "-Dlog.query.path=/.*/drill/log/sqlline_queries\\.json", "-XX:MaxPermSize=512M",
                "sqlline\\.SqlLine", "-d", "org\\.apache\\.drill\\.jdbc\\.Driver", "--maxWidth=10000",
                "--color=true" };

        // Special handling if this machine happens to have the default
        // /var/log/drill log location.

        if (new File("/var/log/drill").exists()) {
            args[0] = "-Dlog\\.path=/var/log/drill/sqlline\\.log";
            args[1] = "-Dlog\\.query\\.path=/var/log/drill/sqlline_queries\\.json";
        }
        return args;
    }

    public static final boolean USE_SOURCE = true;
    public static final String TEMP_DIR = "/tmp";
    public static boolean useSource = USE_SOURCE;

    private ScriptUtils() {
        String drillScriptsDir = System.getProperty("drillScriptDir");
        assertNotNull(drillScriptsDir);
        distribDir = new File(drillScriptsDir);
        javaHome = new File(System.getProperty("java.home"));
    }

    public static ScriptUtils instance() {
        return instance;
    }

    public ScriptUtils fromSource(String sourceDir) {
        useSource = true;
        return this;
    }

    public ScriptUtils fromDistrib(String distrib) {
        distribDir = new File(distrib);
        useSource = false;
        return this;
    }

    /**
     * Out-of-the-box command-line arguments when launching Drill.
     * Order is not important here (though it is to Java.)
     */

    public static String stdArgs[] = buildStdArgs();

    private static String[] buildStdArgs() {
        String args[] = { "-Xms4G", "-Xmx4G", "-XX:MaxDirectMemorySize=8G", "-XX:MaxPermSize=512M",
                "-XX:ReservedCodeCacheSize=1G",
                // Removed in Drill 1.8
                //      "-Ddrill\\.exec\\.enable-epoll=true",
                "-XX:\\+CMSClassUnloadingEnabled", "-XX:\\+UseG1GC",
                "org\\.apache\\.drill\\.exec\\.server\\.Drillbit",
                "-Dlog\\.path=/.*/script-test/drill/log/drillbit\\.log",
                "-Dlog\\.query\\.path=/.*/script-test/drill/log/drillbit_queries\\.json", };

        // Special handling if this machine happens to have the default
        // /var/log/drill log location.

        if (new File("/var/log/drill").exists()) {
            args[args.length - 2] = "-Dlog\\.path=/var/log/drill/drillbit\\.log";
            args[args.length - 1] = "-Dlog\\.query\\.path=/var/log/drill/drillbit_queries\\.json";
        }
        return args;
    };

    /**
     * Out-of-the-box class-path before any custom additions.
     */

    static String stdCp[] = { "conf", "jars/*", "jars/ext/*", "jars/3rdparty/*", "jars/classb/*" };

    /**
     * Directories to create to simulate a Drill distribution.
     */

    static String distribDirs[] = { "bin", "jars", "jars/3rdparty", "jars/ext", "conf" };

    /**
     * Out-of-the-box Jar directories.
     */

    static String jarDirs[] = { "jars", "jars/3rdparty", "jars/ext", };

    /**
     * Scripts we must copy from the source tree to create a simulated
     * Drill bin directory.
     */

    public static String scripts[] = { "drill-config.sh", "drill-embedded", "drill-localhost", "drill-on-yarn.sh",
            "drillbit.sh", "drill-conf",
            //dumpcat
            //hadoop-excludes.txt
            "runbit", "sqlline",
            //sqlline.bat
            //submit_plan
            "yarn-drillbit.sh" };

    /**
     * Create the basic test directory. Tests add or remove details.
     */

    public void initialSetup() throws IOException {
        File tempDir = new File(TEMP_DIR);
        testDir = new File(tempDir, "script-test");
        testDrillHome = new File(testDir, "drill");
        testSiteDir = new File(testDir, "site");
        File varLogDrill = new File("/var/log/drill");
        if (varLogDrill.exists()) {
            testLogDir = varLogDrill;
            externalLogDir = true;
        } else {
            testLogDir = new File(testDrillHome, "log");
        }
        if (testDir.exists()) {
            FileUtils.forceDelete(testDir);
        }
        testDir.mkdirs();
        testSiteDir.mkdir();
        testLogDir.mkdir();
    }

    public void createMockDistrib() throws IOException {
        if (ScriptUtils.useSource) {
            buildFromSource();
        } else {
            buildFromDistrib();
        }
    }

    /**
     * Build the Drill distribution directory directly from sources.
     */

    private void buildFromSource() throws IOException {
        createMockDirs();
        copyScripts(ScriptUtils.instance().distribDir);
    }

    /**
     * Build the shell of a Drill distribution directory by creating the required
     * directory structure.
     */

    private void createMockDirs() throws IOException {
        if (testDrillHome.exists()) {
            FileUtils.forceDelete(testDrillHome);
        }
        testDrillHome.mkdir();
        for (String path : ScriptUtils.distribDirs) {
            File subDir = new File(testDrillHome, path);
            subDir.mkdirs();
        }
        for (String path : ScriptUtils.jarDirs) {
            makeDummyJar(new File(testDrillHome, path), "dist");
        }
    }

    /**
     * The tests should not require jar files, but we simulate them to be a bit
     * more realistic. Since we dont' run Java, the jar files can be simulated.
     */

    public File makeDummyJar(File dir, String prefix) throws IOException {
        String jarName = "";
        if (prefix != null) {
            jarName += prefix + "-";
        }
        jarName += dir.getName() + ".jar";
        File jarFile = new File(dir, jarName);
        writeFile(jarFile, "Dummy jar");
        return jarFile;
    }

    /**
     * Create a simple text file with the given contents.
     */

    public void writeFile(File file, String contents) throws IOException {
        try (PrintWriter out = new PrintWriter(new FileWriter(file))) {
            out.println(contents);
        }
    }

    /**
     * Create a drill-env.sh or distrib-env.sh file with the given environment in
     * the recommended format.
     */

    public void createEnvFile(File file, Map<String, String> env) throws IOException {
        try (PrintWriter out = new PrintWriter(new FileWriter(file))) {
            out.println("#!/usr/bin/env bash");
            for (String key : env.keySet()) {
                String value = env.get(key);
                out.print("export ");
                out.print(key);
                out.print("=${");
                out.print(key);
                out.print(":-\"");
                out.print(value);
                out.println("\"}");
            }
        }
    }

    /**
     * Copy the standard scripts from source location to the mock distribution
     * directory.
     */

    private void copyScripts(File sourceDir) throws IOException {
        File binDir = new File(testDrillHome, "bin");
        for (String script : ScriptUtils.scripts) {
            File source = new File(sourceDir, script);
            File dest = new File(binDir, script);
            copyFile(source, dest);
            dest.setExecutable(true);
        }

        // Create the "magic" wrapper script that simulates the Drillbit and
        // captures the output we need for testing.

        String wrapper = "wrapper.sh";
        File dest = new File(binDir, wrapper);
        try (InputStream is = getClass().getResourceAsStream("/" + wrapper)) {
            Files.copy(is, dest.toPath(), StandardCopyOption.REPLACE_EXISTING);
        }
        dest.setExecutable(true);
    }

    private void buildFromDistrib() {
        // TODO Auto-generated method stub

    }

    /**
     * Consume the input from a stream, specifically the stderr or stdout stream
     * from a process.
     *
     * @see http://stackoverflow.com/questions/14165517/processbuilder-forwarding-stdout-and-stderr-of-started-processes-without-blocki
     */

    private static class StreamGobbler extends Thread {
        InputStream is;
        public StringBuilder buf = new StringBuilder();

        private StreamGobbler(InputStream is) {
            this.is = is;
        }

        @Override
        public void run() {
            try {
                InputStreamReader isr = new InputStreamReader(is);
                BufferedReader br = new BufferedReader(isr);
                String line = null;
                while ((line = br.readLine()) != null) {
                    buf.append(line);
                    buf.append("\n");
                }
            } catch (IOException ioe) {
                ioe.printStackTrace();
            }
        }
    }

    /**
     * Handy run result class to capture the information we need for testing and
     * to do various kinds of validation on it.
     */

    public static class RunResult {
        File logDir;
        File logFile;
        String stdout;
        String stderr;
        List<String> echoArgs;
        int returnCode;
        String classPath[];
        String libPath[];
        String log;
        public File pidFile;
        public File outFile;
        String out;

        /**
         * Split the class path into strings for easier validation.
         */

        public void analyze() {
            if (echoArgs == null) {
                return;
            }
            for (int i = 0; i < echoArgs.size(); i++) {
                String arg = echoArgs.get(i);
                if (arg.equals("-cp")) {
                    classPath = Pattern.compile(":").split((echoArgs.get(i + 1)));
                    break;
                }
            }
            String probe = "-Djava.library.path=";
            for (int i = 0; i < echoArgs.size(); i++) {
                String arg = echoArgs.get(i);
                if (arg.startsWith(probe)) {
                    assertNull(libPath);
                    libPath = Pattern.compile(":").split((arg.substring(probe.length())));
                    break;
                }
            }
        }

        /**
         * Read the log file, if any, generated by the process.
         */

        public void loadLog() throws IOException {
            log = loadFile(logFile);
        }

        private String loadFile(File file) throws IOException {
            StringBuilder buf = new StringBuilder();
            try (BufferedReader reader = new BufferedReader(new FileReader(file))) {
                String line;
                while ((line = reader.readLine()) != null) {
                    buf.append(line);
                    buf.append("\n");
                }
                return buf.toString();
            } catch (FileNotFoundException e) {
                return null;
            }
        }

        /**
         * Validate that the first argument invokes Java correctly.
         */

        public void validateJava() {
            assertNotNull(echoArgs);
            String java = instance.javaHome + "/bin/java";
            List<String> actual = echoArgs;
            assertEquals(java, actual.get(0));
        }

        public boolean containsArg(String arg) {
            for (String actual : echoArgs) {
                if (actual.equals(arg)) {
                    return true;
                }
            }
            return false;
        }

        public void validateStockArgs() {
            for (String arg : ScriptUtils.stdArgs) {
                assertTrue("Argument not found: " + arg + " in " + echoArgs, containsArgRegex(arg));
            }
        }

        public void validateArg(String arg) {
            validateArgs(Collections.singletonList(arg));
        }

        public void validateArgs(String args[]) {
            validateArgs(Arrays.asList(args));
        }

        public void validateArgs(List<String> args) {
            validateJava();
            for (String arg : args) {
                assertTrue(containsArg(arg));
            }
        }

        public void validateArgRegex(String arg) {
            assertTrue(containsArgRegex(arg));
        }

        public void validateArgsRegex(List<String> args) {
            assertTrue(containsArgsRegex(args));
        }

        public boolean containsArgsRegex(List<String> args) {
            for (String arg : args) {
                if (!containsArgRegex(arg)) {
                    return false;
                }
            }
            return true;
        }

        public boolean containsArgsRegex(String args[]) {
            for (String arg : args) {
                if (!containsArgRegex(arg)) {
                    return false;
                }
            }
            return true;
        }

        public boolean containsArgRegex(String arg) {
            for (String actual : echoArgs) {
                if (actual.matches(arg)) {
                    return true;
                }
            }
            return false;
        }

        public void validateClassPath(String expectedCP) {
            assertTrue(classPathContains(expectedCP));
        }

        public void validateClassPath(String expectedCP[]) {
            assertTrue(classPathContains(expectedCP));
        }

        public boolean classPathContains(String expectedCP[]) {
            for (String entry : expectedCP) {
                if (!classPathContains(entry)) {
                    return false;
                }
            }
            return true;
        }

        public boolean classPathContains(String expectedCP) {
            if (classPath == null) {
                fail("No classpath returned");
            }
            String tail = "/" + instance.testDir.getName() + "/" + instance.testDrillHome.getName() + "/";
            String expectedPath;
            if (expectedCP.startsWith("/")) {
                expectedPath = expectedCP;
            } else {
                expectedPath = tail + expectedCP;
            }
            for (String entry : classPath) {
                if (entry.endsWith(expectedPath)) {
                    return true;
                }
            }
            return false;
        }

        public void loadOut() throws IOException {
            out = loadFile(outFile);
        }

        /**
         * Ensure that the Drill log file contains at least the sample message
         * written by the wrapper.
         */

        public void validateDrillLog() {
            assertNotNull(log);
            assertTrue(log.contains("Drill Log Message"));
        }

        /**
         * Validate that the stdout contained the expected message.
         */

        public void validateStdOut() {
            assertTrue(stdout.contains("Starting drillbit on"));
        }

        /**
         * Validate that the stderr contained the sample error message from the
         * wrapper.
         */

        public void validateStdErr() {
            assertTrue(stderr.contains("Stderr Message"));
        }

        public int getPid() throws IOException {
            try (BufferedReader reader = new BufferedReader(new FileReader(pidFile))) {
                return Integer.parseInt(reader.readLine());
            } finally {
            }
        }

    }

    /**
     * The "business end" of the tests: runs drillbit.sh and captures results.
     */

    public static class ScriptRunner {
        // Drillbit commands

        public static String DRILLBIT_RUN = "run";
        public static String DRILLBIT_START = "start";
        public static String DRILLBIT_STATUS = "status";
        public static String DRILLBIT_STOP = "stop";
        public static String DRILLBIT_RESTART = "restart";

        public File cwd = instance.testDir;
        public File drillHome = instance.testDrillHome;
        public String script;
        public List<String> args = new ArrayList<>();
        public Map<String, String> env = new HashMap<>();
        public File logDir;
        public File pidFile;
        public File outputFile;
        public boolean preserveLogs;

        public ScriptRunner(String script) {
            this.script = script;
        }

        public ScriptRunner(String script, String cmd) {
            this(script);
            args.add(cmd);
        }

        public ScriptRunner(String script, String cmdArgs[]) {
            this(script);
            for (String arg : cmdArgs) {
                args.add(arg);
            }
        }

        public ScriptRunner withArg(String arg) {
            args.add(arg);
            return this;
        }

        public ScriptRunner withSite(File siteDir) {
            if (siteDir != null) {
                args.add("--site");
                args.add(siteDir.getAbsolutePath());
            }
            return this;
        }

        public ScriptRunner withEnvironment(Map<String, String> env) {
            if (env != null) {
                this.env.putAll(env);
            }
            return this;
        }

        public ScriptRunner addEnv(String key, String value) {
            env.put(key, value);
            return this;
        }

        public ScriptRunner withLogDir(File logDir) {
            this.logDir = logDir;
            return this;
        }

        public ScriptRunner preserveLogs() {
            preserveLogs = true;
            return this;
        }

        public RunResult run() throws IOException {
            File binDir = new File(drillHome, "bin");
            File scriptFile = new File(binDir, script);
            assertTrue(scriptFile.exists());
            outputFile = new File(instance.testDir, "output.txt");
            outputFile.delete();
            if (logDir == null) {
                logDir = instance.testLogDir;
            }
            if (!preserveLogs) {
                cleanLogs(logDir);
            }

            Process proc = startProcess(scriptFile);
            RunResult result = runProcess(proc);
            if (result.returnCode == 0) {
                captureOutput(result);
                captureLog(result);
            }
            return result;
        }

        private void cleanLogs(File logDir) throws IOException {
            if (logDir == instance.testLogDir && instance.externalLogDir) {
                return;
            }
            if (logDir.exists()) {
                FileUtils.forceDelete(logDir);
            }
        }

        private Process startProcess(File scriptFile) throws IOException {
            outputFile.delete();
            List<String> cmd = new ArrayList<>();
            cmd.add(scriptFile.getAbsolutePath());
            cmd.addAll(args);
            ProcessBuilder pb = new ProcessBuilder().command(cmd).directory(cwd);
            Map<String, String> pbEnv = pb.environment();
            pbEnv.clear();
            pbEnv.putAll(env);
            File binDir = new File(drillHome, "bin");
            File wrapperCmd = new File(binDir, "wrapper.sh");

            // Set the magic wrapper to capture output.

            pbEnv.put("_DRILL_WRAPPER_", wrapperCmd.getAbsolutePath());
            pbEnv.put("JAVA_HOME", instance.javaHome.getAbsolutePath());
            return pb.start();
        }

        private RunResult runProcess(Process proc) {
            StreamGobbler errorGobbler = new StreamGobbler(proc.getErrorStream());
            StreamGobbler outputGobbler = new StreamGobbler(proc.getInputStream());
            outputGobbler.start();
            errorGobbler.start();

            try {
                proc.waitFor();
            } catch (InterruptedException e) {
                // Won't occur.
            }

            RunResult result = new RunResult();
            result.stderr = errorGobbler.buf.toString();
            result.stdout = outputGobbler.buf.toString();
            result.returnCode = proc.exitValue();
            return result;
        }

        private void captureOutput(RunResult result) throws IOException {
            // Capture the Java arguments which the wrapper script wrote to a file.

            try (BufferedReader reader = new BufferedReader(new FileReader(outputFile))) {
                result.echoArgs = new ArrayList<>();
                String line;
                while ((line = reader.readLine()) != null) {
                    result.echoArgs.add(line);
                }
                result.analyze();
            } catch (FileNotFoundException e) {
                ;
            }
        }

        private void captureLog(RunResult result) throws IOException {
            result.logDir = logDir;
            result.logFile = new File(logDir, "drillbit.log");
            if (result.logFile.exists()) {
                result.loadLog();
            } else {
                result.logFile = null;
            }
        }
    }

    public static class DrillbitRun extends ScriptRunner {
        public File pidDir;

        public DrillbitRun() {
            super("drillbit.sh");
        }

        public DrillbitRun(String cmd) {
            super("drillbit.sh", cmd);
        }

        public DrillbitRun withPidDir(File pidDir) {
            this.pidDir = pidDir;
            return this;
        }

        public DrillbitRun asDaemon() {
            addEnv("KEEP_RUNNING", "1");
            return this;
        }

        public RunResult start() throws IOException {
            if (pidDir == null) {
                pidDir = drillHome;
            }
            pidFile = new File(pidDir, "drillbit.pid");
            // pidFile.delete();
            asDaemon();
            RunResult result = run();
            if (result.returnCode == 0) {
                capturePidFile(result);
                captureDrillOut(result);
            }
            return result;
        }

        private void capturePidFile(RunResult result) {
            assertTrue(pidFile.exists());
            result.pidFile = pidFile;
        }

        private void captureDrillOut(RunResult result) throws IOException {
            // Drillbit.out

            result.outFile = new File(result.logDir, "drillbit.out");
            if (result.outFile.exists()) {
                result.loadOut();
            } else {
                result.outFile = null;
            }
        }

    }

    /**
     * Build a "starter" conf or site directory by creating a mock
     * drill-override.conf file.
     */

    public void createMockConf(File siteDir) throws IOException {
        createDir(siteDir);
        File override = new File(siteDir, "drill-override.conf");
        writeFile(override, "# Dummy override");
    }

    public void removeDir(File dir) throws IOException {
        if (dir.exists()) {
            FileUtils.forceDelete(dir);
        }
    }

    /**
     * Remove, then create a directory.
     */

    public File createDir(File dir) throws IOException {
        removeDir(dir);
        dir.mkdirs();
        return dir;
    }

    public void copyFile(File source, File dest) throws IOException {
        Files.copy(source.toPath(), dest.toPath(), StandardCopyOption.REPLACE_EXISTING);
    }

}