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

@@ -175,6 +175,31 @@ 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
+     */
+    @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-client/src/main/java/com/palantir/atlasdb/transaction/impl/CachingTransaction.java

@@ -124,7 +124,28 @@ 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) {
+            throw Throwables.rewrapAndThrowUncheckedException(e.getCause());
+        }
+    }
+
+    @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 -> value != PtBytes.EMPTY_BYTE_ARRAY)
+                                .count();
+                        long numberOfCellsExpectingValuePostCache =
+                                expectedNumberOfPresentCells - nonEmptyValuesInCache;
+
+                        return Futures.immediateFuture(super.getWithExpectedNumberOfCells(
+                                tableReference, toRead, numberOfCellsExpectingValuePostCache));
+                    })
                     .get();
         } catch (InterruptedException | ExecutionException e) {
             throw Throwables.rewrapAndThrowUncheckedException(e.getCause());
@@ -133,7 +154,7 @@ public Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells) {
 
     @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 +177,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 +292,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

@@ -83,6 +83,12 @@ public Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells) {
         return delegate().get(tableRef, cells);
     }
 
+    @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);

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

@@ -64,6 +64,13 @@ public Map<Cell, byte[]> get(TableReference tableRef, Set<Cell> cells) {
         return delegate().get(tableRef, cells);
     }
 
+    @Override
+    public Map<Cell, byte[]> getWithExpectedNumberOfCells(
+            TableReference tableRef, Set<Cell> cells, long expectedNumberOfPresentCells) {
+        checkTableName(tableRef);
+        return delegate.getWithExpectedNumberOfCells(tableRef, cells, expectedNumberOfPresentCells);
+    }
+
     @Override
     public BatchingVisitable<RowResult<byte[]>> getRange(TableReference tableRef, RangeRequest rangeRequest) {
         checkTableName(tableRef);

atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/keyvalue/api/cache/TransactionScopedCacheImpl.java

@@ -43,10 +43,9 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.concurrent.ThreadSafe;
-import org.immutables.value.Value;
 
 @ThreadSafe
-final class TransactionScopedCacheImpl implements TransactionScopedCache {
+public final class TransactionScopedCacheImpl implements TransactionScopedCache {
     private final TransactionCacheValueStore valueStore;
     private final CacheMetrics metrics;
     private volatile boolean finalised = false;
@@ -56,7 +55,7 @@ private TransactionScopedCacheImpl(TransactionCacheValueStore valueStore, CacheM
         this.metrics = metrics;
     }
 
-    static TransactionScopedCache create(ValueCacheSnapshot snapshot, CacheMetrics metrics) {
+    public static TransactionScopedCache create(ValueCacheSnapshot snapshot, CacheMetrics metrics) {
         return new TransactionScopedCacheImpl(new TransactionCacheValueStoreImpl(snapshot), metrics);
     }
 
@@ -81,15 +80,34 @@ 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) {
+        ensureNotFinalised();
+        return AtlasFutures.getUnchecked(getAsyncWithCachedRef(tableReference, cells, valueLoader));
+    }
+
     @Override
     public ListenableFuture<Map<Cell, byte[]>> getAsync(
             TableReference tableReference,
             Set<Cell> cells,
             Function<Set<Cell>, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
         ensureNotFinalised();
+        return getAsyncWithCachedRef(
+                tableReference, cells, cacheLookupResult -> valueLoader.apply(cacheLookupResult.missedCells()));
+    }
+
+    @Override
+    public ListenableFuture<Map<Cell, byte[]>> getAsyncWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        ensureNotFinalised();
         // Short-cut all the logic below if the table is not watched.
         if (!valueStore.isWatched(tableReference)) {
-            return valueLoader.apply(cells);
+            return valueLoader.apply(CacheLookupResult.of(Map.of(), cells));
         }
 
         CacheLookupResult cacheLookup = cacheLookup(tableReference, cells);
@@ -98,7 +116,7 @@ public ListenableFuture<Map<Cell, byte[]>> getAsync(
             return Futures.immediateFuture(filterEmptyValues(cacheLookup.cacheHits()));
         } else {
             return Futures.transform(
-                    valueLoader.apply(cacheLookup.missedCells()),
+                    valueLoader.apply(cacheLookup),
                     uncachedValues -> processUncachedCells(
                             tableReference, cacheLookup.cacheHits(), cacheLookup.missedCells(), uncachedValues),
                     MoreExecutors.directExecutor());
@@ -246,18 +264,4 @@ private static Map<Cell, byte[]> filterEmptyValues(Map<Cell, CacheValue> snapsho
                 .map(value -> value.value().get())
                 .collectToMap();
     }
-
-    @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-impl-shared/src/main/java/com/palantir/atlasdb/keyvalue/api/cache/ValidatingTransactionScopedCache.java

@@ -89,13 +89,30 @@ 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,
             Set<Cell> cells,
             Function<Set<Cell>, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
+        return getAsyncWithCachedRef(
+                tableReference, cells, cacheLookupResult -> valueLoader.apply(cacheLookupResult.missedCells()));
+    }
+
+    @Override
+    public ListenableFuture<Map<Cell, byte[]>> getAsyncWithCachedRef(
+            TableReference tableReference,
+            Set<Cell> cells,
+            Function<CacheLookupResult, ListenableFuture<Map<Cell, byte[]>>> valueLoader) {
         if (shouldValidate()) {
-            ListenableFuture<Map<Cell, byte[]>> remoteReads = valueLoader.apply(cells);
+            ListenableFuture<Map<Cell, byte[]>> remoteReads = valueLoader.apply(CacheLookupResult.of(Map.of(), cells));
             ListenableFuture<Map<Cell, byte[]>> cacheReads = delegate.getAsync(
                     tableReference,
                     cells,
@@ -110,7 +127,7 @@ public ListenableFuture<Map<Cell, byte[]>> getAsync(
                     },
                     MoreExecutors.directExecutor());
         } else {
-            return delegate.getAsync(tableReference, cells, valueLoader);
+            return delegate.getAsyncWithCachedRef(tableReference, cells, valueLoader);
         }
     }