[Named min timestamp leases] Extract exclusive lock collection to its…

… own class (#7330)
palantir · Oct 8, 2024 · df27e8d · df27e8d
1 parent 8505982
commit df27e8d
Show file tree

Hide file tree

Showing 5 changed files with 73 additions and 43 deletions.
diff --git a/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/AsyncLockService.java b/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/AsyncLockService.java
@@ -155,12 +155,12 @@ private AsyncResult<HeldLocks> acquireLocks(
             Set<LockDescriptor> lockDescriptors,
             TimeLimit timeout,
             Optional<LockRequestMetadata> metadata) {
-        OrderedLocks orderedLocks = locks.getAll(lockDescriptors);
+        OrderedLocks orderedLocks = locks.getAllExclusiveLocks(lockDescriptors);
         return lockAcquirer.acquireLocks(requestId, orderedLocks, timeout, metadata);
     }
 
     private AsyncResult<Void> awaitLocks(UUID requestId, Set<LockDescriptor> lockDescriptors, TimeLimit timeout) {
-        OrderedLocks orderedLocks = locks.getAll(lockDescriptors);
+        OrderedLocks orderedLocks = locks.getAllExclusiveLocks(lockDescriptors);
         return lockAcquirer.waitForLocks(requestId, orderedLocks, timeout);
     }
 

diff --git a/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/ExclusiveLockCollection.java b/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/ExclusiveLockCollection.java
@@ -0,0 +1,55 @@
+/*
+ * (c) Copyright 2024 Palantir Technologies Inc. All rights reserved.
+ *
+ * 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 com.palantir.atlasdb.timelock.lock;
+
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.github.benmanes.caffeine.cache.LoadingCache;
+import com.google.common.collect.Lists;
+import com.palantir.lock.LockDescriptor;
+import java.util.ArrayList;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Set;
+
+final class ExclusiveLockCollection {
+    private final LoadingCache<LockDescriptor, AsyncLock> locksById;
+
+    ExclusiveLockCollection() {
+        locksById = Caffeine.newBuilder().weakValues().build(ExclusiveLock::new);
+    }
+
+    OrderedLocks getAll(Set<LockDescriptor> descriptors) {
+        List<LockDescriptor> orderedDescriptors = sort(descriptors);
+
+        List<AsyncLock> locks = Lists.newArrayListWithExpectedSize(descriptors.size());
+        for (LockDescriptor descriptor : orderedDescriptors) {
+            locks.add(getLock(descriptor));
+        }
+
+        return OrderedLocks.fromOrderedList(locks);
+    }
+
+    private static List<LockDescriptor> sort(Set<LockDescriptor> descriptors) {
+        List<LockDescriptor> orderedDescriptors = new ArrayList<>(descriptors);
+        orderedDescriptors.sort(Comparator.naturalOrder());
+        return orderedDescriptors;
+    }
+
+    private AsyncLock getLock(LockDescriptor descriptor) {
+        return locksById.get(descriptor);
+    }
+}
diff --git a/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/LockCollection.java b/timelock-impl/src/main/java/com/palantir/atlasdb/timelock/lock/LockCollection.java
@@ -15,41 +15,13 @@
  */
 package com.palantir.atlasdb.timelock.lock;
 
-import com.github.benmanes.caffeine.cache.Caffeine;
-import com.github.benmanes.caffeine.cache.LoadingCache;
-import com.google.common.collect.Lists;
 import com.palantir.lock.LockDescriptor;
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.List;
 import java.util.Set;
 
-public class LockCollection {
+final class LockCollection {
+    private final ExclusiveLockCollection exclusiveLocks = new ExclusiveLockCollection();
 
-    private final LoadingCache<LockDescriptor, AsyncLock> locksById;
-
-    public LockCollection() {
-        locksById = Caffeine.newBuilder().weakValues().build(ExclusiveLock::new);
-    }
-
-    public OrderedLocks getAll(Set<LockDescriptor> descriptors) {
-        List<LockDescriptor> orderedDescriptors = sort(descriptors);
-
-        List<AsyncLock> locks = Lists.newArrayListWithExpectedSize(descriptors.size());
-        for (LockDescriptor descriptor : orderedDescriptors) {
-            locks.add(getLock(descriptor));
-        }
-
-        return OrderedLocks.fromOrderedList(locks);
-    }
-
-    private static List<LockDescriptor> sort(Set<LockDescriptor> descriptors) {
-        List<LockDescriptor> orderedDescriptors = new ArrayList<>(descriptors);
-        orderedDescriptors.sort(Comparator.naturalOrder());
-        return orderedDescriptors;
-    }
-
-    private AsyncLock getLock(LockDescriptor descriptor) {
-        return locksById.get(descriptor);
+    OrderedLocks getAllExclusiveLocks(Set<LockDescriptor> descriptors) {
+        return exclusiveLocks.getAll(descriptors);
     }
 }
diff --git a/timelock-server/src/test/java/com/palantir/atlasdb/timelock/lock/AsyncLockServiceTest.java b/timelock-server/src/test/java/com/palantir/atlasdb/timelock/lock/AsyncLockServiceTest.java
@@ -75,7 +75,7 @@ public void before() {
         when(acquirer.acquireLocks(any(), any(), any())).thenReturn(new AsyncResult<>());
         when(acquirer.acquireLocks(any(), any(), any(), any())).thenReturn(new AsyncResult<>());
         when(acquirer.waitForLocks(any(), any(), any())).thenReturn(new AsyncResult<>());
-        when(locks.getAll(any())).thenReturn(OrderedLocks.fromSingleLock(newLock()));
+        when(locks.getAllExclusiveLocks(any())).thenReturn(OrderedLocks.fromSingleLock(newLock()));
         when(immutableTimestampTracker.getImmutableTimestamp()).thenReturn(Optional.empty());
         when(immutableTimestampTracker.getLockFor(anyLong())).thenReturn(newLock());
     }
@@ -84,7 +84,7 @@ public void before() {
     public void passesOrderedLocksToAcquirer() {
         OrderedLocks expected = orderedLocks(newLock(), newLock());
         Set<LockDescriptor> descriptors = descriptors(LOCK_A, LOCK_B);
-        when(locks.getAll(descriptors)).thenReturn(expected);
+        when(locks.getAllExclusiveLocks(descriptors)).thenReturn(expected);
 
         lockService.lock(REQUEST_ID, descriptors, DEADLINE);
 
@@ -95,7 +95,7 @@ public void passesOrderedLocksToAcquirer() {
     public void passesOrderedLocksToAcquirerWhenWaitingForLocks() {
         OrderedLocks expected = orderedLocks(newLock(), newLock());
         Set<LockDescriptor> descriptors = descriptors(LOCK_A, LOCK_B);
-        when(locks.getAll(descriptors)).thenReturn(expected);
+        when(locks.getAllExclusiveLocks(descriptors)).thenReturn(expected);
 
         lockService.waitForLocks(REQUEST_ID, descriptors, DEADLINE);
 

diff --git a/timelock-server/src/test/java/com/palantir/atlasdb/timelock/lock/LockCollectionTest.java b/timelock-server/src/test/java/com/palantir/atlasdb/timelock/lock/LockCollectionTest.java
@@ -36,7 +36,7 @@ public class LockCollectionTest {
     public void createsLocksOnDemand() {
         Set<LockDescriptor> descriptors = descriptors("foo", "bar");
 
-        List<AsyncLock> locks = lockCollection.getAll(descriptors).get();
+        List<AsyncLock> locks = lockCollection.getAllExclusiveLocks(descriptors).get();
 
         assertThat(locks).hasSize(2);
         assertThat(ImmutableSet.copyOf(locks)).hasSize(2);
@@ -46,8 +46,10 @@ public void createsLocksOnDemand() {
     public void returnsSameLockForMultipleRequests() {
         Set<LockDescriptor> descriptors = descriptors("foo", "bar");
 
-        List<AsyncLock> locks1 = lockCollection.getAll(descriptors).get();
-        List<AsyncLock> locks2 = lockCollection.getAll(descriptors).get();
+        List<AsyncLock> locks1 =
+                lockCollection.getAllExclusiveLocks(descriptors).get();
+        List<AsyncLock> locks2 =
+                lockCollection.getAllExclusiveLocks(descriptors).get();
 
         assertThat(locks1).containsExactlyElementsOf(locks2);
     }
@@ -60,12 +62,13 @@ public void returnsLocksInOrder() {
                 .sorted()
                 .collect(Collectors.toList());
         List<AsyncLock> expectedOrder = orderedDescriptors.stream()
-                .map(descriptor -> lockCollection.getAll(ImmutableSet.of(descriptor)))
+                .map(descriptor -> lockCollection.getAllExclusiveLocks(ImmutableSet.of(descriptor)))
                 .map(orderedLocks -> orderedLocks.get().get(0))
                 .collect(Collectors.toList());
 
-        List<AsyncLock> actualOrder =
-                lockCollection.getAll(ImmutableSet.copyOf(orderedDescriptors)).get();
+        List<AsyncLock> actualOrder = lockCollection
+                .getAllExclusiveLocks(ImmutableSet.copyOf(orderedDescriptors))
+                .get();
 
         assertThat(actualOrder).containsExactlyElementsOf(expectedOrder);
     }