From 41c5355c0ccd0352dc0cbd5aee7d4fcea2d34d44 Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Thu, 7 Dec 2023 18:06:19 -0500 Subject: [PATCH] chore: add gax package private classes --- .../gaxx/retrying/AttemptCallable.java | 84 ++++ .../gaxx/retrying/RetryingCallable.java | 58 +++ .../RetryingServerStreamingCallable.java | 99 +++++ .../ServerStreamingAttemptCallable.java | 366 ++++++++++++++++++ 4 files changed, 607 insertions(+) create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/AttemptCallable.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingCallable.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingServerStreamingCallable.java create mode 100644 google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ServerStreamingAttemptCallable.java diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/AttemptCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/AttemptCallable.java new file mode 100644 index 0000000000..5bb8847e60 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/AttemptCallable.java @@ -0,0 +1,84 @@ +/* +* Copyright 2023 Google LLC +* +* 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 +* +* https://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.google.cloud.bigtable.gaxx.retrying; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.InternalApi; +import com.google.api.gax.retrying.NonCancellableFuture; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.common.base.Preconditions; +import java.util.concurrent.Callable; +import org.threeten.bp.Duration; + +// TODO: remove this once ApiResultRetryAlgorithm is added to gax. +/** + * A callable representing an attempt to make an RPC call. This class is used from {@link + * com.google.cloud.bigtable.gaxx.retrying.RetryingCallable}. + * + * @param request type + * @param response type + */ +@InternalApi +public class AttemptCallable implements Callable { + private final UnaryCallable callable; + private final RequestT request; + private final ApiCallContext originalCallContext; + + private volatile RetryingFuture externalFuture; + + AttemptCallable( + UnaryCallable callable, RequestT request, ApiCallContext callContext) { + this.callable = Preconditions.checkNotNull(callable); + this.request = Preconditions.checkNotNull(request); + this.originalCallContext = Preconditions.checkNotNull(callContext); + } + + public void setExternalFuture(RetryingFuture externalFuture) { + this.externalFuture = Preconditions.checkNotNull(externalFuture); + } + + @Override + public ResponseT call() { + ApiCallContext callContext = originalCallContext; + + try { + // Set the RPC timeout if the caller did not provide their own. + Duration rpcTimeout = externalFuture.getAttemptSettings().getRpcTimeout(); + if (!rpcTimeout.isZero() && callContext.getTimeout() == null) { + callContext = callContext.withTimeout(rpcTimeout); + } + + externalFuture.setAttemptFuture(new NonCancellableFuture()); + if (externalFuture.isDone()) { + return null; + } + + callContext + .getTracer() + .attemptStarted(request, externalFuture.getAttemptSettings().getOverallAttemptCount()); + + ApiFuture internalFuture = callable.futureCall(request, callContext); + externalFuture.setAttemptFuture(internalFuture); + } catch (Throwable e) { + externalFuture.setAttemptFuture(ApiFutures.immediateFailedFuture(e)); + } + + return null; + } +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingCallable.java new file mode 100644 index 0000000000..bfb5541624 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingCallable.java @@ -0,0 +1,58 @@ +/* + * Copyright 2023 Google LLC + * + * 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 + * + * https://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.google.cloud.bigtable.gaxx.retrying; + +import com.google.api.core.InternalApi; +import com.google.api.gax.retrying.RetryingExecutorWithContext; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.UnaryCallable; +import com.google.common.base.Preconditions; + +// TODO: remove this once ApiResultRetryAlgorithm is added to gax. +/** + * A UnaryCallable that will keep issuing calls to an inner callable until it succeeds or times out. + */ +@InternalApi +public class RetryingCallable extends UnaryCallable { + private final ApiCallContext callContextPrototype; + private final UnaryCallable callable; + private final RetryingExecutorWithContext executor; + + public RetryingCallable( + ApiCallContext callContextPrototype, + UnaryCallable callable, + RetryingExecutorWithContext executor) { + this.callContextPrototype = (ApiCallContext) Preconditions.checkNotNull(callContextPrototype); + this.callable = (UnaryCallable) Preconditions.checkNotNull(callable); + this.executor = (RetryingExecutorWithContext) Preconditions.checkNotNull(executor); + } + + public RetryingFuture futureCall(RequestT request, ApiCallContext inputContext) { + ApiCallContext context = this.callContextPrototype.nullToSelf(inputContext); + AttemptCallable retryCallable = + new AttemptCallable(this.callable, request, context); + RetryingFuture retryingFuture = + this.executor.createFuture(retryCallable, inputContext); + retryCallable.setExternalFuture(retryingFuture); + retryCallable.call(); + return retryingFuture; + } + + public String toString() { + return String.format("retrying(%s)", this.callable); + } +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingServerStreamingCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingServerStreamingCallable.java new file mode 100644 index 0000000000..5b693d9075 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/RetryingServerStreamingCallable.java @@ -0,0 +1,99 @@ +/* + * Copyright 2023 Google LLC + * + * 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 + * + * https://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.google.cloud.bigtable.gaxx.retrying; + +import static com.google.common.util.concurrent.MoreExecutors.directExecutor; + +import com.google.api.core.ApiFutureCallback; +import com.google.api.core.ApiFutures; +import com.google.api.core.InternalApi; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.retrying.ScheduledRetryingExecutor; +import com.google.api.gax.retrying.ServerStreamingAttemptException; +import com.google.api.gax.retrying.StreamResumptionStrategy; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; + +// TODO: remove this once ApiResultRetryAlgorithm is added to gax. +/** + * A ServerStreamingCallable that implements resumable retries. + * + *

Wraps a request, a {@link ResponseObserver} and an inner {@link ServerStreamingCallable} and + * coordinates retries between them. When the inner callable throws an error, this class will + * schedule retries using the configured {@link ScheduledRetryingExecutor}. + * + *

Streams can be resumed using a {@link StreamResumptionStrategy}. The {@link + * StreamResumptionStrategy} is notified of incoming responses and is expected to track the progress + * of the stream. Upon receiving an error, the {@link StreamResumptionStrategy} is asked to modify + * the original request to resume the stream. + */ +@InternalApi +public final class RetryingServerStreamingCallable + extends ServerStreamingCallable { + + private final ServerStreamingCallable innerCallable; + private final ScheduledRetryingExecutor executor; + private final StreamResumptionStrategy resumptionStrategyPrototype; + + public RetryingServerStreamingCallable( + ServerStreamingCallable innerCallable, + ScheduledRetryingExecutor executor, + StreamResumptionStrategy resumptionStrategyPrototype) { + this.innerCallable = innerCallable; + this.executor = executor; + this.resumptionStrategyPrototype = resumptionStrategyPrototype; + } + + @Override + public void call( + RequestT request, + final ResponseObserver responseObserver, + ApiCallContext context) { + + ServerStreamingAttemptCallable attemptCallable = + new ServerStreamingAttemptCallable<>( + innerCallable, + resumptionStrategyPrototype.createNew(), + request, + context, + responseObserver); + + RetryingFuture retryingFuture = executor.createFuture(attemptCallable, context); + attemptCallable.setExternalFuture(retryingFuture); + attemptCallable.start(); + + // Bridge the future result back to the external responseObserver + ApiFutures.addCallback( + retryingFuture, + new ApiFutureCallback() { + @Override + public void onFailure(Throwable throwable) { + // Make sure to unwrap the underlying ApiException + if (throwable instanceof ServerStreamingAttemptException) { + throwable = throwable.getCause(); + } + responseObserver.onError(throwable); + } + + @Override + public void onSuccess(Void ignored) { + responseObserver.onComplete(); + } + }, + directExecutor()); + } +} \ No newline at end of file diff --git a/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ServerStreamingAttemptCallable.java b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ServerStreamingAttemptCallable.java new file mode 100644 index 0000000000..2b8e7d39a5 --- /dev/null +++ b/google-cloud-bigtable/src/main/java/com/google/cloud/bigtable/gaxx/retrying/ServerStreamingAttemptCallable.java @@ -0,0 +1,366 @@ +/* + * Copyright 2023 Google LLC + * + * 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 + * + * https://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.google.cloud.bigtable.gaxx.retrying; + +import com.google.api.core.InternalApi; +import com.google.api.core.SettableApiFuture; +import com.google.api.gax.retrying.RetryingFuture; +import com.google.api.gax.retrying.ServerStreamingAttemptException; +import com.google.api.gax.retrying.StreamResumptionStrategy; +import com.google.api.gax.rpc.ApiCallContext; +import com.google.api.gax.rpc.ResponseObserver; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.api.gax.rpc.StateCheckingResponseObserver; +import com.google.api.gax.rpc.StreamController; +import com.google.common.base.Preconditions; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; + +// TODO: remove this once ApiResultRetryAlgorithm is added to gax. +/** + * A callable that generates Server Streaming attempts. At any one time, it is responsible for at + * most a single outstanding attempt. During an attempt, it proxies all incoming message to the + * outer {@link ResponseObserver} and the {@link StreamResumptionStrategy}. Once the attempt + * completes, the external {@link RetryingFuture} future is notified. If the {@link RetryingFuture} + * decides to retry the attempt, it will invoke {@link #call()}. + * + *

The lifecycle of this class is: + * + *

    + *
  1. The caller instantiates this class. + *
  2. The caller sets the {@link RetryingFuture} via {@link #setExternalFuture(RetryingFuture)}. + * The {@link RetryingFuture} will be responsible for scheduling future attempts. + *
  3. The caller calls {@link #start()}. This notifies the outer {@link ResponseObserver} that + * call is about to start. + *
  4. The outer {@link ResponseObserver} configures inbound flow control via the {@link + * StreamController} that it received in {@link ResponseObserver#onStart(StreamController)}. + *
  5. The attempt call is sent via the inner/upstream {@link ServerStreamingCallable}. + *
  6. A future representing the end state of the inner attempt is passed to the outer {@link + * RetryingFuture}. + *
  7. All messages received from the inner {@link ServerStreamingCallable} are recorded by the + * {@link StreamResumptionStrategy}. + *
  8. All messages received from the inner {@link ServerStreamingCallable} are forwarded to the + * outer {@link ResponseObserver}. + *
  9. Upon attempt completion (either success or failure) are communicated to the outer {@link + * RetryingFuture}. + *
  10. If the {@link RetryingFuture} decides to resume the RPC, it will invoke {@link #call()}, + * which will consult the {@link StreamResumptionStrategy} for the resuming request and + * restart the process at step 5. + *
  11. Once the {@link RetryingFuture} decides to stop the retry loop, it will notify the outer + * {@link ResponseObserver}. + *
+ * + *

This class is meant to be used as middleware between an outer {@link ResponseObserver} and an + * inner {@link ServerStreamingCallable}. As such it follows the general threading model of {@link + * ServerStreamingCallable}s: + * + *

    + *
  • {@code onStart} must be called in the same thread that invoked {@code call()} + *
  • The outer {@link ResponseObserver} can call {@code request()} and {@code cancel()} on this + * class' {@link StreamController} from any thread + *
  • The inner callable will serialize calls to {@code onResponse()}, {@code onError()} and + * {@code onComplete} + *
+ * + *

With this model in mind, this class only needs to synchronize access data that is shared + * between: the outer {@link ResponseObserver} (via this class' {@link StreamController}) and the + * inner {@link ServerStreamingCallable}: pendingRequests, cancellationCause and the current + * innerController. + * + * @param request type + * @param response type + */ +@InternalApi +public final class ServerStreamingAttemptCallable implements Callable { + private final Object lock = new Object(); + + private final ServerStreamingCallable innerCallable; + private final StreamResumptionStrategy resumptionStrategy; + private final RequestT initialRequest; + private ApiCallContext context; + private final ResponseObserver outerObserver; + + // Start state + private boolean autoFlowControl = true; + private boolean isStarted; + + // Outer state + private Throwable cancellationCause; + + private int pendingRequests; + + private RetryingFuture outerRetryingFuture; + + // Internal retry state + private int numAttempts; + + private StreamController innerController; + + private boolean seenSuccessSinceLastError; + private SettableApiFuture innerAttemptFuture; + + public ServerStreamingAttemptCallable( + ServerStreamingCallable innerCallable, + StreamResumptionStrategy resumptionStrategy, + RequestT initialRequest, + ApiCallContext context, + ResponseObserver outerObserver) { + this.innerCallable = innerCallable; + this.resumptionStrategy = resumptionStrategy; + this.initialRequest = initialRequest; + this.context = context; + this.outerObserver = outerObserver; + } + + /** Sets controlling {@link RetryingFuture}. Must be called be before {@link #start()}. */ + void setExternalFuture(RetryingFuture retryingFuture) { + Preconditions.checkState(!isStarted, "Can't change the RetryingFuture once the call has start"); + Preconditions.checkNotNull(retryingFuture, "RetryingFuture can't be null"); + + this.outerRetryingFuture = retryingFuture; + } + + /** + * 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; + + // Call the inner callable + call(); + } + + /** + * Sends the actual RPC. The request being sent will first be transformed by the {@link + * StreamResumptionStrategy}. + * + *

This method expects to be called by one thread at a time. Furthermore, it expects that the + * current RPC finished before the next time it's called. + */ + @Override + public Void call() { + Preconditions.checkState(isStarted, "Must be started first"); + + RequestT request = + (++numAttempts == 1) ? initialRequest : resumptionStrategy.getResumeRequest(initialRequest); + + // Should never happen. onAttemptError will check if ResumptionStrategy can create a resume + // request, + // which the RetryingFuture/StreamResumptionStrategy should respect. + Preconditions.checkState(request != null, "ResumptionStrategy returned a null request."); + + innerAttemptFuture = SettableApiFuture.create(); + seenSuccessSinceLastError = false; + + ApiCallContext attemptContext = context; + + if (!outerRetryingFuture.getAttemptSettings().getRpcTimeout().isZero() + && attemptContext.getTimeout() == null) { + attemptContext = + attemptContext.withTimeout(outerRetryingFuture.getAttemptSettings().getRpcTimeout()); + } + + attemptContext + .getTracer() + .attemptStarted(request, outerRetryingFuture.getAttemptSettings().getOverallAttemptCount()); + + innerCallable.call( + request, + new StateCheckingResponseObserver() { + @Override + public void onStartImpl(StreamController controller) { + onAttemptStart(controller); + } + + @Override + public void onResponseImpl(ResponseT response) { + onAttemptResponse(response); + } + + @Override + public void onErrorImpl(Throwable t) { + onAttemptError(t); + } + + @Override + public void onCompleteImpl() { + onAttemptComplete(); + } + }, + attemptContext); + + outerRetryingFuture.setAttemptFuture(innerAttemptFuture); + + return null; + } + + /** + * Called by the inner {@link ServerStreamingCallable} when the call is about to start. This will + * transfer unfinished state from the previous attempt. + * + * @see ResponseObserver#onStart(StreamController) + */ + private void onAttemptStart(StreamController controller) { + if (!autoFlowControl) { + controller.disableAutoInboundFlowControl(); + } + + Throwable localCancellationCause; + int numToRequest = 0; + + synchronized (lock) { + innerController = controller; + + localCancellationCause = this.cancellationCause; + + if (!autoFlowControl) { + numToRequest = pendingRequests; + } + } + + if (localCancellationCause != null) { + controller.cancel(); + } else if (numToRequest > 0) { + controller.request(numToRequest); + } + } + + /** + * 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(); + } + } + + /** + * 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; + } + + // Note: there is a race condition here where the count might go to the previous attempt's + // StreamController after it failed. But it doesn't matter, because the controller will just + // ignore it and the current controller will pick it up onStart. + if (localInnerController != null) { + localInnerController.request(count); + } + } + + /** Called when the inner callable has responses to deliver. */ + private void onAttemptResponse(ResponseT message) { + if (!autoFlowControl) { + synchronized (lock) { + pendingRequests--; + } + } + // Update local state to allow for future resume. + seenSuccessSinceLastError = true; + message = resumptionStrategy.processResponse(message); + // Notify the outer observer. + outerObserver.onResponse(message); + } + + /** + * Called when the current RPC fails. The error will be bubbled up to the outer {@link + * RetryingFuture} via the {@link #innerAttemptFuture}. + */ + private void onAttemptError(Throwable throwable) { + Throwable localCancellationCause; + synchronized (lock) { + localCancellationCause = cancellationCause; + } + + if (localCancellationCause != null) { + // Take special care to preserve the cancellation's stack trace. + innerAttemptFuture.setException(localCancellationCause); + } else { + // Wrap the original exception and provide more context for StreamingRetryAlgorithm. + innerAttemptFuture.setException( + new ServerStreamingAttemptException( + throwable, resumptionStrategy.canResume(), seenSuccessSinceLastError)); + } + } + + /** + * Called when the current RPC successfully completes. Notifies the outer {@link RetryingFuture} + * via {@link #innerAttemptFuture}. + */ + private void onAttemptComplete() { + innerAttemptFuture.set(null); + } +} \ No newline at end of file