ISPN-3803 Stale locks during state transfer in non-tx caches

Clean up SingleKeyNonTxInvocationContext and allow the key to be
locked/unlocked independently of whether the entry was looked-up or not.

core/src/main/java/org/infinispan/context/SingleKeyNonTxInvocationContext.java

@@ -57,21 +57,26 @@ public Object getLockOwner() {
 
    @Override
    public Set<Object> getLockedKeys() {
-      return isLocked && key != null ?
-            Collections.singleton(key) : InfinispanCollections.emptySet();
+      return isLocked ? Collections.singleton(key) : InfinispanCollections.emptySet();
    }
 
    @Override
    public void clearLockedKeys() {
-      key = null;
-      cacheEntry = null;
       isLocked = false;
+      // TODO Dan: this shouldn't be necessary, but we don't clear the looked up keys
+      // when retrying a non-tx command because of a topology change
+      cacheEntry = null;
    }
 
    @Override
    public void addLockedKey(final Object key) {
-      if (cacheEntry != null && !keyEquivalence.equals(key, this.key))
+      if (this.key == null) {
+         // Set the key here
+         this.key = key;
+      } else if (!keyEquivalence.equals(key, this.key)) {
          throw illegalStateException();
+      }
+
       isLocked = true;
    }
 
@@ -81,27 +86,34 @@ private IllegalStateException illegalStateException() {
 
    @Override
    public CacheEntry lookupEntry(final Object key) {
-      if (key != null && this.key !=null && keyEquivalence.equals(key, this.key))
+      if (key != null && this.key != null && keyEquivalence.equals(key, this.key))
          return cacheEntry;
 
       return null;
    }
 
    @Override
    public Map<Object, CacheEntry> getLookedUpEntries() {
-      return key == null ? InfinispanCollections.<Object, CacheEntry>emptyMap() : Collections.singletonMap(key, cacheEntry);
+      return cacheEntry == null ? InfinispanCollections.<Object, CacheEntry>emptyMap() : Collections.singletonMap(key, cacheEntry);
    }
 
    @Override
    public void putLookedUpEntry(final Object key, final CacheEntry e) {
-      this.key = key;
+      if (this.key == null) {
+         // Set the key here
+         this.key = key;
+      } else if (!keyEquivalence.equals(key, this.key)) {
+         throw illegalStateException();
+      }
+
       this.cacheEntry = e;
    }
 
    @Override
    public void removeLookedUpEntry(final Object key) {
-      if (keyEquivalence.equals(key, this.key))
-         clearLockedKeys();
+      if (keyEquivalence.equals(key, this.key)) {
+         this.cacheEntry = null;
+      }
    }
 
    public Object getKey() {

core/src/test/java/org/infinispan/distribution/rehash/NonTxBackupOwnerBecomingPrimaryOwnerTest.java

@@ -37,10 +37,7 @@
 import static org.mockito.Matchers.anyInt;
 import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.doAnswer;
-import static org.testng.AssertJUnit.assertEquals;
-import static org.testng.AssertJUnit.assertNotNull;
-import static org.testng.AssertJUnit.assertNull;
-import static org.testng.AssertJUnit.assertTrue;
+import static org.testng.AssertJUnit.*;
 
 /**
  * Tests data loss during state transfer a backup owner of a key becomes the primary owner
@@ -185,6 +182,11 @@ public Object call() throws Exception {
       assertEquals("v", cache0.get(key));
       assertEquals("v", cache1.get(key));
       assertEquals("v", cache2.get(key));
+
+      // Check that there are no leaked locks
+      assertFalse(cache0.getAdvancedCache().getLockManager().isLocked(key));
+      assertFalse(cache1.getAdvancedCache().getLockManager().isLocked(key));
+      assertFalse(cache2.getAdvancedCache().getLockManager().isLocked(key));
    }
 
    private static class CustomConsistentHashFactory extends SingleSegmentConsistentHashFactory {

core/src/test/java/org/infinispan/distribution/rehash/NonTxPutIfAbsentDuringJoinStressTest.java

@@ -7,6 +7,7 @@
 import org.infinispan.test.MultipleCacheManagersTest;
 import org.infinispan.test.fwk.CleanupAfterMethod;
 import org.infinispan.transaction.TransactionMode;
+import org.infinispan.util.concurrent.locks.LockManager;
 import org.testng.annotations.Test;
 
 import java.util.concurrent.Callable;
@@ -15,7 +16,7 @@
 import java.util.concurrent.TimeUnit;
 
 import static org.testng.AssertJUnit.assertEquals;
-import static org.testng.AssertJUnit.assertNull;
+import static org.testng.AssertJUnit.assertFalse;
 import static org.testng.AssertJUnit.assertTrue;
 
 /**
@@ -98,5 +99,14 @@ public Object call() throws Exception {
             }
          }
       }
+
+      for (int i = 0; i < caches().size(); i++) {
+         LockManager lockManager = advancedCache(i).getLockManager();
+         assertEquals(0, lockManager.getNumberOfLocksHeld());
+         for (int j = 0; j < NUM_KEYS; j++) {
+            String key = "key_" + j;
+            assertFalse(lockManager.isLocked(key));
+         }
+      }
    }
 }