Add shutdown hook for closing CLI commands

This commit enables CLI commands to be closeable and installs a runtime
shutdown hook to ensure that if the JVM shuts down (as opposed to
aborting) the close method is called.

It is not enough to wrap uses of commands in main methods in
try-with-resources blocks as these will not run if, say, the virtual
machine is terminated in response to SIGINT, or system shutdown event.

Relates #22126

core/src/main/java/org/elasticsearch/cli/Command.java

@@ -24,17 +24,21 @@ import joptsimple.OptionParser;
 import joptsimple.OptionSet;
 import joptsimple.OptionSpec;
 import org.apache.logging.log4j.Level;
+import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.common.SuppressForbidden;
 import org.elasticsearch.common.logging.LogConfigurator;
 import org.elasticsearch.common.settings.Settings;
 
+import java.io.Closeable;
 import java.io.IOException;
+import java.io.PrintWriter;
+import java.io.StringWriter;
 import java.util.Arrays;
 
 /**
  * An action to execute within a cli.
  */
-public abstract class Command {
+public abstract class Command implements Closeable {
 
     /** A description of the command, used in the help output. */
     protected final String description;
@@ -44,15 +48,37 @@ public abstract class Command {
 
     private final OptionSpec<Void> helpOption = parser.acceptsAll(Arrays.asList("h", "help"), "show help").forHelp();
     private final OptionSpec<Void> silentOption = parser.acceptsAll(Arrays.asList("s", "silent"), "show minimal output");
-    private final OptionSpec<Void> verboseOption = parser.acceptsAll(Arrays.asList("v", "verbose"), "show verbose output")
-            .availableUnless(silentOption);
+    private final OptionSpec<Void> verboseOption =
+        parser.acceptsAll(Arrays.asList("v", "verbose"), "show verbose output").availableUnless(silentOption);
 
     public Command(String description) {
         this.description = description;
     }
 
+    final SetOnce<Thread> shutdownHookThread = new SetOnce<>();
+
     /** Parses options for this command from args and executes it. */
     public final int main(String[] args, Terminal terminal) throws Exception {
+        if (addShutdownHook()) {
+            shutdownHookThread.set(new Thread(() -> {
+                try {
+                    this.close();
+                } catch (final IOException e) {
+                    try (
+                        final StringWriter sw = new StringWriter();
+                        final PrintWriter pw = new PrintWriter(sw)) {
+                        e.printStackTrace(pw);
+                        terminal.println(sw.toString());
+                    } catch (final IOException impossible) {
+                        // StringWriter#close declares a checked IOException from the Closeable interface but the Javadocs for StringWriter
+                        // say that an exception here is impossible
+                        throw new AssertionError(impossible);
+                    }
+                }
+            }));
+            Runtime.getRuntime().addShutdownHook(shutdownHookThread.get());
+        }
+
         // initialize default for es.logger.level because we will not read the log4j2.properties
         final String loggerLevel = System.getProperty("es.logger.level", Level.INFO.name());
         final Settings settings = Settings.builder().put("logger.level", loggerLevel).build();
@@ -118,4 +144,19 @@ public abstract class Command {
      * Any runtime user errors (like an input file that does not exist), should throw a {@link UserException}. */
     protected abstract void execute(Terminal terminal, OptionSet options) throws Exception;
 
+    /**
+     * Return whether or not to install the shutdown hook to cleanup resources on exit. This method should only be overridden in test
+     * classes.
+     *
+     * @return whether or not to install the shutdown hook
+     */
+    protected boolean addShutdownHook() {
+        return true;
+    }
+
+    @Override
+    public void close() throws IOException {
+
+    }
+
 }

core/src/main/java/org/elasticsearch/plugins/InstallPluginCommand.java

@@ -320,7 +320,7 @@ class InstallPluginCommand extends SettingCommand {
     /** Downloads a zip from the url, as well as a SHA1 checksum, and checks the checksum. */
     private Path downloadZipAndChecksum(Terminal terminal, String urlString, Path tmpDir) throws Exception {
         Path zip = downloadZip(terminal, urlString, tmpDir);
-
+        pathsToDeleteOnShutdown.add(zip);
         URL checksumUrl = new URL(urlString + ".sha1");
         final String expectedChecksum;
         try (InputStream in = checksumUrl.openStream()) {
@@ -344,9 +344,9 @@ class InstallPluginCommand extends SettingCommand {
         // unzip plugin to a staging temp dir
 
         final Path target = stagingDirectory(pluginsDir);
+        pathsToDeleteOnShutdown.add(target);
 
         boolean hasEsDir = false;
-        // TODO: we should wrap this in a try/catch and try deleting the target dir on failure?
         try (ZipInputStream zipInput = new ZipInputStream(Files.newInputStream(zip))) {
             ZipEntry entry;
             byte[] buffer = new byte[8192];
@@ -605,4 +605,12 @@ class InstallPluginCommand extends SettingCommand {
             Files.setPosixFilePermissions(path, permissions);
         }
     }
+
+    private final List<Path> pathsToDeleteOnShutdown = new ArrayList<>();
+
+    @Override
+    public void close() throws IOException {
+        IOUtils.rm(pathsToDeleteOnShutdown.toArray(new Path[pathsToDeleteOnShutdown.size()]));
+    }
+
 }

core/src/main/java/org/elasticsearch/plugins/PluginCli.java

@@ -19,27 +19,37 @@
 
 package org.elasticsearch.plugins;
 
+import org.apache.lucene.util.IOUtils;
+import org.elasticsearch.cli.Command;
 import org.elasticsearch.cli.MultiCommand;
 import org.elasticsearch.cli.Terminal;
-import org.elasticsearch.common.logging.LogConfigurator;
-import org.elasticsearch.common.settings.Settings;
-import org.elasticsearch.env.Environment;
-import org.elasticsearch.node.internal.InternalSettingsPreparer;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
 
 /**
  * A cli tool for adding, removing and listing plugins for elasticsearch.
  */
 public class PluginCli extends MultiCommand {
 
+    private final Collection<Command> commands;
+
     private PluginCli() {
         super("A tool for managing installed elasticsearch plugins");
         subcommands.put("list", new ListPluginsCommand());
         subcommands.put("install", new InstallPluginCommand());
         subcommands.put("remove", new RemovePluginCommand());
+        commands = Collections.unmodifiableCollection(subcommands.values());
     }
 
     public static void main(String[] args) throws Exception {
         exit(new PluginCli().main(args, Terminal.DEFAULT));
     }
 
+    @Override
+    public void close() throws IOException {
+        IOUtils.close(commands);
+    }
+
 }

core/src/main/java/org/elasticsearch/plugins/RemovePluginCommand.java

@@ -44,7 +44,7 @@ import static org.elasticsearch.cli.Terminal.Verbosity.VERBOSE;
 /**
  * A command for the plugin cli to remove a plugin from elasticsearch.
  */
-final class RemovePluginCommand extends SettingCommand {
+class RemovePluginCommand extends SettingCommand {
 
     private final OptionSpec<String> arguments;
 

core/src/test/java/org/elasticsearch/cli/CommandTests.java

@@ -26,30 +26,49 @@ import org.elasticsearch.test.ESTestCase;
 public class CommandTests extends ESTestCase {
 
     static class UserErrorCommand extends Command {
+
         UserErrorCommand() {
             super("Throws a user error");
         }
+
         @Override
         protected void execute(Terminal terminal, OptionSet options) throws Exception {
             throw new UserException(ExitCodes.DATA_ERROR, "Bad input");
         }
+
+        @Override
+        protected boolean addShutdownHook() {
+            return false;
+        }
+
     }
 
     static class UsageErrorCommand extends Command {
+
         UsageErrorCommand() {
             super("Throws a usage error");
         }
+
         @Override
         protected void execute(Terminal terminal, OptionSet options) throws Exception {
             throw new UserException(ExitCodes.USAGE, "something was no good");
         }
+
+        @Override
+        protected boolean addShutdownHook() {
+            return false;
+        }
+
     }
 
     static class NoopCommand extends Command {
+
         boolean executed = false;
+
         NoopCommand() {
             super("Does nothing");
         }
+
         @Override
         protected void execute(Terminal terminal, OptionSet options) throws Exception {
             terminal.println("Normal output");
@@ -57,10 +76,17 @@ public class CommandTests extends ESTestCase {
             terminal.println(Terminal.Verbosity.VERBOSE, "Verbose output");
             executed = true;
         }
+
         @Override
         protected void printAdditionalHelp(Terminal terminal) {
             terminal.println("Some extra help");
         }
+
+        @Override
+        protected boolean addShutdownHook() {
+            return false;
+        }
+
     }
 
     public void testHelp() throws Exception {
@@ -92,7 +118,7 @@ public class CommandTests extends ESTestCase {
             command.mainWithoutErrorHandling(args, terminal);
         });
         assertTrue(e.getMessage(),
-                e.getMessage().contains("Option(s) [v/verbose] are unavailable given other options on the command line"));
+            e.getMessage().contains("Option(s) [v/verbose] are unavailable given other options on the command line"));
     }
 
     public void testSilentVerbosity() throws Exception {
@@ -143,4 +169,5 @@ public class CommandTests extends ESTestCase {
         assertTrue(output, output.contains("Throws a usage error"));
         assertTrue(output, output.contains("ERROR: something was no good"));
     }
+
 }

qa/evil-tests/src/test/java/org/elasticsearch/cli/EvilCommandTests.java

@@ -0,0 +1,69 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.cli;
+
+import joptsimple.OptionSet;
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.IOException;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.Matchers.isEmptyString;
+
+public class EvilCommandTests extends ESTestCase {
+
+    public void testCommandShutdownHook() throws Exception {
+        final AtomicBoolean closed = new AtomicBoolean();
+        final boolean shouldThrow = randomBoolean();
+        final Command command = new Command("test-command-shutdown-hook") {
+            @Override
+            protected void execute(Terminal terminal, OptionSet options) throws Exception {
+
+            }
+
+            @Override
+            public void close() throws IOException {
+                closed.set(true);
+                if (shouldThrow) {
+                    throw new IOException("fail");
+                }
+            }
+        };
+        final MockTerminal terminal = new MockTerminal();
+        command.main(new String[0], terminal);
+        assertNotNull(command.shutdownHookThread.get());
+        // successful removal here asserts that the runtime hook was installed in Command#main
+        assertTrue(Runtime.getRuntime().removeShutdownHook(command.shutdownHookThread.get()));
+        command.shutdownHookThread.get().run();
+        command.shutdownHookThread.get().join();
+        assertTrue(closed.get());
+        final String output = terminal.getOutput();
+        if (shouldThrow) {
+            // ensure that we dump the exception
+            assertThat(output, containsString("java.io.IOException: fail"));
+            // ensure that we dump the stack trace too
+            assertThat(output, containsString("\tat org.elasticsearch.cli.EvilCommandTests$1.close"));
+        } else {
+            assertThat(output, isEmptyString());
+        }
+    }
+
+}

qa/evil-tests/src/test/java/org/elasticsearch/plugins/InstallPluginCommandTests.java

@@ -603,7 +603,12 @@ public class InstallPluginCommandTests extends ESTestCase {
 
     public void testOfficialPluginsHelpSorted() throws Exception {
         MockTerminal terminal = new MockTerminal();
-        new InstallPluginCommand().main(new String[] { "--help" }, terminal);
+        new InstallPluginCommand() {
+            @Override
+            protected boolean addShutdownHook() {
+                return false;
+            }
+        }.main(new String[] { "--help" }, terminal);
         try (BufferedReader reader = new BufferedReader(new StringReader(terminal.getOutput()))) {
             String line = reader.readLine();
 
@@ -625,7 +630,12 @@ public class InstallPluginCommandTests extends ESTestCase {
 
     public void testOfficialPluginsIncludesXpack() throws Exception {
         MockTerminal terminal = new MockTerminal();
-        new InstallPluginCommand().main(new String[] { "--help" }, terminal);
+        new InstallPluginCommand() {
+            @Override
+            protected boolean addShutdownHook() {
+                return false;
+            }
+        }.main(new String[] { "--help" }, terminal);
         assertTrue(terminal.getOutput(), terminal.getOutput().contains("x-pack"));
     }
 

qa/evil-tests/src/test/java/org/elasticsearch/plugins/ListPluginsCommandTests.java

@@ -59,21 +59,26 @@ public class ListPluginsCommandTests extends ESTestCase {
     static MockTerminal listPlugins(Path home) throws Exception {
         return listPlugins(home, new String[0]);
     }
-    
+
     static MockTerminal listPlugins(Path home, String[] args) throws Exception {
         String[] argsAndHome = new String[args.length + 1];
         System.arraycopy(args, 0, argsAndHome, 0, args.length);
         argsAndHome[args.length] = "-Epath.home=" + home;
         MockTerminal terminal = new MockTerminal();
-        int status = new ListPluginsCommand().main(argsAndHome, terminal);
+        int status = new ListPluginsCommand() {
+            @Override
+            protected boolean addShutdownHook() {
+                return false;
+            }
+        }.main(argsAndHome, terminal);
         assertEquals(ExitCodes.OK, status);
         return terminal;
     }
-    
+
     static String buildMultiline(String... args){
         return Arrays.asList(args).stream().collect(Collectors.joining("\n", "", "\n"));
     }
-    
+
     static void buildFakePlugin(Environment env, String description, String name, String classname) throws IOException {
         PluginTestUtil.writeProperties(env.pluginsFile().resolve(name),
                 "description", description,
@@ -108,7 +113,7 @@ public class ListPluginsCommandTests extends ESTestCase {
         MockTerminal terminal = listPlugins(home);
         assertEquals(terminal.getOutput(), buildMultiline("fake1", "fake2"));
     }
-    
+
     public void testPluginWithVerbose() throws Exception {
         buildFakePlugin(env, "fake desc", "fake_plugin", "org.fake");
         String[] params = { "-v" };
@@ -116,7 +121,7 @@ public class ListPluginsCommandTests extends ESTestCase {
         assertEquals(terminal.getOutput(), buildMultiline("Plugins directory: " + env.pluginsFile(), "fake_plugin",
                 "- Plugin information:", "Name: fake_plugin", "Description: fake desc", "Version: 1.0", " * Classname: org.fake"));
     }
-    
+
     public void testPluginWithVerboseMultiplePlugins() throws Exception {
         buildFakePlugin(env, "fake desc 1", "fake_plugin1", "org.fake");
         buildFakePlugin(env, "fake desc 2", "fake_plugin2", "org.fake2");
@@ -127,7 +132,7 @@ public class ListPluginsCommandTests extends ESTestCase {
                 " * Classname: org.fake", "fake_plugin2", "- Plugin information:", "Name: fake_plugin2",
                 "Description: fake desc 2", "Version: 1.0", " * Classname: org.fake2"));
     }
-        
+
     public void testPluginWithoutVerboseMultiplePlugins() throws Exception {
         buildFakePlugin(env, "fake desc 1", "fake_plugin1", "org.fake");
         buildFakePlugin(env, "fake desc 2", "fake_plugin2", "org.fake2");
@@ -135,13 +140,13 @@ public class ListPluginsCommandTests extends ESTestCase {
         String output = terminal.getOutput();
         assertEquals(output, buildMultiline("fake_plugin1", "fake_plugin2"));
     }
-    
+
     public void testPluginWithoutDescriptorFile() throws Exception{
         Files.createDirectories(env.pluginsFile().resolve("fake1"));
         NoSuchFileException e = expectThrows(NoSuchFileException.class, () -> listPlugins(home));
         assertEquals(e.getFile(), env.pluginsFile().resolve("fake1").resolve(PluginInfo.ES_PLUGIN_PROPERTIES).toString());
     }
-    
+
     public void testPluginWithWrongDescriptorFile() throws Exception{
         PluginTestUtil.writeProperties(env.pluginsFile().resolve("fake1"),
                 "description", "fake desc");
@@ -149,5 +154,5 @@ public class ListPluginsCommandTests extends ESTestCase {
         assertEquals(e.getMessage(), "Property [name] is missing in [" +
                 env.pluginsFile().resolve("fake1").resolve(PluginInfo.ES_PLUGIN_PROPERTIES).toString() + "]");
     }
-    
+
 }

qa/evil-tests/src/test/java/org/elasticsearch/plugins/RemovePluginCommandTests.java

@@ -140,7 +140,12 @@ public class RemovePluginCommandTests extends ESTestCase {
         assertEquals("plugin fake not found; run 'elasticsearch-plugin list' to get list of installed plugins", e.getMessage());
 
         MockTerminal terminal = new MockTerminal();
-        new RemovePluginCommand().main(new String[] { "-Epath.home=" + home, "fake" }, terminal);
+        new RemovePluginCommand() {
+            @Override
+            protected boolean addShutdownHook() {
+                return false;
+            }
+        }.main(new String[] { "-Epath.home=" + home, "fake" }, terminal);
         try (BufferedReader reader = new BufferedReader(new StringReader(terminal.getOutput()))) {
             assertEquals("-> Removing fake...", reader.readLine());
             assertEquals("ERROR: plugin fake not found; run 'elasticsearch-plugin list' to get list of installed plugins",

test/framework/src/main/java/org/elasticsearch/bootstrap/ESElasticsearchCliTestCase.java

@@ -50,6 +50,11 @@ abstract class ESElasticsearchCliTestCase extends ESTestCase {
                     init.set(true);
                     initConsumer.accept(!daemonize, pidFile, quiet, esSettings);
                 }
+
+                @Override
+                protected boolean addShutdownHook() {
+                    return false;
+                }
             }, terminal);
             assertThat(status, equalTo(expectedStatus));
             assertThat(init.get(), equalTo(expectedInit));