Lift blocking server streaming calls into gax

Author: igorbernstein2

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcDirectServerStreamingCallable.java

@@ -36,8 +36,6 @@
 importcom.google.common.base.Preconditions;
 importio.grpc.ClientCall;
 importio.grpc.MethodDescriptor;
-importio.grpc.stub.ClientCalls;
-importjava.util.Iterator;
 
 /**
  * {@code GrpcDirectServerStreamingCallable} creates server-streaming gRPC calls.
@@ -69,11 +67,4 @@ public void call(
 newGrpcDirectStreamController<>(call, responseObserver);
 controller.start(request);
  }
-
-@Override
-publicIterator<ResponseT> blockingServerStreamingCall(RequestTrequest, ApiCallContextcontext) {
-Preconditions.checkNotNull(request);
-ClientCall<RequestT, ResponseT> call = GrpcClientCalls.newCall(descriptor, context);
-returnClientCalls.blockingServerStreamingCall(call, request);
- }
 }

gax-grpc/src/main/java/com/google/api/gax/grpc/GrpcServerStreamingRequestParamCallable.java

@@ -35,7 +35,6 @@
 importcom.google.api.gax.rpc.ResponseObserver;
 importcom.google.api.gax.rpc.ServerStreamingCallable;
 importcom.google.common.base.Preconditions;
-importjava.util.Iterator;
 
 /**
  * A {@code ServerStreamingCallable} that extracts values from the fields of the request and inserts
@@ -62,12 +61,6 @@ public void call(
 callable.call(request, responseObserver, contextWithParamsEncoder(request, context));
  }
 
-@Override
-publicIterator<ResponseT> blockingServerStreamingCall(RequestTrequest, ApiCallContextcontext) {
-returncallable.blockingServerStreamingCall(
-request, contextWithParamsEncoder(request, context));
- }
-
 privateApiCallContextcontextWithParamsEncoder(RequestTrequest, ApiCallContextinputContext) {
 returnGrpcCallContext.createDefault()
  .nullToSelf(inputContext)

gax-grpc/src/test/java/com/google/api/gax/grpc/GrpcDirectServerStreamingCallableTest.java

@@ -36,10 +36,11 @@
 importcom.google.api.gax.grpc.testing.InProcessServer;
 importcom.google.api.gax.rpc.ClientContext;
 importcom.google.api.gax.rpc.ResponseObserver;
+importcom.google.api.gax.rpc.ServerStream;
 importcom.google.api.gax.rpc.ServerStreamingCallable;
 importcom.google.api.gax.rpc.StreamController;
 importcom.google.api.gax.rpc.testing.FakeCallContext;
-importcom.google.common.collect.Iterators;
+importcom.google.common.collect.Lists;
 importcom.google.common.truth.Truth;
 importcom.google.type.Color;
 importcom.google.type.Money;
@@ -49,8 +50,6 @@
 importio.grpc.StatusRuntimeException;
 importio.grpc.inprocess.InProcessChannelBuilder;
 importjava.io.IOException;
-importjava.util.ArrayList;
-importjava.util.Iterator;
 importjava.util.List;
 importjava.util.concurrent.CancellationException;
 importjava.util.concurrent.CountDownLatch;
@@ -227,9 +226,8 @@ public void onComplete() {
 @Test
 publicvoidtestBlockingServerStreaming() throwsException {
 Colorrequest = Color.newBuilder().setRed(0.5f).build();
-Iterator<Money> response = streamingCallable.blockingServerStreamingCall(request);
-List<Money> responseData = newArrayList<>();
-Iterators.addAll(responseData, response);
+ServerStream<Money> response = streamingCallable.call(request);
+List<Money> responseData = Lists.newArrayList(response);
 
 Moneyexpected = Money.newBuilder().setCurrencyCode("USD").setUnits(127).build();
 Truth.assertThat(responseData).containsExactly(expected);

gax/src/main/java/com/google/api/gax/rpc/ServerStream.java

@@ -0,0 +1,287 @@
+/*
+ * Copyright 2017, Google LLC All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ * * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ * * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ * * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+packagecom.google.api.gax.rpc;
+
+importcom.google.api.core.BetaApi;
+importcom.google.api.core.InternalApi;
+importcom.google.common.base.Throwables;
+importcom.google.common.collect.Queues;
+importjava.util.Iterator;
+importjava.util.NoSuchElementException;
+importjava.util.concurrent.BlockingQueue;
+importjavax.annotation.Nonnull;
+
+/**
+ * A blocking Iterable-style wrapper around server stream responses.
+ *
+ * <p>This class asynchronously pulls responses from upstream via {@link
+ * StreamController#request(int)} and exposes them via its Iterator. The implementation is back
+ * pressure aware and uses a constant buffer of 1 item.
+ *
+ * <p>Please note that the stream can only be consumed once and must either be fully consumed or be
+ * canceled.
+ *
+ * <p>This class is not thread safe.
+ *
+ * <p>Example usage:
+ *
+ * <pre>{@code
+ * ServerStream<Item> stream = ...;
+ *
+ * for (Item item : stream) {
+ * System.out.println(item.id());
+ *
+ * // Allow for early termination
+ * if (item.id().equals("needle")) {
+ * stream.cancel();
+ * }
+ * }
+ * }</pre>
+ *
+ * @param <V> The type of each response.
+ */
+@BetaApi("The surface for streaming is not stable yet and may change in the future.")
+publicfinalclassServerStream<V> implementsIterable<V> {
+privatestaticfinalObjectEOF_MARKER = newObject();
+
+privatefinalQueuingResponseObserver<V> observer = newQueuingResponseObserver<>();
+privatefinalServerStreamIterator<V> iterator = newServerStreamIterator<>(observer);
+privatebooleanconsumed;
+
+@InternalApi("For use by ServerStreamingCallable only.")
+ServerStream() {}
+
+@InternalApi("For use by ServerStreamingCallable only.")
+ResponseObserver<V> observer() {
+returnobserver;
+ }
+
+/** {@inheritDoc} */
+@Override
+@Nonnull
+publicIterator<V> iterator() {
+if (consumed) {
+thrownewIllegalStateException("Iterator already consumed");
+ }
+consumed = true;
+
+returniterator;
+ }
+
+/**
+ * Returns true if the next call to the iterator's hasNext() or next() is guaranteed to be
+ * nonblocking.
+ *
+ * @return If the call on any of the iterator's methods is guaranteed to be nonblocking.
+ */
+publicbooleanisReady() {
+returniterator.isReady();
+ }
+
+/**
+ * Cleanly cancels a partially consumed stream. The associated iterator will return false for the
+ * hasNext() in the next iteration. This maintains the contract that an observed true from
+ * hasNext() will yield an item in next(), but afterwards will return false.
+ */
+publicvoidcancel() {
+observer.cancel();
+ }
+
+/**
+ * Internal implementation of a blocking Iterator, which will coordinate with the
+ * QueuingResponseObserver fetch new items from upstream. The Iterator expects the observer to
+ * request the first item, afterwards, new items will be requested when the current ones are
+ * consumed by next().
+ *
+ * @param <V> The type of items to be Iterated over.
+ */
+privatestaticfinalclassServerStreamIterator<V> implementsIterator<V> {
+privatefinalQueuingResponseObserver<V> observer;
+privateObjectlast;
+
+ServerStreamIterator(QueuingResponseObserver<V> observer) {
+this.observer = observer;
+ }
+
+/**
+ * Checks if the next call to {@code hasNext()} (and {@code next()}) will block.
+ *
+ * @return true if the next call is guaranteed to be nonblocking.
+ */
+booleanisReady() {
+returnlast != null || observer.isReady();
+ }
+
+/**
+ * Consumes the next response and asynchronously request the next response from the observer.
+ *
+ * @return The next response.
+ * @throws NoSuchElementException If the stream has been consumed or cancelled.
+ */
+@Override
+publicVnext() {
+if (!hasNext()) {
+thrownewNoSuchElementException();
+ }
+try {
+observer.request();
+@SuppressWarnings("unchecked")
+Vtmp = (V) last;
+returntmp;
+ } finally {
+last = null;
+ }
+ }
+
+/**
+ * Checks if the stream has been fully consumed or cancelled. This method will block until the
+ * observer enqueues another event (response or completion event). If the observer encountered
+ * an error, this method will propagte the error and put itself into an error where it will
+ * always return the same error.
+ *
+ * @return true If iterator has more responses.
+ */
+@Override
+publicbooleanhasNext() {
+if (last == null) {
+try {
+last = observer.getNext();
+ } catch (InterruptedExceptione) {
+Thread.currentThread().interrupt();
+thrownewRuntimeException(e);
+ }
+ }
+if (lastinstanceofThrowable) {
+Throwablethrowable = (Throwable) this.last;
+
+Throwables.throwIfUnchecked(throwable);
+thrownewRuntimeException(throwable);
+ }
+returnlast != EOF_MARKER;
+ }
+
+/** Unsupported. */
+@Override
+publicvoidremove() {
+thrownewUnsupportedOperationException();
+ }
+ }
+
+/**
+ * A back pressure aware bridge from a {@link ResponseObserver} to a {@link BlockingQueue}. The
+ * queue size is fixed to 1 item & a close signal. The observer will manage it's own flow control
+ * keeping the queue in one of 3 states:
+ *
+ * <ul>
+ * <li>empty: a item has been requested and we are awaiting the next item
+ * <li>1 item: an in progress stream with 1 item buffered
+ * <li>1 control signal: either a Throwable or an EOF_MARKER means that the stream is closed
+ * <li>1 item & 1 control signal: this is the last item of the stream
+ * </ul>
+ *
+ * The observer can also be abruptly cancelled, which cancels the underlying call and always
+ * returns an EOF_MARKER.
+ *
+ * @param <V> The item type.
+ */
+privatestaticfinalclassQueuingResponseObserver<V> implementsResponseObserver<V> {
+privatefinalBlockingQueue<Object> buffer = Queues.newArrayBlockingQueue(2);
+privateStreamControllercontroller;
+privatebooleanisCancelled;
+
+voidrequest() {
+controller.request(1);
+ }
+
+ObjectgetNext() throwsInterruptedException {
+if (isCancelled) {
+returnEOF_MARKER;
+ }
+returnbuffer.take();
+ }
+
+booleanisReady() {
+returnisCancelled || !buffer.isEmpty();
+ }
+
+/**
+ * Cancels the underlying RPC and causes getNext to always return EOF_MARKER. This can only be
+ * called after starting the underlying call.
+ */
+voidcancel() {
+isCancelled = true;
+controller.cancel();
+ }
+
+/**
+ * Before starting the RPC, disable automatic flow control and retain a reference to the
+ * controller.
+ *
+ * @param controller The controller for the stream.
+ */
+@Override
+publicvoidonStart(StreamControllercontroller) {
+this.controller = controller;
+controller.disableAutoInboundFlowControl();
+controller.request(1);
+ }
+
+/**
+ * Buffer the response. There should be at most 1 response in the buffer
+ *
+ * @param response The received response.
+ */
+@Override
+publicvoidonResponse(Vresponse) {
+buffer.add(response);
+ }
+
+/**
+ * Enqueue the error to be thrown later on. The error might occur w/o a request so the queue
+ * might grow to 2 elements, in that case the previous response will be consumed first.
+ *
+ * @param t The error occurred on the stream
+ */
+@Override
+publicvoidonError(Throwablet) {
+buffer.add(t);
+ }
+
+/**
+ * Enqueue a marker to notify the consumer that the stream is finished. In most situations this
+ * will cause the queue to grow to 2 elements: the requested response and an unsolicited
+ * completion marker.
+ */
+@Override
+publicvoidonComplete() {
+buffer.add(EOF_MARKER);
+ }
+ }
+}