BackoffHelper.java

/**
 * 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 org.apache.aurora.common.util;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;

import org.apache.aurora.common.base.ExceptionalSupplier;
import org.apache.aurora.common.quantity.Amount;
import org.apache.aurora.common.quantity.Time;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A utility for dealing with backoffs of retryable actions.
 *
 * <p>TODO(John Sirois): investigate synergies with BackoffDecider.
 *
 * @author John Sirois
 */
public class BackoffHelper {
  private static final Logger LOG = LoggerFactory.getLogger(BackoffHelper.class);

  private static final Amount<Long,Time> DEFAULT_INITIAL_BACKOFF = Amount.of(1L, Time.SECONDS);
  private static final Amount<Long,Time> DEFAULT_MAX_BACKOFF = Amount.of(1L, Time.MINUTES);

  private final Clock clock;
  private final BackoffStrategy backoffStrategy;

  /**
   * Creates a new BackoffHelper that uses truncated binary backoff starting at a 1 second backoff
   * and maxing out at a 1 minute backoff.
   */
  public BackoffHelper() {
    this(DEFAULT_INITIAL_BACKOFF, DEFAULT_MAX_BACKOFF);
  }

  /**
   * Creates a new BackoffHelper that uses truncated binary backoff starting at the given
   * {@code initialBackoff} and maxing out at the given {@code maxBackoff}.
   *
   * @param initialBackoff the initial amount of time to back off
   * @param maxBackoff the maximum amount of time to back off
   */
  public BackoffHelper(Amount<Long, Time> initialBackoff, Amount<Long, Time> maxBackoff) {
    this(new TruncatedBinaryBackoff(initialBackoff, maxBackoff));
  }

  /**
   * Creates a BackoffHelper that uses the given {@code backoffStrategy} to calculate backoffs
   * between retries.
   *
   * @param backoffStrategy the backoff strategy to use
   */
  public BackoffHelper(BackoffStrategy backoffStrategy) {
    this(Clock.SYSTEM_CLOCK, backoffStrategy);
  }

  @VisibleForTesting BackoffHelper(Clock clock, BackoffStrategy backoffStrategy) {
    this.clock = Preconditions.checkNotNull(clock);
    this.backoffStrategy = Preconditions.checkNotNull(backoffStrategy);
  }

  /**
   * Executes the given task using the configured backoff strategy until the task succeeds as
   * indicated by returning {@code true}.
   *
   * @param task the retryable task to execute until success
   * @throws InterruptedException if interrupted while waiting for the task to execute successfully
   * @throws BackoffStoppedException if the backoff stopped unsuccessfully
   * @throws E if the task throws
   */
  public <E extends Exception> void doUntilSuccess(final ExceptionalSupplier<Boolean, E> task)
      throws InterruptedException, BackoffStoppedException, E {
    doUntilResult(() -> {
      Boolean result = task.get();
      return Boolean.TRUE.equals(result) ? result : null;
    });
  }

  /**
   * Gets {@link BackoffStrategy} instance the BackoffHelper is initialized with.
   * @return instance of {@link BackoffStrategy} used by BackoffHelper.
   */
  public BackoffStrategy getBackoffStrategy() {
    return backoffStrategy;
  }

  /**
   * Executes the given task using the configured backoff strategy until the task succeeds as
   * indicated by returning a non-null value.
   *
   * @param task the retryable task to execute until success
   * @return the result of the successfully executed task
   * @throws InterruptedException if interrupted while waiting for the task to execute successfully
   * @throws BackoffStoppedException if the backoff stopped unsuccessfully
   * @throws E if the task throws
   */
  public <T, E extends Exception> T doUntilResult(ExceptionalSupplier<T, E> task)
      throws InterruptedException, BackoffStoppedException, E {
    T result = task.get(); // give an immediate try
    return (result != null) ? result : retryWork(task);
  }

  private <T, E extends Exception> T retryWork(ExceptionalSupplier<T, E> work)
      throws E, InterruptedException, BackoffStoppedException {
    long currentBackoffMs = 0;
    while (backoffStrategy.shouldContinue(currentBackoffMs)) {
      currentBackoffMs = backoffStrategy.calculateBackoffMs(currentBackoffMs);
      LOG.debug("Operation failed, backing off for " + currentBackoffMs + "ms");
      clock.waitFor(currentBackoffMs);

      T result = work.get();
      if (result != null) {
        return result;
      }
    }
    throw new BackoffStoppedException(String.format("Backoff stopped without succeeding."));
  }

  /**
   * Occurs after the backoff strategy should stop.
   */
  public static class BackoffStoppedException extends RuntimeException {
    public BackoffStoppedException(String msg) {
      super(msg);
    }
  }
}