Streaming results for List Keys/Buckets, 2i, and MR

src/main/java/com/basho/riak/client/api/RiakClient.java

@@ -402,6 +402,28 @@ public <T,S> RiakFuture<T,S> executeAsync(RiakCommand<T,S> command)
         return command.executeAsync(cluster);
     }
 
+    /**
+     * Execute a StreamableRiakCommand asynchronously, and stream the results back before the command is complete.
+     * <p>
+     * Calling this method  causes the client to execute the provided StreamableRiakCommand
+     * asynchronously. It will immediately return a RiakFuture that contains an immediately available result that data
+     * will be streamed to. The RiakFuture will also keep track of the overall operation's progress with the
+     * {@see RiakFuture#isDone}, etc methods.
+     *
+     * @param <I> StreamableRiakCommand's immediate return type, available before the command/operation is complete.
+     * @param <S> The RiakCommand's query info type.
+     * @param command The RiakCommand to execute.
+     * @param timeoutMS The loading timeout in milliseconds for each result chunk.
+     *                  If the timeout is reached a {@see null} will be returned from the result's iterator,
+     *                  instead of blocking indefinitely.
+     * @return a RiakFuture for the operation
+     * @see RiakFuture
+     */
+    public <I,S> RiakFuture<I,S> executeAsyncStreaming(StreamableRiakCommand<I, ?, S> command, int timeoutMS)
+    {
+        return command.executeAsyncStreaming(cluster, timeoutMS);
+    }
+
     /**
      * Shut down the client and the underlying RiakCluster.
      * <p>

src/main/java/com/basho/riak/client/api/StreamableRiakCommand.java

@@ -0,0 +1,22 @@
+package com.basho.riak.client.api;
+
+import com.basho.riak.client.core.RiakCluster;
+import com.basho.riak.client.core.RiakFuture;
+
+/*
+ * The base class for all Streamable Riak Commands.
+ * Allows the user to either use {@link RiakCommand#executeAsync} and return a "batch-mode" result
+ * that is only available after the command is complete, or
+ * use {@link StreamableRiakCommand#executeAsyncStreaming} and return a "immediate" or "stream-mode" result
+ * that data will flow into.
+ * @param <S> The response type returned by "streaming mode" {@link executeAsyncStreaming}
+ * @param <R> The response type returned by the "batch mode" @{link executeAsync}
+ * @param <Q> The query info type
+ * @author Dave Rusek
+ * @author Brian Roach <roach at basho.com>
+ * @since 2.0
+ */
+public abstract class StreamableRiakCommand<S, R, Q> extends RiakCommand<R, Q>
+{
+    protected abstract RiakFuture<S, Q> executeAsyncStreaming(RiakCluster cluster, int timeout);
+}

src/main/java/com/basho/riak/client/api/commands/ChunkedResponseIterator.java

@@ -0,0 +1,122 @@
+package com.basho.riak.client.api.commands;
+
+import com.basho.riak.client.core.StreamingRiakFuture;
+import com.basho.riak.client.core.util.BinaryValue;
+
+import java.util.Iterator;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TransferQueue;
+import java.util.function.Function;
+
+public class ChunkedResponseIterator<FinalT, ChunkT extends Iterable<CoreT>, CoreT> implements Iterator<FinalT>
+{
+    private final int timeout;
+    private volatile BinaryValue continuation = null;
+    private final StreamingRiakFuture<ChunkT, ?> coreFuture;
+    private final TransferQueue<ChunkT> chunkQueue;
+    private final Function<CoreT, FinalT> createNext;
+    private final Function<ChunkT, Iterator<CoreT>> getNextIterator;
+    private final Function<ChunkT, BinaryValue> getContinuationFn;
+
+    private Iterator<CoreT> currentIterator = null;
+
+    public ChunkedResponseIterator(StreamingRiakFuture<ChunkT, ?> coreFuture,
+                                   int pollTimeout,
+                                   Function<CoreT, FinalT> createNextFn,
+                                   Function<ChunkT, Iterator<CoreT>> getNextIteratorFn)
+    {
+        this(coreFuture, pollTimeout, createNextFn, getNextIteratorFn, (x) -> null);
+    }
+
+    public ChunkedResponseIterator(StreamingRiakFuture<ChunkT, ?> coreFuture,
+                                   int pollTimeout,
+                                   Function<CoreT, FinalT> createNextFn,
+                                   Function<ChunkT, Iterator<CoreT>> getNextIteratorFn,
+                                   Function<ChunkT, BinaryValue> getContinuationFn)
+    {
+        this.timeout = pollTimeout;
+        this.coreFuture = coreFuture;
+        this.chunkQueue = coreFuture.getResultsQueue();
+        this.createNext = createNextFn;
+        this.getNextIterator = getNextIteratorFn;
+        this.getContinuationFn = getContinuationFn;
+        loadNextChunkIterator();
+    }
+
+    @Override
+    public boolean hasNext()
+    {
+        if(currentIteratorHasNext())
+        {
+            return true;
+        }
+
+        loadNextChunkIterator();
+
+        return currentIteratorHasNext();
+    }
+
+    private boolean currentIteratorHasNext()
+    {
+        return currentIterator != null && currentIterator.hasNext();
+    }
+
+    private boolean possibleChunksRemaining()
+    {
+        // Chunks may remain if :
+        // Core Operation Not Done OR items still in chunk Queue
+        return !coreFuture.isDone() || !chunkQueue.isEmpty();
+    }
+
+    public boolean hasContinuation()
+    {
+        return continuation != null || possibleChunksRemaining();
+    }
+
+    public BinaryValue getContinuation()
+    {
+        return continuation;
+    }
+
+    @Override
+    public synchronized FinalT next()
+    {
+        return createNext.apply(currentIterator.next());
+    }
+
+    private void loadNextChunkIterator()
+    {
+        this.currentIterator = null;
+        boolean populatedChunkLoaded = false;
+
+        try
+        {
+            while (!populatedChunkLoaded && possibleChunksRemaining())
+            {
+                final ChunkT nextChunk = chunkQueue.poll(timeout, TimeUnit.MILLISECONDS);
+
+                if (nextChunk != null)
+                {
+                    this.currentIterator = getNextIterator.apply(nextChunk);
+                    populatedChunkLoaded = currentIteratorHasNext();
+
+                    loadContinuation(nextChunk);
+                }
+            }
+        }
+        catch (InterruptedException e)
+        {
+            Thread.currentThread().interrupt();
+            throw new RuntimeException(e);
+        }
+    }
+
+    private void loadContinuation(ChunkT nextChunk)
+    {
+        final BinaryValue fetchedContinuation = getContinuationFn.apply(nextChunk);
+        if(this.continuation == null && fetchedContinuation != null)
+        {
+            this.continuation = fetchedContinuation;
+        }
+    }
+}

src/main/java/com/basho/riak/client/api/commands/ImmediateCoreFutureAdapter.java

@@ -0,0 +1,83 @@
+/*
+ * Copyright 2014 Brian Roach <roach at basho dot com>.
+ *
+ * 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.basho.riak.client.api.commands;
+
+import com.basho.riak.client.core.RiakFuture;
+
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+
+/**
+ * Used when the converted response is available before the core future is complete.
+ *
+ * @param <T> The core response type.
+ * @param <S> The core query info type.
+ * @param <T2> The converted response type.
+ * @param <S2> The converted query info type.
+ *
+ * @author Alex Moore <amoore at basho dot com>
+ * @since 2.1
+ */
+public abstract class ImmediateCoreFutureAdapter<T2,S2,T,S> extends CoreFutureAdapter<T2,S2,T,S>
+{
+    private final T2 immediateResponse;
+
+    protected ImmediateCoreFutureAdapter(RiakFuture<T, S> coreFuture, T2 immediateResponse)
+    {
+        super(coreFuture);
+        this.immediateResponse = immediateResponse;
+    }
+
+    @Override
+    public T2 get() throws InterruptedException, ExecutionException
+    {
+        return immediateResponse;
+    }
+
+    @Override
+    public T2 get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException
+    {
+        return immediateResponse;
+    }
+
+    @Override
+    public T2 getNow()
+    {
+        return immediateResponse;
+    }
+
+    @Override
+    protected T2 convertResponse(T unused) { return null; }
+
+    @Override
+    protected abstract S2 convertQueryInfo(S coreQueryInfo);
+
+    public static abstract class SameQueryInfo<T2,S,T> extends ImmediateCoreFutureAdapter<T2,S,T,S>
+    {
+        protected SameQueryInfo(RiakFuture<T, S> coreFuture, T2 immediateResponse)
+        {
+            super(coreFuture, immediateResponse);
+        }
+
+        @Override
+        protected S convertQueryInfo(S coreQueryInfo)
+        {
+            return coreQueryInfo;
+        }
+    }
+}

src/main/java/com/basho/riak/client/api/commands/buckets/ListBuckets.java

@@ -15,10 +15,13 @@
  */
 package com.basho.riak.client.api.commands.buckets;
 
-import com.basho.riak.client.api.RiakCommand;
+import com.basho.riak.client.api.commands.ChunkedResponseIterator;
 import com.basho.riak.client.api.commands.CoreFutureAdapter;
+import com.basho.riak.client.api.StreamableRiakCommand;
+import com.basho.riak.client.api.commands.ImmediateCoreFutureAdapter;
 import com.basho.riak.client.core.RiakCluster;
 import com.basho.riak.client.core.RiakFuture;
+import com.basho.riak.client.core.StreamingRiakFuture;
 import com.basho.riak.client.core.operations.ListBucketsOperation;
 import com.basho.riak.client.core.query.Namespace;
 import com.basho.riak.client.core.util.BinaryValue;
@@ -42,7 +45,7 @@
  * @author Dave Rusek <drusek at basho dot com>
  * @since 2.0
  */
-public final class ListBuckets extends RiakCommand<ListBuckets.Response, BinaryValue>
+public final class ListBuckets extends StreamableRiakCommand<ListBuckets.StreamingResponse, ListBuckets.Response, BinaryValue>
 {
     private final int timeout;
     private final BinaryValue type;
@@ -57,7 +60,7 @@ public final class ListBuckets extends RiakCommand<ListBuckets.Response, BinaryV
     protected RiakFuture<Response, BinaryValue> executeAsync(RiakCluster cluster)
     {
         RiakFuture<ListBucketsOperation.Response, BinaryValue> coreFuture =
-            cluster.execute(buildCoreOperation());
+            cluster.execute(buildCoreOperation(false));
 
         CoreFutureAdapter<ListBuckets.Response, BinaryValue, ListBucketsOperation.Response, BinaryValue> future =
             new CoreFutureAdapter<ListBuckets.Response, BinaryValue, ListBucketsOperation.Response, BinaryValue>(coreFuture)
@@ -78,7 +81,22 @@ protected BinaryValue convertQueryInfo(BinaryValue coreQueryInfo)
         return future;
     }
 
-    private ListBucketsOperation buildCoreOperation()
+    protected RiakFuture<StreamingResponse, BinaryValue> executeAsyncStreaming(RiakCluster cluster, int timeout)
+    {
+        StreamingRiakFuture<ListBucketsOperation.Response, BinaryValue> coreFuture =
+                cluster.execute(buildCoreOperation(true));
+
+        final StreamingResponse streamingResponse = new StreamingResponse(type, timeout, coreFuture);
+
+        ImmediateCoreFutureAdapter.SameQueryInfo<StreamingResponse, BinaryValue, ListBucketsOperation.Response> future =
+                new ImmediateCoreFutureAdapter.SameQueryInfo<StreamingResponse, BinaryValue, ListBucketsOperation.Response>
+                        (coreFuture, streamingResponse) {};
+
+        coreFuture.addListener(future);
+        return future;
+    }
+
+    private ListBucketsOperation buildCoreOperation(boolean streamResults)
     {
         ListBucketsOperation.Builder builder = new ListBucketsOperation.Builder();
         if (timeout > 0)
@@ -91,6 +109,8 @@ private ListBucketsOperation buildCoreOperation()
             builder.withBucketType(type);
         }
 
+        builder.streamResults(streamResults);
+
         return builder.build();
     }
 
@@ -126,6 +146,29 @@ public Iterator<Namespace> iterator()
         }
     }
 
+    public static class StreamingResponse extends Response
+    {
+        private final ChunkedResponseIterator<Namespace, ListBucketsOperation.Response, BinaryValue>
+                chunkedResponseIterator;
+
+        StreamingResponse(BinaryValue type,
+                          int pollTimeout,
+                          StreamingRiakFuture<ListBucketsOperation.Response, BinaryValue> coreFuture)
+        {
+            super(type, null);
+            chunkedResponseIterator = new ChunkedResponseIterator<>(coreFuture,
+                                                                    pollTimeout,
+                                                                    (bucketName) -> new Namespace(type, bucketName),
+                                                                    (response) -> response.getBuckets().iterator());
+        }
+
+        @Override
+        public Iterator<Namespace> iterator()
+        {
+            return chunkedResponseIterator;
+        }
+    }
+
     private static class Itr implements Iterator<Namespace>
     {
         private final Iterator<BinaryValue> iterator;
@@ -207,4 +250,5 @@ public ListBuckets build()
             return new ListBuckets(this);
         }
     }
+
 }