Remove VanillaChronicleHash from OnExitHook on close()

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

@@ -0,0 +1,78 @@
+/*
+ *      Copyright (C) 2016 Roman Leventov
+ *
+ *      This program is free software: you can redistribute it and/or modify
+ *      it under the terms of the GNU Lesser General Public License as published by
+ *      the Free Software Foundation, either version 3 of the License.
+ *
+ *      This program is distributed in the hope that it will be useful,
+ *      but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *      GNU Lesser General Public License for more details.
+ *
+ *      You should have received a copy of the GNU Lesser General Public License
+ *      along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package net.openhft.chronicle.hash.impl;
+
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.TreeMap;
+import java.util.WeakHashMap;
+
+final class ChronicleHashCloseOnExitHook {
+    private static final Logger LOG = LoggerFactory.getLogger(ChronicleHashCloseOnExitHook.class);
+
+    private static WeakHashMap<VanillaChronicleHash.Identity, Long> maps = new WeakHashMap<>();
+    private static long order = 0;
+
+    static {
+        Runtime.getRuntime().addShutdownHook(new Thread() {
+            @Override
+            public void run() {
+                closeAll();
+            }
+        });
+    }
+
+    static synchronized void add(VanillaChronicleHash hash) {
+        if (maps == null)
+            throw new IllegalStateException("Shutdown in progress");
+        maps.put(hash.identity, order++);
+    }
+
+    static synchronized void remove(VanillaChronicleHash hash) {
+        if (maps == null)
+            return; // we are already in shutdown
+        maps.remove(hash.identity);
+    }
+
+    private static void closeAll() {
+        try {
+            WeakHashMap<VanillaChronicleHash.Identity, Long> maps;
+            synchronized (ChronicleHashCloseOnExitHook.class) {
+                maps = ChronicleHashCloseOnExitHook.maps;
+                ChronicleHashCloseOnExitHook.maps = null;
+            }
+
+            TreeMap<Long, VanillaChronicleHash> orderedMaps = new TreeMap<>();
+            maps.forEach((identity, order) -> orderedMaps.put(order, identity.hash()));
+            // close later added maps first
+            orderedMaps.descendingMap().values().forEach(h -> {
+                try {
+                    h.close();
+                } catch (RuntimeException | AssertionError e) {
+                    LOG.error("", e);
+                }
+            });
+        } catch (Throwable t) {
+            LOG.error("", t);
+            throw t;
+        }
+    }
+
+    private ChronicleHashCloseOnExitHook() {}
+}

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

@@ -164,9 +164,8 @@ public abstract class VanillaChronicleHash<K,
     private transient ArrayList<WeakReference<ChainingInterface>> allContexts;
 
     /**
-     * {@link net.openhft.chronicle.map.ChronicleMapCloseOnExitHook} needs to use
-     * {@code VanillaChronicleHash}es as WeakHashMap keys, but with identity comparison, not Map's
-     * equals() and hashCode().
+     * {@link ChronicleHashCloseOnExitHook} needs to use {@code VanillaChronicleHash}es as
+     * WeakHashMap keys, but with identity comparison, not Map's equals() and hashCode().
      */
     public class Identity {
         public VanillaChronicleHash hash() {
@@ -489,6 +488,10 @@ public abstract class VanillaChronicleHash<K,
         this.cleaner = Cleaner.create(this, resourceReleaser);
     }
 
+    public void addToOnExitHook() {
+        ChronicleHashCloseOnExitHook.add(this);
+    }
+
     public final void createMappedStoreAndSegments(ChronicleHashResourceReleaser resourceReleaser)
             throws IOException {
         this.resourceReleaser = resourceReleaser;
@@ -678,6 +681,7 @@ public abstract class VanillaChronicleHash<K,
         // Releases nothing after resourceReleaser.releaseManually(), only removes the cleaner
         // from the internal linked list of all cleaners.
         cleaner.clean();
+        ChronicleHashCloseOnExitHook.remove(this);
     }
 
     @Override