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/BigSegmentHeader.java

@@ -52,7 +52,7 @@ public final class BigSegmentHeader implements SegmentHeader {
      * Previously this value was 2 seconds, but GC pauses often take more time that shouldn't
      * result to IllegalStateException.
      */
-    private static final int LOCK_TIMEOUT_SECONDS = 60;
+    public static final int LOCK_TIMEOUT_SECONDS = 60;
 
     private static RuntimeException deadLock() {
         return new RuntimeException("Failed to acquire the lock in " + LOCK_TIMEOUT_SECONDS +

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

src/main/java/net/openhft/chronicle/hash/impl/stage/entry/SegmentStages.java

@@ -259,6 +259,12 @@ public abstract class SegmentStages implements SegmentLock, LocksInterface {
     public abstract boolean locksInit();
 
     void initLocks() {
+        // This is a dummy check needed to make the Locks stage depend on the chaining.Used stage,
+        // to move the chaining.Used stage up in topological order, to make it closed later in
+        // the global context close() method, to ensure the context is unlocked (this is done in
+        // the chaining.closeUsed() method) after all other stages, potentially accessing Map's
+        // off-heap memory, are closed.
+        assert chaining.used;
         // Ensure SegmentHeader is init. This method initLocks() doesn't trigger SegmentHeader
         // initialization otherwise, so it could remain uninit on the beginning of locking methods.
         // But Locks still anyway depends on SegmentHeader stage (at least via

src/main/java/net/openhft/chronicle/hash/impl/stage/hash/Chaining.java

@@ -16,6 +16,7 @@
 
 package net.openhft.chronicle.hash.impl.stage.hash;
 
+import net.openhft.sg.Stage;
 import net.openhft.sg.Staged;
 
 import java.util.ArrayList;
@@ -27,6 +28,9 @@ public abstract class Chaining extends ChainingInterface {
 
     public final List<ChainingInterface> contextChain;
     public final int indexInContextChain;
+    /**
+     * First context, ever created in this thread. rootContextInThisThread === contextChain.get(0).
+     */
     public final ChainingInterface rootContextInThisThread;
 
     public Chaining() {
@@ -53,20 +57,30 @@ public abstract class Chaining extends ChainingInterface {
         return (T) contextChain.get(index);
     }
 
-    boolean used;
+
+    @Stage("Used") public boolean used;
+    @Stage("Used") private boolean firstContextLockedInThisThread;
 
     @Override
     public boolean usedInit() {
         return used;
     }
 
+    /**
+     * Init method parameterized to avoid automatic initialization. {@code used} argument should
+     * always be {@code true}.
+     */
     @Override
     public void initUsed(boolean used) {
-        this.used = used;
+        assert used;
+        firstContextLockedInThisThread = rootContextInThisThread.lockContextLocally();
+        this.used = true;
     }
     
     void closeUsed() {
         used = false;
+        if (firstContextLockedInThisThread)
+            rootContextInThisThread.unlockContextLocally();
     }
 
     @Override

src/main/java/net/openhft/chronicle/hash/impl/stage/hash/CheckOnEachPublicOperation.java

@@ -16,7 +16,6 @@
 
 package net.openhft.chronicle.hash.impl.stage.hash;
 
-import net.openhft.chronicle.hash.impl.VanillaChronicleHashHolder;
 import net.openhft.chronicle.hash.impl.stage.entry.ReadLock;
 import net.openhft.chronicle.hash.impl.stage.entry.UpdateLock;
 import net.openhft.chronicle.hash.impl.stage.entry.WriteLock;
@@ -28,7 +27,6 @@ import net.openhft.sg.Staged;
 public abstract class CheckOnEachPublicOperation {
     
     @StageRef OwnerThreadHolder holder;
-    @StageRef VanillaChronicleHashHolder<?> hh;
     
     public void checkOnEachPublicOperation() {
         checkOnEachLockOperation();
@@ -42,7 +40,5 @@ public abstract class CheckOnEachPublicOperation {
      */
     public void checkOnEachLockOperation() {
         holder.checkAccessingFromOwnerThread();
-        if (!hh.h().isOpen())
-            throw new IllegalStateException("Access to Chronicle Hash after close()");
     }
 }

src/main/java/net/openhft/chronicle/hash/impl/stage/hash/ThreadLocalState.java

@@ -16,8 +16,114 @@
 
 package net.openhft.chronicle.hash.impl.stage.hash;
 
+import net.openhft.chronicle.core.Memory;
+import net.openhft.chronicle.core.OS;
+import net.openhft.chronicle.hash.ChronicleHash;
+
+import java.lang.reflect.Field;
+import java.util.Arrays;
+import java.util.concurrent.TimeUnit;
+
+import static net.openhft.chronicle.hash.impl.BigSegmentHeader.LOCK_TIMEOUT_SECONDS;
+
 public abstract class ThreadLocalState {
     public boolean iterationContextLockedInThisThread;
 
+    private static final Memory MEMORY = OS.memory();
+    private static final long CONTEXT_LOCK_OFFSET;
+    private static final int CONTEXT_UNLOCKED = 0;
+    private static final int CONTEXT_LOCKED_LOCALLY = 1;
+    private static final int CONTEXT_CLOSED = 2;
+
+    static {
+        try {
+            Field contextLockField =
+                    ThreadLocalState.class.getDeclaredField("contextLock");
+            contextLockField.setAccessible(true);
+            CONTEXT_LOCK_OFFSET = MEMORY.getFieldOffset(contextLockField);
+        } catch (NoSuchFieldException | SecurityException e) {
+            throw new AssertionError(e);
+        }
+    }
+
+    private volatile int contextLock = CONTEXT_UNLOCKED;
+
+    /**
+     * Returns {@code true} if this is the outer context lock in this thread, {@code false} if this
+     * is a nested context.
+     */
+    public boolean lockContextLocally() {
+        // h().isOpen() check guarantees no starvation of a thread calling chMap.close() and trying
+        // to close this context by closeContext() method below, while the thread owning this
+        // context frequently locks and unlocks it (e. g. in a loop). This is also the only check
+        // for chMap openness during the whole context usage lifecycle.
+        if (hash().isOpen() && MEMORY.compareAndSwapInt(this, CONTEXT_LOCK_OFFSET,
+                CONTEXT_UNLOCKED, CONTEXT_LOCKED_LOCALLY)) {
+            return true;
+        } else {
+            if (contextLock == CONTEXT_LOCKED_LOCALLY)
+                return false;
+            if (contextLock == CONTEXT_CLOSED)
+                throw new IllegalStateException("Access to ChronicleHash after close()");
+            throw new AssertionError("Unknown context lock state: " + contextLock);
+        }
+    }
+
+    public void unlockContextLocally() {
+        // Ensure all reads from mapped memory are done before thread calling chronicleMap.close()
+        // frees resources potentially nulling some pointer used in those reads.
+        MEMORY.loadFence();
+        // Avoid volatile write to avoid expensive store-load barrier
+        MEMORY.writeOrderedInt(this, CONTEXT_LOCK_OFFSET, CONTEXT_UNLOCKED);
+    }
+
+    public void closeContext() {
+        if (tryCloseContext())
+            return;
+        if (owner() == Thread.currentThread())
+            throw new IllegalStateException("Attempt to close a Chronicle Hash in the context " +
+                    "of not yet finished query or iteration");
+        // Double the current timeout for segment locks "without timeout", that effectively
+        // specifies maximum lock (hence context) holding time
+        long timeoutMillis = TimeUnit.SECONDS.toMillis(LOCK_TIMEOUT_SECONDS) * 2;
+        long lastTime = System.currentTimeMillis();
+        do {
+            if (tryCloseContext())
+                return;
+            Thread.yield();
+            long now = System.currentTimeMillis();
+            if (now != lastTime) {
+                lastTime = now;
+                timeoutMillis--;
+            }
+        } while (timeoutMillis >= 0);
+        throw new RuntimeException("Failed to close a context, belonging to the thread\n" +
+                owner() + ", in the state: " + owner().getState() + ", current stack trace:\n" +
+                Arrays.toString(owner().getStackTrace()) + "\n" +
+                "Possible reasons:\n" +
+                "- The context owner thread exited before closing this context. Ensure that you\n" +
+                "always close opened Chronicle Map's contexts, the best way to do this is to use\n" +
+                "try-with-resources blocks." +
+                "- The context owner thread runs some context operation (e. g. a query) for\n" +
+                "unexpectedly long time (at least " + LOCK_TIMEOUT_SECONDS + " seconds).\n" +
+                "You should either redesign your logic to spend less time in Chronicle Map\n" +
+                "contexts (recommended) or synchronize map.close() with queries externally,\n" +
+                "so that close() is called only after all query operations finished.\n" +
+                "- Iteration over a large Chronicle Map takes more than " + LOCK_TIMEOUT_SECONDS +
+                " seconds.\n" +
+                "In this case you should synchronize map.close() with iterations over the map\n" +
+                "externally, so that close() is called only after all iterations are finished.\n" +
+                "- This is a dead lock involving the context owner thread and this thread (from\n" +
+                "which map.close() method is called. Make sure you always close Chronicle Map\n" +
+                "contexts, preferably using try-with-resources blocks.");
+    }
+
+    private boolean tryCloseContext() {
+        return MEMORY.compareAndSwapInt(this, CONTEXT_LOCK_OFFSET,
+                CONTEXT_UNLOCKED, CONTEXT_CLOSED);
+    }
+
     public abstract Thread owner();
+
+    public abstract ChronicleHash<?, ?, ?, ?> hash();
 }

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

@@ -563,6 +563,7 @@ public class ReplicatedChronicleMap<K, V, R> extends VanillaChronicleMap<K, V, R
         queryContext = cxt.get();
         if (queryContext == null) {
             queryContext = new CompiledReplicatedMapQueryContext<>(ReplicatedChronicleMap.this);
+            addContext(queryContext);
             cxt.set(queryContext);
         }
         return queryContext;
@@ -597,6 +598,7 @@ public class ReplicatedChronicleMap<K, V, R> extends VanillaChronicleMap<K, V, R
         iterContext = cxt.get();
         if (iterContext == null) {
             iterContext = new CompiledReplicatedMapIterationContext<>(ReplicatedChronicleMap.this);
+            addContext(iterContext);
             cxt.set(iterContext);
         }
         return iterContext;

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

@@ -34,9 +34,7 @@ import net.openhft.chronicle.wire.WireIn;
 import net.openhft.chronicle.wire.WireOut;
 import org.jetbrains.annotations.NotNull;
 
-import java.io.File;
 import java.io.IOException;
-import java.io.RandomAccessFile;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
 import java.util.Set;
@@ -248,6 +246,7 @@ public class VanillaChronicleMap<K, V, R>
         queryContext = cxt.get();
         if (queryContext == null) {
             queryContext = new CompiledMapQueryContext<>(VanillaChronicleMap.this);
+            addContext(queryContext);
             cxt.set(queryContext);
         }
         return queryContext;
@@ -263,6 +262,7 @@ public class VanillaChronicleMap<K, V, R>
         iterContext = cxt.get();
         if (iterContext == null) {
             iterContext = new CompiledMapIterationContext<>(VanillaChronicleMap.this);
+            addContext(iterContext);
             cxt.set(iterContext);
         }
         return iterContext;

src/main/java/net/openhft/chronicle/map/impl/stage/map/ReplicatedChronicleMapHolderImpl.java

@@ -57,6 +57,7 @@ public abstract class ReplicatedChronicleMapHolderImpl<K, V, R> extends Chaining
         return m.chronicleSet;
     }
 
+    @Override
     public ChronicleHash<K, ?, ?, ?> hash() {
         return set() != null ? set() : map();
     }

src/main/java/net/openhft/chronicle/map/impl/stage/map/VanillaChronicleMapHolderImpl.java

@@ -57,6 +57,7 @@ public abstract class VanillaChronicleMapHolderImpl<K, V, R> extends Chaining
         return m.chronicleSet;
     }
 
+    @Override
     public ChronicleHash<K, ?, ?, ?> hash() {
         return set() != null ? set() : map();
     }