feature: Expose API to let client define number of expected values in a Get.

.palantir/revapi.yml

@@ -129,3 +129,28 @@ acceptedBreaks:
       new: "parameter <T> long com.palantir.atlasdb.keyvalue.cassandra.thrift.ThriftObjectSizeUtils::getCollectionSize(java.util.Collection<T>,\
         \ ===java.util.function.ToLongFunction<T>===)"
       justification: "Avoid autoboxing long to Long"
+  "0.884.0":
+    com.palantir.atlasdb:atlasdb-api:
+    - code: "java.method.addedToInterface"
+      new: "method java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell, byte[]> com.palantir.atlasdb.transaction.api.Transaction::getWithExpectedNumberOfCells(com.palantir.atlasdb.keyvalue.api.TableReference,\
+        \ java.util.Set<com.palantir.atlasdb.keyvalue.api.Cell>, int)"
+      justification: "Just introducing new Get API"
+  "0.885.0":
+    com.palantir.atlasdb:atlasdb-api:
+    - code: "java.method.addedToInterface"
+      new: "method com.google.common.util.concurrent.ListenableFuture<java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell,\
+        \ byte[]>> com.palantir.atlasdb.keyvalue.api.cache.TransactionScopedCache::getAsyncWithCachedRef(com.palantir.atlasdb.keyvalue.api.TableReference,\
+        \ java.util.Set<com.palantir.atlasdb.keyvalue.api.Cell>, java.util.function.Function<com.palantir.atlasdb.keyvalue.api.cache.TransactionScopedCache.CacheLookupResult,\
+        \ com.google.common.util.concurrent.ListenableFuture<java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell,\
+        \ byte[]>>>)"
+      justification: "just adding a new get method"
+    - code: "java.method.addedToInterface"
+      new: "method java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell, byte[]> com.palantir.atlasdb.keyvalue.api.cache.TransactionScopedCache::getWithCachedRef(com.palantir.atlasdb.keyvalue.api.TableReference,\
+        \ java.util.Set<com.palantir.atlasdb.keyvalue.api.Cell>, java.util.function.Function<com.palantir.atlasdb.keyvalue.api.cache.TransactionScopedCache.CacheLookupResult,\
+        \ com.google.common.util.concurrent.ListenableFuture<java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell,\
+        \ byte[]>>>)"
+      justification: "just adding a new get method"
+    - code: "java.method.addedToInterface"
+      new: "method java.util.Map<com.palantir.atlasdb.keyvalue.api.Cell, byte[]> com.palantir.atlasdb.transaction.api.Transaction::getWithExpectedNumberOfCells(com.palantir.atlasdb.keyvalue.api.TableReference,\
+        \ java.util.Set<com.palantir.atlasdb.keyvalue.api.Cell>, long)"
+      justification: "just adding a new get method"

atlasdb-api/src/main/java/com/palantir/atlasdb/keyvalue/api/cache/NoOpTransactionScopedCache.java

@@ -51,6 +51,14 @@ public Map<Cell, byte[]> get(
         return AtlasFutures.getUnchecked(getAsync(tableReference, cells, valueLoader));
     }
 
+    @Override
+    public Map<Cell, byte[]> getWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        return AtlasFutures.getUnchecked(getAsyncWithCachedRef(tableReference, cells, valueLoader));
+    }
+
     @Override
     public ListenableFuture<Map<Cell, byte[]>> getAsync(
             TableReference tableReference,
@@ -59,6 +67,14 @@ public ListenableFuture<Map<Cell, byte[]>> getAsync(
         return valueLoader.apply(cells);
     }
 
+    @Override
+    public ListenableFuture<Map<Cell, byte[]>> getAsyncWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        return valueLoader.apply(CacheLookupResult.of(Map.of(), cells));
+    }
+
     @Override
     public NavigableMap<byte[], RowResult<byte[]>> getRows(
             TableReference tableRef,

atlasdb-api/src/main/java/com/palantir/atlasdb/keyvalue/api/cache/ReadOnlyTransactionScopedCache.java

@@ -56,6 +56,14 @@ public Map<Cell, byte[]> get(
         return delegate.get(tableReference, cells, valueLoader);
     }
 
+    @Override
+    public Map<Cell, byte[]> getWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        return delegate.getWithCachedRef(tableReference, cells, valueLoader);
+    }
+
     @Override
     public ListenableFuture<Map<Cell, byte[]>> getAsync(
             TableReference tableReference,
@@ -64,6 +72,14 @@ public ListenableFuture<Map<Cell, byte[]>> getAsync(
         return delegate.getAsync(tableReference, cells, valueLoader);
     }
 
+    @Override
+    public ListenableFuture<Map<Cell, byte[]>> getAsyncWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        return delegate.getAsyncWithCachedRef(tableReference, cells, valueLoader);
+    }
+
     @Override
     public NavigableMap<byte[], RowResult<byte[]>> getRows(
             TableReference tableRef,

atlasdb-api/src/main/java/com/palantir/atlasdb/keyvalue/api/cache/TransactionScopedCache.java

@@ -27,6 +27,7 @@
 import java.util.NavigableMap;
 import java.util.Set;
 import java.util.function.Function;
+import org.immutables.value.Value;
 
 /**
  * The {@link LockWatchValueScopingCache} will provide one of these to every (relevant) transaction, and this will
@@ -60,11 +61,21 @@ Map<Cell, byte[]> get(
             Set<Cell> cells,
             Function<Set<Cell>, ListenableFuture<Map<Cell, byte[]>>> valueLoader);
 
+    Map<Cell, byte[]> getWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader);
+
     ListenableFuture<Map<Cell, byte[]>> getAsync(
             TableReference tableReference,
             Set<Cell> cells,
             Function<Set<Cell>, ListenableFuture<Map<Cell, byte[]>>> valueLoader);
 
+    ListenableFuture<Map<Cell, byte[]>> getAsyncWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader);
+
     /**
      * The cache will try to fulfil as much of the request as possible with cached values. In the case where some of the
      * columns are present in the cache for a row, the cellLoader will be used to read the remaining cells remotely.
@@ -93,4 +104,18 @@ NavigableMap<byte[], RowResult<byte[]>> getRows(
     HitDigest getHitDigest();
 
     TransactionScopedCache createReadOnlyCache(CommitUpdate commitUpdate);
+
+    @Value.Immutable
+    interface CacheLookupResult {
+        Map<Cell, CacheValue> cacheHits();
+
+        Set<Cell> missedCells();
+
+        static CacheLookupResult of(Map<Cell, CacheValue> cachedValues, Set<Cell> missedCells) {
+            return ImmutableCacheLookupResult.builder()
+                    .cacheHits(cachedValues)
+                    .missedCells(missedCells)
+                    .build();
+        }
+    }
 }

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/Transaction.java

@@ -16,6 +16,7 @@
 package com.palantir.atlasdb.transaction.api;
 
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.errorprone.annotations.RestrictedApi;
 import com.palantir.atlasdb.keyvalue.api.BatchColumnRangeSelection;
 import com.palantir.atlasdb.keyvalue.api.Cell;
 import com.palantir.atlasdb.keyvalue.api.ColumnRangeSelection;
@@ -24,6 +25,7 @@
 import com.palantir.atlasdb.keyvalue.api.RowResult;
 import com.palantir.atlasdb.keyvalue.api.TableReference;
 import com.palantir.atlasdb.spi.KeyValueServiceConfig;
+import com.palantir.atlasdb.transaction.api.annotations.ReviewedRestrictedApiUsage;
 import com.palantir.atlasdb.transaction.service.TransactionService;
 import com.palantir.common.annotation.Idempotent;
 import com.palantir.common.base.BatchingVisitable;
@@ -175,6 +177,41 @@ Iterator<Map.Entry<Cell, byte[]>> getSortedColumns(
     @Idempotent
     Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells);
 
+    /**
+     * Similar to {@link #get(TableReference, Set)}, but allows the caller to specify the number of expected present
+     * cells. This is important to allow clients that might have schemas that guarantees only a subset of the columns
+     * are present to still benefit from being able to skip the immutable timestamp lock check on non-empty reads of
+     * thoroughly swept tables.
+     *
+     * If we find values for the exact number of expected present cells, we'll skip the immutable timestamp lock check.
+     * If we find less values, we'll perform the immutable timestamp lock check.
+     * If we find more values, we'll throw an exception, as it means that either:
+     *  1 - There is a bug in the implementation of this method.
+     *  2 - The client is making an incorrect assumption about the maximum number of values that will be present, and
+     *      we could have returned empty values in the past when we should have not.
+     *
+     * WARNING: This method should only be used if you REALLY know what you're doing. Otherwise, you could have
+     * correctness issues by reading empty values when one is actually present if you don't use this method correctly.
+     *
+     * @param tableRef the table from which to get the values
+     * @param cells the cells for which we want to get the values
+     * @param expectedNumberOfPresentCells the maximum number of cells that are expected to be present.
+     * @return a {@link Map} from {@link Cell} to {@code byte[]} representing cell/value pairs
+     */
+    @RestrictedApi(
+            explanation = "This API is only meant to be used by AtlasDb proxies that want to make usage of the "
+                    + "performance improvement of skipping the immutable timestamp lock check on non-empty reads of "
+                    + "thoroughly swept tables, but their schema won't allow them to have non empty reads due to "
+                    + "columns being mutually exclusive, for example. So this API gives them a good way to specifying "
+                    + "the max number of values they'd ever expect in a get, which when met would allow them to skip "
+                    + "the lock check. Wrong usage of this API can cause correctness issues, so it's restricted.",
+            link = "https://github.com/palantir/atlasdb/pull/6655",
+            allowedOnPath = ".*/src/test/.*", // Unsafe behavior in tests is ok.
+            allowlistAnnotations = {ReviewedRestrictedApiUsage.class})
+    @Idempotent
+    Map<Cell, byte[]> getWithExpectedNumberOfCells(
+            TableReference tableRef, Set<Cell> cells, long expectedNumberOfPresentCells);
+
     /**
      * Gets the values associated for each cell in {@code cells} from table specified by {@code tableRef}. It is not
      * guaranteed that the actual implementations are in fact asynchronous.

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/annotations/ReviewedRestrictedApiUsage.java

@@ -0,0 +1,23 @@
+/*
+ * (c) Copyright 2023 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.transaction.api.annotations;
+
+/**
+ * Annotation to be used to indicate that usage of a restricted method is intentional and all the nuances about it
+ * are understood.
+ */
+public @interface ReviewedRestrictedApiUsage {}

atlasdb-api/src/main/java/com/palantir/atlasdb/transaction/api/exceptions/MoreCellsPresentThanExpectedException.java

@@ -0,0 +1,67 @@
+/*
+ * (c) Copyright 2023 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.transaction.api.exceptions;
+
+import com.palantir.atlasdb.keyvalue.api.Cell;
+import com.palantir.logsafe.Arg;
+import com.palantir.logsafe.Safe;
+import com.palantir.logsafe.SafeArg;
+import com.palantir.logsafe.SafeLoggable;
+import com.palantir.logsafe.Unsafe;
+import com.palantir.logsafe.UnsafeArg;
+import com.palantir.logsafe.exceptions.SafeExceptions;
+import java.util.List;
+import java.util.Map;
+
+public final class MoreCellsPresentThanExpectedException extends IllegalStateException implements SafeLoggable {
+
+    private static final String MESSAGE =
+            "KeyValueService returned more results than Get expected. This means there is a bug"
+                    + "either in the SnapshotTransaction implementation or in how the client is "
+                    + "using such method.";
+    private final List<Arg<?>> arguments;
+    private final Map<Cell, byte[]> fetchedCells;
+
+    public MoreCellsPresentThanExpectedException(Map<Cell, byte[]> fetchedCells, long expectedNumberOfCells) {
+        super(SafeExceptions.renderMessage(
+                MESSAGE, argsFrom(fetchedCells, expectedNumberOfCells).toArray(Arg<?>[]::new)));
+        this.fetchedCells = fetchedCells;
+        this.arguments = argsFrom(fetchedCells, expectedNumberOfCells);
+    }
+
+    public Map<Cell, byte[]> getFetchedCells() {
+        return fetchedCells;
+    }
+
+    @Override
+    public @Safe String getLogMessage() {
+        return MESSAGE;
+    }
+
+    @Override
+    public List<Arg<?>> getArgs() {
+        return arguments;
+    }
+
+    @Unsafe
+    private static List<Arg<?>> argsFrom(Map<Cell, byte[]> retrievedCells, long expectedNumberOfCells) {
+        return List.of(
+                SafeArg.of("expectedNumberOfCells", expectedNumberOfCells),
+                SafeArg.of("numberOfCellsRetrieved", retrievedCells.size()),
+                UnsafeArg.of("retrievedCells", retrievedCells));
+    }
+}

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/CachingTransaction.java

@@ -33,12 +33,14 @@
 import com.palantir.atlasdb.keyvalue.impl.Cells;
 import com.palantir.atlasdb.transaction.api.Transaction;
 import com.palantir.atlasdb.transaction.api.TransactionFailedException;
+import com.palantir.atlasdb.transaction.api.annotations.ReviewedRestrictedApiUsage;
 import com.palantir.atlasdb.transaction.service.TransactionService;
 import com.palantir.common.base.Throwables;
 import com.palantir.logsafe.SafeArg;
 import com.palantir.logsafe.logger.SafeLogger;
 import com.palantir.logsafe.logger.SafeLoggerFactory;
 import com.palantir.util.Pair;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -124,16 +126,42 @@ public Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells) {
             return getWithLoader(
                             tableRef,
                             cells,
-                            (tableReference, toRead) -> Futures.immediateFuture(super.get(tableReference, toRead)))
+                            (tableReference, _cachedCells, toRead) ->
+                                    Futures.immediateFuture(super.get(tableReference, toRead)))
                     .get();
-        } catch (InterruptedException | ExecutionException e) {
+        } catch (InterruptedException e) {
+            throw Throwables.rewrapAndThrowUncheckedException(e);
+        } catch (ExecutionException e) {
+            throw Throwables.rewrapAndThrowUncheckedException(e.getCause());
+        }
+    }
+
+    @ReviewedRestrictedApiUsage
+    @Override
+    public Map<Cell, byte[]> getWithExpectedNumberOfCells(
+            TableReference tableRef, Set<Cell> cells, long expectedNumberOfPresentCells) {
+        try {
+            return getWithLoader(tableRef, cells, (tableReference, cachedCells, toRead) -> {
+                        long nonEmptyValuesInCache = cachedCells.values().stream()
+                                .filter(value -> !Arrays.equals(value, PtBytes.EMPTY_BYTE_ARRAY))
+                                .count();
+                        long numberOfCellsExpectingValuePostCache =
+                                expectedNumberOfPresentCells - nonEmptyValuesInCache;
+
+                        return Futures.immediateFuture(super.getWithExpectedNumberOfCells(
+                                tableReference, toRead, numberOfCellsExpectingValuePostCache));
+                    })
+                    .get();
+        } catch (InterruptedException e) {
+            throw Throwables.rewrapAndThrowUncheckedException(e);
+        } catch (ExecutionException e) {
             throw Throwables.rewrapAndThrowUncheckedException(e.getCause());
         }
     }
 
     @Override
     public ListenableFuture<Map<Cell, byte[]>> getAsync(TableReference tableRef, Set<Cell> cells) {
-        return getWithLoader(tableRef, cells, super::getAsync);
+        return getWithLoader(tableRef, cells, (table, _cacheCells, cellsToLoad) -> super.getAsync(table, cellsToLoad));
     }
 
     private ListenableFuture<Map<Cell, byte[]>> getWithLoader(
@@ -156,7 +184,7 @@ private ListenableFuture<Map<Cell, byte[]>> getWithLoader(
         }
 
         return Futures.transform(
-                cellLoader.load(tableRef, toLoad),
+                cellLoader.load(tableRef, cacheHit, toLoad),
                 loadedCells -> {
                     cacheLoadedCells(tableRef, toLoad, loadedCells);
                     cacheHit.putAll(loadedCells);
@@ -271,6 +299,7 @@ public void abort() {
 
     @FunctionalInterface
     private interface CellLoader {
-        ListenableFuture<Map<Cell, byte[]>> load(TableReference tableReference, Set<Cell> toRead);
+        ListenableFuture<Map<Cell, byte[]>> load(
+                TableReference tableReference, Map<Cell, byte[]> cachedCells, Set<Cell> toRead);
     }
 }

atlasdb-client/src/main/java/com/palantir/atlasdb/transaction/impl/ForwardingTransaction.java

@@ -29,6 +29,7 @@
 import com.palantir.atlasdb.transaction.api.Transaction;
 import com.palantir.atlasdb.transaction.api.TransactionFailedException;
 import com.palantir.atlasdb.transaction.api.TransactionReadSentinelBehavior;
+import com.palantir.atlasdb.transaction.api.annotations.ReviewedRestrictedApiUsage;
 import com.palantir.atlasdb.transaction.service.TransactionService;
 import com.palantir.common.base.BatchingVisitable;
 import java.util.Iterator;
@@ -83,6 +84,13 @@ public Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells) {
         return delegate().get(tableRef, cells);
     }
 
+    @ReviewedRestrictedApiUsage
+    @Override
+    public Map<Cell, byte[]> getWithExpectedNumberOfCells(
+            TableReference tableRef, Set<Cell> cells, long expectedNumberOfPresentCells) {
+        return delegate().getWithExpectedNumberOfCells(tableRef, cells, expectedNumberOfPresentCells);
+    }
+
     @Override
     public BatchingVisitable<RowResult<byte[]>> getRange(TableReference tableRef, RangeRequest rangeRequest) {
         return delegate().getRange(tableRef, rangeRequest);