Make OutputEncapsulator a CommandRunner

This is in preparation for having other "command changing things" like
redirecting to file and maybe variable substitutions in the arguments.

"On the way" added a
RUNNING "some_command" "--some_flag" "some_value"
to the output encapsulator so that if we run multiple commands, we know what
command was called where.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=219837452

java/google/registry/tools/ShellCommand.java

@@ -28,6 +28,8 @@ import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Ascii;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableTable;
+import com.google.common.escape.Escaper;
+import com.google.common.escape.Escapers;
 import google.registry.util.Clock;
 import google.registry.util.SystemClock;
 import java.io.BufferedReader;
@@ -70,10 +72,25 @@ public class ShellCommand implements Command {
   private static final String ALERT_COLOR = "\u001b[1;41;97m"; // red background
   private static final Duration IDLE_THRESHOLD = Duration.standardHours(1);
   private static final String SUCCESS = "SUCCESS";
-  private static final String FAILURE = "FAILURE ";
+  private static final String FAILURE = "FAILURE";
+  private static final String RUNNING = "RUNNING";
+  private static final Escaper STRING_ESCAPER =
+      Escapers.builder()
+          .addEscape('\\', "\\\\")
+          .addEscape('"', "\\\"")
+          .addEscape('\n', "\\n")
+          .addEscape('\r', "\\r")
+          .addEscape('\t', "\\t")
+          .build();
 
+  /**
+   * The runner we received in the constructor.
+   *
+   * <p>We might want to update this runner based on flags (e.g. --encapsulate_output), but these
+   * flags aren't available in the constructor so we have to do it in the {@link #run} function.
+   */
+  private final CommandRunner originalRunner;
 
-  private final CommandRunner runner;
   private final BufferedReader lineReader;
   private final ConsoleReader consoleReader;
   private final Clock clock;
@@ -96,7 +113,7 @@ public class ShellCommand implements Command {
   boolean encapsulateOutput = false;
 
   public ShellCommand(CommandRunner runner) throws IOException {
-    this.runner = runner;
+    this.originalRunner = runner;
     InputStream in = System.in;
     if (System.console() != null) {
       consoleReader = new ConsoleReader();
@@ -114,7 +131,7 @@ public class ShellCommand implements Command {
 
   @VisibleForTesting
   ShellCommand(BufferedReader bufferedReader, Clock clock, CommandRunner runner) {
-    this.runner = runner;
+    this.originalRunner = runner;
     this.lineReader = bufferedReader;
     this.clock = clock;
     this.consoleReader = null;
@@ -145,42 +162,11 @@ public class ShellCommand implements Command {
     return this;
   }
 
-  private static class OutputEncapsulator {
+  private static class OutputEncapsulator implements CommandRunner {
-    private PrintStream orgStdout;
+    private final CommandRunner runner;
-    private PrintStream orgStderr;
 
-    private EncapsulatingOutputStream encapsulatedOutputStream = null;
+    private OutputEncapsulator(CommandRunner runner) {
-    private EncapsulatingOutputStream encapsulatedErrorStream = null;
+      this.runner = runner;
-
-    private Exception error;
-
-    private OutputEncapsulator() {
-      orgStdout = System.out;
-      orgStderr = System.err;
-      encapsulatedOutputStream = new EncapsulatingOutputStream(System.out, "out: ");
-      encapsulatedErrorStream = new EncapsulatingOutputStream(System.out, "err: ");
-      System.setOut(new PrintStream(encapsulatedOutputStream));
-      System.setErr(new PrintStream(encapsulatedErrorStream));
-    }
-
-    void setError(Exception e) {
-      error = e;
-    }
-
-    private void restoreOriginalStreams() {
-      try {
-        encapsulatedOutputStream.dumpLastLine();
-        encapsulatedErrorStream.dumpLastLine();
-        System.setOut(orgStdout);
-        System.setErr(orgStderr);
-        if (error != null) {
-          emitFailure(error);
-        } else {
-          emitSuccess();
-        }
-      } catch (IOException e) {
-        throw new RuntimeException(e);
-      }
     }
 
     /**
@@ -188,7 +174,7 @@ public class ShellCommand implements Command {
      *
      * <p>Dumps the last line of output prior to doing this.
      */
-    private void emitSuccess() {
+    private static void emitSuccess() {
       System.out.println(SUCCESS);
       System.out.flush();
     }
@@ -198,23 +184,45 @@ public class ShellCommand implements Command {
      *
      * <p>Dumps the last line of output prior to doing this.
      */
-    private void emitFailure(Throwable e) {
+    private static void emitFailure(Throwable e) {
-      System.out.println(
+      System.out.format(
-          FAILURE
+          "%s %s %s\n", FAILURE, e.getClass().getName(), STRING_ESCAPER.escape(e.getMessage()));
-              + e.getClass().getName()
+      System.out.flush();
-              + " "
+    }
-              + e.getMessage().replace("\\", "\\\\").replace("\n", "\\n"));
+
+    private static void emitArguments(String[] args) {
+      System.out.print(RUNNING);
+      Arrays.stream(args).forEach(arg -> System.out.format(" \"%s\"", STRING_ESCAPER.escape(arg)));
+      System.out.println();
+      System.out.flush();
+    }
+
+    private void encapsulatedRun(String[] args) throws Exception {
+      PrintStream orgOut = System.out;
+      PrintStream orgErr = System.err;
+      try (PrintStream newOut =
+              new PrintStream(new EncapsulatingOutputStream(System.out, "out: "));
+          PrintStream newErr =
+              new PrintStream(new EncapsulatingOutputStream(System.out, "err: "))) {
+        System.setOut(newOut);
+        System.setErr(newErr);
+        runner.run(args);
+      } finally {
+        System.setOut(orgOut);
+        System.setErr(orgErr);
+      }
     }
 
     /** Run "func" with output encapsulation. */
-    static void run(CommandRunner runner, String[] args) {
+    @Override
-      OutputEncapsulator encapsulator = new OutputEncapsulator();
+    public void run(String[] args) {
+
       try {
-        runner.run(args);
+        emitArguments(args);
+        encapsulatedRun(args);
+        emitSuccess();
       } catch (Exception e) {
-        encapsulator.setError(e);
+        emitFailure(e);
-      } finally {
-        encapsulator.restoreOriginalStreams();
       }
     }
   }
@@ -222,6 +230,10 @@ public class ShellCommand implements Command {
   /** Run the shell until the user presses "Ctrl-D". */
   @Override
   public void run() {
+    // Wrap standard output and error if requested. We have to do so here in run because the flags
+    // haven't been processed in the constructor.
+    CommandRunner runner =
+        encapsulateOutput ? new OutputEncapsulator(originalRunner) : originalRunner;
     // On Production we want to be extra careful - to prevent accidental use.
     boolean beExtraCareful = (RegistryToolEnvironment.get() == RegistryToolEnvironment.PRODUCTION);
     setPrompt(RegistryToolEnvironment.get(), beExtraCareful);
@@ -242,16 +254,11 @@ public class ShellCommand implements Command {
         continue;
       }
 
-      // Wrap standard output and error if requested. We have to do so here in run because the flags
+      try {
-      // haven't been processed in the constructor.
+        runner.run(lineArgs);
-      if (encapsulateOutput) {
+      } catch (Exception e) {
-        OutputEncapsulator.run(runner, lineArgs);
+        System.err.println("Got an exception:\n" + e);
-      } else {
+        e.printStackTrace();
-        try {
-          runner.run(lineArgs);
-        } catch (Exception e) {
-          System.err.println("Got an exception:\n" + e);
-        }
       }
     }
     if (!encapsulateOutput) {
@@ -566,6 +573,14 @@ public class ShellCommand implements Command {
     @Override
     public void flush() throws IOException {
       dumpLastLine();
+      super.flush();
+    }
+
+    @Override
+    public void close() throws IOException {
+      dumpLastLine();
+      // We do NOT want to call super.close as that would close the original outputStream
+      // (System.out)
     }
 
     /** Dump the accumulated last line of output, if there was one. */

javatests/google/registry/tools/ShellCommandTest.java

@@ -253,10 +253,11 @@ public class ShellCommandTest {
   @Test
   public void testEncapsulatedOutputStream_basicFuncionality() {
     ByteArrayOutputStream backing = new ByteArrayOutputStream();
-    PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
+    try (PrintStream out =
-    out.println("first line");
+        new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "))) {
-    out.print("second line\ntrailing data");
+      out.println("first line");
-    out.flush();
+      out.print("second line\ntrailing data");
+    }
     assertThat(backing.toString())
         .isEqualTo("out: first line\nout: second line\nout: trailing data\n");
   }
@@ -264,8 +265,8 @@ public class ShellCommandTest {
   @Test
   public void testEncapsulatedOutputStream_emptyStream() {
     ByteArrayOutputStream backing = new ByteArrayOutputStream();
-    PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
+    try (PrintStream out =
-    out.flush();
+        new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "))) {}
     assertThat(backing.toString()).isEqualTo("");
   }
 
@@ -288,10 +289,31 @@ public class ShellCommandTest {
     assertThat(stderr.toString()).isEmpty();
     assertThat(stdout.toString())
         .isEqualTo(
-            "out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n"
+            "RUNNING \"command1\"\n"
+                + "out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n"
                 + "SUCCESS\n");
   }
 
+  @Test
+  public void testEncapsulatedOutput_throws() throws Exception {
+    RegistryToolEnvironment.ALPHA.setup();
+    captureOutput();
+    ShellCommand shellCommand =
+        new ShellCommand(
+            args -> {
+              System.out.println("first line");
+              throw new Exception("some error!");
+            });
+    shellCommand.encapsulateOutput = true;
+    shellCommand.run();
+    assertThat(stderr.toString()).isEmpty();
+    assertThat(stdout.toString())
+        .isEqualTo(
+            "RUNNING \"command1\"\n"
+                + "out: first line\n"
+                + "FAILURE java.lang.Exception some error!\n");
+  }
+
   @Test
   public void testEncapsulatedOutput_noCommand() throws Exception {
     captureOutput();
@@ -307,7 +329,7 @@ public class ShellCommandTest {
     shellCommand.run();
     assertThat(stderr.toString()).isEmpty();
     assertThat(stdout.toString())
-        .isEqualTo("out: first line\nSUCCESS\n");
+        .isEqualTo("RUNNING \"do\" \"something\"\nout: first line\nSUCCESS\n");
   }
 
   void captureOutput() {