SnapshotTransaction stream refactor proof-of-concept

atlasdb-api/src/main/java/com/palantir/atlasdb/cell/api/TransactionKeyValueService.java

@@ -18,6 +18,7 @@
 
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.errorprone.annotations.MustBeClosed;
+import com.palantir.atlasdb.futures.AtlasFutures;
 import com.palantir.atlasdb.keyvalue.api.BatchColumnRangeSelection;
 import com.palantir.atlasdb.keyvalue.api.Cell;
 import com.palantir.atlasdb.keyvalue.api.ColumnRangeSelection;
@@ -30,8 +31,10 @@
 import com.palantir.atlasdb.keyvalue.api.Value;
 import com.palantir.common.base.ClosableIterator;
 import com.palantir.processors.AutoDelegate;
+import com.palantir.processors.DoNotDelegate;
 import com.palantir.util.paging.TokenBackedBasicResultsPage;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 
 /**
  * Key-Value API to be used with user data tables.
@@ -67,6 +70,12 @@ RowColumnRangeIterator getRowsColumnRange(
 
     ListenableFuture<Map<Cell, Value>> getAsync(TableReference tableRef, Map<Cell, Long> timestampByCell);
 
+    @DoNotDelegate
+    default CompletableFuture<Map<Cell, Value>> getCompletableAsync(
+            TableReference tableRef, Map<Cell, Long> timestampByCell) {
+        return AtlasFutures.toCompletableFuture(getAsync(tableRef, timestampByCell));
+    }
+
     Map<Cell, Long> getLatestTimestamps(TableReference tableRef, Map<Cell, Long> timestampByCell);
 
     void multiPut(Map<TableReference, ? extends Map<Cell, byte[]>> valuesByTable, long timestamp)

atlasdb-commons/src/main/java/com/palantir/atlasdb/futures/AtlasFutures.java

@@ -16,18 +16,23 @@
 
 package com.palantir.atlasdb.futures;
 
+import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.MoreExecutors;
 import com.palantir.common.base.Throwables;
 import com.palantir.common.streams.KeyedStream;
 import com.palantir.tracing.DeferredTracer;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public final class AtlasFutures {
     private AtlasFutures() {}
@@ -109,6 +114,67 @@ public static <R> R getUnchecked(Future<R> future) {
         }
     }
 
+    public static <R> CompletableFuture<R> toCompletableFuture(ListenableFuture<R> listenableFuture) {
+        CompletableFuture<R> completableFuture = new CompletableFuture<R>() {
+            @Override
+            public boolean cancel(boolean mayInterruptIfRunning) {
+                boolean result = listenableFuture.cancel(mayInterruptIfRunning);
+                super.cancel(mayInterruptIfRunning);
+                return result;
+            }
+        };
+        Futures.addCallback(
+                listenableFuture,
+                new FutureCallback<R>() {
+                    @Override
+                    public void onSuccess(R result) {
+                        completableFuture.complete(result);
+                    }
+
+                    @Override
+                    public void onFailure(Throwable t) {
+                        completableFuture.completeExceptionally(t);
+                    }
+                },
+                MoreExecutors.directExecutor());
+        return completableFuture;
+    }
+
+    public static <R> ListenableFuture<R> toListenableFuture(CompletableFuture<R> completableFuture) {
+        return new ListenableFuture<R>() {
+            @Override
+            public void addListener(Runnable listener, Executor executor) {
+                completableFuture.thenRunAsync(listener, executor);
+            }
+
+            @Override
+            public boolean cancel(boolean mayInterruptIfRunning) {
+                return completableFuture.cancel(mayInterruptIfRunning);
+            }
+
+            @Override
+            public boolean isCancelled() {
+                return completableFuture.isCancelled();
+            }
+
+            @Override
+            public boolean isDone() {
+                return completableFuture.isDone();
+            }
+
+            @Override
+            public R get() throws InterruptedException, ExecutionException {
+                return completableFuture.get();
+            }
+
+            @Override
+            public R get(long timeout, TimeUnit unit)
+                    throws InterruptedException, ExecutionException, TimeoutException {
+                return completableFuture.get(timeout, unit);
+            }
+        };
+    }
+
     private static Executor traceRestoringExecutor(Executor executor, String operation) {
         DeferredTracer deferredTracer = new DeferredTracer(operation);
         return command -> executor.execute(() -> deferredTracer.withTrace(() -> {

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

@@ -27,7 +27,6 @@
 import com.palantir.atlasdb.keyvalue.api.Value;
 import com.palantir.common.base.ClosableIterator;
 import com.palantir.common.base.ClosableIterators;
-import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.Map;
@@ -110,6 +109,6 @@ protected Iterator<Map.Entry<Cell, byte[]>> computeNext() {
 
     @FunctionalInterface
     interface PostFilterer {
-        Collection<Map.Entry<Cell, byte[]>> postFilter(Map<Cell, Value> rawResults);
+        Map<Cell, byte[]> postFilter(Map<Cell, Value> rawResults);
     }
 }