Remove VanillaChronicleHash from OnExitHook on close()

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

@@ -14,18 +14,17 @@
  *      along with this program.  If not, see <http://www.gnu.org/licenses/>.
  */
 
-package net.openhft.chronicle.map;
+package net.openhft.chronicle.hash.impl;
 
 
-import net.openhft.chronicle.hash.impl.VanillaChronicleHash;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.util.TreeMap;
 import java.util.WeakHashMap;
 
-final class ChronicleMapCloseOnExitHook {
-    private static final Logger LOG = LoggerFactory.getLogger(ChronicleMapCloseOnExitHook.class);
+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;
@@ -45,12 +44,18 @@ final class ChronicleMapCloseOnExitHook {
         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 (ChronicleMapCloseOnExitHook.class) {
-                maps = ChronicleMapCloseOnExitHook.maps;
-                ChronicleMapCloseOnExitHook.maps = null;
+            synchronized (ChronicleHashCloseOnExitHook.class) {
+                maps = ChronicleHashCloseOnExitHook.maps;
+                ChronicleHashCloseOnExitHook.maps = null;
             }
 
             TreeMap<Long, VanillaChronicleHash> orderedMaps = new TreeMap<>();
@@ -69,5 +74,5 @@ final class ChronicleMapCloseOnExitHook {
         }
     }
 
-    private ChronicleMapCloseOnExitHook() {}
+    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

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

@@ -21,8 +21,6 @@ import net.openhft.chronicle.algo.MemoryUnit;
 import net.openhft.chronicle.algo.hashing.LongHashFunction;
 import net.openhft.chronicle.bytes.Byteable;
 import net.openhft.chronicle.bytes.Bytes;
-import net.openhft.chronicle.bytes.BytesStore;
-import net.openhft.chronicle.bytes.NativeBytesStore;
 import net.openhft.chronicle.core.Jvm;
 import net.openhft.chronicle.core.OS;
 import net.openhft.chronicle.hash.ChronicleHashBuilder;
@@ -1524,7 +1522,7 @@ public final class ChronicleMapBuilder<K, V> implements
         map.registerCleaner();
         // Ensure safe publication of a ChronicleMap
         OS.memory().storeFence();
-        ChronicleMapCloseOnExitHook.add(map);
+        map.addToOnExitHook();
     }
 
     /**