CLI: Close subcommands in MultiCommand

server/cli/src/main/java/org/elasticsearch/cli/MultiCommand.java

@@ -19,6 +19,8 @@
 
 package org.elasticsearch.cli;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -74,4 +76,48 @@ protected void execute(Terminal terminal, OptionSet options) throws Exception {
         }
         subcommand.mainWithoutErrorHandling(Arrays.copyOfRange(args, 1, args.length), terminal);
     }
+
+    @Override
+    public void close() throws IOException, RuntimeException {
+        Throwable th = null;
+        for (Closeable object : subcommands.values()) {
+            try {
+                if (object != null) {
+                    object.close();
+                }
+            } catch (Throwable t) {
+                addSuppressed(th, t);
+                if (th == null) {
+                    th = t;
+                }
+            }
+        }
+        if (th != null) {
+            throw reThrowAlways(th);
+        }
+    }
+
+    // Following methods are similar to IOUtils,
+    // avoiding lucene dependency in CLI.
+    private static void addSuppressed(Throwable t, Throwable suppressed) {
+        if (t != null && suppressed != null) {
+            t.addSuppressed(suppressed);
+        }
+    }
+
+    private static Error reThrowAlways(Throwable th) throws IOException, RuntimeException {
+        if (th instanceof IOException) {
+            throw (IOException) th;
+        }
+
+        if (th instanceof RuntimeException) {
+            throw (RuntimeException) th;
+        }
+
+        if (th instanceof Error) {
+            throw (Error) th;
+        }
+
+        throw new RuntimeException(th);
+    }
 }

server/src/test/java/org/elasticsearch/cli/MultiCommandTests.java

@@ -21,6 +21,14 @@
 
 import joptsimple.OptionSet;
 import org.junit.Before;
+import org.mockito.Mockito;
+
+import java.io.IOException;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
 
 public class MultiCommandTests extends CommandTestCase {
 
@@ -102,4 +110,32 @@ public void testSubcommandArguments() throws Exception {
         assertFalse(output, output.contains("command1"));
         assertTrue(output, output.contains("Arguments: [foo, bar]"));
     }
+
+    public void testClose() throws Exception {
+        Command spySubCommand1 = spy(new DummySubCommand());
+        Command spySubCommand2 = spy(new DummySubCommand());
+        multiCommand.subcommands.put("command1", spySubCommand1);
+        multiCommand.subcommands.put("command2", spySubCommand2);
+        multiCommand.close();
+        verify(spySubCommand1, times(1)).close();
+        verify(spySubCommand2, times(1)).close();
+    }
+
+    public void testCloseWhenSubCommandCloseThrowsException() throws Exception {
+        Command subCommand1 = mock(DummySubCommand.class);
+        Mockito.doThrow(new IOException()).when(subCommand1).close();
+        Command spySubCommand2 = spy(new DummySubCommand());
+        multiCommand.subcommands.put("command1", subCommand1);
+        multiCommand.subcommands.put("command2", spySubCommand2);
+        IOException ioe = null;
+        try {
+            multiCommand.close();
+        } catch (IOException e) {
+            ioe = e;
+        }
+        // verify exception is thrown, as well as other non failed sub-commands closed
+        // properly.
+        assertNotNull("Expected IOException", ioe);
+        verify(spySubCommand2, times(1)).close();
+    }
 }