/*
 * Copyright 2019 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
 *
 *       http://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.spanner;

import com.google.api.core.NanoClock;
import com.google.api.gax.retrying.ResultRetryAlgorithm;
import com.google.api.gax.retrying.RetrySettings;
import com.google.api.gax.retrying.TimedAttemptSettings;
import com.google.cloud.RetryHelper;
import com.google.cloud.RetryHelper.RetryHelperException;
import com.google.common.base.Throwables;
import io.grpc.Context;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import org.threeten.bp.Duration;

/**
 * Util class for retrying aborted transactions. This class is a wrapper around {@link RetryHelper}
 * that uses specific settings to only retry on aborted transactions, without a timeout and without
 * a cap on the number of retries.
 */
class SpannerRetryHelper {
  private static final RetrySettings txRetrySettings =
      RetrySettings.newBuilder()
          .setInitialRetryDelay(Duration.ofMillis(1000L))
          .setMaxRetryDelay(Duration.ofMillis(32000L))
          .setTotalTimeout(Duration.ofMillis(Integer.MAX_VALUE))
          .build();

  /** Executes the {@link Callable} and retries if it fails with an {@link AbortedException}. */
  static <T> T runTxWithRetriesOnAborted(Callable<T> callable) {
    try {
      return RetryHelper.runWithRetries(
          callable, txRetrySettings, new TxRetryAlgorithm<>(), NanoClock.getDefaultClock());
    } catch (RetryHelperException e) {
      if (e.getCause() != null) {
        Throwables.throwIfUnchecked(e.getCause());
      }
      throw e;
    }
  }

  private static class TxRetryAlgorithm<T> implements ResultRetryAlgorithm<T> {
    @Override
    public TimedAttemptSettings createNextAttempt(
        Throwable prevThrowable, T prevResponse, TimedAttemptSettings prevSettings) {
      if (prevThrowable != null) {
        long retryDelay = SpannerException.extractRetryDelay(prevThrowable);
        if (retryDelay > -1L) {
          return prevSettings
              .toBuilder()
              .setRandomizedRetryDelay(Duration.ofMillis(retryDelay))
              .build();
        }
      }
      return null;
    }

    @Override
    public boolean shouldRetry(Throwable prevThrowable, T prevResponse)
        throws CancellationException {
      if (Context.current().isCancelled()) {
        throw SpannerExceptionFactory.newSpannerExceptionForCancellation(Context.current(), null);
      }
      return prevThrowable != null
          && (prevThrowable instanceof AbortedException
              || prevThrowable instanceof com.google.api.gax.rpc.AbortedException);
    }
  }
}

@@ -945,29 +945,189 @@
 le.c
-ommon.base.Throwables
+loud.spanner.v1.stub.SpannerStub;%0Aimport com.google.cloud.spanner.v1.stub.SpannerStubSettings;%0Aimport com.google.common.base.Throwables;%0Aimport com.google.spanner.v1.RollbackRequest
 ;%0Aim
@@ -1545,290 +1545,699 @@
 r %7B%0A
+%0A
   
-private static final RetrySettings txRetrySettings =%0A      R
+/**%0A   * Use the same %7B@link RetrySettings%7D for retrying an aborted transaction as for retrying a %7B@link%0A   * RollbackRequest%7D. The %7B@link RollbackRequest%7D automatically uses the default r
 etry
-S
+ s
 ettings
-.newBuilder()%0A          .setInitialRetryDelay(Duration.ofMillis(1000L))%0A          .setMaxRetryDelay(Duration.ofMillis(32000L))%0A          .setTotalTimeout(Duration.ofMillis(Integer.MAX_VALUE))%0A          .build
+%0A   * defined for the %7B@link SpannerStub%7D. By referencing these settings, the retry settings for%0A   * retrying aborted transactions will also automatically be updated if the default retry settings%0A   * are updated.%0A   *%0A   * %3Cp%3EThese default %7B@link RetrySettings%7D are only used if no retry information is returned by the%0A   * %7B@link AbortedException%7D.%0A   */%0A  private static final RetrySettings txRetrySettings =%0A      SpannerStubSettings.newBuilder().rollbackSettings().getRetrySettings
 ();%0A

use default retry settings for aborted tx (#48)