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();
+    }
 }