Add RetryingContext

gax/src/main/java/com/google/api/gax/retrying/BasicRetryingFuture.java

@@ -34,6 +34,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.core.InternalApi;
 import com.google.common.util.concurrent.AbstractFuture;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.util.concurrent.Callable;
@@ -57,15 +58,20 @@ class BasicRetryingFuture<ResponseT> extends AbstractFuture<ResponseT>
   private final Callable<ResponseT> callable;
 
   private final RetryAlgorithm<ResponseT> retryAlgorithm;
+  private final RetryingContext retryingContext;
 
   private volatile TimedAttemptSettings attemptSettings;
 
   private volatile ApiFuture<ResponseT> latestCompletedAttemptResult;
   private volatile ApiFuture<ResponseT> attemptResult;
 
-  BasicRetryingFuture(Callable<ResponseT> callable, RetryAlgorithm<ResponseT> retryAlgorithm) {
+  BasicRetryingFuture(
+      Callable<ResponseT> callable,
+      RetryAlgorithm<ResponseT> retryAlgorithm,
+      RetryingContext retryingContext) {
     this.callable = checkNotNull(callable);
     this.retryAlgorithm = checkNotNull(retryAlgorithm);
+    this.retryingContext = retryingContext;
 
     this.attemptSettings = retryAlgorithm.createFirstAttempt();
 
@@ -77,6 +83,11 @@ class BasicRetryingFuture<ResponseT> extends AbstractFuture<ResponseT>
     super.addListener(new CompletionListener(), MoreExecutors.directExecutor());
   }
 
+  @InternalApi("VisibleForTesting")
+  RetryingContext getRetryingContext() {
+    return retryingContext;
+  }
+
   @Override
   public void setAttemptFuture(ApiFuture<ResponseT> attemptFuture) {
     try {

gax/src/main/java/com/google/api/gax/retrying/CallbackChainRetryingFuture.java

@@ -52,14 +52,15 @@
  * <p>This class is thread-safe.
  */
 class CallbackChainRetryingFuture<ResponseT> extends BasicRetryingFuture<ResponseT> {
-  private final RetryingExecutor<ResponseT> retryingExecutor;
+  private final ScheduledRetryingExecutor<ResponseT> retryingExecutor;
   private volatile AttemptCompletionListener attemptFutureCompletionListener;
 
   CallbackChainRetryingFuture(
       Callable<ResponseT> callable,
       RetryAlgorithm<ResponseT> retryAlgorithm,
-      RetryingExecutor<ResponseT> retryingExecutor) {
-    super(callable, retryAlgorithm);
+      ScheduledRetryingExecutor<ResponseT> retryingExecutor,
+      RetryingContext retryingContext) {
+    super(callable, retryAlgorithm, retryingContext);
     this.retryingExecutor = checkNotNull(retryingExecutor);
   }
 

gax/src/main/java/com/google/api/gax/retrying/DirectRetryingExecutor.java

@@ -33,9 +33,11 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.core.BetaApi;
 import java.io.InterruptedIOException;
 import java.nio.channels.ClosedByInterruptException;
 import java.util.concurrent.Callable;
+import javax.annotation.Nonnull;
 import org.threeten.bp.Duration;
 
 /**
@@ -46,7 +48,7 @@
  *
  * @param <ResponseT> response type
  */
-public class DirectRetryingExecutor<ResponseT> implements RetryingExecutor<ResponseT> {
+public class DirectRetryingExecutor<ResponseT> implements RetryingExecutorWithContext<ResponseT> {
 
   private final RetryAlgorithm<ResponseT> retryAlgorithm;
 
@@ -70,15 +72,30 @@ public DirectRetryingExecutor(RetryAlgorithm<ResponseT> retryAlgorithm) {
    */
   @Override
   public RetryingFuture<ResponseT> createFuture(Callable<ResponseT> callable) {
-    return new BasicRetryingFuture<>(callable, retryAlgorithm);
+    return createFuture(callable, RetryingContext.createDefault());
+  }
+
+  /**
+   * Creates a {@link RetryingFuture}, which is a facade, returned to the client code to wait for
+   * any retriable operation to complete. The future is bounded to {@code this} executor instance.
+   *
+   * @param callable the actual callable, which should be executed in a retriable context
+   * @return retrying future facade
+   */
+  @BetaApi("The surface for passing per operation state is not yet stable")
+  @Override
+  public RetryingFuture<ResponseT> createFuture(
+      @Nonnull Callable<ResponseT> callable, @Nonnull RetryingContext retryingContext) {
+    return new BasicRetryingFuture<>(callable, retryAlgorithm, retryingContext);
   }
 
   /**
    * Submits an attempt for execution in the current thread, causing the current thread to sleep for
    * the specified by the {@link RetryingFuture#getAttemptSettings()} amount of time. As result,
    * this method completes execution only after the specified {@code retryingFuture} completes.
    *
-   * @param retryingFuture the future previously returned by {@link #createFuture(Callable)}
+   * @param retryingFuture the future previously returned by {@link #createFuture(Callable,
+   *     RetryingContext)}
    * @return returns completed {@code retryingFuture}
    */
   @Override

gax/src/main/java/com/google/api/gax/retrying/RetryingContext.java

@@ -0,0 +1,70 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
+import com.google.api.gax.rpc.ApiCallContext;
+import com.google.auto.value.AutoValue;
+import javax.annotation.Nullable;
+
+/**
+ * Context for a retryable operation.
+ *
+ * <p>It provides state to individual {@link RetryingFuture}s via the {@link RetryingExecutor}.
+ */
+@BetaApi
+@InternalExtensionOnly
+@AutoValue
+public abstract class RetryingContext {
+  public static Builder newBuilder() {
+    return new AutoValue_RetryingContext.Builder();
+  }
+
+  public static RetryingContext createDefault() {
+    return newBuilder().build();
+  }
+
+  public static RetryingContext fromCallContext(@Nullable ApiCallContext callContext) {
+    Builder builder = newBuilder();
+
+    if (callContext != null) {
+      // TODO: transfer relevant state from the ApiCallContext
+    }
+    return builder.build();
+  }
+
+  @BetaApi
+  @InternalExtensionOnly
+  @AutoValue.Builder
+  public abstract static class Builder {
+    public abstract RetryingContext build();
+  }
+}

gax/src/main/java/com/google/api/gax/retrying/RetryingExecutorWithContext.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2018 Google LLC
+ *
+ * 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
+import java.util.concurrent.Callable;
+import javax.annotation.Nonnull;
+
+/**
+ * A {@link RetryingExecutor} that accepts a per-operation context.
+ *
+ * @see RetryingExecutor
+ */
+// TODO(igorbernstein2): Consider replacing this with a default implementation in RetryingExecutor
+// once support for java 7 is dropped
+@BetaApi("The surface for per invocation state is unstable and will probably change in the future")
+@InternalExtensionOnly
+public interface RetryingExecutorWithContext<ResponseT> extends RetryingExecutor<ResponseT> {
+  RetryingFuture<ResponseT> createFuture(
+      @Nonnull Callable<ResponseT> callable, @Nonnull RetryingContext retryingContext);
+}

gax/src/main/java/com/google/api/gax/retrying/ScheduledRetryingExecutor.java

@@ -31,6 +31,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
+import com.google.api.core.BetaApi;
 import com.google.api.core.ListenableFutureToApiFuture;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningScheduledExecutorService;
@@ -39,6 +40,7 @@
 import java.util.concurrent.RejectedExecutionException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
+import javax.annotation.Nonnull;
 
 /**
  * The retry executor which uses {@link ScheduledExecutorService} to schedule an attempt tasks.
@@ -53,7 +55,8 @@
  *
  * @param <ResponseT> response type
  */
-public class ScheduledRetryingExecutor<ResponseT> implements RetryingExecutor<ResponseT> {
+public class ScheduledRetryingExecutor<ResponseT>
+    implements RetryingExecutorWithContext<ResponseT> {
 
   private final RetryAlgorithm<ResponseT> retryAlgorithm;
   private final ListeningScheduledExecutorService scheduler;
@@ -81,13 +84,30 @@ public ScheduledRetryingExecutor(
    */
   @Override
   public RetryingFuture<ResponseT> createFuture(Callable<ResponseT> callable) {
-    return new CallbackChainRetryingFuture<>(callable, retryAlgorithm, this);
+    return createFuture(callable, RetryingContext.createDefault());
+  }
+
+  /**
+   * Creates a {@link RetryingFuture}, which is a facade, returned to the client code to wait for
+   * any retriable operation to complete. The returned future is bounded to {@code this} executor
+   * instance.
+   *
+   * @param callable the actual callable, which should be executed in a retriable context
+   * @param retryingContext the context for this operation
+   * @return retrying future facade
+   */
+  @BetaApi("The surface for passing per operation state is not yet stable")
+  @Override
+  public RetryingFuture<ResponseT> createFuture(
+      @Nonnull Callable<ResponseT> callable, @Nonnull RetryingContext retryingContext) {
+    return new CallbackChainRetryingFuture<>(callable, retryAlgorithm, this, retryingContext);
   }
 
   /**
    * Submits an attempt for execution in a different thread.
    *
-   * @param retryingFuture the future previously returned by {@link #createFuture(Callable)}
+   * @param retryingFuture the future previously returned by {@link #createFuture(Callable,
+   *     RetryingContext)}
    * @return submitted attempt future
    */
   @Override

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

@@ -35,7 +35,6 @@
 import com.google.api.gax.longrunning.OperationSnapshot;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
 import com.google.api.gax.retrying.RetryAlgorithm;
-import com.google.api.gax.retrying.RetryingExecutor;
 import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.api.gax.retrying.StreamingRetryAlgorithm;
 
@@ -64,7 +63,7 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> retrying(
             new ApiResultRetryAlgorithm<ResponseT>(),
             new ExponentialRetryAlgorithm(
                 callSettings.getRetrySettings(), clientContext.getClock()));
-    RetryingExecutor<ResponseT> retryingExecutor =
+    ScheduledRetryingExecutor<ResponseT> retryingExecutor =
         new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
     return new RetryingCallable<>(
         clientContext.getDefaultCallContext(), innerCallable, retryingExecutor);

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

@@ -36,7 +36,7 @@
 import com.google.api.gax.longrunning.OperationFuture;
 import com.google.api.gax.longrunning.OperationFutureImpl;
 import com.google.api.gax.longrunning.OperationSnapshot;
-import com.google.api.gax.retrying.RetryingExecutor;
+import com.google.api.gax.retrying.RetryingExecutorWithContext;
 import com.google.api.gax.retrying.RetryingFuture;
 
 /**
@@ -50,14 +50,14 @@ class OperationCallableImpl<RequestT, ResponseT, MetadataT>
     extends OperationCallable<RequestT, ResponseT, MetadataT> {
 
   private final UnaryCallable<RequestT, OperationSnapshot> initialCallable;
-  private final RetryingExecutor<OperationSnapshot> executor;
+  private final RetryingExecutorWithContext<OperationSnapshot> executor;
   private final LongRunningClient longRunningClient;
   private final ApiFunction<OperationSnapshot, ResponseT> responseTransformer;
   private final ApiFunction<OperationSnapshot, MetadataT> metadataTransformer;
 
   OperationCallableImpl(
       UnaryCallable<RequestT, OperationSnapshot> initialCallable,
-      RetryingExecutor<OperationSnapshot> executor,
+      RetryingExecutorWithContext<OperationSnapshot> executor,
       LongRunningClient longRunningClient,
       OperationCallSettings<RequestT, ResponseT, MetadataT> operationCallSettings) {
     this.initialCallable = checkNotNull(initialCallable);

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

@@ -29,7 +29,8 @@
  */
 package com.google.api.gax.rpc;
 
-import com.google.api.gax.retrying.RetryingExecutor;
+import com.google.api.gax.retrying.RetryingContext;
+import com.google.api.gax.retrying.RetryingExecutorWithContext;
 import com.google.api.gax.retrying.RetryingFuture;
 import com.google.common.base.Preconditions;
 
@@ -43,10 +44,11 @@
  */
 class RecheckingCallable<RequestT, ResponseT> extends UnaryCallable<RequestT, ResponseT> {
   private final UnaryCallable<RequestT, ResponseT> callable;
-  private final RetryingExecutor<ResponseT> executor;
+  private final RetryingExecutorWithContext<ResponseT> executor;
 
   RecheckingCallable(
-      UnaryCallable<RequestT, ResponseT> callable, RetryingExecutor<ResponseT> executor) {
+      UnaryCallable<RequestT, ResponseT> callable,
+      RetryingExecutorWithContext<ResponseT> executor) {
     this.callable = Preconditions.checkNotNull(callable);
     this.executor = Preconditions.checkNotNull(executor);
   }
@@ -56,7 +58,10 @@ public RetryingFuture<ResponseT> futureCall(RequestT request, ApiCallContext inp
     CheckingAttemptCallable<RequestT, ResponseT> checkingAttemptCallable =
         new CheckingAttemptCallable<>(callable);
 
-    RetryingFuture<ResponseT> retryingFuture = executor.createFuture(checkingAttemptCallable);
+    RetryingContext retryingContext = RetryingContext.fromCallContext(inputContext);
+
+    RetryingFuture<ResponseT> retryingFuture =
+        executor.createFuture(checkingAttemptCallable, retryingContext);
     checkingAttemptCallable.setExternalFuture(retryingFuture);
     checkingAttemptCallable.call();