[Named min timestamp Leases] Modify client API (#7386)

atlasdb-impl-shared/build.gradle

@@ -23,6 +23,7 @@ dependencies {
   implementation project(":atlasdb-client")
   implementation project(":atlasdb-commons")
   implementation project(":atlasdb-coordination-impl")
+  implementation project(':timelock-api')
   implementation project(":timestamp-api")
 
   implementation 'com.github.ben-manes.caffeine:caffeine'

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/transaction/impl/InstrumentedTimelockService.java

@@ -17,21 +17,23 @@
 
 import com.codahale.metrics.Meter;
 import com.palantir.atlasdb.AtlasDbMetricNames;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.atlasdb.util.MetricsManager;
 import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
-import com.palantir.lock.v2.AcquireNamedMinTimestampLeaseResult;
 import com.palantir.lock.v2.ClientLockingOptions;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
 import com.palantir.lock.v2.LockRequest;
 import com.palantir.lock.v2.LockResponse;
 import com.palantir.lock.v2.LockToken;
 import com.palantir.lock.v2.StartIdentifiedAtlasDbTransactionResponse;
 import com.palantir.lock.v2.TimelockService;
+import com.palantir.lock.v2.TimestampLeaseResults;
 import com.palantir.lock.v2.WaitForLocksRequest;
 import com.palantir.lock.v2.WaitForLocksResponse;
 import com.palantir.timestamp.TimestampRange;
 import com.palantir.tritium.metrics.registry.MetricName;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.function.Supplier;
 
@@ -135,16 +137,14 @@ public long currentTimeMillis() {
 
     @ReviewedRestrictedApiUsage
     @Override
-    public AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(
-            String timestampName, int numFreshTimestamps) {
-        return executeWithRecord(
-                () -> timelockService.acquireNamedMinTimestampLease(timestampName, numFreshTimestamps));
+    public TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests) {
+        return executeWithRecord(() -> timelockService.acquireTimestampLeases(requests));
     }
 
     @ReviewedRestrictedApiUsage
     @Override
-    public long getMinLeasedTimestampForName(String timestampName) {
-        return executeWithRecord(() -> timelockService.getMinLeasedTimestampForName(timestampName));
+    public Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames) {
+        return executeWithRecord(() -> timelockService.getMinLeasedTimestamps(timestampNames));
     }
 
     private <T> T executeWithRecord(Supplier<T> method) {

lock-api-objects/build.gradle

@@ -4,6 +4,7 @@ apply from: "../gradle/shared.gradle"
 libsDirName = file('build/artifacts')
 
 dependencies {
+    api project(":timelock-api")
     api project(":timestamp-api")
     api project(":timestamp-client")
     implementation project(":atlasdb-commons")

lock-api-objects/src/main/java/com/palantir/lock/v2/TimelockService.java

@@ -16,13 +16,15 @@
 package com.palantir.lock.v2;
 
 import com.google.errorprone.annotations.RestrictedApi;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
 import com.palantir.logsafe.Safe;
 import com.palantir.processors.AutoDelegate;
 import com.palantir.processors.DoNotDelegate;
 import com.palantir.timestamp.TimestampRange;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -114,19 +116,21 @@ default List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTr
     long currentTimeMillis();
 
     /**
-     * Acquires a lease on a named timestamp. The lease is taken out with a new fresh timestamp.
-     * The returned timestamps are fresh timestamps obtained strictly after the lease is taken out.
+     * Acquires a lease on named timestamps. The lease is taken out with a new fresh timestamp.
+     * The timestamps supplied are fresh timestamps obtained strictly after the lease is taken out.
+     * The supplier returns exactly the number of timestamps requested and throws on any additional
+     * interactions.
      */
     @RestrictedApi(
             explanation =
                     "This method is for internal Atlas and internal library use only. Clients MUST NOT use it unless"
                         + " given explicit approval. Mis-use can result in SEVERE DATA CORRUPTION and the API contract"
                         + " is subject to change at any time.",
             allowlistAnnotations = ReviewedRestrictedApiUsage.class)
-    AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(String timestampName, int numFreshTimestamps);
+    TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests);
 
     /**
-     * Returns the smallest leased timestamp in the associated named collection at the time of the call.
+     * Returns the smallest leased timestamp for each of the associated named collections at the time of the call.
      * If there are no active leases, a fresh timestamp is obtained and returned.
      */
     @RestrictedApi(
@@ -135,5 +139,5 @@ default List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTr
                         + " given explicit approval. Mis-use can result in SEVERE DATA CORRUPTION and the API contract"
                         + " is subject to change at any time.",
             allowlistAnnotations = ReviewedRestrictedApiUsage.class)
-    long getMinLeasedTimestampForName(String timestampName);
+    Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames);
 }

lock-api-objects/src/main/java/com/palantir/lock/v2/AcquireNamedMinTimestampLeaseResult.java → lock-api-objects/src/main/java/com/palantir/lock/v2/TimestampLeaseResult.java

@@ -16,17 +16,14 @@
 
 package com.palantir.lock.v2;
 
-import java.util.List;
+import java.util.function.LongSupplier;
 import org.immutables.value.Value;
 
 @Value.Immutable
-public interface AcquireNamedMinTimestampLeaseResult {
+public interface TimestampLeaseResult {
     @Value.Parameter
     long minLeasedTimestamp();
 
     @Value.Parameter
-    LockToken lock();
-
-    @Value.Parameter
-    List<Long> freshTimestamps();
+    LongSupplier freshTimestampsSupplier();
 }

lock-api-objects/src/main/java/com/palantir/lock/v2/TimestampLeaseResults.java

@@ -0,0 +1,30 @@
+/*
+ * (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.lock.v2;
+
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
+import java.util.Map;
+import org.immutables.value.Value;
+
+@Value.Immutable
+public interface TimestampLeaseResults {
+    @Value.Parameter
+    LockToken lock();
+
+    @Value.Parameter
+    Map<TimestampLeaseName, TimestampLeaseResult> results();
+}

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

@@ -19,16 +19,17 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Stopwatch;
 import com.google.common.util.concurrent.RateLimiter;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.common.base.Throwables;
 import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
-import com.palantir.lock.v2.AcquireNamedMinTimestampLeaseResult;
 import com.palantir.lock.v2.ClientLockingOptions;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
 import com.palantir.lock.v2.LockRequest;
 import com.palantir.lock.v2.LockResponse;
 import com.palantir.lock.v2.LockToken;
 import com.palantir.lock.v2.StartIdentifiedAtlasDbTransactionResponse;
 import com.palantir.lock.v2.TimelockService;
+import com.palantir.lock.v2.TimestampLeaseResults;
 import com.palantir.lock.v2.WaitForLocksRequest;
 import com.palantir.lock.v2.WaitForLocksResponse;
 import com.palantir.logsafe.SafeArg;
@@ -39,6 +40,7 @@
 import com.palantir.timestamp.TimestampRange;
 import java.time.Duration;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
@@ -173,17 +175,14 @@ public long currentTimeMillis() {
 
     @ReviewedRestrictedApiUsage
     @Override
-    public AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(
-            String timestampName, int numFreshTimestamps) {
-        return runTaskTimed(
-                "acquireNamedMinTimestampLease",
-                () -> delegate.acquireNamedMinTimestampLease(timestampName, numFreshTimestamps));
+    public TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests) {
+        return runTaskTimed("acquireTimestampLeases", () -> delegate.acquireTimestampLeases(requests));
     }
 
     @ReviewedRestrictedApiUsage
     @Override
-    public long getMinLeasedTimestampForName(String timestampName) {
-        return runTaskTimed("getMinLeasedNamedTimestamp", () -> delegate.getMinLeasedTimestampForName(timestampName));
+    public Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames) {
+        return runTaskTimed("getMinLeasedTimestamps", () -> delegate.getMinLeasedTimestamps(timestampNames));
     }
 
     private <T> T runTaskTimed(String actionName, Supplier<T> action) {

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

@@ -20,8 +20,7 @@
 import com.palantir.atlasdb.timelock.api.ConjureGetFreshTimestampsResponseV2;
 import com.palantir.atlasdb.timelock.api.ConjureTimestampRange;
 import com.palantir.atlasdb.timelock.api.Namespace;
-import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
-import com.palantir.lock.v2.AcquireNamedMinTimestampLeaseResult;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.lock.v2.ClientLockingOptions;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
 import com.palantir.lock.v2.LockRequest;
@@ -30,6 +29,7 @@
 import com.palantir.lock.v2.NamespacedTimelockRpcClient;
 import com.palantir.lock.v2.StartIdentifiedAtlasDbTransactionResponse;
 import com.palantir.lock.v2.TimelockService;
+import com.palantir.lock.v2.TimestampLeaseResults;
 import com.palantir.lock.v2.WaitForLocksRequest;
 import com.palantir.lock.v2.WaitForLocksResponse;
 import com.palantir.lock.watch.LockWatchCache;
@@ -38,6 +38,7 @@
 import com.palantir.logsafe.logger.SafeLoggerFactory;
 import com.palantir.timestamp.TimestampRange;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -170,17 +171,14 @@ public long currentTimeMillis() {
         return rpcClient.currentTimeMillis();
     }
 
-    @ReviewedRestrictedApiUsage
     @Override
-    public AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(
-            String timestampName, int numFreshTimestamps) {
+    public TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests) {
         // TODO(aalouane): implement!
         throw new UnsupportedOperationException("Not implemented yet!");
     }
 
-    @ReviewedRestrictedApiUsage
     @Override
-    public long getMinLeasedTimestampForName(String timestampName) {
+    public Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames) {
         // TODO(aalouane): implement!
         throw new UnsupportedOperationException("Not implemented yet!");
     }

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

@@ -17,12 +17,12 @@
 
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.common.base.Throwables;
 import com.palantir.common.concurrent.NamedThreadFactory;
 import com.palantir.common.concurrent.PTExecutors;
 import com.palantir.leader.NotCurrentLeaderException;
 import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
-import com.palantir.lock.v2.AcquireNamedMinTimestampLeaseResult;
 import com.palantir.lock.v2.ClientLockingOptions;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
 import com.palantir.lock.v2.LockLeaseRefresher;
@@ -31,6 +31,7 @@
 import com.palantir.lock.v2.LockToken;
 import com.palantir.lock.v2.StartIdentifiedAtlasDbTransactionResponse;
 import com.palantir.lock.v2.TimelockService;
+import com.palantir.lock.v2.TimestampLeaseResults;
 import com.palantir.lock.v2.WaitForLocksRequest;
 import com.palantir.lock.v2.WaitForLocksResponse;
 import com.palantir.timestamp.CloseableTimestampService;
@@ -39,6 +40,7 @@
 import java.net.ConnectException;
 import java.net.UnknownHostException;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ScheduledExecutorService;
@@ -184,15 +186,14 @@ public long currentTimeMillis() {
 
     @ReviewedRestrictedApiUsage
     @Override
-    public AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(
-            String timestampName, int numFreshTimestamps) {
-        return delegate.acquireNamedMinTimestampLease(timestampName, numFreshTimestamps);
+    public TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests) {
+        return executeOnTimeLock(() -> delegate.acquireTimestampLeases(requests));
     }
 
     @ReviewedRestrictedApiUsage
     @Override
-    public long getMinLeasedTimestampForName(String timestampName) {
-        return delegate.getMinLeasedTimestampForName(timestampName);
+    public Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames) {
+        return executeOnTimeLock(() -> delegate.getMinLeasedTimestamps(timestampNames));
     }
 
     private static <T> T executeOnTimeLock(Callable<T> callable) {

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

@@ -19,22 +19,24 @@
 import com.google.common.collect.Sets;
 import com.palantir.atlasdb.buggify.api.BuggifyFactory;
 import com.palantir.atlasdb.buggify.impl.DefaultBuggifyFactory;
+import com.palantir.atlasdb.timelock.api.TimestampLeaseName;
 import com.palantir.lock.annotations.ReviewedRestrictedApiUsage;
-import com.palantir.lock.v2.AcquireNamedMinTimestampLeaseResult;
 import com.palantir.lock.v2.ClientLockingOptions;
 import com.palantir.lock.v2.LockImmutableTimestampResponse;
 import com.palantir.lock.v2.LockRequest;
 import com.palantir.lock.v2.LockResponse;
 import com.palantir.lock.v2.LockToken;
 import com.palantir.lock.v2.StartIdentifiedAtlasDbTransactionResponse;
 import com.palantir.lock.v2.TimelockService;
+import com.palantir.lock.v2.TimestampLeaseResults;
 import com.palantir.lock.v2.WaitForLocksRequest;
 import com.palantir.lock.v2.WaitForLocksResponse;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.logger.SafeLogger;
 import com.palantir.logsafe.logger.SafeLoggerFactory;
 import com.palantir.timestamp.TimestampRange;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -159,15 +161,14 @@ public long currentTimeMillis() {
 
     @ReviewedRestrictedApiUsage
     @Override
-    public AcquireNamedMinTimestampLeaseResult acquireNamedMinTimestampLease(
-            String timestampName, int numFreshTimestamps) {
-        return delegate.acquireNamedMinTimestampLease(timestampName, numFreshTimestamps);
+    public TimestampLeaseResults acquireTimestampLeases(Map<TimestampLeaseName, Integer> requests) {
+        return delegate.acquireTimestampLeases(requests);
     }
 
     @ReviewedRestrictedApiUsage
     @Override
-    public long getMinLeasedTimestampForName(String timestampName) {
-        return delegate.getMinLeasedTimestampForName(timestampName);
+    public Map<TimestampLeaseName, Long> getMinLeasedTimestamps(Set<TimestampLeaseName> timestampNames) {
+        return delegate.getMinLeasedTimestamps(timestampNames);
     }
 
     private void maybeRandomlyIncreaseTimestamp() {

lock-impl/build.gradle

@@ -19,6 +19,7 @@ dependencies {
   implementation project(":atlasdb-commons")
   implementation project(':commons-executors')
   implementation project(':lock-api-objects')
+  implementation project(':timelock-api')
   implementation project(':timestamp-api')
 
   testImplementation 'com.fasterxml.jackson.core:jackson-databind'