new scripting extension: shutdown; scripted shutdown hooks

engine-extensions/src/main/java/io/nosqlbench/engine/shutdown/ShutdownHookPlugin.java

@@ -0,0 +1,32 @@
+package io.nosqlbench.engine.shutdown;
+
+import com.codahale.metrics.MetricRegistry;
+import org.apache.logging.log4j.Logger;
+
+import javax.script.ScriptContext;
+import java.util.function.Function;
+
+public class ShutdownHookPlugin {
+    private final Logger logger;
+    private final MetricRegistry metricRegistry;
+    private final ScriptContext scriptContext;
+
+    public ShutdownHookPlugin(Logger logger, MetricRegistry metricRegistry, ScriptContext scriptContext) {
+
+        this.logger = logger;
+        this.metricRegistry = metricRegistry;
+        this.scriptContext = scriptContext;
+    }
+
+    public void addShutdownHook(String name, Object f) {
+        if (!(f instanceof Function)) {
+            throw new RuntimeException("The object provided to the shutdown hook plugin was not recognized as a function.");
+        }
+        String shutdownName = "shutdown-function-" + name;
+        Thread runnable = new ShutdownRunnableFunction(logger, name, (Function<?,?>)f);
+        runnable.setName(shutdownName);
+        Runtime.getRuntime().addShutdownHook(runnable);
+        logger.info("Registered shutdown hook to run under name '" + shutdownName + "'");
+
+    }
+}

engine-extensions/src/main/java/io/nosqlbench/engine/shutdown/ShutdownHookPluginMetadata.java

@@ -0,0 +1,22 @@
+package io.nosqlbench.engine.shutdown;
+
+import com.codahale.metrics.MetricRegistry;
+import io.nosqlbench.engine.api.extensions.ScriptingPluginInfo;
+import io.nosqlbench.nb.annotations.Service;
+import org.apache.logging.log4j.Logger;
+
+import javax.script.ScriptContext;
+
+@Service(value=ScriptingPluginInfo.class,selector = "shutdown")
+public class ShutdownHookPluginMetadata implements ScriptingPluginInfo<ShutdownHookPlugin> {
+
+    @Override
+    public String getDescription() {
+        return "Register shutdown hooks in the form of javascript functions.";
+    }
+
+    @Override
+    public ShutdownHookPlugin getExtensionObject(Logger logger, MetricRegistry metricRegistry, ScriptContext scriptContext) {
+        return new ShutdownHookPlugin(logger,metricRegistry,scriptContext);
+    }
+}

engine-extensions/src/main/java/io/nosqlbench/engine/shutdown/ShutdownRunnableFunction.java

@@ -0,0 +1,32 @@
+package io.nosqlbench.engine.shutdown;
+
+import org.apache.logging.log4j.Logger;
+
+import java.util.function.Function;
+
+
+public class ShutdownRunnableFunction extends Thread {
+    private final String name;
+    private final Function<Object[],Object> function;
+    private final Logger logger;
+
+    public ShutdownRunnableFunction(Logger logger, String name, Function<?, ?> function) {
+        this.logger = logger;
+        this.name = name;
+        this.function = (Function<Object[],Object>)function;
+    }
+
+    @Override
+    public void run() {
+        logger.info("Running shutdown hook '" + name + "'...");
+        try {
+            Object result = function.apply(new Object[0]);
+            if (result instanceof CharSequence) {
+                logger.info("shutdown hook returned output:\n" + ((CharSequence) result).toString());
+            }
+            logger.info("Completed shutdown hook '" + name + "'...");
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

nb/src/test/java/io/nosqlbench/engine/core/script/AsyncScriptIntegrationTests.java

@@ -237,6 +237,14 @@ public class AsyncScriptIntegrationTests {
         assertThat(scenarioResult.getIOLog()).contains("count: ");
     }
 
+    @Test
+    public void testShutdownHook() {
+        ScenarioResult scenarioResult = runScenario("extension_shutdown_hook");
+        assertThat(scenarioResult.getIOLog()).doesNotContain("shutdown hook running").describedAs(
+            "shutdown hooks should not run in the same IO context as the main scenario"
+        );
+    }
+
     @Test
     public void testExceptionPropagationFromMotorThread() {
         ScenarioResult scenarioResult = runScenario("activityerror");