Make VanillaChronicleHash.close() to wait until concurrent threads accessing the hash finish their work and don't permit new threads to come while close() is in progress and after that

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

@@ -23,6 +23,7 @@ import net.openhft.chronicle.bytes.NativeBytesStore;
 import net.openhft.chronicle.core.Maths;
 import net.openhft.chronicle.core.OS;
 import net.openhft.chronicle.hash.*;
+import net.openhft.chronicle.hash.impl.stage.hash.ChainingInterface;
 import net.openhft.chronicle.hash.impl.util.BuildVersion;
 import net.openhft.chronicle.hash.impl.util.CanonicalRandomAccessFiles;
 import net.openhft.chronicle.hash.impl.util.jna.PosixMsync;
@@ -44,10 +45,10 @@ import sun.misc.Cleaner;
 import java.io.File;
 import java.io.IOException;
 import java.io.RandomAccessFile;
+import java.lang.ref.WeakReference;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.*;
 import java.util.concurrent.TimeUnit;
 
 import static java.lang.Long.numberOfTrailingZeros;
@@ -156,12 +157,14 @@ public abstract class VanillaChronicleHash<K,
     public transient CompactOffHeapLinearHashTable hashLookup;
 
     protected transient volatile boolean closed;
+    private transient Object closeLock;
 
     private transient VanillaGlobalMutableState globalMutableState;
 
+    private transient ArrayList<WeakReference<ChainingInterface>> allContexts;
+
     public VanillaChronicleHash(ChronicleMapBuilder<K, ?> builder) {
         createdOrInMemory = true;
-        closed = false;
 
         // Version
         dataFileVersion = BuildVersion.version();
@@ -227,12 +230,10 @@ public abstract class VanillaChronicleHash<K,
     }
 
     protected void readMarshallableFields(@NotNull WireIn wireIn) {
-        // Previously assignment of these values was done in default field initializers, but
-        // with Wire serialization VanillaChronicleMap instance is created with
-        // unsafe.allocateInstance(), that doesn't guarantee (?) to initialize fields with default
-        // values (false for boolean)
+        // Previously this assignment was done in default field initializer, but with Wire
+        // serialization VanillaChronicleMap instance is created with unsafe.allocateInstance(),
+        // that doesn't guarantee (?) to initialize fields with default values (false for boolean)
         createdOrInMemory = false;
-        closed = false;
 
         dataFileVersion = wireIn.read(() -> "dataFileVersion").text();
 
@@ -370,6 +371,8 @@ public abstract class VanillaChronicleHash<K,
     }
 
     private void initOwnTransients() {
+        closed = false;
+        closeLock = new Object();
         globalMutableState = createGlobalMutableState();
         tierBulkOffsets = new ArrayList<>();
         if (tierHashLookupSlotSize == 4) {
@@ -380,6 +383,7 @@ public abstract class VanillaChronicleHash<K,
             throw new AssertionError("hash lookup slot size could be 4 or 8, " +
                     tierHashLookupSlotSize + " observed");
         }
+        allContexts = new ArrayList<>();
     }
 
     public final void initBeforeMapping(
@@ -633,20 +637,41 @@ public abstract class VanillaChronicleHash<K,
     }
 
     @Override
-    public synchronized void close() {
-        if (closed)
-            return;
-        closed = true;
-        bs.release();
-        assert bs.refCount() == 0;
-        for (TierBulkData bulkData : tierBulkOffsets) {
-            if (bulkData.bytesStore != bs) {
-                bulkData.bytesStore.release();
-                assert bulkData.bytesStore.refCount() == 0;
+    public void close() {
+        // Mutually exclude concurrent close() calls and ensure close procedure is done only once
+        synchronized (closeLock) {
+            if (closed)
+                return;
+            // Synchronization on allContexts here in conjunction with checkOpen() in addContext()
+            // ensures threads using this chronicleHash for the first time concurrently with close()
+            // will not be able to access chronicleHash, or either will be visible during iteration
+            // over allContexts below
+            synchronized (allContexts) {
+                closed = true;
             }
+            for (WeakReference<ChainingInterface> contextRef : allContexts) {
+                ChainingInterface context = contextRef.get();
+                if (context != null && context.owner().isAlive()) {
+                    // Ensures that if the thread owning this context will come to access
+                    // chronicleHash concurrently with resource releasing operations below, it will
+                    // fail due to the check in context.lockContextLocally() method.
+                    // If the thread owning this context is currently accessing chronicleHash,
+                    // closeContext() will spin-wait until the end of this access session.
+                    context.closeContext();
+                }
+            }
+
+            bs.release();
+            assert bs.refCount() == 0;
+            for (TierBulkData bulkData : tierBulkOffsets) {
+                if (bulkData.bytesStore != bs) {
+                    bulkData.bytesStore.release();
+                    assert bulkData.bytesStore.refCount() == 0;
+                }
+            }
+            if (rafCleaner != null)
+                rafCleaner.clean();
         }
-        if (rafCleaner != null)
-            rafCleaner.clean();
     }
 
     @Override
@@ -984,4 +1009,42 @@ public abstract class VanillaChronicleHash<K,
                     offsetWithinMapping += tierBulkSizeInBytes));
         }
     }
+
+    private void checkOpen() {
+        if (closed)
+            throw new IllegalStateException("Access to ChronicleHash after close()");
+    }
+
+    protected void addContext(ChainingInterface context) {
+        synchronized (allContexts) {
+            checkOpen();
+            allContexts.removeAll(pseudoCollectionForExpunge);
+            allContexts.add(new WeakReference<>(context));
+        }
+    }
+
+    /** For testing only */
+    public List<WeakReference<ChainingInterface>> allContexts() {
+        return Collections.unmodifiableList(allContexts);
+    }
+
+    private static final Collection<WeakReference<ChainingInterface>> pseudoCollectionForExpunge =
+            new AbstractCollection<WeakReference<ChainingInterface>>() {
+
+                @Override
+                public boolean contains(Object o) {
+                    ChainingInterface context = ((WeakReference<ChainingInterface>) o).get();
+                    return context == null || !context.owner().isAlive();
+                }
+
+                @Override
+                public Iterator<WeakReference<ChainingInterface>> iterator() {
+                    throw new UnsupportedOperationException();
+                }
+
+                @Override
+                public int size() {
+                    throw new UnsupportedOperationException();
+                }
+            };
 }