Merge pull request #99 from moshegood/moshe/ShouldSkipBlockingWait/eventually.should.work

Fix: ShouldSkipBlockingWait should still acquire a dead lock if tried for longer than TTL

Author: shetsa-amzn

src/main/java/com/amazonaws/services/dynamodbv2/AmazonDynamoDBLockClient.java

@@ -233,6 +233,7 @@ public class AmazonDynamoDBLockClient implements Runnable, Closeable {
     private final boolean holdLockOnServiceUnavailable;
     private final String ownerName;
     private final ConcurrentHashMap<String, LockItem> locks;
+    private final ConcurrentHashMap<String, LockItem> notMyLocks = new ConcurrentHashMap<>();
     private final ConcurrentHashMap<String, Thread> sessionMonitors;
     private final Optional<Thread> backgroundThread;
     private final Function<String, ThreadFactory> namedThreadCreator;
@@ -470,11 +471,30 @@ public LockItem acquireLock(final AcquireLockOptions options) throws LockNotGran
                     }
 
                     if (options.shouldSkipBlockingWait() && existingLock.isPresent() && !existingLock.get().isExpired()) {
+                        String id = existingLock.get().getUniqueIdentifier();
+                        // Let's check to see if this existingLock expired based on old data we cached.
+                        // Or cache it if we haven't seen this recordVersion before.
+                        boolean isReallyExpired = false;
+                        if (notMyLocks.containsKey(id) &&
+                              notMyLocks.get(id).getRecordVersionNumber()
+                              .equals(existingLock.get().getRecordVersionNumber())) {
+
+                          isReallyExpired = notMyLocks.get(id).isExpired();
+                          if (isReallyExpired) {
+                              // short circuit the waiting that we normally do.
+                              lockTryingToBeAcquired = notMyLocks.get(id);
+                          }
+                        } else {
+                            notMyLocks.put(id, existingLock.get());
+                        }
+
                         /*
                          * The lock is being held by some one and is still not expired. And the caller explicitly said not to perform a blocking wait;
                          * We will throw back a lock not grant exception, so that the caller can retry if needed.
                          */
-                        throw new LockCurrentlyUnavailableException("The lock being requested is being held by another client.");
+                        if (!isReallyExpired) {
+                            throw new LockCurrentlyUnavailableException("The lock being requested is being held by another client.");
+                        }
                     }
 
                     Optional<ByteBuffer> newLockData = Optional.empty();

src/test/java/com/amazonaws/services/dynamodbv2/AmazonDynamoDBLockClientTest.java

@@ -443,10 +443,45 @@ public void acquireLock_whenLockAlreadyExistsAndIsNotReleased_andSkipBlockingWai
             .thenReturn(GetItemResponse.builder().item(item).build())
             .thenReturn(GetItemResponse.builder().build());
         AcquireLockOptions acquireLockOptions = AcquireLockOptions.builder("customer1")
-                .withShouldSkipBlockingWait(true)
-                .withDeleteLockOnRelease(false).build();
+            .withShouldSkipBlockingWait(true)
+            .withDeleteLockOnRelease(false).build();
         client.acquireLock(acquireLockOptions);
     }
+    /*
+     * Test case for the scenario, where the lock is being held by the first owner and the lock duration has not past
+     * the lease duration. In this case, We should expect a LockAlreadyOwnedException when shouldSkipBlockingWait is set.
+     * But if we try again later, we should get the lock.
+     */
+    @Test
+    public void acquireLock_whenLockAlreadyExistsAndIsNotReleased_andSkipBlockingWait_eventuallyGetsTheLock()
+        throws InterruptedException {
+        UUID uuid = setOwnerNameToUuid();
+        AmazonDynamoDBLockClient client = getLockClient();
+        Map<String, AttributeValue> item = new HashMap<>(5);
+        item.put("customer", AttributeValue.builder().s("customer1").build());
+        item.put("ownerName", AttributeValue.builder().s("foobar").build());
+        item.put("recordVersionNumber", AttributeValue.builder().s(uuid.toString()).build());
+        item.put("leaseDuration", AttributeValue.builder().s("100").build());
+        when(dynamodb.getItem(Mockito.<GetItemRequest>any()))
+            .thenReturn(GetItemResponse.builder().item(item).build())
+            .thenReturn(GetItemResponse.builder().build());
+        AcquireLockOptions acquireLockOptions = AcquireLockOptions.builder("customer1")
+            .withShouldSkipBlockingWait(true)
+            .withDeleteLockOnRelease(false).build();
+
+        try {
+            client.acquireLock(acquireLockOptions);
+        } catch (LockCurrentlyUnavailableException e) {
+            // This is expected
+        } catch (RuntimeException e) {
+            Assert.fail("Expected LockCurrentlyUnavailableException, but got " + e.getClass().getName());
+        }
+
+        // Now wait for the TTL to expire and try to acquire the lock again
+        Thread.sleep(101);
+        LockItem lockItem = client.acquireLock(acquireLockOptions);
+        Assert.assertNotNull("Failed to get lock item, when the lock is not present in the db", lockItem);
+    }
 
     @Test(expected = IllegalArgumentException.class)
     public void sendHeartbeat_whenDeleteDataTrueAndDataNotNull_throwsIllegalArgumentException() {