From a36e1bcfb74aa75067f2e06f82e75ef15367fe2f Mon Sep 17 00:00:00 2001
From: leventov <leventov@ya.ru>
Date: Mon, 19 Jan 2015 09:01:29 +0700
Subject: [PATCH] More bug fixes

---
 .../chronicle/map/AbstractChronicleMap.java       | 14 ++++++++++----
 .../net/openhft/chronicle/map/VanillaContext.java | 15 ++++++++++-----
 .../openhft/chronicle/map/ChronicleMapTest.java   |  5 ++---
 3 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/src/main/java/net/openhft/chronicle/map/AbstractChronicleMap.java b/src/main/java/net/openhft/chronicle/map/AbstractChronicleMap.java
index c9bdaac4b..f4260986b 100644
--- a/src/main/java/net/openhft/chronicle/map/AbstractChronicleMap.java
+++ b/src/main/java/net/openhft/chronicle/map/AbstractChronicleMap.java
@@ -299,12 +299,18 @@ abstract class AbstractChronicleMap<K, V> extends AbstractMap<K, V>
         }
 
         @Override
-        public void accept(long key, long pos) {
+        public void accept(long hash, long pos) {
             context.pos = pos;
             context.initKeyFromPos();
-            if (!context.containsKey()) // for replicated map
-                return;
-            entryBuffer.add(new WriteThroughEntry(context.immutableKey(), context.getUsing(null)));
+            try {
+                if (!context.containsKey()) // for replicated map
+                    return;
+                K key = context.immutableKey();
+                V value = context.getUsing(null);
+                entryBuffer.add(new WriteThroughEntry(key, value));
+            } finally {
+                context.closeKeySearch();
+            }
         }
 
         @Override
diff --git a/src/main/java/net/openhft/chronicle/map/VanillaContext.java b/src/main/java/net/openhft/chronicle/map/VanillaContext.java
index 1002aa531..eada44bd0 100644
--- a/src/main/java/net/openhft/chronicle/map/VanillaContext.java
+++ b/src/main/java/net/openhft/chronicle/map/VanillaContext.java
@@ -647,10 +647,11 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
                 segmentHeader.readUnlock(segmentHeaderAddress);
                 closeKeySearch();
             }
-            rootContextOnThisSegment.totalReadLockCount -= readLockCount;
-            rootContextOnThisSegment.totalUpdateLockCount -= updateLockCount;
-            rootContextOnThisSegment.totalWriteLockCount -= writeLockCount;
         }
+        rootContextOnThisSegment.totalReadLockCount -= readLockCount;
+        rootContextOnThisSegment.totalUpdateLockCount -= updateLockCount;
+        rootContextOnThisSegment.totalWriteLockCount -= writeLockCount;
+        readLockCount = updateLockCount = writeLockCount = 0;
         rootContextOnThisSegment = null;
     }
 
@@ -1089,7 +1090,8 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
     SearchState state;
     long pos;
     long keyOffset;
-    final MultiStoreBytes entry = new MultiStoreBytes();
+    final MultiStoreBytes entryCache = new MultiStoreBytes();
+    MultiStoreBytes entry;
 
     void initKeySearch() {
         if (state != null)
@@ -1141,6 +1143,8 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
 
     void closeKeySearch0() {
         state = null;
+        pos = -1;
+        entry = null;
         hashLookup.closeSearch0();
     }
 
@@ -1626,6 +1630,7 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
     }
 
     void writeNewValueAndSwitch() {
+        entry.position(valueOffset);
         metaValueInterop.write(valueInterop, entry, newValue);
         value = newValue;
         closeNewValue0();
@@ -1693,7 +1698,7 @@ class VanillaContext<K, KI, MKI extends MetaBytesInterop<K, ? super KI>,
     }
 
     final void reuse(long pos) {
-        reuse(entry, pos);
+        entry = reuse(entryCache, pos);
     }
 
     final long entrySize(long keySize, long valueSize) {
diff --git a/src/test/java/net/openhft/chronicle/map/ChronicleMapTest.java b/src/test/java/net/openhft/chronicle/map/ChronicleMapTest.java
index e4807ed1d..7d4701142 100755
--- a/src/test/java/net/openhft/chronicle/map/ChronicleMapTest.java
+++ b/src/test/java/net/openhft/chronicle/map/ChronicleMapTest.java
@@ -1652,8 +1652,7 @@ public class ChronicleMapTest {
         }
 
         // this will remove the entry
-        try (MapKeyContext<LongValue> context = map.acquireContext("one", value)) {
-            assert value == context.get();
+        try (MapKeyContext<LongValue> context = map.context("one")) {
             context.remove();
         }
 
@@ -1789,7 +1788,7 @@ public class ChronicleMapTest {
         }
 
         // this will remove the entry
-        try (MapKeyContext<LongValue> context = map.acquireContext("one", value)) {
+        try (MapKeyContext<LongValue> context = map.context("one")) {
             context.remove();
         }