/*
 * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
 * (the "License"). You may not use this work except in compliance with the License, which is
 * available at www.apache.org/licenses/LICENSE-2.0
 *
 * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied, as more fully set forth in the License.
 *
 * See the NOTICE file distributed with this work for information regarding copyright ownership.
 */

package alluxio;

import alluxio.exception.AlluxioException;
import alluxio.thrift.AlluxioService;
import alluxio.thrift.AlluxioTException;
import alluxio.thrift.ThriftIOException;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import org.apache.thrift.TException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

import javax.annotation.concurrent.ThreadSafe;

/**
 * The base class for clients that use {@link alluxio.network.connection.ThriftClientPool}.
 *
 * @param <C> the Alluxio service type
 */
@ThreadSafe
public abstract class AbstractThriftClient<C extends AlluxioService.Client> {

  private static final Logger LOG = LoggerFactory.getLogger(Constants.LOGGER_TYPE);

  private static final int RPC_MAX_NUM_RETRY = 30;

  /**
   * @return a Thrift service client
   */
  protected abstract C acquireClient() throws IOException;

  /**
   * @param client the client to release
   */
  protected abstract void releaseClient(C client);

  /**
   * The RPC to be executed in {@link #retryRPC(RpcCallable)}.
   *
   * @param <V> the return value of {@link #call(AlluxioService.Client)}
   * @param <C> the Alluxio service type
   */
  protected interface RpcCallable<V, C extends AlluxioService.Client> {
    /**
     * The task where RPC happens.
     *
     * @return RPC result
     * @throws TException when any exception defined in thrift happens
     */
    V call(C client) throws TException;
  }

  /**
   * Same with {@link RpcCallable} except that this RPC call throws {@link AlluxioTException} and
   * is to be executed in {@link #retryRPC(RpcCallableThrowsAlluxioTException)}.
   *
   * @param <V> the return value of {@link #call(AlluxioService.Client)}
   * @param <C> the Alluxio service type
   */
  protected interface RpcCallableThrowsAlluxioTException<V, C extends AlluxioService.Client> {
    /**
     * The task where RPC happens.
     *
     * @return RPC result
     * @throws AlluxioTException when any {@link AlluxioException} happens during RPC and is wrapped
     *         into {@link AlluxioTException}
     * @throws TException when any exception defined in thrift happens
     */
    V call(C client) throws AlluxioTException, TException;
  }

  /**
   * Tries to execute an RPC defined as a {@link RpcCallable}.
   *
   * @param rpc the RPC call to be executed
   * @param <V> type of return value of the RPC call
   * @return the return value of the RPC call
   * @throws IOException when retries exceeds {@link #RPC_MAX_NUM_RETRY} or some server
   *         side IOException occurred.
   */
  protected <V> V retryRPC(RpcCallable<V, C> rpc) throws IOException {
    int countingRetry = 0;
    TException exception = null;
    while ((countingRetry++) <= RPC_MAX_NUM_RETRY) {
      C client = acquireClient();
      try {
        return rpc.call(client);
      } catch (ThriftIOException e) {
        throw new IOException(e);
      } catch (AlluxioTException e) {
        throw Throwables.propagate(AlluxioException.fromThrift(e));
      } catch (TException e) {
        LOG.error(e.getMessage(), e);
        closeClient(client);
        exception = e;
      } finally {
        releaseClient(client);
      }
    }

    LOG.error("Failed after " + countingRetry + " retries.");
    Preconditions.checkNotNull(exception);
    throw new IOException(exception);
  }

  /**
   * Similar to {@link #retryRPC(RpcCallable)} except that the RPC call may throw
   * {@link AlluxioTException} and once it is thrown, it will be transformed into
   * {@link AlluxioException} and be thrown.
   *
   * @param rpc the RPC call to be executed
   * @param <V> type of return value of the RPC call
   * @return the return value of the RPC call
   * @throws AlluxioException when {@link AlluxioTException} is thrown by the RPC call
   * @throws IOException when retries exceeds {@link #RPC_MAX_NUM_RETRY} or some server
   *         side IOException occurred.
   */
  protected <V> V retryRPC(RpcCallableThrowsAlluxioTException<V, C> rpc)
      throws AlluxioException, IOException {
    int countingRetry = 0;
    TException exception = null;
    while ((countingRetry++) <= RPC_MAX_NUM_RETRY) {
      C client = acquireClient();
      try {
        return rpc.call(client);
      } catch (AlluxioTException e) {
        throw AlluxioException.fromThrift(e);
      } catch (ThriftIOException e) {
        throw new IOException(e);
      } catch (TException e) {
        LOG.error(e.getMessage(), e);
        closeClient(client);
        exception = e;
      } finally {
        releaseClient(client);
      }
    }

    LOG.error("Failed after " + countingRetry + " retries.");
    Preconditions.checkNotNull(exception);
    throw new IOException(exception);
  }

  /**
   * Close the given client.
   *
   * @param client the client to close
   */
  private void closeClient(C client) {
    client.getOutputProtocol().getTransport().close();
  }
}

@@ -560,32 +560,68 @@
 luxioException;%0A
+import alluxio.retry.CountingRetry;%0A
 import alluxio.t
@@ -3231,35 +3231,8 @@
 n %7B%0A
-    int countingRetry = 0;%0A
     
@@ -3260,33 +3260,57 @@
 = null;%0A    
-while ((c
+CountingRetry retryPolicy = new C
 ountingRetry
@@ -3301,39 +3301,33 @@
 ew CountingRetry
-++) %3C= 
+(
 RPC_MAX_NUM_RETR
@@ -3320,32 +3320,40 @@
 C_MAX_NUM_RETRY)
+;%0A    do
  %7B%0A      C clien
@@ -3771,32 +3771,68 @@
 );%0A      %7D%0A    %7D
+ while (retryPolicy.attemptRetry());
 %0A%0A    LOG.error(
@@ -3845,37 +3845,51 @@
 d after %22 + 
-countingRetry
+retryPolicy.getRetryCount()
  + %22 retries
@@ -4685,35 +4685,8 @@
 n %7B%0A
-    int countingRetry = 0;%0A
     
@@ -4722,17 +4722,41 @@
     
-while ((c
+CountingRetry retryPolicy = new C
 ount
@@ -4767,15 +4767,9 @@
 etry
-++) %3C= 
+(
 RPC_
@@ -4782,16 +4782,24 @@
 M_RETRY)
+;%0A    do
  %7B%0A     
@@ -5211,16 +5211,52 @@
  %7D%0A    %7D
+ while (retryPolicy.attemptRetry());
 %0A%0A    LO
@@ -5285,21 +5285,35 @@
 %22 + 
-countingRetry
+retryPolicy.getRetryCount()
  + %22

use CountingRetry class