BIGTOP-835. The shell exec method must have variants which have timeout and can run in background

2 files changed, 135 insertions, 9 deletions

diff --git a/bigtop-test-framework/src/main/groovy/org/apache/bigtop/itest/shell/Shell.groovy b/bigtop-test-framework/src/main/groovy/org/apache/bigtop/itest/shell/Shell.groovy
index ae3da68b..4df36422 100644
--- a/bigtop-test-framework/src/main/groovy/org/apache/bigtop/itest/shell/Shell.groovy
+++ b/bigtop-test-framework/src/main/groovy/org/apache/bigtop/itest/shell/Shell.groovy
@@ -62,10 +62,13 @@ class Shell {
    * stdout as getOut() and stderr as getErr(). The script itself can be accessed
    * as getScript()
    * WARNING: it isn't thread safe
+   * @param timeout timeout in milliseconds to wait before killing the script.
+   * If timeout < 0, then this method will wait until the script completes
+   * and will not be killed.
    * @param args shell script split into multiple Strings
    * @return Shell object for chaining
    */
-  Shell exec(Object... args) {
+  Shell execWithTimeout(int timeout, Object... args) {
     def proc = user ? "sudo -u $user PATH=${System.getenv('PATH')} $shell".execute() :
                                     "$shell".execute()
     script = args.join("\n")
@@ -78,20 +81,30 @@ class Shell {
       writer.println(script)
       writer.close()
     }
-    ByteArrayOutputStream baosErr = new ByteArrayOutputStream(4096);
-    proc.consumeProcessErrorStream(baosErr);
-    out = proc.in.readLines()
+    ByteArrayOutputStream outStream = new ByteArrayOutputStream(4096)
+    ByteArrayOutputStream errStream = new ByteArrayOutputStream(4096)
+    Thread.start {
+      proc.consumeProcessOutput(outStream, errStream)
+    }
+    if (timeout >= 0) {
+      proc.waitForOrKill(timeout)
+    } else {
+      proc.waitFor()
+    }
 
     // Possibly a bug in String.split as it generates a 1-element array on an
     // empty String
-    if (baosErr.size() != 0) {
-      err = baosErr.toString().split('\n');
+    if (outStream.size() != 0) {
+      out = outStream.toString().split('\n')
+    } else {
+      out = Collections.EMPTY_LIST
+    }
+    if (errStream.size() != 0) {
+      err = errStream.toString().split('\n')
     }
     else {
-      err = new ArrayList<String>();
+      err = Collections.EMPTY_LIST
     }
-
-    proc.waitFor()
     ret = proc.exitValue()
 
     if (LOG.isTraceEnabled()) {
@@ -105,7 +118,74 @@ class Shell {
            LOG.trace("\n<stderr>\n${err.join('\n')}\n</stderr>");
         }
     }
+    return this
+  }
+
+  /**
+   * Execute shell script consisting of as many Strings as we have arguments,
+   * possibly under an explicit username (requires sudoers privileges).
+   * NOTE: individual strings are concatenated into a single script as though
+   * they were delimited with new line character. All quoting rules are exactly
+   * what one would expect in standalone shell script.
+   *
+   * After executing the script its return code can be accessed as getRet(),
+   * stdout as getOut() and stderr as getErr(). The script itself can be accessed
+   * as getScript()
+   * WARNING: it isn't thread safe
+   * @param timeout timeout in milliseconds to wait before killing the script
+   * . If timeout < 0, then this method will wait until the script completes
+   * and will not be killed.
+   * @param args shell script split into multiple Strings
+   * @return Shell object for chaining
+   */
+  Shell exec(Object... args) {
+    return execWithTimeout(-1, args)
+  }
 
+  /**
+   * Executes a shell script consisting of as many strings as we have args,
+   * under an explicit user name. This method does the same job as
+   * {@linkplain #exec(java.lang.Object[])}, but will return immediately,
+   * with the process continuing execution in the background. If this method
+   * is called, the output stream and error stream of this script  will be
+   * available in the {@linkplain #out} and {@linkplain #err} lists.
+   * WARNING: it isn't thread safe
+   * <strong>CAUTION:</strong>
+   * If this shell object is used to run other script while a script is
+   * being executed in the background, then the output stream and error
+   * stream of the script executed later will be what is available,
+   * and the output and error streams of this script may be lost.
+   * @param args
+   * @return Shell object for chaining
+   */
+  Shell fork(Object... args) {
+    forkWithTimeout(-1, args)
+    return this
+  }
+
+  /**
+   * Executes a shell script consisting of as many strings as we have args,
+   * under an explicit user name. This method does the same job as
+   * {@linkplain #execWithTimeout(int, java.lang.Object[])}, but will return immediately,
+   * with the process continuing execution in the background for timeout
+   * milliseconds (or until the script completes , whichever is earlier). If
+   * this method
+   * is called, the output stream and error stream of this script will be
+   * available in the {@linkplain #out} and {@linkplain #err} lists.
+   * WARNING: it isn't thread safe
+   * <strong>CAUTION:</strong>
+   * If this shell object is used to run other script while a script is
+   * being executed in the background, then the output stream and error
+   * stream of the script executed later will be what is available,
+   * and the output and error streams of this script may be lost.
+   * @param timeout The timoeut in milliseconds before the script is killed
+   * @param args
+   * @return Shell object for chaining
+   */
+  Shell forkWithTimeout(int timeout, Object... args) {
+    Thread.start {
+      execWithTimeout(timeout, args)
+    }
     return this
   }
 }
diff --git a/bigtop-test-framework/src/test/groovy/org/apache/bigtop/itest/shell/ShellTest.groovy b/bigtop-test-framework/src/test/groovy/org/apache/bigtop/itest/shell/ShellTest.groovy
index 1571e100..63ca7ee3 100644
--- a/bigtop-test-framework/src/test/groovy/org/apache/bigtop/itest/shell/ShellTest.groovy
+++ b/bigtop-test-framework/src/test/groovy/org/apache/bigtop/itest/shell/ShellTest.groovy
@@ -23,6 +23,7 @@ import org.junit.Test
 
 import static org.junit.Assert.assertEquals
 import static org.junit.Assert.assertFalse
+import static org.junit.Assert.assertTrue
 
 class ShellTest {
   @Test
@@ -48,4 +49,49 @@ class ShellTest {
     assertEquals("got extra stdout ${sh.out}", 0, sh.out.size())
     assertEquals("got extra stderr ${sh.err}", 0, sh.err.size())
   }
+
+  @Test
+  void testRegularShellWithTimeout() {
+    Shell sh = new Shell("/bin/bash -s")
+    def preTime = System.currentTimeMillis()
+    sh.execWithTimeout(500, 'A=a ; sleep 30 ; r() { return $1; } ; echo $A ; ' +
+      'r `id -u`')
+    assertTrue(System.currentTimeMillis() - preTime < 30000)
+  }
+
+  @Test
+  void testBackgroundExecWithTimeoutFailure() {
+    Shell sh = new Shell("/bin/bash -s")
+    def preTime = System.currentTimeMillis()
+    sh.forkWithTimeout(500, 'A=a ; sleep 30 ; r() { return $1; } ' +
+      '; echo $A ; r `id -u`')
+    //Wait for script to get killed
+    Thread.sleep(750)
+    assertTrue("got wrong stdout ${sh.out}", sh.out.isEmpty())
+    assertTrue("got wrong srderr ${sh.out}", sh.err.isEmpty())
+  }
+
+  @Test
+  void testForkWithTimeoutSuccess() {
+    Shell sh = new Shell("/bin/bash -s")
+    def preTime = System.currentTimeMillis()
+    sh.forkWithTimeout(30000,
+      'A=a ; r() { return $1; } ; echo $A ; r `id -u`')
+    //Wait for the script to complete
+    Thread.sleep(500)
+    assertFalse("${sh.script} exited with a non-zero status", sh.ret == 0)
+    assertEquals("got wrong stdout ${sh.out}", "a", sh.out[0])
+    assertEquals("got extra stderr ${sh.err}", 0, sh.err.size())
+  }
+
+  @Test
+  void testForkWithoutTimeoutSuccess() {
+    Shell sh = new Shell("/bin/bash -s")
+    sh.fork('A=a ; r() { return $1; } ; echo $A ; r `id -u`')
+    //Wait for the script to complete
+    Thread.sleep(550)
+    assertFalse("${sh.script} exited with a non-zero status", sh.ret == 0)
+    assertEquals("got wrong stdout ${sh.out}", "a", sh.out[0])
+    assertEquals("got extra stderr ${sh.err}", 0, sh.err.size())
+  }
 }