RetryPolicy.java

/*
 * Copyright 2021 TiKV Project Authors.
 *
 * 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.tikv.common.policy;

import com.google.common.collect.ImmutableSet;
import io.grpc.Status;
import io.prometheus.client.Counter;
import io.prometheus.client.Histogram;
import java.util.concurrent.Callable;
import org.tikv.common.exception.GrpcException;
import org.tikv.common.log.SlowLogSpan;
import org.tikv.common.operation.ErrorHandler;
import org.tikv.common.util.BackOffer;
import org.tikv.common.util.ConcreteBackOffer;
import org.tikv.common.util.HistogramUtils;

public abstract class RetryPolicy<RespT> {
  BackOffer backOffer = ConcreteBackOffer.newCopNextMaxBackOff();
  public static final Histogram GRPC_SINGLE_REQUEST_LATENCY =
      HistogramUtils.buildDuration()
          .name("client_java_grpc_single_requests_latency")
          .help("grpc request latency.")
          .labelNames("type", "cluster")
          .register();
  public static final Histogram CALL_WITH_RETRY_DURATION =
      HistogramUtils.buildDuration()
          .name("client_java_call_with_retry_duration")
          .help("callWithRetry duration.")
          .labelNames("type", "cluster")
          .register();
  public static final Counter GRPC_REQUEST_RETRY_NUM =
      Counter.build()
          .name("client_java_grpc_requests_retry_num")
          .help("grpc request retry num.")
          .labelNames("type", "cluster")
          .register();

  // handles PD and TiKV's error.
  private final ErrorHandler<RespT> handler;

  private final ImmutableSet<Status.Code> unrecoverableStatus =
      ImmutableSet.of(
          Status.Code.ALREADY_EXISTS, Status.Code.PERMISSION_DENIED,
          Status.Code.INVALID_ARGUMENT, Status.Code.NOT_FOUND,
          Status.Code.UNIMPLEMENTED, Status.Code.OUT_OF_RANGE,
          Status.Code.UNAUTHENTICATED, Status.Code.CANCELLED);

  RetryPolicy(ErrorHandler<RespT> handler) {
    this.handler = handler;
  }

  private void rethrowNotRecoverableException(Exception e) {
    Status status = Status.fromThrowable(e);
    if (unrecoverableStatus.contains(status.getCode())) {
      throw new GrpcException(e);
    }
  }

  public RespT callWithRetry(Callable<RespT> proc, String methodName, BackOffer backOffer) {
    String[] labels = new String[] {methodName, backOffer.getClusterId().toString()};
    Histogram.Timer callWithRetryTimer = CALL_WITH_RETRY_DURATION.labels(labels).startTimer();
    SlowLogSpan callWithRetrySlowLogSpan = backOffer.getSlowLog().start("callWithRetry");
    callWithRetrySlowLogSpan.addProperty("method", methodName);
    try {
      while (true) {
        RespT result = null;
        try {
          // add single request duration histogram
          Histogram.Timer requestTimer = GRPC_SINGLE_REQUEST_LATENCY.labels(labels).startTimer();
          SlowLogSpan slowLogSpan = backOffer.getSlowLog().start("gRPC");
          slowLogSpan.addProperty("method", methodName);
          try {
            result = proc.call();
          } finally {
            slowLogSpan.end();
            requestTimer.observeDuration();
          }
        } catch (Exception e) {
          rethrowNotRecoverableException(e);
          // Handle request call error
          backOffer.checkTimeout();
          boolean retry = handler.handleRequestError(backOffer, e);
          if (retry) {
            GRPC_REQUEST_RETRY_NUM.labels(labels).inc();
            continue;
          } else {
            return result;
          }
        }

        // Handle response error
        if (handler != null) {
          boolean retry = handler.handleResponseError(backOffer, result);
          if (retry) {
            GRPC_REQUEST_RETRY_NUM.labels(labels).inc();
            continue;
          }
        }
        return result;
      }
    } finally {
      callWithRetryTimer.observeDuration();
      callWithRetrySlowLogSpan.end();
    }
  }

  public interface Builder<T> {
    RetryPolicy<T> create(ErrorHandler<T> handler);
  }
}