Server streaming retries take 2

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

@@ -238,6 +238,11 @@ ServerStreamingCallable<RequestT, ResponseT> createServerStreamingCallable(
         new GrpcExceptionServerStreamingCallable<>(
             callable, streamingCallSettings.getRetryableCodes());
 
+    if (!streamingCallSettings.getRetryableCodes().isEmpty()
+        && streamingCallSettings.getRetrySettings().getMaxAttempts() > 1) {
+      callable = Callables.retrying(callable, streamingCallSettings, clientContext);
+    }
+
     return callable.withDefaultCallContext(clientContext.getDefaultCallContext());
   }
 

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

@@ -46,15 +46,15 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     return false;
   }
 
-  void cancelPrivately() {
+  public void cancelPrivately() {
     super.cancel(false);
   }
 
-  boolean setPrivately(ResponseT value) {
+  public boolean setPrivately(ResponseT value) {
     return super.set(value);
   }
 
-  boolean setExceptionPrivately(Throwable throwable) {
+  public boolean setExceptionPrivately(Throwable throwable) {
     return super.setException(throwable);
   }
 }

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2018, 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.InternalApi;
+
+/**
+ * A wrapper exception thrown by {@code ServerStreamingAttemptCallable} to communicate additional
+ * context to the {@link StreamingRetryAlgorithm} and to pass the original cancellation stack trace
+ * to {@code RetryingServerStreamingCallable}.
+ *
+ * <p>For internal use only - public for technical reasons.
+ */
+@InternalApi("For internal use only")
+public class ServerStreamingAttemptException extends RuntimeException {
+  private final boolean canResume;
+  private final boolean seenResponses;
+
+  public ServerStreamingAttemptException(
+      Throwable cause, boolean canResume, boolean seenResponses) {
+    super(cause);
+    this.canResume = canResume;
+    this.seenResponses = seenResponses;
+  }
+
+  /** If the {@link StreamResumptionStrategy} supports resuming after this error. */
+  public boolean canResume() {
+    return canResume;
+  }
+
+  /**
+   * If the current RPC attempt has seen any streamed messages. This is used as a signal by {@link
+   * StreamingRetryAlgorithm} to reset timers.
+   */
+  public boolean hasSeenResponses() {
+    return seenResponses;
+  }
+}

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

@@ -0,0 +1,64 @@
+/*
+ * Copyright 2018, 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.BetaApi;
+import com.google.common.base.Preconditions;
+
+/**
+ * Simplest implementation of a {@link StreamResumptionStrategy} which returns the initial request
+ * for unstarted streams.
+ */
+@BetaApi("The surface for streaming is not stable yet and may change in the future.")
+public final class SimpleStreamResumptionStrategy<RequestT, ResponseT>
+    implements StreamResumptionStrategy<RequestT, ResponseT> {
+  private boolean seenFirstResponse;
+
+  @Override
+  public StreamResumptionStrategy<RequestT, ResponseT> createNew() {
+    return new SimpleStreamResumptionStrategy<>();
+  }
+
+  @Override
+  public void onProgress(ResponseT response) {
+    seenFirstResponse = true;
+  }
+
+  @Override
+  public RequestT getResumeRequest(RequestT originalRequest) {
+    Preconditions.checkState(!seenFirstResponse, "Tried to resume an unresumeable stream.");
+    return originalRequest;
+  }
+
+  @Override
+  public boolean canResume() {
+    return !seenFirstResponse;
+  }
+}

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

@@ -0,0 +1,63 @@
+/*
+ * Copyright 2018, 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.BetaApi;
+
+/**
+ * This is part of the server streaming retry api. Its implementers are responsible for tracking the
+ * progress of the stream and calculating a request to resume it in case of an error.
+ *
+ * <p>Implementations don't have to be threadsafe because all of the calls will be serialized.
+ */
+@BetaApi("The surface for streaming is not stable yet and may change in the future.")
+public interface StreamResumptionStrategy<RequestT, ResponseT> {
+
+  /** Creates a new instance of this StreamResumptionStrategy without accumulated state */
+  StreamResumptionStrategy<RequestT, ResponseT> createNew();
+
+  /**
+   * Called by the {@code ServerStreamingAttemptCallable} when a response has been successfully
+   * received.
+   */
+  void onProgress(ResponseT response);
+
+  /**
+   * Called when a stream needs to be restarted, the implementation should generate a request that
+   * will yield a new stream whose first response would come right after the last response received
+   * by onProgress.
+   *
+   * @return A request that can be used to resume the stream.
+   */
+  RequestT getResumeRequest(RequestT originalRequest);
+
+  /** If a resume request can be created. */
+  boolean canResume();
+}

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

@@ -0,0 +1,100 @@
+/*
+ * Copyright 2018, 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.
+ */
+package com.google.api.gax.retrying;
+
+import com.google.api.core.InternalApi;
+import java.util.concurrent.CancellationException;
+
+/**
+ * The streaming retry algorithm, which makes decision based either on the thrown exception and the
+ * execution time settings of the previous attempt. This extends {@link RetryAlgorithm} to take
+ * additional information (provided by {@code ServerStreamingAttemptCallable}) into account.
+ *
+ * <p>This class is thread-safe.
+ *
+ * <p>Internal use only - public for technical reasons.
+ */
+@InternalApi("For internal use only")
+public final class StreamingRetryAlgorithm<ResponseT> extends RetryAlgorithm<ResponseT> {
+  public StreamingRetryAlgorithm(
+      ResultRetryAlgorithm<ResponseT> resultAlgorithm, TimedRetryAlgorithm timedAlgorithm) {
+    super(resultAlgorithm, timedAlgorithm);
+  }
+
+  /**
+   * {@inheritDoc}
+   *
+   * <p>The attempt settings will be reset if the stream attempt produced any messages.
+   */
+  @Override
+  public TimedAttemptSettings createNextAttempt(
+      Throwable prevThrowable, ResponseT prevResponse, TimedAttemptSettings prevSettings) {
+
+    if (prevThrowable instanceof ServerStreamingAttemptException) {
+      ServerStreamingAttemptException wrapper = (ServerStreamingAttemptException) prevThrowable;
+      prevThrowable = prevThrowable.getCause();
+
+      // If we have made progress in the last attempt, then reset the delays
+      if (wrapper.hasSeenResponses()) {
+        prevSettings =
+            createFirstAttempt()
+                .toBuilder()
+                .setFirstAttemptStartTimeNanos(prevSettings.getFirstAttemptStartTimeNanos())
+                .build();
+      }
+    }
+
+    return super.createNextAttempt(prevThrowable, prevResponse, prevSettings);
+  }
+
+  /**
+   * {@inheritDoc}
+   *
+   * <p>Ensures retries are only scheduled if the {@link StreamResumptionStrategy} in the {@code
+   * ServerStreamingAttemptCallable} supports it.
+   */
+  @Override
+  public boolean shouldRetry(
+      Throwable prevThrowable, ResponseT prevResponse, TimedAttemptSettings nextAttemptSettings)
+      throws CancellationException {
+
+    // Unwrap
+    if (prevThrowable instanceof ServerStreamingAttemptException) {
+      ServerStreamingAttemptException wrapper = (ServerStreamingAttemptException) prevThrowable;
+      prevThrowable = prevThrowable.getCause();
+
+      if (!wrapper.canResume()) {
+        return false;
+      }
+    }
+
+    return super.shouldRetry(prevThrowable, prevResponse, nextAttemptSettings);
+  }
+}

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

@@ -37,6 +37,7 @@
 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;
 
 /**
  * Class with utility methods to create callable objects using provided settings.
@@ -65,6 +66,37 @@ public static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> retrying(
         clientContext.getDefaultCallContext(), innerCallable, retryingExecutor);
   }
 
+  @BetaApi("The surface for streaming is not stable yet and may change in the future.")
+  public static <RequestT, ResponseT> ServerStreamingCallable<RequestT, ResponseT> retrying(
+      ServerStreamingCallable<RequestT, ResponseT> innerCallable,
+      ServerStreamingCallSettings<RequestT, ResponseT> callSettings,
+      ClientContext clientContext) {
+
+    StreamingRetryAlgorithm<Void> retryAlgorithm =
+        new StreamingRetryAlgorithm<>(
+            new ApiResultRetryAlgorithm<Void>(),
+            new ExponentialRetryAlgorithm(
+                callSettings.getRetrySettings(), clientContext.getClock()));
+
+    ScheduledRetryingExecutor<Void> retryingExecutor =
+        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
+
+    // NOTE: This creates a Watchdog per streaming API method. Ideally, there should only be a
+    // single Watchdog for the entire process, however that change would be fairly invasive and
+    // the cost of multiple Watchdogs is fairly small, since they all use the same executor. If this
+    // becomes an issue, the watchdog can be moved to ClientContext.
+    Watchdog<ResponseT> watchdog =
+        new Watchdog<>(
+            clientContext.getExecutor(),
+            clientContext.getClock(),
+            callSettings.getTimeoutCheckInterval(),
+            callSettings.getIdleTimeout());
+    watchdog.start();
+
+    return new RetryingServerStreamingCallable<>(
+        watchdog, innerCallable, retryingExecutor, callSettings.getResumptionStrategy());
+  }
+
   /**
    * Create a callable object that represents a batching API method. Designed for use by generated
    * code.