CLI: Close subcommands in MultiCommand

server/cli/build.gradle

@@ -36,6 +36,7 @@ archivesBaseName = 'elasticsearch-cli'
 
 dependencies {
     compile 'net.sf.jopt-simple:jopt-simple:5.0.2'
+    compile "org.elasticsearch:elasticsearch-core:${version}"
 }
 
 test.enabled = false

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

@@ -19,13 +19,17 @@
 
 package org.elasticsearch.cli;
 
+import java.io.Closeable;
+import java.io.IOException;
 import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.Map;
 
 import joptsimple.NonOptionArgumentSpec;
 import joptsimple.OptionSet;
 
+import org.elasticsearch.core.internal.io.IOUtils;
+
 /**
  * A cli tool which is made up of multiple subcommands.
  */
@@ -74,4 +78,10 @@ protected void execute(Terminal terminal, OptionSet options) throws Exception {
         }
         subcommand.mainWithoutErrorHandling(Arrays.copyOfRange(args, 1, args.length), terminal);
     }
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.close(subcommands.values());
+    }
+
 }

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

@@ -22,22 +22,46 @@
 import joptsimple.OptionSet;
 import org.junit.Before;
 
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
 public class MultiCommandTests extends CommandTestCase {
 
     static class DummyMultiCommand extends MultiCommand {
+
+        AtomicBoolean closed = new AtomicBoolean();
         DummyMultiCommand() {
             super("A dummy multi command", () -> {});
         }
+        @Override
+        public void close() throws IOException {
+            super.close();
+            this.closed.compareAndSet(false, true);
+        }
     }
 
     static class DummySubCommand extends Command {
+        AtomicBoolean throwsExceptionOnClose = new AtomicBoolean();
+        AtomicBoolean closed = new AtomicBoolean();
         DummySubCommand() {
             super("A dummy subcommand", () -> {});
         }
+        DummySubCommand(boolean throwsExceptionOnClose) {
+            super("A dummy subcommand", () -> {});
+            this.throwsExceptionOnClose.compareAndSet(false, throwsExceptionOnClose);
+        }
         @Override
         protected void execute(Terminal terminal, OptionSet options) throws Exception {
             terminal.println("Arguments: " + options.nonOptionArguments().toString());
         }
+        @Override
+        public void close() throws IOException {
+            if (throwsExceptionOnClose.get()) {
+                throw new IOException();
+            } else {
+                closed.compareAndSet(false, true);
+            }
+        }
     }
 
     DummyMultiCommand multiCommand;
@@ -102,4 +126,27 @@ public void testSubcommandArguments() throws Exception {
         assertFalse(output, output.contains("command1"));
         assertTrue(output, output.contains("Arguments: [foo, bar]"));
     }
+
+    public void testClose() throws Exception {
+        DummySubCommand subCommand1 = new DummySubCommand();
+        DummySubCommand subCommand2 = new DummySubCommand();
+        multiCommand.subcommands.put("command1", subCommand1);
+        multiCommand.subcommands.put("command2", subCommand2);
+        multiCommand.close();
+        assertTrue("MultiCommand must have been closed when close method is invoked", multiCommand.closed.get());
+        assertTrue("SubCommand1 must have been closed when close method is invoked", subCommand1.closed.get());
+        assertTrue("SubCommand2 must have been closed when close method is invoked", subCommand2.closed.get());
+    }
+
+    public void testCloseWhenSubCommandCloseThrowsException() {
+        boolean throwExceptionWhenClosed = true;
+        DummySubCommand subCommand1 = new DummySubCommand(throwExceptionWhenClosed);
+        DummySubCommand subCommand2 = new DummySubCommand();
+        multiCommand.subcommands.put("command1", subCommand1);
+        multiCommand.subcommands.put("command2", subCommand2);
+        // verify exception is thrown, as well as other non failed sub-commands closed
+        // properly.
+        expectThrows(IOException.class, () -> multiCommand.close());
+        assertTrue("SubCommand2 must have been closed when close method is invoked", subCommand2.closed.get());
+    }
 }