Add SSH and SCP command utility and refactor shell utility

core/common/src/main/java/alluxio/util/ShellUtils.java

@@ -21,6 +21,7 @@
 import java.io.InputStreamReader;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -135,8 +136,8 @@ private static UnixMountInfo.Options parseUnixMountOptions(String line) {
   }
 
   @NotThreadSafe
-  private static final class Command {
-    private String[] mCommand;
+  private static class Command {
+    protected String[] mCommand;
 
     private Command(String[] execString) {
       mCommand = execString.clone();
@@ -148,7 +149,7 @@ private Command(String[] execString) {
      * @return the output
      * @throws ExitCodeException if the command returns a non-zero exit code
      */
-    private String run() throws ExitCodeException, IOException {
+    protected String run() throws ExitCodeException, IOException {
       Process process = new ProcessBuilder(mCommand).redirectErrorStream(true).start();
 
       BufferedReader inReader =
@@ -193,6 +194,101 @@ private String run() throws ExitCodeException, IOException {
         process.destroy();
       }
     }
+
+    /**
+     * Runs a command and returns its stdout, stderr and exit code.
+     * No matter it succeeds or not.
+     *
+     * @return {@link CommandReturn} object representation of stdout, stderr and exit code
+     */
+    protected CommandReturn runTolerateFailure() throws IOException {
+      Process process = new ProcessBuilder(mCommand).start();
+
+      BufferedReader inReader =
+              new BufferedReader(new InputStreamReader(process.getInputStream()));
+      BufferedReader errReader =
+              new BufferedReader(new InputStreamReader(process.getErrorStream()));
+
+      try {
+        // read the output of the command
+        StringBuilder stdout = new StringBuilder();
+        String outLine = inReader.readLine();
+        while (outLine != null) {
+          stdout.append(outLine);
+          stdout.append("\n");
+          outLine = inReader.readLine();
+        }
+
+        // read the stderr of the command
+        StringBuilder stderr = new StringBuilder();
+        String errLine = errReader.readLine();
+        while (errLine != null) {
+          stderr.append(errLine);
+          stderr.append("\n");
+          errLine = inReader.readLine();
+        }
+
+        // wait for the process to finish and check the exit code
+        int exitCode = process.waitFor();
+        if (exitCode != 0) {
+          // log error instead of throwing exception
+          LOG.warn(String.format("Failed to run command %s%nExit Code: %s%nStderr: %s",
+                  Arrays.toString(mCommand), exitCode, stderr.toString()));
+        }
+
+        return new CommandReturn(exitCode, stdout.toString(), stderr.toString());
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new IOException(e);
+      } finally {
+        try {
+          // JDK 7 tries to automatically drain the input streams for us
+          // when the process exits, but since close is not synchronized,
+          // it creates a race if we close the stream first and the same
+          // fd is recycled. the stream draining thread will attempt to
+          // drain that fd!! it may block, OOM, or cause bizarre behavior
+          // see: https://bugs.openjdk.java.net/browse/JDK-8024521
+          // issue is fixed in build 7u60
+          InputStream stdoutStream = process.getInputStream();
+
+          synchronized (stdoutStream) {
+            inReader.close();
+          }
+        } catch (IOException e) {
+          LOG.warn("Error while closing the input stream", e);
+        }
+        try {
+          InputStream stderrStream = process.getErrorStream();
+          synchronized (stderrStream) {
+            errReader.close();
+          }
+        } catch (IOException e) {
+          LOG.warn("Error while closing the error stream", e);
+        }
+        process.destroy();
+      }
+    }
+  }
+
+  @NotThreadSafe
+  private static class SshCommand extends Command {
+    protected String mHostName;
+
+    private SshCommand(String[] execString) {
+      this("localhost", execString);
+    }
+
+    private SshCommand(String hostname, String[] execString) {
+      super(execString);
+      mHostName = hostname;
+      mCommand = new String[]{"bash", "-c",
+              String.format("ssh %s %s %s", ShellUtils.COMMON_SSH_OPTS, hostname,
+                      String.join(" ", execString))};
+    }
+
+    protected String getHostName() {
+      return mHostName;
+    }
   }
 
   /**
@@ -234,14 +330,113 @@ public String toString() {
   }
 
   /**
-   * Static method to execute a shell command.
+   * Object representation of a command execution.
+   */
+  public static class CommandReturn {
+    private int mExitCode;
+    private String mStdOut;
+    private String mStdErr;
+
+    /**
+     * Creates object from the contents.
+     *
+     * @param code exit code
+     * @param stdOut stdout content
+     * @param stdErr stderr content
+     */
+    public CommandReturn(int code, String stdOut, String stdErr) {
+      mExitCode = code;
+      mStdOut = stdOut;
+      mStdErr = stdErr;
+    }
+
+    /**
+     * Gets the stdout content.
+     *
+     * @return stdout content
+     */
+    public String getStdOut() {
+      return mStdOut;
+    }
+
+    /**
+     * Gets the stderr content.
+     *
+     * @return stderr content
+     */
+    public String getStdErr() {
+      return mStdErr;
+    }
+
+    /**
+     * Gets the exit code.
+     *
+     * @return exit code of execution
+     */
+    public int getExitCode() {
+      return mExitCode;
+    }
+
+    /**
+     * Formats the object to more readable format.
+     * This is not done in toString() because stdout and stderr may be long.
+     *
+     * @return pretty formatted output
+     */
+    public String getFormattedOutput() {
+      return String.format("StatusCode:%s%nStdOut:%n%s%nStdErr:%n%s", getExitCode(),
+              getStdOut(), getStdErr());
+    }
+  }
+
+  /**
+   * Static method to execute a shell command. The StdErr is not returned.
+   * If execution failed
    *
    * @param cmd shell command to execute
    * @return the output of the executed command
+   * @throws IOException in various situations:
+   *  1. {@link ExitCodeException} when exit code is non-zero
+   *  2. when the executable is not valid, i.e. running ls in Windows
+   *  3. execution interrupted
+   *  4. other normal reasons for IOException
    */
   public static String execCommand(String... cmd) throws IOException {
     return new Command(cmd).run();
   }
 
+  /**
+   * Static method to execute a shell command and tolerate non-zero exit code.
+   * Preserve exit code, stderr and stdout in object.
+   *
+   * @param cmd shell command to execute
+   * @return the output of the executed command
+   * @throws IOException in various situations:
+   *  1. when the executable is not valid, i.e. running ls in Windows
+   *  2. execution interrupted
+   *  3. other normal reasons for IOException
+   */
+  public static CommandReturn execCommandTolerateFailure(String... cmd) throws IOException {
+    return new Command(cmd).runTolerateFailure();
+  }
+
+  /**
+   * Static method to execute a shell command remotely via ssh.
+   * Preserve exit code, stderr and stdout in object.
+   * SSH must be password-less.
+   *
+   * @param hostname Hostname where the command should execute
+   * @param cmd shell command to execute
+   * @return the output of the executed command
+   * @throws IOException in various situations:
+   *  1. when the executable is not valid, i.e. running ls in Windows
+   *  2. execution interrupted
+   *  3. other normal reasons for IOException
+   */
+  public static CommandReturn sshExecCommandTolerateFailure(String hostname, String... cmd)
+          throws IOException {
+    return new SshCommand(hostname, cmd).runTolerateFailure();
+  }
+
   private ShellUtils() {} // prevent instantiation
 }

core/common/src/test/java/alluxio/util/ShellUtilsTest.java

@@ -16,18 +16,33 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assume.assumeTrue;
 
+import alluxio.AlluxioTestDirectory;
 import alluxio.Constants;
 
 import com.google.common.base.Optional;
+import org.junit.Rule;
 import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Paths;
 import java.util.List;
 
 /**
  * Tests the {@link ShellUtils} class.
  */
 public final class ShellUtilsTest {
 
+  @Rule
+  public ExpectedException mExceptionRule = ExpectedException.none();
+
+  public static File createFileInDir(File dir, String fileName) throws IOException {
+    File newFile = new File(Paths.get(dir.getAbsolutePath(), fileName).toString());
+    newFile.createNewFile();
+    return newFile;
+  }
+
   /**
    * Tests the {@link ShellUtils#execCommand(String...)} method.
    *
@@ -41,6 +56,15 @@ public void execCommand() throws Exception {
     assertEquals(testString + "\n", result);
   }
 
+  @Test
+  public void execCommandFail() throws Exception {
+    String testString = "false";
+    mExceptionRule.expect(ShellUtils.ExitCodeException.class);
+    // run a command that guarantees to fail
+    String result = ShellUtils.execCommand("bash", "-c", " " + testString);
+    assertEquals(testString + "\n", result);
+  }
+
   /**
    * Tests the {@link ShellUtils#execCommand(String...)} method for a group of commands.
    *
@@ -110,4 +134,66 @@ public void getMountInfo() throws Exception {
     List<UnixMountInfo> info = ShellUtils.getUnixMountInfo();
     assertTrue(info.size() > 0);
   }
+
+  @Test
+  public void execCommandTolerateFailure() throws Exception {
+    // create temp file
+    File testDir = AlluxioTestDirectory.createTemporaryDirectory("command");
+
+    File testFile = createFileInDir(testDir, "testFile");
+
+    // ls temp file
+    String[] testCommandSucceed = new String[]{"ls",
+            String.format("%s", testDir.getAbsolutePath())};
+    ShellUtils.CommandReturn crs = ShellUtils.execCommandTolerateFailure(testCommandSucceed);
+    assertEquals(0, crs.getExitCode());
+    assertTrue(crs.getStdOut().contains(testFile.getName()));
+
+    // do sth wrong
+    String[] testCommandFail = new String[]{"ls",
+            String.format("%saaaa", testDir.getAbsolutePath())};
+    ShellUtils.CommandReturn crf = ShellUtils.execCommandTolerateFailure(testCommandFail);
+    System.out.println(crf.getFormattedOutput());
+    assertFalse(crf.getExitCode() == 0);
+
+    // if there's no such command there will be IOException
+    mExceptionRule.expect(IOException.class);
+    mExceptionRule.expectMessage("No such file or directory");
+    String[] testCommandExcept = new String[]{"lsa",
+            String.format("%s", testDir.getAbsolutePath())};
+    // lsa is not a valid executable
+    ShellUtils.execCommandTolerateFailure(testCommandExcept);
+  }
+
+  @Test
+  public void sshExecuteCommandTolerateFailure() throws Exception {
+    // create temp file
+    File testDir = AlluxioTestDirectory.createTemporaryDirectory("command");
+
+    File testFile = createFileInDir(testDir, "testFile");
+
+    // the temp file is found
+    String[] testCommandSucceed = new String[]{"ls",
+            String.format("%s", testDir.getAbsolutePath())};
+    ShellUtils.CommandReturn crs = ShellUtils.sshExecCommandTolerateFailure("localhost",
+            testCommandSucceed);
+    assertEquals(0, crs.getExitCode());
+    assertTrue(crs.getStdOut().contains(testFile.getName()));
+
+    // do sth wrong
+    String[] testCommandFail = new String[]{"ls",
+            String.format("%saaaa", testDir.getAbsolutePath())};
+    ShellUtils.CommandReturn crf = ShellUtils.sshExecCommandTolerateFailure("localhost",
+            testCommandFail);
+    assertFalse(crf.getExitCode() == 0);
+
+    // if there's no such command there will be IOException
+    String[] testCommandExcept = new String[]{"lsa",
+            String.format("%s", testDir.getAbsolutePath())};
+    mExceptionRule.expect(IOException.class);
+    mExceptionRule.expectMessage("No such file or directory");
+    // lsa is not a valid executable
+    ShellUtils.CommandReturn cre = ShellUtils.sshExecCommandTolerateFailure("localhost",
+            testCommandExcept);
+  }
 }