[Cross Client Batching LeaderTimes - 2a] | Refactor coalescing leader…

…TimeGetter (#5120)

* Fix tests + remove redundant tests

* Refactor LeaderTimeGetter

* Configurable bufferSize

* Implement LeaderTimeCoalescingBatcher

* Wire LeaderTimeGetter

* Fix build

* Address comments

* Remove redundant line

* Fix tests

* Add generated changelog entries

* Fix nits

atlasdb-autobatch/src/main/java/com/palantir/atlasdb/autobatch/Autobatchers.java

@@ -23,6 +23,7 @@
 import com.palantir.tracing.Observability;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalInt;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -88,6 +89,7 @@ public static final class AutobatcherBuilder<I, O> {
         private final ImmutableMap.Builder<String, String> safeTags = ImmutableMap.builder();
 
         private Observability observability = Observability.UNDECIDED;
+        private OptionalInt bufferSize = OptionalInt.empty();
 
         @Nullable
         private String purpose;
@@ -111,16 +113,24 @@ public AutobatcherBuilder<I, O> observability(Observability observabilityParam)
             return this;
         }
 
+        public AutobatcherBuilder<I, O> bufferSize(OptionalInt bufferSizeParam) {
+            this.bufferSize = bufferSizeParam;
+            return this;
+        }
+
         public DisruptorAutobatcher<I, O> build() {
             Preconditions.checkArgument(purpose != null, "purpose must be provided");
-            EventHandler<BatchElement<I, O>> handler = this.handlerFactory.apply(DEFAULT_BUFFER_SIZE);
 
-            EventHandler<BatchElement<I, O>> tracingHandler = new TracingEventHandler<>(handler, DEFAULT_BUFFER_SIZE);
+            int bufferSizeValue = bufferSize.orElse(DEFAULT_BUFFER_SIZE);
+
+            EventHandler<BatchElement<I, O>> handler = this.handlerFactory.apply(bufferSizeValue);
+
+            EventHandler<BatchElement<I, O>> tracingHandler = new TracingEventHandler<>(handler, bufferSizeValue);
 
             EventHandler<BatchElement<I, O>> profiledHandler =
                     new ProfilingEventHandler<>(tracingHandler, purpose, safeTags.build());
 
-            return DisruptorAutobatcher.create(profiledHandler, DEFAULT_BUFFER_SIZE, purpose);
+            return DisruptorAutobatcher.create(profiledHandler, bufferSizeValue, purpose);
         }
     }
 }

changelog/@unreleased/pr-5120.v2.yml

@@ -0,0 +1,8 @@
+type: improvement
+improvement:
+  description: '```LeaderTimeGetter``` is an interface responsible for making ```leaderTime```
+    requests to TimeLock with the intend of adding a client side request batching
+    layer. TimeLock adaptor on AtlasDB client now accepts custom implementation of
+    the same.'
+  links:
+  - https://github.com/palantir/atlasdb/pull/5120

lock-api/src/main/java/com/palantir/lock/client/LeaderTimeCoalescingBatcher.java

@@ -0,0 +1,63 @@
+/*
+ * (c) Copyright 2020 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.lock.client;
+
+import com.palantir.atlasdb.autobatch.Autobatchers;
+import com.palantir.atlasdb.autobatch.CoalescingRequestFunction;
+import com.palantir.atlasdb.autobatch.DisruptorAutobatcher;
+import com.palantir.atlasdb.futures.AtlasFutures;
+import com.palantir.atlasdb.timelock.api.MultiClientConjureTimelockService;
+import com.palantir.atlasdb.timelock.api.Namespace;
+import com.palantir.lock.v2.LeaderTime;
+import com.palantir.tokens.auth.AuthHeader;
+import java.util.Map;
+import java.util.OptionalInt;
+import java.util.Set;
+
+public class LeaderTimeCoalescingBatcher implements AutoCloseable {
+    private static final AuthHeader AUTH_HEADER = AuthHeader.valueOf("Bearer omitted");
+    private final DisruptorAutobatcher<Namespace, LeaderTime> batcher;
+
+    public LeaderTimeCoalescingBatcher(MultiClientConjureTimelockService delegate, OptionalInt bufferSize) {
+        this.batcher = Autobatchers.coalescing(new LeaderTimeCoalescingConsumer(delegate))
+                .bufferSize(bufferSize)
+                .safeLoggablePurpose("get-leader-times")
+                .build();
+    }
+
+    public LeaderTime apply(Namespace namespace) {
+        return AtlasFutures.getUnchecked(batcher.apply(namespace));
+    }
+
+    @Override
+    public void close() {
+        batcher.close();
+    }
+
+    static class LeaderTimeCoalescingConsumer implements CoalescingRequestFunction<Namespace, LeaderTime> {
+        private final MultiClientConjureTimelockService delegate;
+
+        public LeaderTimeCoalescingConsumer(MultiClientConjureTimelockService delegate) {
+            this.delegate = delegate;
+        }
+
+        @Override
+        public Map<Namespace, LeaderTime> apply(Set<Namespace> namespaces) {
+            return delegate.leaderTimes(AUTH_HEADER, namespaces).getLeaderTimes();
+        }
+    }
+}

lock-api/src/main/java/com/palantir/lock/client/LeaderTimeGetter.java

@@ -0,0 +1,26 @@
+/*
+ * (c) Copyright 2020 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.lock.client;
+
+import com.palantir.lock.v2.LeaderTime;
+
+public interface LeaderTimeGetter extends AutoCloseable {
+    LeaderTime leaderTime();
+
+    @Override
+    void close();
+}

lock-api/src/main/java/com/palantir/lock/client/LegacyLeaderTimeGetter.java

@@ -0,0 +1,36 @@
+/*
+ * (c) Copyright 2020 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.lock.client;
+
+import com.palantir.common.concurrent.CoalescingSupplier;
+import com.palantir.lock.v2.LeaderTime;
+
+public class LegacyLeaderTimeGetter implements LeaderTimeGetter {
+    private final CoalescingSupplier<LeaderTime> time;
+
+    public LegacyLeaderTimeGetter(NamespacedConjureTimelockService delegate) {
+        this.time = new CoalescingSupplier<>(delegate::leaderTime);
+    }
+
+    @Override
+    public LeaderTime leaderTime() {
+        return time.get();
+    }
+
+    @Override
+    public void close() {}
+}

lock-api/src/main/java/com/palantir/lock/client/LockLeaseService.java

@@ -27,7 +27,6 @@
 import com.palantir.atlasdb.timelock.api.ConjureUnlockRequest;
 import com.palantir.atlasdb.timelock.api.GetCommitTimestampsRequest;
 import com.palantir.atlasdb.timelock.api.GetCommitTimestampsResponse;
-import com.palantir.common.concurrent.CoalescingSupplier;
 import com.palantir.lock.v2.LeaderTime;
 import com.palantir.lock.v2.Lease;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
@@ -47,19 +46,20 @@
 class LockLeaseService {
     private final NamespacedConjureTimelockService delegate;
     private final UUID clientId;
-    private final CoalescingSupplier<LeaderTime> time;
+    private final LeaderTimeGetter leaderTimeGetter;
     private final BlockEnforcingLockService lockService;
 
     @VisibleForTesting
-    LockLeaseService(NamespacedConjureTimelockService delegate, UUID clientId) {
+    LockLeaseService(NamespacedConjureTimelockService delegate, UUID clientId, LeaderTimeGetter leaderTimeGetter) {
         this.delegate = delegate;
         this.clientId = clientId;
-        this.time = new CoalescingSupplier<>(delegate::leaderTime);
+        this.leaderTimeGetter = leaderTimeGetter;
         this.lockService = BlockEnforcingLockService.create(delegate);
     }
 
-    static LockLeaseService create(NamespacedConjureTimelockService conjureTimelock) {
-        return new LockLeaseService(conjureTimelock, UUID.randomUUID());
+    static LockLeaseService create(
+            NamespacedConjureTimelockService conjureTimelock, LeaderTimeGetter leaderTimeGetter) {
+        return new LockLeaseService(conjureTimelock, UUID.randomUUID(), leaderTimeGetter);
     }
 
     LockImmutableTimestampResponse lockImmutableTimestamp() {
@@ -128,7 +128,7 @@ Set<LockToken> refreshLockLeases(Set<LockToken> uncastedTokens) {
             return uncastedTokens;
         }
 
-        LeaderTime leaderTime = time.get();
+        LeaderTime leaderTime = leaderTimeGetter.leaderTime();
         Set<LeasedLockToken> allTokens = leasedTokens(uncastedTokens);
 
         Set<LeasedLockToken> validByLease =

lock-api/src/main/java/com/palantir/lock/client/NamespacedCoalescingLeaderTimeGetter.java

@@ -0,0 +1,47 @@
+/*
+ * (c) Copyright 2020 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.lock.client;
+
+import com.palantir.atlasdb.timelock.api.Namespace;
+import com.palantir.lock.v2.LeaderTime;
+
+/**
+ * This class maintains the context of namespace for a client and directs leaderTime requests to
+ * {@link LeaderTimeCoalescingBatcher} that coalesces and batches requests across clients.
+ *
+ * This is different from {@link LegacyLeaderTimeGetter} in that LegacyLeaderTimeGetter coalesces requests only for a
+ * single namespace.
+ */
+public class NamespacedCoalescingLeaderTimeGetter implements LeaderTimeGetter {
+    private final LeaderTimeCoalescingBatcher batcher;
+    private final Namespace namespace;
+
+    public NamespacedCoalescingLeaderTimeGetter(String namespace, LeaderTimeCoalescingBatcher batcher) {
+        this.namespace = Namespace.of(namespace);
+        this.batcher = batcher;
+    }
+
+    @Override
+    public LeaderTime leaderTime() {
+        return batcher.apply(namespace);
+    }
+
+    @Override
+    public void close() {
+        batcher.close();
+    }
+}

lock-api/src/main/java/com/palantir/lock/client/RemoteTimelockServiceAdapter.java

@@ -42,9 +42,10 @@ public final class RemoteTimelockServiceAdapter implements TimelockService, Auto
     private RemoteTimelockServiceAdapter(
             NamespacedTimelockRpcClient rpcClient,
             NamespacedConjureTimelockService conjureTimelockService,
-            LockWatchEventCache lockWatchEventCache) {
+            LockWatchEventCache lockWatchEventCache,
+            LeaderTimeGetter leaderTimeGetter) {
         this.rpcClient = rpcClient;
-        this.lockLeaseService = LockLeaseService.create(conjureTimelockService);
+        this.lockLeaseService = LockLeaseService.create(conjureTimelockService, leaderTimeGetter);
         this.transactionStarter = TransactionStarter.create(lockLeaseService, lockWatchEventCache);
         this.commitTimestampGetter = CommitTimestampGetter.create(lockLeaseService, lockWatchEventCache);
         this.conjureTimelockService = conjureTimelockService;
@@ -54,7 +55,15 @@ public static RemoteTimelockServiceAdapter create(
             NamespacedTimelockRpcClient rpcClient,
             NamespacedConjureTimelockService conjureClient,
             LockWatchEventCache lockWatchEventCache) {
-        return new RemoteTimelockServiceAdapter(rpcClient, conjureClient, lockWatchEventCache);
+        return create(rpcClient, conjureClient, lockWatchEventCache, new LegacyLeaderTimeGetter(conjureClient));
+    }
+
+    public static RemoteTimelockServiceAdapter create(
+            NamespacedTimelockRpcClient rpcClient,
+            NamespacedConjureTimelockService conjureClient,
+            LockWatchEventCache lockWatchEventCache,
+            LeaderTimeGetter leaderTimeGetter) {
+        return new RemoteTimelockServiceAdapter(rpcClient, conjureClient, lockWatchEventCache, leaderTimeGetter);
     }
 
     @Override

lock-api/src/test/java/com/palantir/lock/client/LockLeaseServiceTest.java

@@ -92,7 +92,7 @@ public void before() {
             ConjureUnlockRequest request = inv.getArgument(0);
             return ConjureUnlockResponse.of(request.getTokens());
         });
-        lockLeaseService = new LockLeaseService(timelock, SERVICE_ID);
+        lockLeaseService = new LockLeaseService(timelock, SERVICE_ID, new LegacyLeaderTimeGetter(timelock));
     }
 
     @Test

timelock-impl/src/test/java/com/palantir/atlasdb/timelock/batch/MultiClientConjureTimelockResourceTest.java

@@ -53,7 +53,6 @@ public class MultiClientConjureTimelockResourceTest {
 
     private Map<String, AsyncTimelockService> namespaces = new HashMap();
     private Map<String, LeadershipId> namespaceToLeaderMap = new HashMap();
-
     private MultiClientConjureTimelockResource resource;
 
     @Before