Merge pull request #159 from udaysagar2177/fix2

add pre-shutdown action to be executed inside ChronicleHashCloseOnExitHook

src/main/java/net/openhft/chronicle/hash/ChronicleHashBuilder.java

@@ -31,6 +31,7 @@ import java.io.IOException;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
 
 /**
  * Base interface for {@link ChronicleMapBuilder} and {@link ChronicleSetBuilder}, i. e. defines
@@ -744,6 +745,22 @@ public interface ChronicleHashBuilder<K, H extends ChronicleHash<K, ?, ?, ?>,
             File file, boolean sameBuilderConfigAndLibraryVersion,
             ChronicleHashCorruption.Listener corruptionListener) throws IOException;
 
+    /**
+     * A {@link ChronicleHash} created using this builder is closed using a JVM shutdown hook.
+     * This method lets you perform an action before the {@link ChronicleHash} is closed.
+     * <p>
+     * <p>The registered action is not executed when JVM is running and a
+     * {@link ChronicleHash#close()} is explicitly called.
+     * <p>
+     * <p>Example usage of this call: To carry out a graceful shutdown and explicitly
+     * control when the {@link ChronicleHash} is closed, the action can be a wait on a
+     * {@link CountDownLatch} that would be released appropriately.
+     *
+     * @param preShutdownAction action to run before closing the {@link ChronicleHash} in a
+     *                          JVM shutdown hook.
+     */
+    void setPreShutdownAction(Runnable preShutdownAction);
+
     /**
      * @deprecated don't use private API in the client code
      */

src/main/java/net/openhft/chronicle/hash/ChronicleHashBuilderPrivateAPI.java

@@ -89,4 +89,12 @@ public interface ChronicleHashBuilderPrivateAPI<K, RO> {
     void removedEntryCleanupTimeout(long removedEntryCleanupTimeout, TimeUnit unit);
 
     void remoteOperations(RO remoteOperations);
+
+    /**
+     * Provides registered action to be executed before closing a {@link ChronicleHash} in a JVM
+     * shutdown hook.
+     *
+     * @return registered pre-shutdown action.
+     */
+    Runnable getPreShutdownAction();
 }

src/main/java/net/openhft/chronicle/hash/impl/ChronicleHashCloseOnExitHook.java

@@ -60,6 +60,20 @@ final class ChronicleHashCloseOnExitHook {
             // close later added maps first
             orderedMaps.descendingMap().values().forEach(h -> {
                 try {
+                    Runnable preShutdownAction = h.getPreShutdownAction();
+                    if (preShutdownAction != null) {
+                        try {
+                            preShutdownAction.run();
+                        } catch (Throwable throwable) {
+                            try {
+                                LOG.error("Error running pre-shutdown action for " + h.toIdentityString() +
+                                        " :", throwable);
+                            } catch (Throwable t2) {
+                                throwable.addSuppressed(t2);
+                                throwable.printStackTrace();
+                            }
+                        }
+                    }
                     h.close();
                 } catch (Throwable throwable) {
                     try {

src/main/java/net/openhft/chronicle/hash/impl/VanillaChronicleHash.java

@@ -122,6 +122,7 @@ public abstract class VanillaChronicleHash<K,
     public transient CompactOffHeapLinearHashTable hashLookup;
     public transient Identity identity;
     protected int log2TiersInBulk;
+    private Runnable preShutdownAction;
     /////////////////////////////////////////////////
     // Bytes Store (essentially, the base address) and serialization-dependent offsets
     protected transient BytesStore bs;
@@ -205,6 +206,8 @@ public abstract class VanillaChronicleHash<K,
         tierBulkSizeInBytes = computeTierBulkBytesSize(tiersInBulk);
 
         checksumEntries = privateAPI.checksumEntries();
+
+        preShutdownAction = privateAPI.getPreShutdownAction();
     }
 
     public static IOException throwRecoveryOrReturnIOException(
@@ -227,6 +230,10 @@ public abstract class VanillaChronicleHash<K,
         initTransients();
     }
 
+    public Runnable getPreShutdownAction() {
+        return preShutdownAction;
+    }
+
     protected void readMarshallableFields(@NotNull WireIn wireIn) {
         dataFileVersion = wireIn.read(() -> "dataFileVersion").text();
 

src/main/java/net/openhft/chronicle/map/ChronicleMapBuilder.java

@@ -180,6 +180,7 @@ public final class ChronicleMapBuilder<K, V> implements
     MapMethods<K, V, ?> methods = DefaultSpi.mapMethods();
     MapEntryOperations<K, V, ?> entryOperations = mapEntryOperations();
     MapRemoteOperations<K, V, ?> remoteOperations = mapRemoteOperations();
+    Runnable preShutdownAction;
     private String name;
     // not final because of cloning
     private ChronicleMapBuilderPrivateAPI<K, V> privateAPI =
@@ -1589,6 +1590,11 @@ public final class ChronicleMapBuilder<K, V> implements
                 corruptionListener);
     }
 
+    @Override
+    public void setPreShutdownAction(Runnable preShutdownAction) {
+        this.preShutdownAction = preShutdownAction;
+    }
+
     @Override
     public ChronicleMap<K, V> create() {
         // clone() to make this builder instance thread-safe, because createWithoutFile() method

src/main/java/net/openhft/chronicle/map/ChronicleMapBuilderPrivateAPI.java

@@ -115,4 +115,9 @@ class ChronicleMapBuilderPrivateAPI<K, V>
     public void remoteOperations(MapRemoteOperations<K, V, ?> remoteOperations) {
         b.remoteOperations(remoteOperations);
     }
+
+    @Override
+    public Runnable getPreShutdownAction() {
+        return b.preShutdownAction;
+    }
 }

src/main/java/net/openhft/chronicle/set/ChronicleSetBuilder.java

@@ -352,6 +352,11 @@ public final class ChronicleSetBuilder<K>
         return new SetFromMap<>((VanillaChronicleMap<K, DummyValue, ?>) map);
     }
 
+    @Override
+    public void setPreShutdownAction(Runnable preShutdownAction) {
+        chronicleMapBuilder.setPreShutdownAction(preShutdownAction);
+    }
+
     /**
      * @deprecated don't use private API in the client code
      */

src/main/java/net/openhft/chronicle/set/ChronicleSetBuilderPrivateAPI.java

@@ -133,4 +133,9 @@ class ChronicleSetBuilderPrivateAPI<K>
             }
         });
     }
+
+    @Override
+    public Runnable getPreShutdownAction() {
+        return mapB.getPreShutdownAction();
+    }
 }

src/test/java/net/openhft/chronicle/map/ExitHookTest.java

@@ -16,6 +16,7 @@
 
 package net.openhft.chronicle.map;
 
+import com.google.common.base.Preconditions;
 import net.openhft.chronicle.core.OS;
 import org.junit.Rule;
 import org.junit.Test;
@@ -24,12 +25,18 @@ import org.junit.rules.TemporaryFolder;
 import java.io.File;
 import java.io.IOException;
 import java.lang.reflect.Field;
+import java.nio.file.Files;
+import java.util.Iterator;
+import java.util.stream.Stream;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 public class ExitHookTest {
 
     private static final int KEY = 1;
+    private static final String PRE_SHUTDOWN_ACTION_EXECUTED = "PRE_SHUTDOWN_ACTION_EXECUTED";
 
     @Rule
     public TemporaryFolder folder = new TemporaryFolder();
@@ -37,18 +44,30 @@ public class ExitHookTest {
     public static void main(String[] args) throws IOException, InterruptedException {
         System.out.println("Other process started");
         File mapFile = new File(args[0]);
-        ChronicleMap<Integer, Integer> map = createMap(mapFile);
+        File preShutdownActionExecutionConfirmationFile = new File(args[1]);
+        ChronicleMapBuilder<Integer, Integer> mapBuilder = createMapBuilder();
+        mapBuilder.setPreShutdownAction(() -> {
+            try {
+                System.out.println("Executing pre-shutdown action");
+                Preconditions.checkState(preShutdownActionExecutionConfirmationFile.exists());
+                Files.write(preShutdownActionExecutionConfirmationFile.toPath(),
+                        PRE_SHUTDOWN_ACTION_EXECUTED.getBytes());
+            } catch (Exception e) {
+                throw new RuntimeException(e);
+            }
+        });
+        ChronicleMap<Integer, Integer> map = mapBuilder.createPersistedTo(mapFile);
         try (ExternalMapQueryContext<Integer, Integer, ?> c = map.queryContext(KEY)) {
             c.writeLock().lock();
             Thread.sleep(30_000);
         }
     }
 
-    private static ChronicleMap<Integer, Integer> createMap(File mapFile) throws IOException {
-        return ChronicleMap
+    private static ChronicleMapBuilder<Integer, Integer> createMapBuilder()
+            throws IOException {
+        return ChronicleMapBuilder
                 .of(Integer.class, Integer.class)
-                .entries(1)
-                .createPersistedTo(mapFile);
+                .entries(1);
     }
 
     // http://stackoverflow.com/a/33171840/648955
@@ -73,9 +92,9 @@ public class ExitHookTest {
         if (!OS.isLinux() && !OS.isMacOSX())
             return; // This test runs only in Unix-like OSes
         File mapFile = folder.newFile();
-        System.out.println("map file: " + mapFile.getAbsolutePath());
+        File preShutdownActionExecutionConfirmationFile = folder.newFile();
         // Create a process which opens the map, acquires the lock and "hangs" for 30 seconds
-        Process process = startOtherProcess(mapFile);
+        Process process = startOtherProcess(mapFile, preShutdownActionExecutionConfirmationFile);
         // Let the other process actually reach the moment when it locks the map
         // (JVM startup and chronicle map creation are not instant)
         Thread.sleep(10_000);
@@ -85,12 +104,19 @@ public class ExitHookTest {
         int actual = process.exitValue();
         if (actual != 0) // clean shutdown
             assertEquals(130, actual); // 130 is exit code for SIGINT (interruption).
-        ChronicleMap<Integer, Integer> map = createMap(mapFile);
+        ChronicleMap<Integer, Integer> map = createMapBuilder().createPersistedTo(mapFile);
         try (ExternalMapQueryContext<Integer, Integer, ?> c = map.queryContext(KEY)) {
             // Test that we are able to lock the segment, i. e. the lock was released in other
             // process thanks to shutdown hooks
             c.writeLock().lock();
         }
+        try (Stream<String> lines = Files.lines(preShutdownActionExecutionConfirmationFile.toPath())) {
+            Iterator<String> lineIterator = lines.iterator();
+            assertTrue(lineIterator.hasNext());
+            String line = lineIterator.next();
+            assertEquals(PRE_SHUTDOWN_ACTION_EXECUTED, line);
+            assertFalse(lineIterator.hasNext());
+        }
     }
 
     // http://stackoverflow.com/a/7835467/648955
@@ -99,7 +125,7 @@ public class ExitHookTest {
     }
 
     // http://stackoverflow.com/a/723914/648955
-    private Process startOtherProcess(File mapFile) throws IOException {
+    private Process startOtherProcess(File mapFile, File outputFile) throws IOException {
         String javaHome = System.getProperty("java.home");
         String javaBin = javaHome +
                 File.separator + "bin" +
@@ -108,8 +134,9 @@ public class ExitHookTest {
         System.out.println("Classpath: " + classpath);
         String className = ExitHookTest.class.getCanonicalName();
 
-        ProcessBuilder builder = new ProcessBuilder(
-                javaBin, "-cp", classpath, className, mapFile.getAbsolutePath());
+        String[] command = new String[] {javaBin, "-cp", classpath, className,
+                mapFile.getAbsolutePath(), outputFile.getAbsolutePath()};
+        ProcessBuilder builder = new ProcessBuilder(command);
         builder.inheritIO();
         return builder.start();
     }