Examples with StreamController com.google.api.gax.rpc.StreamController used on opensource projects

Example 11 with StreamController

use of com.google.api.gax.rpc.StreamController in project java-firestore by googleapis.

the class Query method internalStream.

private void internalStream(final QuerySnapshotObserver documentObserver, final long startTimeNanos, @Nullable final ByteString transactionId, @Nullable final Timestamp readTime) {
    RunQueryRequest.Builder request = RunQueryRequest.newBuilder();
    request.setStructuredQuery(buildQuery()).setParent(options.getParentPath().toString());
    if (transactionId != null) {
        request.setTransaction(transactionId);
    }
    if (readTime != null) {
        request.setReadTime(readTime.toProto());
    }
    Tracing.getTracer().getCurrentSpan().addAnnotation(TraceUtil.SPAN_NAME_RUNQUERY + ": Start", ImmutableMap.of("transactional", AttributeValue.booleanAttributeValue(transactionId != null)));
    final AtomicReference<QueryDocumentSnapshot> lastReceivedDocument = new AtomicReference<>();
    ResponseObserver<RunQueryResponse> observer = new ResponseObserver<RunQueryResponse>() {

        Timestamp readTime;

        boolean firstResponse;

        int numDocuments;

        @Override
        public void onStart(StreamController streamController) {
        }

        @Override
        public void onResponse(RunQueryResponse response) {
            if (!firstResponse) {
                firstResponse = true;
                Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: First response");
            }
            if (response.hasDocument()) {
                numDocuments++;
                if (numDocuments % 100 == 0) {
                    Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Received 100 documents");
                }
                Document document = response.getDocument();
                QueryDocumentSnapshot documentSnapshot = QueryDocumentSnapshot.fromDocument(rpcContext, Timestamp.fromProto(response.getReadTime()), document);
                documentObserver.onNext(documentSnapshot);
                lastReceivedDocument.set(documentSnapshot);
            }
            if (readTime == null) {
                readTime = Timestamp.fromProto(response.getReadTime());
            }
        }

        @Override
        public void onError(Throwable throwable) {
            QueryDocumentSnapshot cursor = lastReceivedDocument.get();
            if (shouldRetry(cursor, throwable)) {
                Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Retryable Error");
                Query.this.startAfter(cursor).internalStream(documentObserver, startTimeNanos, /* transactionId= */
                null, options.getRequireConsistency() ? cursor.getReadTime() : null);
            } else {
                Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Error");
                documentObserver.onError(throwable);
            }
        }

        @Override
        public void onComplete() {
            Tracing.getTracer().getCurrentSpan().addAnnotation("Firestore.Query: Completed", ImmutableMap.of("numDocuments", AttributeValue.longAttributeValue(numDocuments)));
            documentObserver.onCompleted(readTime);
        }

        boolean shouldRetry(DocumentSnapshot lastDocument, Throwable t) {
            if (transactionId != null) {
                // Transactional queries are retried via the transaction runner.
                return false;
            }
            if (lastDocument == null) {
                // failure are handled by Google Gax, which also implements backoff.
                return false;
            }
            if (!isRetryableError(t)) {
                return false;
            }
            if (rpcContext.getTotalRequestTimeout().isZero()) {
                return true;
            }
            Duration duration = Duration.ofNanos(rpcContext.getClock().nanoTime() - startTimeNanos);
            return duration.compareTo(rpcContext.getTotalRequestTimeout()) < 0;
        }
    };
    rpcContext.streamRequest(request.build(), observer, rpcContext.getClient().runQueryCallable());
}

Example 12 with StreamController

use of com.google.api.gax.rpc.StreamController in project gax-java by googleapis.

the class GrpcDirectServerStreamingCallableTest method testObserverErrorCancelsCall.

@Test
public void testObserverErrorCancelsCall() throws Throwable {
    final RuntimeException expectedCause = new RuntimeException("some error");
    final SettableApiFuture<Throwable> actualErrorF = SettableApiFuture.create();
    ResponseObserver<Money> moneyObserver = new StateCheckingResponseObserver<Money>() {

        @Override
        protected void onStartImpl(StreamController controller) {
        }

        @Override
        protected void onResponseImpl(Money response) {
            throw expectedCause;
        }

        @Override
        protected void onErrorImpl(Throwable t) {
            actualErrorF.set(t);
        }

        @Override
        protected void onCompleteImpl() {
            actualErrorF.set(null);
        }
    };
    streamingCallable.call(DEFAULT_REQUEST, moneyObserver);
    Throwable actualError = actualErrorF.get(500, TimeUnit.MILLISECONDS);
    Truth.assertThat(actualError).isInstanceOf(ApiException.class);
    Truth.assertThat(((ApiException) actualError).getStatusCode().getCode()).isEqualTo(StatusCode.Code.CANCELLED);
    // grpc is responsible for the immediate cancellation
    Truth.assertThat(actualError.getCause()).isInstanceOf(StatusRuntimeException.class);
    // and the client error is cause for grpc to cancel it
    Truth.assertThat(actualError.getCause().getCause()).isSameInstanceAs(expectedCause);
}

Example 13 with StreamController

use of com.google.api.gax.rpc.StreamController in project java-bigquerystorage by googleapis.

the class ReadRowsAttemptCallable method onRequest.

/**
 * Called when the outer {@link ResponseObserver} is ready for more data.
 *
 * @see StreamController#request(int)
 */
private void onRequest(int count) {
    Preconditions.checkState(!autoFlowControl, "Automatic flow control is enabled");
    Preconditions.checkArgument(count > 0, "Count must be > 0");
    final StreamController localInnerController;
    synchronized (lock) {
        int maxInc = Integer.MAX_VALUE - pendingRequests;
        count = Math.min(maxInc, count);
        pendingRequests += count;
        localInnerController = this.innerController;
    }
    // ignore it and the current controller will pick it up onStart.
    if (localInnerController != null) {
        localInnerController.request(count);
    }
}

Example 14 with StreamController

use of com.google.api.gax.rpc.StreamController in project java-bigquerystorage by googleapis.

the class ReadRowsAttemptCallable method start.

/**
 * Starts the initial call. The call is attempted on the caller's thread. Further call attempts
 * will be scheduled by the {@link RetryingFuture}.
 */
public void start() {
    Preconditions.checkState(!isStarted, "Already started");
    // Initialize the outer observer
    outerObserver.onStart(new StreamController() {

        @Override
        public void disableAutoInboundFlowControl() {
            Preconditions.checkState(!isStarted, "Can't disable auto flow control once the stream is started");
            autoFlowControl = false;
        }

        @Override
        public void request(int count) {
            onRequest(count);
        }

        @Override
        public void cancel() {
            onCancel();
        }
    });
    if (autoFlowControl) {
        synchronized (lock) {
            pendingRequests = Integer.MAX_VALUE;
        }
    }
    isStarted = true;
    // Propagate the totalTimeout as the overall stream deadline.
    Duration totalTimeout = outerRetryingFuture.getAttemptSettings().getGlobalSettings().getTotalTimeout();
    if (totalTimeout != null && context != null) {
        context = context.withTimeout(totalTimeout);
    }
    // Call the inner callable
    call();
}

Example 15 with StreamController

use of com.google.api.gax.rpc.StreamController in project java-bigquerystorage by googleapis.

the class ReadRowsAttemptCallable method onCancel.

/**
 * Called when the outer {@link ResponseObserver} wants to prematurely cancel the stream.
 *
 * @see StreamController#cancel()
 */
private void onCancel() {
    StreamController localInnerController;
    synchronized (lock) {
        if (cancellationCause != null) {
            return;
        }
        // NOTE: BasicRetryingFuture will replace j.u.c.CancellationExceptions with it's own,
        // which will not have the current stacktrace, so a special wrapper has be used here.
        cancellationCause = new ServerStreamingAttemptException(new CancellationException("User cancelled stream"), resumptionStrategy.canResume(), seenSuccessSinceLastError);
        localInnerController = innerController;
    }
    if (localInnerController != null) {
        localInnerController.cancel();
    }
}