HCOLL-353 Implement nested contexts

README.md

@@ -1556,8 +1556,8 @@ enough to
  - Backup all changes to ChronicleMap to some durable storage, e. g. SQL database
  - Perform multi-ChronicleMap operations correctly in concurrent environment,
    by acquiring locks on all ChronicleMaps before updating them.
- - <s>Perform multi-key operations on a single ChronicleMap correctly in concurrent environment,
-   by acquiring locks on all keys before updating the entries</s> SOON
+ - Perform multi-key operations on a single ChronicleMap correctly in concurrent environment,
+   by acquiring locks on all keys before updating the entries
  - Define own replication/reconciliation logic for distributed ChronicleMaps
 
 ## Examples
@@ -1917,6 +1917,125 @@ ChronicleMap<Integer, Set<Integer>> map1 = ChronicleMapBuilder
         .create();
 ```
 
+### Undirected graph
+
+Using nested contexts:
+
+```java
+    public static boolean addEdge(
+            ChronicleMap<Integer, Set<Integer>> graph, int source, int target) {
+        if (source == target)
+            throw new IllegalArgumentException("loops are forbidden");
+
+        try (ExternalMapQueryContext<Integer, Set<Integer>, ?> sc = graph.queryContext(source)) {
+            sc.updateLock().lock();
+            try (ExternalMapQueryContext<Integer, Set<Integer>, ?> tc =
+                         graph.queryContext(target)) {
+                tc.updateLock().lock();
+                MapEntry<Integer, Set<Integer>> sEntry = sc.entry();
+                if (sEntry != null) {
+                    MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                    if (tEntry != null) {
+                        return addEdgeBothPresent(sc, sEntry, source, tc, tEntry, target);
+                    } else {
+                        addEdgePresentAbsent(sc, sEntry, source, tc, target);
+                        return true;
+                    }
+                } else {
+                    MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                    if (tEntry != null) {
+                        addEdgePresentAbsent(tc, tEntry, target, sc, source);
+                    } else {
+                        addEdgeBothAbsent(sc, source, tc, target);
+                    }
+                    return true;
+                }
+            }
+        }
+    }
+
+    private static boolean addEdgeBothPresent(
+            MapQueryContext<Integer, Set<Integer>, ?> sc,
+            @NotNull MapEntry<Integer, Set<Integer>> sEntry, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc,
+            @NotNull MapEntry<Integer, Set<Integer>> tEntry, int target) {
+        Set<Integer> sNeighbours = sEntry.value().get();
+        if (sNeighbours.add(target)) {
+            Set<Integer> tNeighbours = tEntry.value().get();
+            boolean added = tNeighbours.add(source);
+            assert added;
+            sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+            tEntry.doReplaceValue(tc.wrapValueAsData(tNeighbours));
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    private static void addEdgePresentAbsent(
+            MapQueryContext<Integer, Set<Integer>, ?> sc,
+            @NotNull MapEntry<Integer, Set<Integer>> sEntry, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc, int target) {
+        Set<Integer> sNeighbours = sEntry.value().get();
+        boolean added = sNeighbours.add(target);
+        assert added;
+        sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+
+        addEdgeOneSide(tc, source);
+    }
+
+    private static void addEdgeBothAbsent(MapQueryContext<Integer, Set<Integer>, ?> sc, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc, int target) {
+        addEdgeOneSide(sc, target);
+        addEdgeOneSide(tc, source);
+    }
+
+    private static void addEdgeOneSide(MapQueryContext<Integer, Set<Integer>, ?> tc, int source) {
+        Set<Integer> tNeighbours = new HashSet<>();
+        tNeighbours.add(source);
+        MapAbsentEntry<Integer, Set<Integer>> tAbsentEntry = tc.absentEntry();
+        assert tAbsentEntry != null;
+        tAbsentEntry.doInsert(tc.wrapValueAsData(tNeighbours));
+    }
+
+    public static boolean removeEdge(
+            ChronicleMap<Integer, Set<Integer>> graph, int source, int target) {
+        try (ExternalMapQueryContext<Integer, Set<Integer>, ?> sc = graph.queryContext(source)) {
+            sc.updateLock().lock();
+            MapEntry<Integer, Set<Integer>> sEntry = sc.entry();
+            if (sEntry == null)
+                return false;
+            Set<Integer> sNeighbours = sEntry.value().get();
+            if (!sNeighbours.remove(target))
+                return false;
+            try (ExternalMapQueryContext<Integer, Set<Integer>, ?> tc =
+                         graph.queryContext(target)) {
+                tc.updateLock().lock();
+                MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                if (tEntry == null)
+                    throw new IllegalStateException("target node should be present in the graph");
+                Set<Integer> tNeighbours = tEntry.value().get();
+                if (!tNeighbours.remove(source))
+                    throw new IllegalStateException("the target node have an edge to the source");
+                sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+                tEntry.doReplaceValue(tc.wrapValueAsData(tNeighbours));
+                return true;
+            }
+        }
+    }
+```
+
+Usage:
+```java
+ChronicleMap<Integer, Set<Integer>> graph = ChronicleMapBuilder
+        .of(Integer.class, (Class<Set<Integer>>) (Class) Set.class)
+        .entries(100)
+        .averageValueSize(1000)
+        .create();
+
+addEdge(graph, 1, 2);
+removeEdge(graph, 1, 2);
+```
 
 # Performance Topics
 

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

@@ -44,7 +44,9 @@ public class ReadLock implements InterProcessLock {
     @Override
     public void lock() {
         if (s.localLockState == UNLOCKED) {
-            s.segmentHeader.readLock(s.segmentHeaderAddress);
+            if (s.readZero() && s.updateZero() && s.writeZero())
+                s.segmentHeader.readLock(s.segmentHeaderAddress);
+            s.incrementRead();
             s.setLocalLockState(READ_LOCKED);
         }
     }
@@ -52,7 +54,9 @@ public class ReadLock implements InterProcessLock {
     @Override
     public void lockInterruptibly() throws InterruptedException {
         if (s.localLockState == UNLOCKED) {
-            s.segmentHeader.readLockInterruptibly(s.segmentHeaderAddress);
+            if (s.readZero() && s.updateZero() && s.writeZero())
+                s.segmentHeader.readLockInterruptibly(s.segmentHeaderAddress);
+            s.incrementRead();
             s.setLocalLockState(READ_LOCKED);
         }
     }
@@ -60,7 +64,9 @@ public class ReadLock implements InterProcessLock {
     @Override
     public boolean tryLock() {
         if (s.localLockState == UNLOCKED) {
-            if (s.segmentHeader.tryReadLock(s.segmentHeaderAddress)) {
+            if (!s.readZero() || !s.updateZero() || !s.writeZero() ||
+                    s.segmentHeader.tryReadLock(s.segmentHeaderAddress)) {
+                s.incrementRead();
                 s.setLocalLockState(READ_LOCKED);
                 return true;
             } else {
@@ -74,7 +80,9 @@ public class ReadLock implements InterProcessLock {
     @Override
     public boolean tryLock(long time, @NotNull TimeUnit unit) throws InterruptedException {
         if (s.localLockState == UNLOCKED) {
-            if (s.segmentHeader.tryReadLock(s.segmentHeaderAddress, time, unit)) {
+            if (!s.readZero() || !s.updateZero() || !s.writeZero() ||
+                    s.segmentHeader.tryReadLock(s.segmentHeaderAddress, time, unit)) {
+                s.incrementRead();
                 s.setLocalLockState(READ_LOCKED);
                 return true;
             } else {
@@ -87,18 +95,7 @@ public class ReadLock implements InterProcessLock {
 
     @Override
     public void unlock() {
-        switch (s.localLockState) {
-            case UNLOCKED:
-                return;
-            case READ_LOCKED:
-                s.segmentHeader.readUnlock(s.segmentHeaderAddress);
-                break;
-            case UPDATE_LOCKED:
-                s.segmentHeader.updateUnlock(s.segmentHeaderAddress);
-                break;
-            case WRITE_LOCKED:
-                s.segmentHeader.writeUnlock(s.segmentHeaderAddress);
-        }
+        s.readUnlockAndDecrementCount();
         s.setLocalLockState(UNLOCKED);
         // TODO what should close here?
         hlp.closeHashLookupPos();

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

@@ -51,9 +51,13 @@ public abstract class SegmentStages implements InterProcessReadWriteUpdateLock {
         
     }
 
+    // "SegHeader" because stage-generator doesn't support stage names - one is a prefix of another
+    // there is a stage named "Segment"
     @Stage("SegHeader") long segmentHeaderAddress;
     @Stage("SegHeader") SegmentHeader segmentHeader = null;
 
+    abstract boolean segHeaderInit();
+
     private void initSegHeader() {
         segmentHeaderAddress = hh.h().ms.address() + hh.h().segmentHeaderOffset(segmentIndex);
         segmentHeader = BigSegmentHeader.INSTANCE;
@@ -104,8 +108,7 @@ public abstract class SegmentStages implements InterProcessReadWriteUpdateLock {
 
     @Stage("Locks")
     public void incrementModCount() {
-        contextModCount = rootContextOnThisSegment.latestSameThreadSegmentModCount =
-                rootContextOnThisSegment.latestSameThreadSegmentModCount + 1;
+        contextModCount = ++rootContextOnThisSegment.latestSameThreadSegmentModCount;
     }
 
     // chain
@@ -116,6 +119,51 @@ public abstract class SegmentStages implements InterProcessReadWriteUpdateLock {
     @Stage("Locks") int totalUpdateLockCount;
     @Stage("Locks") int totalWriteLockCount;
 
+    @Stage("Locks")
+    public boolean readZero() {
+        return rootContextOnThisSegment.totalReadLockCount == 0;
+    }
+
+    @Stage("Locks")
+    public boolean updateZero() {
+        return rootContextOnThisSegment.totalUpdateLockCount == 0;
+    }
+
+    @Stage("Locks")
+    public boolean writeZero() {
+        return rootContextOnThisSegment.totalWriteLockCount == 0;
+    }
+
+    @Stage("Locks")
+    public int decrementRead() {
+        return rootContextOnThisSegment.totalReadLockCount -= 1;
+    }
+
+    @Stage("Locks")
+    public int decrementUpdate() {
+        return rootContextOnThisSegment.totalUpdateLockCount -= 1;
+    }
+
+    @Stage("Locks")
+    public int decrementWrite() {
+        return rootContextOnThisSegment.totalWriteLockCount -= 1;
+    }
+
+    @Stage("Locks")
+    public void incrementRead() {
+        rootContextOnThisSegment.totalReadLockCount += 1;
+    }
+
+    @Stage("Locks")
+    public void incrementUpdate() {
+        rootContextOnThisSegment.totalUpdateLockCount += 1;
+    }
+
+    @Stage("Locks")
+    public void incrementWrite() {
+        rootContextOnThisSegment.totalWriteLockCount += 1;
+    }
+
     public abstract boolean locksInit();
 
     void initLocks() {
@@ -131,27 +179,34 @@ public abstract class SegmentStages implements InterProcessReadWriteUpdateLock {
         }
         rootContextOnThisSegment = this;
         concurrentSameThreadContexts = false;
+
         latestSameThreadSegmentModCount = 0;
         contextModCount = 0;
-        nextNode = null;
+
         totalReadLockCount = 0;
         totalUpdateLockCount = 0;
         totalWriteLockCount = 0;
     }
 
+    @Stage("Locks")
     boolean tryFindInitLocksOfThisSegment(Object thisContext, int index) {
         SegmentStages c = chaining.contextAtIndexInChain(index);
-        if (c.segmentHeader != null && c.segmentHeaderAddress == segmentHeaderAddress &&
-                c.rootContextOnThisSegment != null) {
-            //TODO implement nested contexts. Hardest part - bodies of Read/Update/WriteLock classes
-            throw new IllegalStateException("Nested context not implemented yet");
+        if (c.segHeaderInit() &&
+                c.segmentHeaderAddress == segmentHeaderAddress &&
+                c.locksInit()) {
+            SegmentStages root = c.rootContextOnThisSegment;
+            this.rootContextOnThisSegment = root;
+            root.concurrentSameThreadContexts = true;
+            this.concurrentSameThreadContexts = true;
+            this.contextModCount = root.latestSameThreadSegmentModCount;
+            linkToSegmentContextsChain();
+            return true;
         } else {
             return false;
         }
     }
 
     void closeLocks() {
-        // TODO should throw ISE on attempt to close root context when children not closed
         if (rootContextOnThisSegment == this) {
             closeRootLocks();
         } else {
@@ -164,71 +219,87 @@ public abstract class SegmentStages implements InterProcessReadWriteUpdateLock {
     @Stage("Locks")
     private void closeNestedLocks() {
         unlinkFromSegmentContextsChain();
-        
+        readUnlockAndDecrementCount();
+    }
+
+    @Stage("Locks")
+    public void readUnlockAndDecrementCount() {
         switch (localLockState) {
             case UNLOCKED:
-                break;
+                return;
             case READ_LOCKED:
-                int newTotalReadLockCount = this.rootContextOnThisSegment.totalReadLockCount -= 1;
+                int newTotalReadLockCount = decrementRead();
                 if (newTotalReadLockCount == 0) {
-                    if (this.rootContextOnThisSegment.totalUpdateLockCount == 0 &&
-                            this.rootContextOnThisSegment.totalWriteLockCount == 0) {
+                    if (updateZero() && writeZero())
                         segmentHeader.readUnlock(segmentHeaderAddress);
-                    }
-                } else if (newTotalReadLockCount < 0) {
-                    throw new IllegalStateException("read underflow");
+                } else {
+                    assert newTotalReadLockCount > 0 : "read underflow";
                 }
-                break;
+                return;
             case UPDATE_LOCKED:
-                int newTotalUpdateLockCount =
-                        this.rootContextOnThisSegment.totalUpdateLockCount -= 1;
+                int newTotalUpdateLockCount = decrementUpdate();
                 if (newTotalUpdateLockCount == 0) {
-                    if (this.rootContextOnThisSegment.totalWriteLockCount == 0) {
-                        if (this.rootContextOnThisSegment.totalReadLockCount == 0) {
+                    if (writeZero()) {
+                        if (readZero()) {
                             segmentHeader.updateUnlock(segmentHeaderAddress);
                         } else {
                             segmentHeader.downgradeUpdateToReadLock(segmentHeaderAddress);
                         }
                     }
-                } else if (newTotalUpdateLockCount < 0) {
-                    throw new IllegalStateException("update underflow");
+                } else {
+                    assert newTotalUpdateLockCount > 0 : "update underflow";
                 }
-                break;
+                return;
             case WRITE_LOCKED:
-                int newTotalWriteLockCount = this.rootContextOnThisSegment.totalWriteLockCount -= 1;
+                int newTotalWriteLockCount = decrementWrite();
                 if (newTotalWriteLockCount == 0) {
-                    if (this.rootContextOnThisSegment.totalUpdateLockCount > 0) {
+                    if (!updateZero()) {
                         segmentHeader.downgradeWriteToUpdateLock(segmentHeaderAddress);
                     } else {
-                        if (this.rootContextOnThisSegment.totalReadLockCount > 0) {
+                        if (!readZero()) {
                             segmentHeader.downgradeWriteToReadLock(segmentHeaderAddress);
                         } else {
                             segmentHeader.writeUnlock(segmentHeaderAddress);
                         }
                     }
+                } else {
+                    assert newTotalWriteLockCount > 0 : "write underflow";
                 }
-                break;
         }
     }
 
+    @Stage("Locks")
+    private void linkToSegmentContextsChain() {
+        SegmentStages innermostContextOnThisSegment = rootContextOnThisSegment;
+        while (innermostContextOnThisSegment.nextNode != null)
+            innermostContextOnThisSegment = innermostContextOnThisSegment.nextNode;
+        innermostContextOnThisSegment.nextNode = this;
+    }
+
     @Stage("Locks")
     private void unlinkFromSegmentContextsChain() {
-        SegmentStages prevContext = this.rootContextOnThisSegment;
+        SegmentStages prevContext = rootContextOnThisSegment;
         while (true) {
             assert prevContext.nextNode != null;
             if (prevContext.nextNode == this)
                 break;
             prevContext = prevContext.nextNode;
         }
-        assert nextNode == null;
+        // i. e. structured unlocking
+        verifyInnermostContext();
         prevContext.nextNode = null;
     }
 
+    @Stage("Locks")
+    private void verifyInnermostContext() {
+        if (nextNode != null)
+            throw new IllegalStateException("Attempt to close contexts not structurally");
+    }
+
     @Stage("Locks")
     private void closeRootLocks() {
-        assert nextNode == null;
-        
-        
+        verifyInnermostContext();
+
         switch (localLockState) {
             case UNLOCKED:
                 return;

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

@@ -42,7 +42,12 @@ public class UpdateLock implements InterProcessLock {
     public void lock() {
         switch (s.localLockState) {
             case UNLOCKED:
-                s.segmentHeader.updateLock(s.segmentHeaderAddress);
+                if (s.updateZero() && s.writeZero()) {
+                    if (!s.readZero())
+                        throw forbiddenUpdateLockWhenOuterContextReadLocked();
+                    s.segmentHeader.updateLock(s.segmentHeaderAddress);
+                }
+                s.incrementUpdate();
                 s.setLocalLockState(UPDATE_LOCKED);
                 return;
             case READ_LOCKED:
@@ -61,11 +66,26 @@ public class UpdateLock implements InterProcessLock {
         return new IllegalMonitorStateException("Cannot upgrade from read to update lock");
     }
 
+    /**
+     * Non-static because after compilation it becomes inner class which forbids static methods
+     */
+    @NotNull
+    private IllegalStateException forbiddenUpdateLockWhenOuterContextReadLocked() {
+        return new IllegalStateException("Cannot acquire update lock, because outer context " +
+                "holds read lock. In this case you should acquire update lock in the outer " +
+                "context up front");
+    }
+
     @Override
     public void lockInterruptibly() throws InterruptedException {
         switch (s.localLockState) {
             case UNLOCKED:
-                s.segmentHeader.updateLockInterruptibly(s.segmentHeaderAddress);
+                if (s.updateZero() && s.writeZero()) {
+                    if (!s.readZero())
+                        throw forbiddenUpdateLockWhenOuterContextReadLocked();
+                    s.segmentHeader.updateLockInterruptibly(s.segmentHeaderAddress);
+                }
+                s.incrementUpdate();
                 s.setLocalLockState(UPDATE_LOCKED);
                 return;
             case READ_LOCKED:
@@ -80,11 +100,20 @@ public class UpdateLock implements InterProcessLock {
     public boolean tryLock() {
         switch (s.localLockState) {
             case UNLOCKED:
-                if (s.segmentHeader.tryUpdateLock(s.segmentHeaderAddress)) {
+                if (s.updateZero() && s.writeZero()) {
+                    if (!s.readZero())
+                        throw forbiddenUpdateLockWhenOuterContextReadLocked();
+                    if (s.segmentHeader.tryUpdateLock(s.segmentHeaderAddress)) {
+                        s.incrementUpdate();
+                        s.setLocalLockState(UPDATE_LOCKED);
+                        return true;
+                    } else {
+                        return false;
+                    }
+                } else {
+                    s.incrementUpdate();
                     s.setLocalLockState(UPDATE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case READ_LOCKED:
                 throw forbiddenUpgrade();
@@ -99,11 +128,20 @@ public class UpdateLock implements InterProcessLock {
     public boolean tryLock(long time, @NotNull TimeUnit unit) throws InterruptedException {
         switch (s.localLockState) {
             case UNLOCKED:
-                if (s.segmentHeader.tryUpdateLock(s.segmentHeaderAddress, time, unit)) {
+                if (s.updateZero() && s.writeZero()) {
+                    if (!s.readZero())
+                        throw forbiddenUpdateLockWhenOuterContextReadLocked();
+                    if (s.segmentHeader.tryUpdateLock(s.segmentHeaderAddress, time, unit)) {
+                        s.incrementUpdate();
+                        s.setLocalLockState(UPDATE_LOCKED);
+                        return true;
+                    } else {
+                        return false;
+                    }
+                } else {
+                    s.incrementUpdate();
                     s.setLocalLockState(UPDATE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case READ_LOCKED:
                 throw forbiddenUpgrade();
@@ -121,11 +159,27 @@ public class UpdateLock implements InterProcessLock {
             case READ_LOCKED:
                 return;
             case UPDATE_LOCKED:
-                s.segmentHeader.downgradeUpdateToReadLock(s.segmentHeaderAddress);
+                int newTotalUpdateLockCount = s.decrementUpdate();
+                if (newTotalUpdateLockCount == 0) {
+                    if (s.writeZero())
+                        s.segmentHeader.downgradeUpdateToReadLock(s.segmentHeaderAddress);
+                } else {
+                    assert newTotalUpdateLockCount > 0 : "update underflow";
+                }
                 break;
             case WRITE_LOCKED:
-                s.segmentHeader.downgradeWriteToReadLock(s.segmentHeaderAddress);
+                int newTotalWriteLockCount = s.decrementWrite();
+                if (newTotalWriteLockCount == 0) {
+                    if (!s.updateZero()) {
+                        s.segmentHeader.downgradeWriteToUpdateLock(s.segmentHeaderAddress);
+                    } else {
+                        s.segmentHeader.downgradeWriteToReadLock(s.segmentHeaderAddress);
+                    }
+                } else {
+                    assert newTotalWriteLockCount > 0 : "write underflow";
+                }
         }
+        s.incrementRead();
         s.setLocalLockState(READ_LOCKED);
     }
 }

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

@@ -42,13 +42,27 @@ public class WriteLock implements InterProcessLock {
     public void lock() {
         switch (s.localLockState) {
             case UNLOCKED:
-                s.segmentHeader.writeLock(s.segmentHeaderAddress);
+                if (s.writeZero()) {
+                    if (!s.updateZero()) {
+                        s.segmentHeader.upgradeUpdateToWriteLock(s.segmentHeaderAddress);
+                    } else {
+                        if (!s.readZero())
+                            throw forbiddenWriteLockWhenOuterContextReadLocked();
+                        s.segmentHeader.writeLock(s.segmentHeaderAddress);
+                    }
+                }
+                s.incrementWrite();
                 s.setLocalLockState(WRITE_LOCKED);
                 return;
             case READ_LOCKED:
                 throw forbiddenUpgrade();
             case UPDATE_LOCKED:
-                s.segmentHeader.upgradeUpdateToWriteLock(s.segmentHeaderAddress);
+                if (s.writeZero()) {
+                    assert !s.updateZero();
+                    s.segmentHeader.upgradeUpdateToWriteLock(s.segmentHeaderAddress);
+                }
+                s.decrementUpdate();
+                s.incrementWrite();
                 s.setLocalLockState(WRITE_LOCKED);
             case WRITE_LOCKED:
                 // do nothing
@@ -63,17 +77,42 @@ public class WriteLock implements InterProcessLock {
         return new IllegalMonitorStateException("Cannot upgrade from read to write lock");
     }
 
+    /**
+     * Non-static because after compilation it becomes inner class which forbids static methods
+     */
+    @NotNull
+    private IllegalStateException forbiddenWriteLockWhenOuterContextReadLocked() {
+        return new IllegalStateException("Cannot acquire write lock, because outer context " +
+                "holds read lock. In this case you should acquire update lock in the outer " +
+                "context up front");
+    }
+
     @Override
     public void lockInterruptibly() throws InterruptedException {
         switch (s.localLockState) {
             case UNLOCKED:
-                s.segmentHeader.writeLockInterruptibly(s.segmentHeaderAddress);
+                if (s.writeZero()) {
+                    if (!s.updateZero()) {
+                        s.segmentHeader.upgradeUpdateToWriteLockInterruptibly(
+                                s.segmentHeaderAddress);
+                    } else {
+                        if (!s.readZero())
+                            throw forbiddenWriteLockWhenOuterContextReadLocked();
+                        s.segmentHeader.writeLockInterruptibly(s.segmentHeaderAddress);
+                    }
+                }
+                s.incrementWrite();
                 s.setLocalLockState(WRITE_LOCKED);
                 return;
             case READ_LOCKED:
                 throw forbiddenUpgrade();
             case UPDATE_LOCKED:
-                s.segmentHeader.upgradeUpdateToWriteLockInterruptibly(s.segmentHeaderAddress);
+                if (s.writeZero()) {
+                    assert !s.updateZero();
+                    s.segmentHeader.upgradeUpdateToWriteLockInterruptibly(s.segmentHeaderAddress);
+                }
+                s.decrementUpdate();
+                s.incrementWrite();
                 s.setLocalLockState(WRITE_LOCKED);
             case WRITE_LOCKED:
                 // do nothing
@@ -84,20 +123,49 @@ public class WriteLock implements InterProcessLock {
     public boolean tryLock() {
         switch (s.localLockState) {
             case UNLOCKED:
-                if (s.segmentHeader.tryWriteLock(s.segmentHeaderAddress)) {
+                if (s.writeZero()) {
+                    if (!s.updateZero()) {
+                        if (s.segmentHeader.tryUpgradeUpdateToWriteLock(s.segmentHeaderAddress)) {
+                            s.incrementWrite();
+                            s.setLocalLockState(WRITE_LOCKED);
+                            return true;
+                        } else {
+                            return false;
+                        }
+                    } else {
+                        if (!s.readZero())
+                            throw forbiddenWriteLockWhenOuterContextReadLocked();
+                        if (s.segmentHeader.tryWriteLock(s.segmentHeaderAddress)) {
+                            s.incrementWrite();
+                            s.setLocalLockState(WRITE_LOCKED);
+                            return true;
+                        } else {
+                            return false;
+                        }
+                    }
+                } else {
+                    s.incrementWrite();
                     s.setLocalLockState(WRITE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case READ_LOCKED:
                 throw forbiddenUpgrade();
             case UPDATE_LOCKED:
-                if (s.segmentHeader.tryUpgradeUpdateToWriteLock(s.segmentHeaderAddress)) {
+                if (s.writeZero()) {
+                    assert !s.updateZero();
+                    if (s.segmentHeader.tryUpgradeUpdateToWriteLock(s.segmentHeaderAddress)) {
+                        s.decrementUpdate();
+                        s.incrementWrite();
+                        s.setLocalLockState(WRITE_LOCKED);
+                        return true;
+                    } else {
+                        return false;
+                    }
+                } else {
+                    s.decrementUpdate();
+                    s.incrementWrite();
                     s.setLocalLockState(WRITE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case WRITE_LOCKED:
                 return true;
@@ -109,21 +177,51 @@ public class WriteLock implements InterProcessLock {
     public boolean tryLock(long time, @NotNull TimeUnit unit) throws InterruptedException {
         switch (s.localLockState) {
             case UNLOCKED:
-                if (s.segmentHeader.tryWriteLock(s.segmentHeaderAddress, time, unit)) {
+                if (s.writeZero()) {
+                    if (!s.updateZero()) {
+                        if (s.segmentHeader.tryUpgradeUpdateToWriteLock(
+                                s.segmentHeaderAddress, time, unit)) {
+                            s.incrementWrite();
+                            s.setLocalLockState(WRITE_LOCKED);
+                            return true;
+                        } else {
+                            return false;
+                        }
+                    } else {
+                        if (!s.readZero())
+                            throw forbiddenWriteLockWhenOuterContextReadLocked();
+                        if (s.segmentHeader.tryWriteLock(s.segmentHeaderAddress, time, unit)) {
+                            s.incrementWrite();
+                            s.setLocalLockState(WRITE_LOCKED);
+                            return true;
+                        } else {
+                            return false;
+                        }
+                    }
+                } else {
+                    s.incrementWrite();
                     s.setLocalLockState(WRITE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case READ_LOCKED:
                 throw forbiddenUpgrade();
             case UPDATE_LOCKED:
-                if (s.segmentHeader.tryUpgradeUpdateToWriteLock(
-                        s.segmentHeaderAddress, time, unit)) {
+                if (s.writeZero()) {
+                    assert !s.updateZero();
+                    if (s.segmentHeader.tryUpgradeUpdateToWriteLock(
+                            s.segmentHeaderAddress, time, unit)) {
+                        s.decrementUpdate();
+                        s.incrementWrite();
+                        s.setLocalLockState(WRITE_LOCKED);
+                        return true;
+                    } else {
+                        return false;
+                    }
+                } else {
+                    s.decrementUpdate();
+                    s.incrementWrite();
                     s.setLocalLockState(WRITE_LOCKED);
                     return true;
-                } else {
-                    return false;
                 }
             case WRITE_LOCKED:
                 return true;
@@ -139,8 +237,14 @@ public class WriteLock implements InterProcessLock {
             case UPDATE_LOCKED:
                 return;
             case WRITE_LOCKED:
-                s.segmentHeader.downgradeWriteToUpdateLock(s.segmentHeaderAddress);
+                int newTotalWriteLockCount = s.decrementWrite();
+                if (newTotalWriteLockCount == 0) {
+                    s.segmentHeader.downgradeWriteToUpdateLock(s.segmentHeaderAddress);
+                } else {
+                    assert newTotalWriteLockCount > 0 : "write underflow";
+                }
         }
+        s.incrementUpdate();
         s.setLocalLockState(UPDATE_LOCKED);
     }
 }

src/test/java/net/openhft/chronicle/map/NestedContextsTest.java

@@ -0,0 +1,181 @@
+/*
+ * Copyright 2015 Higher Frequency Trading
+ *
+ *  http://www.higherfrequencytrading.com
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package net.openhft.chronicle.map;
+
+import com.google.common.collect.ImmutableSet;
+import org.jetbrains.annotations.NotNull;
+import org.junit.Test;
+
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.ThreadLocalRandom;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+public class NestedContextsTest {
+
+    @Test
+    public void nestedContextsTest() throws ExecutionException, InterruptedException {
+        ChronicleMap<Integer, Set<Integer>> graph = ChronicleMapBuilder
+                .of(Integer.class, (Class<Set<Integer>>) (Class) Set.class)
+                .actualSegments(1)
+                .create();
+
+        addEdge(graph, 1, 2);
+        addEdge(graph, 2, 3);
+        addEdge(graph, 1, 3);
+
+        assertEquals(ImmutableSet.of(2, 3), graph.get(1));
+        assertEquals(ImmutableSet.of(1, 3), graph.get(2));
+        assertEquals(ImmutableSet.of(1, 2), graph.get(3));
+
+        verifyGraphConsistent(graph);
+
+        ForkJoinPool pool = new ForkJoinPool(8);
+        try {
+            pool.submit(() -> {
+                ThreadLocalRandom.current().ints().limit(10_000).parallel().forEach(i -> {
+                    int sourceNode = Math.abs(i % 10);
+                    int targetNode;
+                    do {
+                        targetNode = ThreadLocalRandom.current().nextInt(10);
+                    } while (targetNode == sourceNode);
+                    if (i % 2 == 0) {
+                        addEdge(graph, sourceNode, targetNode);
+                    } else {
+                        removeEdge(graph, sourceNode, targetNode);
+                    }
+                });
+            }).get();
+            verifyGraphConsistent(graph);
+        } finally {
+            pool.shutdownNow();
+        }
+    }
+
+    private static void verifyGraphConsistent(ChronicleMap<Integer, Set<Integer>> graph) {
+        graph.forEach((node, neighbours) ->
+                neighbours.forEach(neighbour -> assertTrue(graph.get(neighbour).contains(node))));
+    }
+
+    public static boolean addEdge(
+            ChronicleMap<Integer, Set<Integer>> graph, int source, int target) {
+        if (source == target)
+            throw new IllegalArgumentException("loops are forbidden");
+
+        try (ExternalMapQueryContext<Integer, Set<Integer>, ?> sc = graph.queryContext(source)) {
+            sc.updateLock().lock();
+            try (ExternalMapQueryContext<Integer, Set<Integer>, ?> tc =
+                         graph.queryContext(target)) {
+                tc.updateLock().lock();
+                MapEntry<Integer, Set<Integer>> sEntry = sc.entry();
+                if (sEntry != null) {
+                    MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                    if (tEntry != null) {
+                        return addEdgeBothPresent(sc, sEntry, source, tc, tEntry, target);
+                    } else {
+                        addEdgePresentAbsent(sc, sEntry, source, tc, target);
+                        return true;
+                    }
+                } else {
+                    MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                    if (tEntry != null) {
+                        addEdgePresentAbsent(tc, tEntry, target, sc, source);
+                    } else {
+                        addEdgeBothAbsent(sc, source, tc, target);
+                    }
+                    return true;
+                }
+            }
+        }
+    }
+
+    private static boolean addEdgeBothPresent(
+            MapQueryContext<Integer, Set<Integer>, ?> sc,
+            @NotNull MapEntry<Integer, Set<Integer>> sEntry, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc,
+            @NotNull MapEntry<Integer, Set<Integer>> tEntry, int target) {
+        Set<Integer> sNeighbours = sEntry.value().get();
+        if (sNeighbours.add(target)) {
+            Set<Integer> tNeighbours = tEntry.value().get();
+            boolean added = tNeighbours.add(source);
+            assert added;
+            sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+            tEntry.doReplaceValue(tc.wrapValueAsData(tNeighbours));
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    private static void addEdgePresentAbsent(
+            MapQueryContext<Integer, Set<Integer>, ?> sc,
+            @NotNull MapEntry<Integer, Set<Integer>> sEntry, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc, int target) {
+        Set<Integer> sNeighbours = sEntry.value().get();
+        boolean added = sNeighbours.add(target);
+        assert added;
+        sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+
+        addEdgeOneSide(tc, source);
+    }
+
+    private static void addEdgeBothAbsent(MapQueryContext<Integer, Set<Integer>, ?> sc, int source,
+            MapQueryContext<Integer, Set<Integer>, ?> tc, int target) {
+        addEdgeOneSide(sc, target);
+        addEdgeOneSide(tc, source);
+    }
+
+    private static void addEdgeOneSide(MapQueryContext<Integer, Set<Integer>, ?> tc, int source) {
+        Set<Integer> tNeighbours = new HashSet<>();
+        tNeighbours.add(source);
+        MapAbsentEntry<Integer, Set<Integer>> tAbsentEntry = tc.absentEntry();
+        assert tAbsentEntry != null;
+        tAbsentEntry.doInsert(tc.wrapValueAsData(tNeighbours));
+    }
+
+    public static boolean removeEdge(
+            ChronicleMap<Integer, Set<Integer>> graph, int source, int target) {
+        try (ExternalMapQueryContext<Integer, Set<Integer>, ?> sc = graph.queryContext(source)) {
+            sc.updateLock().lock();
+            MapEntry<Integer, Set<Integer>> sEntry = sc.entry();
+            if (sEntry == null)
+                return false;
+            Set<Integer> sNeighbours = sEntry.value().get();
+            if (!sNeighbours.remove(target))
+                return false;
+            try (ExternalMapQueryContext<Integer, Set<Integer>, ?> tc =
+                         graph.queryContext(target)) {
+                tc.updateLock().lock();
+                MapEntry<Integer, Set<Integer>> tEntry = tc.entry();
+                if (tEntry == null)
+                    throw new IllegalStateException("target node should be present in the graph");
+                Set<Integer> tNeighbours = tEntry.value().get();
+                if (!tNeighbours.remove(source))
+                    throw new IllegalStateException("the target node have an edge to the source");
+                sEntry.doReplaceValue(sc.wrapValueAsData(sNeighbours));
+                tEntry.doReplaceValue(tc.wrapValueAsData(tNeighbours));
+                return true;
+            }
+        }
+    }
+}